Looped transformers finally beat standard models when per-token depth is learned jointly

The core finding here is a diagnostic before it is a solution. Looped transformers — architectures that reuse the same layers across multiple iterations before emitting each token — do scale better with test-time compute than standard single-pass models: their accuracy-per-compute-doubling slope is steeper. But steeper slope from a lower starting point still means worse absolute accuracy at any given compute budget. That gap is what this paper sets out to close.

The diagnosis is precise. Examining token-level loss changes across iterations in two existing looped architectures (Ouro and Huginn), the authors find that the majority of tokens cluster near zero improvement, and a substantial fraction actually get worse after additional iterations. Fixed-depth looping wastes compute on tokens that don't need it, and Ouro's learned exit gate helps but not enough to beat the non-looped baseline at matched FLOPs.

TaH2 attacks this by learning, per token, whether another iteration is worth running. A lightweight iteration decider is trained jointly with the backbone — not in a separate frozen stage — using what the paper calls lookahead depth supervision. During training, a no-gradient extra iteration is computed for tokens that stopped early, and the actual change in prediction loss determines whether the decider should have said "continue." Labels are derived online from the current backbone state, not from offline proxies, which matters: ablations show that replacing gain-based labels with simpler mismatch labels (does the top-1 prediction match the target?) costs several accuracy points.

The numbers are specific. At 1.7B parameters on AIME24–26, TaH2 improves the accuracy-compute slope by 53% over the non-looped baseline (2.74 vs. 1.79 points per doubling of decoding FLOPs). When the output budget is extended to 32K tokens, TaH2 exceeds the baseline's peak accuracy by roughly 3.4 points at matched compute. Crucially, as the maximum iteration depth ceiling rises from 2 to 8, existing looped models plateau while TaH2's gain over the baseline grows from 2.8 to 3.9 points. The gains hold at 4B and 8B and transfer to code, QA, and tool-use benchmarks.

The serving story is honest. TaH2 adds about 22% more decoding FLOPs per token compared to the standard model and roughly 30-34% more end-to-end latency, depending on batch size. The fixed-depth looped variant costs twice the FLOPs and more than twice the latency. So the adaptive depth mechanism is doing real work, not just reporting FLOP counts that never materialize as wall-clock savings.

Two limitations are stated plainly: TaH2 requires more training FLOPs than standard supervised fine-tuning (though far less than pretraining), and it has only been studied under SFT, not reinforcement learning or on-policy distillation. The token-depth visualizations are a useful sanity check — mathematical expressions and final code use fewer iterations than the natural-language reasoning that precedes them, which is exactly what you'd want.