# Improving Test-Time Scaling with Adaptive Looped Transformers (TaH2)

> TaH2 post-trains looped transformers with token-adaptive depth supervision, improving accuracy-compute scaling by 53% and beating non-looped baselines at matched test-time compute.

- **Source:** [arXiv](https://arxiv.org/abs/2609.35748)
- **Published:** 2026-10-12
- **Permalink:** https://picx.dev/p/00CeNy
- **Whiteboard:** https://picx.dev/p/00CeNy/image

## Summary

# 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

| Domain | Benchmark | Std. (M=1) | Huginn (M=3) | Ouro (M=2) | TaH2-fixed (M=2) | **TaH2 (M=2)** | **TaH2 (M=4)** | **TaH2 (M=8)** |
|---|---|---|---|---|---|---|---|---|
| math | AIME24 | 11.0 | 13.9 | 10.5 | 15.7 | 15.7 | 16.3 | 16.3 |
| math | AIME25 | 13.0 | 14.2 | 13.9 | 15.0 | 16.5 | 16.9 | 17.9 |
| math | AIME26 | 11.9 | 11.8 | 11.0 | 13.2 | 14.3 | 14.2 | 14.5 |
| math | AMC23 | 45.3 | 45.8 | 45.9 | 51.1 | 50.7 | 50.9 | 52.8 |
| math | MATH500 | 75.2 | 74.4 | 75.1 | 76.1 | 78.1 | 78.6 | 79.3 |
| math | Olympiad | 40.4 | 40.6 | 37.9 | 42.1 | 43.8 | 44.1 | 47.4 |
| code | HE | 63.6 | 61.7 | 61.2 | 63.8 | 66.8 | 65.2 | 72.6 |
| code | MBPP | 70.2 | 70.3 | 69.6 | 70.8 | 71.8 | 72.3 | 74.4 |
| QA | GPQA | 31.5 | 33.7 | 34.0 | 31.9 | 33.1 | 33.0 | 32.0 |
| tool | BFCL | 15.3 | 15.6 | 14.3 | 14.6 | 15.3 | 17.4 | 17.9 |
| **Avg.** | | **37.7** | **38.2** | **37.3** | **39.4** | **40.6** | **40.9** | **42.5** |
| **FLOPs/token** | | **1.00×** | **1.90×** | **1.91×** | **1.99×** | **1.21×** | **1.47×** | **2.37×** |

**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

| Metric | Std. (BS=1) | TaH2-fixed (BS=1) | TaH2 (BS=1) | Std. (BS=4) | TaH2-fixed (BS=4) | TaH2 (BS=4) |
|---|---|---|---|---|---|---|
| GFLOPs/tok | 7.04 | 14.09 | 8.59 | 7.04 | 14.09 | 8.59 |
| vs. Std. | 1.00× | 2.00× | **1.22×** | 1.00× | 2.00× | **1.22×** |
| Latency (s) | 139.8 | 328.7 | 187.0 | 220.3 | 483.3 | 285.4 |
| vs. Std. | 1.00× | 2.35× | **1.34×** | 1.00× | 2.19× | **1.30×** |

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

### Scaling to Larger Models

| Domain | Benchmark | 4B Std. | 4B TaH2 | 8B Std. | 8B TaH2 |
|---|---|---|---|---|---|
| math | AIME24 | 52.0 | 57.7 | 66.4 | 70.8 |
| math | AIME25 | 40.3 | 43.1 | 52.0 | 55.9 |
| math | AIME26 | 45.5 | 52.4 | 60.9 | 63.4 |
| math | AMC23 | 88.6 | 89.2 | 93.5 | 96.6 |
| math | Olymp. | 67.0 | 68.4 | 73.3 | 74.4 |
| math | IMO-AB | 31.2 | 33.8 | 40.3 | 42.0 |
| code | LCB | 35.4 | 37.3 | 42.1 | 44.4 |
| QA | SGPQA | 28.6 | 30.5 | 38.0 | 39.2 |
| tool | BFCL | 29.2 | 33.6 | 38.0 | 39.4 |
| **Avg.** | | **46.4** | **49.6 (+3.2)** | **56.1** | **58.5 (+2.4)** |

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

### Ablation Study (Design Choices)

| Variant | AIME24 | AIME25 | AIME26 | Avg. |
|---|---|---|---|---|
| **TaH2 (default)** | **15.7** | **16.5** | **14.3** | **15.5** |
| Top-1 mismatch labels | 14.1 | 16.5 | 12.6 | 14.4 (−1.1) |
| Uniform decider weights | 12.9 | 11.6 | 11.3 | 11.9 (−3.6) |
| Coverage ρ = 1 | 14.0 | 13.4 | 13.5 | 13.6 (−1.9) |
| Sampled decisions | 14.4 | 16.3 | 12.0 | 14.2 (−1.3) |
| Top-100 embedding updater | 13.3 | 14.1 | 12.7 | 13.4 (−2.1) |
| Final-iteration output | 15.4 | 15.4 | 12.7 | 14.5 (−1.0) |

**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).

---

_Markdown view of https://picx.dev/p/00CeNy, served by PicX — AI-generated visual whiteboard summaries of research papers._