AI/ ai · machine learning · transformers · research

Looped Transformers Learn When to Stop

Researchers found that randomizing loop counts during training cuts out-of-distribution variance and makes looped transformers generalize more reliably.

A new technique makes looped transformers dramatically more stable when solving problems longer than anything they were trained on.

Looped transformers reuse a single shared block repeatedly — a design that suits variable-length tasks like adding binary numbers or matching parentheses. The problem: they tend to latch onto a spurious correlation between sequence length and the number of loops, which makes their behavior brittle on out-of-distribution inputs. Researchers traced this instability to that shortcut and tested whether injecting randomness into the loop count during training would break it. It does. They also evaluated a learned approach called RL-Halting, which trains the model to decide when to stop rather than relying on a fixed or randomly chosen schedule, and found it generally improved the accuracy-stability trade-off across four benchmark tasks.

The finding reframes a design decision that has largely been treated as an afterthought. Most work on looped architectures treats stopping as an inference-time knob — something you tune after the model is already trained. Treating it as a training-time choice opens a path to models that generalize length without the brittleness that has made looped designs hard to deploy.

Length generalization has been a known weak spot for transformers since the architecture debuted, and looped variants were supposed to help — this research suggests they can, but only if you teach them to quit on cue.

TR

The Revision

Written by an AI system from the public sources credited above. How we write →