Improving Test-Time Scaling with Adaptive Looped Transformers

What it claims

Post-trained looped transformers scale better per compute doubling than a non-looped model, yet they are still less accurate at any given compute budget. TaH2 closes that gap. It learns per token whether to run another full pass of the backbone. On mean AIME24–26 at 1.7B it reaches 2.74 accuracy points per doubling of decoding FLOPs, against 1.79 for the non-looped Standard. With evaluation extended to 32K output tokens, it beats Standard's peak by about 3.4 points at matched decoding FLOPs. The gain over Standard grows with the depth ceiling M: +2.8 points at M=2 and +3.9 at M=8. Existing looped models mostly plateau as M grows.

The diagnosis that motivates it

All variants were post-trained from Qwen3-1.7B-Base on the same data with 16K context. The slopes per compute doubling were:

Despite the steeper slopes, all three looped models were less accurate than Standard across the overlapping compute range.

The authors then measured per-token loss change on validation data:

They also identify a cause for Ouro's weak adaptive mode. Its gate is trained afterwards on a frozen backbone that only ever saw full depth, so early exit creates a train–inference mismatch.

How it works

• Backbone: all L layers are looped. A small normalised MLP "updater" reinjects the token embedding together with the previous final-layer state before each pass.
• Decider: after each iteration, the decider reads the embedding, the hidden state and the top-K of the current prediction, and outputs a continue probability g. A token stops at the first depth where g < τ_exit, or else runs to M.
• Output: the continue probabilities become stopping weights that mix the per-iteration predictions. The last executed iteration takes the remaining weight.
• Size: the updater and decider add under 3% parameters at every scale tested.
• Attention: "extended duo-causal attention" lets a token attend to executed KV states at earlier positions and lower-or-equal depths.
• On-policy training: depth follows the current decider's choices, using the same rule as at inference.
• Lookahead supervision: tokens that stop early get one extra no-gradient iteration, so the gain from one more pass can be measured. The gain is δ = ℓ(m) − ℓ(m+1).
• Labels: at each depth, positive gains are ranked and the largest are kept until they cover a fraction ρ of the total gain. Tokens above that cutoff get label "continue". Once a token is labelled "stop", it stays stopped.
• Loss: next-token cross-entropy on the mixed prediction, plus α_D × a cost-sensitive BCE on the decider. The BCE is weighted by the distance from the cutoff, |δ − δ_cut|.
• Settings: ρ = 0.99 and α_D = 0.05.

flowchart TD
  A[Embed token] --> B[Updater + backbone pass]
  B --> C[Head: prediction q]
  C --> D{Decider g < threshold?}
  D -- no, continue --> B
  D -- yes, stop --> E[Weighted mix of q's]
  D -. training only .-> F[No-grad lookahead pass]
  F --> G[Loss gain δ → continue label]
  G --> H[Cost-sensitive BCE on decider]

The new element against prior work is the source of the depth labels. TaH used offline mismatch labels in separate training stages. Ouro trained its gate with the backbone frozen. TaH2's labels are measured online on the backbone as it trains.

Evidence

• Setting: accuracy is the mean over AIME24–26 with 32 samples per problem. Output cutoffs were swept from 4K to 16K in 2K steps. Accuracy was fitted linearly against log₂ of decoding FLOPs per response.
• Baselines: Standard (M=1), TaH2-fixed (no decider, always runs M), Ouro (M=2, 4) and Huginn (M=3, 7). Huginn loops 14 of 28 layers, which matches 2 and 4 effective full passes.
• Absolute accuracy is low: the figure axes span roughly 0.06–0.15. The point gains are real but are measured on a hard benchmark with a small model.
• Other claims not shown in the fetched text: the paper says gains hold at 4B and 8B and extend to code, QA and tool use (its Table 3). Those numbers were cut off before that point.

How to use it

The fetched text gives no code, weights, licence or run instructions. Reproducing it means post-training Qwen3-{1.7B, 4B, 8B}-Base with the recipe in the paper's appendices. Training costs an extra no-grad iteration for early-stopped tokens. Inference drops the lookahead.

Limits

• Post-training only: the authors note that recurrence works well when built in from pretraining, as with Ouro and Huginn. Their finding that looped models underperform applies to post-training only.
• Main results are 1.7B on AIME: the larger-scale and non-math results were not in the fetched excerpt.
• The 3.4-point win needs 32K-token evaluation: it is measured where Standard has already plateaued.
• Serving not covered here: batching with variable per-token depth under duo-causal KV is deferred to the paper's Appendix E, which is not in the excerpt.