Improving Test-Time Scaling with Adaptive Looped Transformers

AI summary

Overview

• Research area: Natural Language Processing — specifically efficient LLM inference, test-time scaling, and looped (recurrent-depth) transformer architectures.
• Technical level: Intermediate. The paper assumes familiarity with transformer decoding, FLOPs-based compute accounting, and post-training (supervised fine-tuning), but its core argument is stated in plain terms.
• Scope (one sentence): The paper post-trains "looped" transformers that reuse their own layers for extra latent computation, and shows that teaching the model which tokens deserve extra iterations (TaH2) beats both fixed-depth looping and the non-looped baseline on accuracy per unit of test-time compute.

What This Paper Is About

Looped transformers gain effective depth by applying the same layers repeatedly before emitting each token, so they can add "thinking" without adding parameters — but every extra iteration still costs decoding compute. Prior work compared looped and non-looped models at matched parameters or matched per-token FLOPs, leaving open whether looping actually improves scaling as generated outputs get longer. This paper post-trains looped and non-looped models from the same checkpoint and measures the accuracy–compute slope, finding that existing looped models scale faster but are still less accurate at matched compute — because fixed-depth looping wastes iterations on tokens that do not benefit. The goal is to post-train a looped LLM that beats the non-looped baseline at the same test-time compute.

Key Contributions

1. A study of test-time scaling in looped models. The authors measure the accuracy–compute slope (accuracy gain per doubling of decoding FLOPs) for post-trained Ouro (full-stack recurrence, M=2), Huginn (middle-block recurrence, M=3) and a non-looped Standard baseline, and show that the looped models have steeper slopes (2.12, 2.27, and 2.26 vs. 1.79) yet remain less accurate over the overlapping compute range.

2. TaH2, an adaptive looped post-training method. TaH2 jointly post-trains the backbone and a token-level iteration decider using lookahead depth supervision, in which depth labels are derived online from whether a further iteration actually reduces prediction loss, and each depth decision is supervised with a cost-sensitive loss.

3. Improved test-time and iteration-depth scaling. On AIME, TaH2 improves the accuracy–compute slope by 53% (2.74 vs. 1.79) over Standard, exceeds Standard's peak accuracy by about 3.4 points at matched compute, and its gain over Standard grows from +2.8 points at depth 2 to +3.9 points at depth 8.

4. Generality across scales and domains. The gains persist at 4B and 8B and extend beyond math to code, QA, and tool-use benchmarks.

Main Findings

• Existing looped models scale faster but lose on accuracy. Fixed-depth Ouro, adaptive Ouro, and Huginn yield slopes of 2.12, 2.27, and 2.26 accuracy points per compute doubling, versus 1.79 for Standard — yet all three remain less accurate than Standard over the overlapping compute range, and Ouro's exit gate lowers decoding cost without closing the gap.

• Most tokens gain little or nothing from extra iterations. Comparing next-token loss after the first and final iterations, for Ouro and Huginn respectively 52.3% and 33.5% of tokens change by at most 10⁻³, while 21.7% and 15.9% become worse by more than this threshold.

• The train–inference mismatch limits Ouro's gate. Ouro's backbone is trained only at full depth, so early exit at inference creates a mismatch; this motivates learning token-dependent depth jointly with the backbone, supervised by measured loss reductions.

• TaH2 improves the accuracy–compute slope by 53%. On AIME24–26 at 1.7B, TaH2 reaches 2.74 points per doubling of decoding FLOPs versus 1.79 for Standard.

• TaH2 exceeds Standard's peak at matched compute. Extending evaluation to 32K tokens, Standard saturates at 12.0% accuracy with 191.7 TFLOPs per response, while TaH2 reaches 15.4% at the same compute — 3.4 points higher.

• Depth scaling keeps helping TaH2 but not other looped models. On validation loss, Ouro and Huginn remain at or above Standard at every depth ceiling, whereas adaptive TaH2 reaches the lowest loss (−0.0106 versus Standard at M=8). Across ten benchmarks, TaH2's average gain over Standard grows from +2.9 points at M=2 to +4.8 at M=8, while Huginn and Ouro stay close to Standard.

• TaH2 also improves parallel scaling. AIME24–26 cons@32 reaches 27.3–29.3% for M=2–8, versus 21.9% for Standard.

• The decider is aligned with actual gains. Grouping validation tokens by continue probability shows tokens near zero probability have negative or negligible loss reduction, and mean gain rises with continue probability; at the 0.5 decision threshold, the corresponding loss reduction is near zero.

• Cost is real but modest. TaH2 adds 22% decoding FLOPs per token and 30–34% end-to-end latency relative to Standard, versus 2.00× GFLOPs/token for TaH2-fixed, and still achieves better test-time scaling and higher attainable accuracy in actual serving (7.04 GFLOPs/token for Standard vs. 8.59 for TaH2).

• Gains transfer to larger models and other domains. At 4B and 8B with M=2, TaH2 raises average accuracy by 3.2 points and 2.4 points respectively, with AIME gains of up to 6.9 points at 4B and 4.4 points at 8B.

• Design choices matter. Compared with TaH2's default 15.5 average AIME accuracy, uniform decider-loss weights cost 3.6 points, a top-100 embedding updater costs 2.1 points, keeping all positive gains (ρ=1) costs 1.9 points, sampled rather than thresholded training decisions cost 1.3 points, top-1 mismatch depth labels cost 1.1 points, and using only the final iteration's prediction costs 1.0 point.

• Depth is allocated unevenly across token types. In math and code examples, mathematical expressions and final code use fewer iterations than preceding natural-language reasoning, while the QA example maintains greater depth throughout.

• Attention heads diverge in iteration preference. In three representative heads on validation sequences, heads attend primarily to first-iteration states, later-iteration states, or both.

Methodology in Plain English

The authors start from a shared pretrained checkpoint (Qwen3-1.7B-Base, later 4B and 8B) and post-train several model variants on the same data so comparisons are apples-to-apples. The non-looped "Standard" model runs the backbone once per token. Looped variants run it multiple times: Ouro repeats the whole stack, Huginn repeats only a middle block (layers 8–21 of 28) with input injection.

To diagnose waste, they compare per-token prediction loss after the first iteration versus after the final one on a held-out validation set, and find that a large share of tokens barely change or get worse. That motivates TaH2, which has three parts: the shared backbone; a small learned "updater" that fuses the original token embedding back into the hidden state between iterations; and a lightweight "decider" that outputs a continue probability at each iteration. A token stops at the first iteration where the continue probability falls below a threshold (0.5), and the predictions from the executed iterations are combined into a weighted mixture, with earlier iterations weighted by the probability of stopping there.

Training uses lookahead depth supervision. At each iteration the authors measure the actual change in loss when one more iteration is applied. Positive gains suggest continuing, negative gains suggest stopping. To avoid noisy labels, they keep the largest positive gains until they cover ρ = 0.99 of the total positive gain; the smallest retained gain defines a per-iteration cutoff, and tokens above it get a "continue" label. Tokens that stop early still receive a no-gradient lookahead iteration purely to generate a supervision target (omitted at inference). The decider is trained with a cost-sensitive binary cross-entropy weighted by how far each token's gain is from the cutoff, and training is on-policy with respect to the depth choices the decider actually makes.

Architecturally, each iteration keeps its own KV cache and uses an "extended duo-causal" attention rule: a query at token t and depth m can attend to executed states at positions ≤ t and depths ≤ m, which concatenates into a block-structured mask so training and prefill parallelize. Serving uses an extended Mini-SGLang engine that batches requests at different iteration depths into a shared forward pass.

Why This Matters

Impact on research. The paper reframes the looped-transformer comparison: steep scaling slopes alone do not mean a model is better, and matching parameter counts or per-token FLOPs can hide the fact that fixed-depth recurrence spends compute on tokens that do not need it. It also adds a concrete data point to the debate about whether recurrence should be introduced at pretraining time or through post-training, and provides a supervision recipe (online, gain-based depth labels) that differs from offline labels (TaH) and frozen-backbone gate training (Ouro).

Real-world applications (plausible from the paper's benchmark set):

• Mathematical reasoning assistants that need higher accuracy per unit of inference spend (AIME24–26, AMC23, MATH500, OlympiadBench).
• Code generation and completion, where the paper's token-level analysis shows final code segments can use fewer iterations (HumanEval, MBPP, LiveCodeBench v6).
• Scientific and general question answering where depth demand varies across a response (GPQA, SuperGPQA).
• Agentic tool-use systems, where latency budgets are tight and the paper reports gains on BFCL v3.

Industry relevance. TaH2 adds fewer than 3% parameters at every scale studied and 22% decoding FLOPs per token relative to Standard, so it slots into an existing serving stack without changing model size. The paper's serving measurements (latency 139.8s, 187.0s, 328.7s at batch size 1 for Standard, TaH2, and TaH2-fixed) give deployment teams a concrete trade-off between accuracy and end-to-end latency rather than a purely theoretical one.

Future Directions

• Extend beyond supervised fine-tuning. The paper explicitly leaves extension to on-policy distillation and reinforcement learning for future work; both could supply better depth supervision than loss-reduction labels computed on a fixed data distribution.
• Reduce the training overhead. TaH2 incurs more training FLOPs than standard SFT (reported in Appendix A.4.3); the authors argue post-training is still far cheaper than pretraining, but reducing this cost is an open engineering question.
• Test whether looping should be introduced earlier. The paper's central diagnosis is a post-training gap for looped architectures; whether a jointly trained-from-scratch looped-and-adaptive model shows the same or larger benefits is not established here.
• Generalize the depth policy beyond the studied settings. The method is evaluated with Qwen3-{1.7B, 4B, 8B} backbones, M up to 8, and a specific serving engine; scaling behavior at larger model sizes, longer depth ceilings, and other backbone families is not reported.

Target Audience

Researchers and engineers working on efficient LLM inference, test-time compute allocation, and adaptive or recurrent-depth architectures. It is also useful for practitioners who care about accuracy-per-FLOP or accuracy-per-latency trade-offs in deployed reasoning systems, and for readers already familiar with Huginn, Ouro, or TaH-style looped post-training who want to understand where fixed-depth recurrence falls short.

Authors’ abstract

Looped transformers have demonstrated promising parameter efficiency by reusing layers for latent computation. Prior studies compare looped and non-looped models at matched parameters or per-token FLOPs. However, to the best of our knowledge, whether looping improves test-time scaling as outputs grow longer remains underexplored. Through post-training looped transformers, we study the accuracy-compute slope, measured as the accuracy gain per doubling of test-time decoding FLOPs. We find that existing looped transformers often yield steeper slopes than their non-looped baseline, yet underperform it at matched compute. While fixed-depth looping spends extra iterations on every token, our analysis shows that many tokens do not benefit from extra iterations. We therefore propose TaH2, which enables the model to focus extra iterations on the tokens that benefit from looping. It jointly post-trains the backbone and an iteration decider through lookahead depth supervision, which uses online labels indicating whether further iteration improves the prediction. TaH2 improves both the efficiency and attainable accuracy of test-time scaling. On challenging AIME benchmarks, TaH2 improves the accuracy-compute slope by 53% (2.74 vs. 1.79) over the non-looped baseline, exceeding the baseline's peak accuracy by about 3.4 points at matched test-time compute. As the maximum iteration depth increases, existing looped models largely plateau, while TaH2's gain over the non-looped baseline continues to grow from +2.8 points at depth 2 to +3.9 points at depth 8. Our code is available at https://github.com/thu-nics/TaH.