Where is the ball? Why AI models lose track, and how to make them stop

Julien Brandoit, Arthur Fyon, Thomas Braipson, Tom Clara, Florent De Geeter, Pierre Sacré, Damien Ernst & Guillaume Drion

Link: https://hdl.handle.net/2268/349312

A sequence model is a program that reads its input one piece at a time (a word, a move, a line of code) and has to make sense of the whole. One of the most basic things it must do is keep track of something. Here is a small puzzle:

Three cups on a table, a ball under the middle one. Swap the left and the middle. Swap the middle and the right. Swap the left and the middle again. [...] Where is the ball?

To answer it, you carry one fact forward (which cup hides the ball) and update it every time a new swap comes in. The same skill lets you follow which variable holds which value through a hundred lines of code, or work out who “she” refers to after three pages of dialogue.

Recurrent networks were built for exactly this. A recurrent network keeps a small memory, a list of numbers, and rewrites it after every new piece of input. The catch is that it has to read the sequence strictly in order, one step at a time, which makes training slow. (Training is how a model learns: it sees many examples and adjusts its internal numbers, called weights, until its answers come out right.)

Newer recurrent designs have largely replaced classical recurrent networks: state space models such as Mamba, gated linear recurrent networks, and linear attention. The details differ, but they share one trick: their memory update is affine. An affine update follows a fixed recipe: multiply the old memory by some numbers, then add something new. Nothing more elaborate happens in between. Because of that, a computer can handle the whole sequence at once, in parallel, on a GPU (a chip, originally made for graphics, that performs thousands of calculations at the same time). On paper, affine models lose nothing in exchange, since the theory says they can keep track of any situation with a fixed, finite number of possibilities, like the three positions of the ball. In practice they can’t. Train one on sequences of sixty-four swaps and it follows the ball perfectly up to sixty-four, starts getting answers wrong somewhere past a hundred, and is guessing soon after.

The natural reaction is to blame training and reach for more data, more training steps, or a better curriculum (the order in which examples are shown, usually easy ones first). Training explains some of the failure. Part of it is arithmetic, though, and that part never goes away. A computer stores every number with a fixed number of digits, so it rounds the result of every calculation. Each step lands slightly off target, and these small errors pile up from one step to the next.

We show that staying correct at every length asks each update to do two opposite things at once. It has to pull the memory back together, so that this step rounding error does not carry over into the next one. It also has to keep genuinely different states far apart (a state is one of the possible situations, such as “ball under the left cup”), so that rounding never makes two of them look alike. An affine update has a single rate at which it squeezes or stretches the memory, and one rate cannot do both jobs. Set it to squeeze, and distinct states collapse onto each other. Set it to stretch, and the errors grow along with the states.

Our main result is that making the model bigger does not help. Whatever fixed number of digits the computer uses, affine models can only track states that depend on the last handful of inputs, however wide (more numbers in each processing stage) or deep (more stages, called layers) you build them. Training longer or buying more GPUs will not push that limit back. In effect, the model has a hidden context window. A model context window is how much past input it can take into account, and it is normally stated up front. Here, old information sinks below the level of rounding noise at some length nobody tells you, and the model gives no warning when that happens.

The technique for processing a sequence in parallel is called a parallel scan. It splits the sequence into chunks, summarizes what each chunk does to the memory, then merges neighbouring summaries two at a time until a single one remains. A parallel scan does not need the update to be affine. All it needs is a compact summary of what a stretch of input does to the memory, and a cheap way to merge two summaries into one. An affine update is one way to get those, and there are others. If the memory only ever takes a handful of clean, well-separated values, then the effect of a stretch of input fits in a lookup table: “if the memory was A, it is now C; if it was B, it is now A”, and so on. Two such tables merge by following the arrows from the first into the second.

We built the Neural Finite-State Machine (NFSM) on that idea. A finite-state machine is about the simplest machine there is: it is always in one of a fixed list of states, and a rulebook says which state comes next for each possible input. At every step, the NFSM snaps its memory back onto one of its clean values, which erases that step rounding error before it can spread. A table that the model learns during training decides which value comes next. Merging tables means following positions in a list instead of adding up numbers, so nothing gets rounded. The parallel run gives exactly the same result as the step-by-step run, down to the last digit, and running in parallel adds no error of its own.

The table below puts numbers on all of this. There are four tasks. Counting modulo 16 works like a clock with 16 positions, where 0 comes again after 15. Two tasks shuffle four or five objects around. The last one is the cup game written in plain English: people swap objects, and at the end the model is asked who is holding what. We trained every model on short sequences, tested it on much longer ones, and recorded the length at which it first got an answer wrong. The affine models fail a few hundred steps past their training length, which is where the theory says they must. A single NFSM layer is still perfect a million steps in.

On the three tasks that follow an exact mathematical rule (counting and the two shuffles), we can also read the learned table out of the model and compare it with the correct one. They matched exactly in every training run (we trained each model several times, starting from different random weights). So we are not guessing from test results that the model will keep working at longer lengths: we checked directly that the rules it learned are the right ones.

Task	Existing parallel model	NFSM (ours)
Counting modulo 16	135	≥ 1,000,000
Shuffling four objects	131	≥ 1,000,000
Shuffling five objects	187	≥ 1,000,000
Who-has-what, in English	1,922	≥ 1,000,000

The middle column gives the best result among three recent affine designs: Mamba with negative eigenvalues (a version of Mamba whose memory update is allowed to multiply by negative numbers), AUSSM and PD-SSM (two other recent state space models). Each of them used four layers, against one for the NFSM. “≥ 1,000,000” means no error up to the longest sequence that fits on a single GPU.

What you get is a recurrent block, a building piece you can plug into a larger model and then stop worrying about. It trains in parallel as fast as the affine designs it replaces: the number of rounds the computation needs grows with the logarithm of the sequence length, which means that doubling the length adds only one round. On these tasks, one layer of it did what four layers of the affine models could not. The NFSM applies the same transitions at step ten as at step ten million, so no hidden length limit is waiting to show up once the model is in real use. Where correctness matters, you don’t need to take our word for it or keep running longer tests. You can read the state machine the model has learned (its list of states and its rulebook) straight from its weights and check it against the one you asked for.

So, where is the ball? An affine model can tell you for the first hundred or so swaps, and after that it starts guessing without saying so. The NFSM gives the right answer at swap ten, at swap ten million, and at every swap in between, and you can open it up and check the rules it follows. If you build models that need to keep track of something over long inputs, the full paper has the proofs and the experiments.