Improving Test-Time Scaling with Adaptive Looped Transformers

Authors: Yichen You*, Tianyu Fu*, Aosong Feng, Xingtai Lv, Xuefei Ning, Ning Ding, Yu Wang
Affiliations: Tsinghua University, Yale University

---

Summary (Overview)

• Core Question: Does looping (recurrent layer application) in transformers improve test-time scaling as output-token budgets grow, and how can looped models be post-trained to outperform non-looped baselines at matched compute?

• Key Finding 1 (Problem Identification): Post-trained looped transformers (Ouro, Huginn) yield steeper accuracy–compute slopes than non-looped baselines (2.12–2.27 vs. 1.79 points per compute doubling) but underperform at matched test-time compute.

• Key Finding 2 (Token-Level Analysis): Fixed-depth looping wastes compute on tokens that gain little or worsen with extra iterations — 52.3% (Ouro) and 33.5% (Huginn) of tokens change by ≤10⁻³ in loss, while 21.7% and 15.9% become worse.

• Key Finding 3 (Proposed Method — TaH2): A new post-training approach that jointly trains the backbone and an iteration decider via lookahead depth supervision, using online labels derived from measured iteration gains.

• Key Finding 4 (Results): On AIME24–26, TaH2 improves the accuracy–compute slope by 53% (2.74 vs. 1.79) over the non-looped baseline, exceeds baseline peak accuracy by ~3.4 points at matched compute, and its gains grow from +2.8 to +3.9 points as max depth increases from 2 to 8.

---

Introduction and Theoretical Foundation

Background and Motivation

Test-time scaling improves language model reasoning by spending additional inference compute, either through:

• Token-space methods: Longer chains of thought (CoT) with budget forcing or RL
• Latent-space methods: Repeated application of shared layers in looped transformers

Looped transformers achieve parameter efficiency by reusing layers for latent computation. However, prior scaling studies compare looped and non-looped models at matched parameters or matched per-token FLOPs, which does not establish whether looping improves accuracy–compute scaling as output-token budgets increase.

Central Research Question

How to post-train looped LLMs to outperform non-looped LLMs at the same test-time compute?

Theoretical Foundation

The recurrent computation of looped transformers is formalized as follows. For a shared Transformer backbone $\mathcal{F}_\theta$ with embedding $\mathbf{e}_t$ at token position $t$:

$$\mathbf{h}_t^{(m)} = \mathcal{F}_\theta\left(\mathbf{h}_t^{(m-1)}\right), \qquad \mathbf{q}_t^{(m)} = \mathrm{softmax}\left(\mathbf{W}_{\mathrm{out}}\mathbf{h}_t^{(m)}\right), \tag{1}$$

where $m$ indexes the iteration, $\mathbf{h}_t^{(0)} = \mathbf{e}_t$, and $\mathbf{W}_{\mathrm{out}}$ is the output projection. Key architectures:

• Ouro: Loops all $L$ layers (full-stack recurrence)
• Huginn: Loops a middle block with input injection

Key Empirical Finding

The paper's motivating observation: "Post-trained looped models gain accuracy faster as test-time compute increases, yet underperform the non-looped baseline at matched compute."

This is attributed to fixed-depth looping spending extra iterations on tokens that gain little or worsen, motivating the need for token-adaptive depth.

---

Methodology

TaH2 Architecture

TaH2 comprises three components sharing parameters across iterations:

1. Shared Transformer Backbone

Uses extended duo-causal attention: at depth $m$, token $t$ attends to executed KV states at positions $s \leq t$ and depths $j \leq m$. During training, stopped tokens take an additional no-gradient lookahead iteration for supervision.

2. Learned Input-Injection Updater

Modifies the recurrence from Equation 1 to:

$$\mathbf{h}_t^{(m+1)} = \mathcal{F}_\theta\left(\mathcal{U}_\psi\left(\mathbf{e}_t, \mathbf{h}_t^{(m)}\right)\right), \tag{2}$$

where $\mathcal{U}_\psi$ is a small normalized MLP that reinjects the input embedding while mapping the previous final-layer state back to the backbone input space.

3. Iteration Decider

After iteration $m < M$, a lightweight decider predicts a conditional continue probability:

$$g_t^{(m)} = \mathcal{D}_\phi\left(\mathbf{e}_t, \mathbf{h}_t^{(m)}, \operatorname{TopK}\left(\mathbf{q}_t^{(m)}\right)\right) \in (0, 1). \tag{3}$$

With exit threshold $\tau_{\mathrm{exit}}$, token $t$ stops at the first depth where $g_t^{(m)} < \tau_{\mathrm{exit}}$. The continue probabilities induce stopping weights over executed depths:

$$\omega_t^{(m)} = \begin{cases} (1 - g_t^{(m)})\prod_{j < m} g_t^{(j)}, & m < m_t, \\ \prod_{j < m_t} g_t^{(j)}, & m = m_t, \end{cases} \quad \mathbf{q}_t = \sum_{m=1}^{m_t} \omega_t^{(m)} \mathbf{q}_t^{(m)}. \tag{4}$$

Training with Lookahead Depth Supervision

Online Continuation Labels

For each supervised token, the prediction loss and iteration gain are:

$$\ell_t^{(m)} = -\log \mathbf{q}_t^{(m)}[y_t^*], \qquad \delta_t^{(m)} = \ell_t^{(m)} - \ell_t^{(m+1)}. \tag{5}$$

Positive gains favor continuing; negative gains favor stopping. The target continue label is:

$$c_t^{(m)} = a_t^{(m)} c_t^{(m-1)} \mathbf{1}\left[\delta_t^{(m)} \geq \delta_{\mathrm{cut}}^{(m)}\right]. \tag{6}$$

where $\delta_{\mathrm{cut}}^{(m)}$ is computed by ranking positive gains and selecting the largest until they account for fraction $\rho$ of the total (coverage).

Joint Objective

$$\mathcal{L}_{\mathrm{SFT}} = \underbrace{\frac{1}{N}\sum_t -\log \mathbf{q}_t[y_t^*]}_{\text{next-token prediction loss}} + \alpha_D \underbrace{\frac{1}{N}\sum_{m=1}^{M-1}\sum_{t: a_t^{(m)}=1} w_t^{(m)} \operatorname{BCE}\left(g_t^{(m)}, c_t^{(m)}\right)}_{\text{cost-sensitive decider loss}}, \tag{7}$$

where $N$ is the number of supervised tokens, $\mathbf{q}_t$ is the stopping-weighted mixture from Equation 4, and $w_t^{(m)}$ are cost-sensitive weights computed from $|\delta_t^{(m)} - \delta_{\mathrm{cut}}^{(m)}|$. Default settings: coverage $\rho = 0.99$, decider coefficient $\alpha_D = 0.05$.

Training Setup

• Backbones: Qwen3-{1.7B, 4B, 8B}-Base
• Data: 273K prompts from AM-Qwen3-Distilled (math, code, QA) + Nemotron-Agentic-v1 (tool use), responses regenerated by Qwen3-8B
• Training: 3 epochs, 16,384-token context, 3.4B total tokens for 1.7B experiments; fixed TPP = 2 tokens/parameter for larger models
• Evaluation: Zero-shot CoT, temperature 0.6, top-p 0.95, top-k 20, up to 32K generation length

---

Empirical Validation / Results

Performance at 1.7B Scale

Table 1: Accuracy (%) of Qwen3-1.7B models across ten benchmarks. TaH2's average gain over Standard grows from +2.9 points (M=2) to +4.8 points (M=8).

Test-Time Scaling Results

• Slope improvement: TaH2 at M=2 achieves 2.74 points per doubling of decoding FLOPs vs. 1.79 for Standard (53% improvement)
• Peak accuracy: Standard saturates at 12.0% with 191.7 TFLOPs per response; TaH2 reaches 15.4% at the same compute (+3.4 points)
• Depth scaling: TaH2's gain over Standard grows from +2.8 points (M=2) to +3.9 points (M=8), while existing looped models plateau

Runtime Efficiency

Table 2: Runtime efficiency on AIME26. TaH2 adds only 22% decoding FLOPs and 30–34% latency vs. Standard.

Scaling to Larger Models

Table 3: Accuracy (%) at 4B and 8B. TaH2 outperforms Standard on every benchmark at both sizes.

Ablation Study (Design Choices)

Table 4: Design choices at 1.7B. The largest drop (−3.6 points) comes from uniform decider loss weights, highlighting the importance of cost-sensitive supervision.

---

Theoretical and Practical Implications

Theoretical Implications

1. Token-Adaptive Depth is Essential: The analysis reveals that fixed-depth looping wastes compute on tokens that gain little or worsen. The success of TaH2 demonstrates that learning which tokens benefit from additional iterations is more important than simply increasing depth.

2. Online Labels Outperform Offline Labels: Gain-based online labels (measuring actual iteration improvement) outperform top-1 mismatch labels by 1.1 points, showing that supervision should reflect whether further iteration improves prediction, not just whether the current prediction is correct.

3. Cost-Sensitive Supervision Matters: Weighting decider loss by distance from the gain cutoff (−3.6 points vs. uniform) shows that the magnitude of iteration gains carries important signal for depth decisions.

4. Train–Inference Consistency: Deterministic (threshold-based) training decisions outperform sampled decisions (−1.3 points), suggesting that consistent token selection across training and inference helps each depth specialize.

Practical Implications

1. Test-Time Scaling Efficiency: TaH2 achieves higher accuracy at lower compute than fixed-depth looped models (1.21× vs. 1.99× FLOPs at M=2), making adaptive looping a practical way to improve reasoning without prohibitive inference costs.

2. Post-Training Applicability: The method works by post-training existing pretrained models (Qwen3-Base), avoiding the need for expensive looped pretraining from scratch.

3. Scalability: Gains persist at 4B and 8B scales and generalize beyond math to code, QA, and tool use, suggesting broad applicability.

4. Serving Efficiency: With batched serving that groups requests at different iteration depths, TaH2 achieves better test-time scaling and higher attainable accuracy in real-world serving scenarios.

---

Conclusion

Main Takeaways

This paper studies the test-time scaling of looped transformers introduced through post-training. Key findings:

1. Problem identified: Existing post-trained looped models (Ouro, Huginn) yield steeper accuracy–compute slopes than non-looped baselines but underperform at matched compute.

2. Method proposed: TaH2 jointly post-trains a backbone and iteration decider through lookahead depth supervision, using online labels from measured iteration gains to allocate depth to tokens that benefit most.

3. Results achieved:

   • 53% improvement in accuracy–compute slope on AIME (2.74 vs. 1.79)
   • ~3.4 points higher peak accuracy than Standard at matched compute
   • Gains grow from +2.8 to +3.9 points as max depth increases from 2 to 8
   • Consistent gains at 4B and 8B scales across math, code, QA, and tool use

Limitations and Future Directions

1. Training cost: TaH2 incurs more training FLOPs than standard SFT, though post-training is substantially cheaper than pretraining.

2. Scope: The method is studied only under SFT; extension to on-policy distillation and reinforcement learning is left for future work.

3. Potential directions: Exploring TaH2 with RL-based training, extending to even larger models, and investigating the interaction between adaptive depth and other test-time scaling techniques (e.g., majority voting, tree search).