Summary of "Scaling to Tens of Thousands of Test-Time Iterations with Loop-Native Attention Residuals"

Summary (Overview)

  • Core Problem Identified: The paper identifies the "carry-last" rule (feeding the latest block output directly as the next recurrent state) as a key bottleneck in looped Transformers. This rule causes performance degradation at scale—noisy updates overwrite correct intermediate deductions and destabilize solved states.
  • Main Contribution: Introduces InfiLoop, a loop-native residual connection that replaces carry-last with learned attention over recurrent outputs, combining content-based scoring with exponential temporal decay. This admits an exact streaming recurrence requiring only O(d)O(d) persistent memory per token, independent of loop count.
  • Key Results: A 7M-parameter InfiLoop model achieves 97.9% exact accuracy on Sudoku-Extreme, 87.4% on Maze-Hard, 49.1% pass@2 on ARC-AGI-1, and 13.6% pass@2 on ARC-AGI-2, outperforming strong baselines like FPRM, TRM, and HRM.
  • Emergent Behavior: The learned connection exhibits self-verification: it suppresses outputs during unsolved search (step size ~0.40), surges to 0.91 when the solution appears, and freezes the state once solved, preventing solution loss at extreme depths.
  • Scalability: InfiLoop continues to improve with test-time looping beyond 20,000 effective steps, translating added depth directly into stronger reasoning.

Introduction and Theoretical Foundation

Background and Motivation

Looped Transformers address complex reasoning by repeatedly applying a shared Transformer block, adding computational depth without new parameters. The standard carry-last rule connects iterations by feeding the final output of one loop directly into the next:

zt=fθ(ht−1;x),ht=zt(1)z_t = f_\theta(h_{t-1}; x), \qquad h_t = z_t \tag{1}

This makes the unrolled computation equivalent to a deep network with shared weights, reintroducing the "curse of depth" where deeper layers contribute progressively less.

Key Insight

The authors argue that state replacement is the fundamental bottleneck. Under carry-last:

  • Earlier computations are accessible only through information preserved in the latest state
  • A noisy or undesirable iteration can degrade useful representations throughout propagation
  • The model cannot explicitly retain and select among outputs from earlier loops

Theoretical Foundation

The paper proposes that the connection between iterations should be learned rather than fixed. The key observation is that "a newer output is not necessarily a better output"—yet carry-last passes every new block output to the next iteration in full.

Methodology

The InfiLoop Architecture

InfiLoop replaces carry-last with a learned attention-based aggregation over recurrent outputs. For candidate states ziz_i and recurrent state hth_t:

ht=∑i=0tαt,izi,αt,i=exp⁡(si−λ(t−i))∑j=0texp⁡(sj−λ(t−j))(2)h_t = \sum_{i=0}^{t} \alpha_{t,i} z_i, \qquad \alpha_{t,i} = \frac{\exp(s_i - \lambda(t-i))}{\sum_{j=0}^{t} \exp(s_j - \lambda(t-j))} \tag{2}

where:

  • si=q⊤RMSNorm(zi)s_i = q^\top \text{RMSNorm}(z_i) is a content score from a shared learned pseudo-query qq
  • λ>0\lambda > 0 is a learned decay rate
  • β=e−λ∈(0,1)\beta = e^{-\lambda} \in (0,1) discounts older candidates

Exact Streaming Recurrence

The aggregation admits an exact constant-memory recurrence. Defining ei=exp⁡(si)e_i = \exp(s_i) and β=e−λ\beta = e^{-\lambda}:

Nt=∑i=0t−1βt−ieizi+etzt=βNt−1+etzt,(4)N_t = \sum_{i=0}^{t-1} \beta^{t-i} e_i z_i + e_t z_t = \beta N_{t-1} + e_t z_t, \tag{4} Dt=∑i=0t−1βt−iei+et=βDt−1+etD_t = \sum_{i=0}^{t-1} \beta^{t-i} e_i + e_t = \beta D_{t-1} + e_t

with ht=Nt/Dth_t = N_t / D_t. This requires only O(d)O(d) memory per token (a vector NtN_t and scalar DtD_t), independent of loop count.

Interpretation as Adaptive Residual Update

Substituting Nt−1=Dt−1ht−1N_{t-1} = D_{t-1} h_{t-1} into Equation 4 yields:

ht=ht−1+γt(zt−ht−1),γt=etβDt−1+et∈(0,1)(5)h_t = h_{t-1} + \gamma_t (z_t - h_{t-1}), \quad \gamma_t = \frac{e_t}{\beta D_{t-1} + e_t} \in (0,1) \tag{5}

The step size γt\gamma_t compares the candidate's content weight ete_t with the retained historical mass βDt−1\beta D_{t-1}, allowing the model to suppress low-scoring proposals while accepting high-scoring ones.

Architecture Details

  • Model: 7M-parameter nested looped Transformer with two layers per block, hidden dimension 512, 8 attention heads, SwiGLU activations, and depthwise 3×3 convolution
  • Connection: Applied independently to both inner and outer recurrent axes, adding only d+1d+1 learned parameters per axis (513 total per axis)
  • Training: Deep supervision with H=3 outer loops and L=6 inner updates for Sudoku; AdamW optimizer with peak learning rate 10−410^{-4} (Sudoku) or 3×10−43 \times 10^{-4} (Maze/ARC)

Empirical Validation / Results

Benchmark Performance

Table 1: Reasoning benchmark results (%) — Best results in bold; best among 7M-parameter models underlined.

MethodParamsSudoku-Ext.Maze-HardARC-1ARC-2
Large language models
DeepSeek R1671B0.00.015.81.3
Claude 3.7 Sonnet-0.00.028.60.7
o3-mini-high-0.00.034.53.0
Gemini 2.5 Pro (32K)---37.04.9
Recursive reasoning models
HRM27M55.074.540.35.0
Attractor Model27M91.493.1--
URM14M77.6-≥53.8‡≥16.0‡
TRM7M74.785.344.67.8
EqR7M93.0---
FPRM7M94.287.047.56.2
InfiLoop (ours)7M97.987.449.113.6

Depth Scaling Results

  • At 1,872 layer evaluations: InfiLoop achieves 90.9% exact accuracy on Sudoku-Extreme, exceeding FPRM's 90.1% at 3,744 layers (half the computation)
  • At 24,960 layers: InfiLoop reaches 92.74%, vs. 91.63% for FPRM and 83.65% for TRM
  • Key finding: InfiLoop at 936 layers (86.5%) is never matched by TRM even at 24,960 layers (83.65%)—a 26× depth advantage

Difficulty Analysis

At 60 empty cells (hardest Sudoku difficulty), InfiLoop maintains 59.5% accuracy vs. 56.5% for FPRM and 36.0% for TRM—a +23.4 percentage point margin over TRM.

Emergent Self-Verification

The learned step size γt\gamma_t reveals three distinct phases:

  1. During search: Suppressed at ~0.40 (below the equal-weight baseline of 0.56)
  2. At solution: Surges to 0.91, giving dominant weight to the newly found solution
  3. After solving: Returns to 0.56 baseline; state freezes

The content score sts_t correlates with solution correctness without explicit supervision: median score of 7.25 for fully correct boards, dropping to −0.58 at 33 wrong cells.

Ablation Results

Table 2: Component and initialization ablations (Sudoku-Extreme, exact accuracy %)

ConfigurationH3 L6H3 L12H3 L96H32 L6
Inner axis alone72.675.277.279.4
Outer axis alone75.176.376.280.0
Both axes (ours)82.185.387.489.7

Temporal decay ablations (both axes, H3 L6): Reverse decay (β=1.82\beta=1.82): 2.0%; Uniform (β=1\beta=1): 38.4%; Forward (β=0.55\beta=0.55): 74.0%; Learned from β0=0.55\beta_0=0.55: 82.1%

Perturbation Stability

Under injected perturbations (ϵ∈[10−4,0.3]\epsilon \in [10^{-4}, 0.3]), InfiLoop retains 100% of solved puzzles, while TRM loses 15-20% even without perturbation. TRM trajectories diverge to O(1) relative distance within 12 updates; InfiLoop constrains drift throughout.

Theoretical and Practical Implications

Theoretical Contributions

  1. Constant-memory aggregation: The exact streaming recurrence reduces persistent memory from O(Td)O(Td) to O(d)O(d) and total aggregation cost from O(T2d)O(T^2d) to O(Td)O(Td) compared to Full Attention Residuals
  2. Convex hull guarantee: Since attention weights are positive and sum to one, ∥ht∥≤max⁡0≤i≤t∥zi∥\|h_t\| \leq \max_{0 \leq i \leq t} \|z_i\|—aggregation never amplifies representation norm
  3. Adaptive residual update: Equation 5 provides a principled framework for understanding how the model balances historical state vs. new proposals

Practical Implications

  • Efficient test-time scaling: InfiLoop achieves higher accuracy with fewer layers, making it practical for deployment where compute is constrained
  • Architecture-agnostic: The framework applies to both standard looped Transformers and nested recurrent reasoners (TRM, HRM), with separate scorer-decay pairs per recurrent axis
  • Solution retention: The learned connection prevents catastrophic forgetting of correct solutions, a critical property for long-horizon reasoning

Conclusion

The paper establishes that state replacement under carry-last is a primary bottleneck in deep looped Transformers. The proposed InfiLoop connection:

  1. Preserves solutions indefinitely: Unlike TRM, which loses 19.7% of solved puzzles within 36 updates, InfiLoop retains all solved states
  2. Scales effectively to extreme depths: Continues improving beyond 20,000 effective steps
  3. Learns self-verification: Emergent behavior suppresses bad proposals and accepts good ones without explicit correctness supervision
  4. Achieves state-of-the-art results: Outperforms all 7M-parameter baselines across four reasoning benchmarks

Future directions suggested by the work include:

  • Extending the framework to larger model scales and more complex reasoning tasks
  • Investigating the interaction between the learned connection and other stabilization techniques (e.g., fixed-point methods)
  • Exploring whether the emergent self-verification behavior can be leveraged for interpretability or verification-aware training

The code is publicly available at https://github.com/pixeli99/InfiLoop.

Related papers