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

AI summary

Overview

Research area: Deep learning architecture and inference-time computation scaling for reasoning models (looped/recurrent Transformers, latent reasoning, test-time compute).

Technical level: Intermediate to Advanced. The motivation and empirical results are accessible, but the core contribution involves attention-weight math, a streaming recurrence, and complexity analysis.

Scope: This paper proposes InfiLoop, a learned residual connection that replaces the "carry-last" state-replacement rule in looped Transformers, and shows that it enables useful reasoning depth to scale to tens of thousands of test-time iterations.

What This Paper Is About

Looped Transformers reuse one shared block many times, which adds computational depth without adding parameters. The standard way to connect iterations is to pass each new block output directly into the next, a rule the authors call "carry-last." The paper argues that this rule is a bottleneck: a noisy new output overwrites correct intermediate deductions, so later loops cannot recover, and accuracy falls as depth grows. The goal is a different connection rule that lets the model learn which earlier computations to keep.

Key Contributions

1. A diagnosis of carry-last failure. The authors measure how a carry-last recursive reasoner (self-attention TRM) behaves as depth grows and show that it overwrites correct predictions, loses already-solved puzzles, and its state keeps moving after solving.
2. InfiLoop, a loop-native attention residual connection. Instead of h_t = z_t, the state becomes a content-weighted, temporally decayed sum over all earlier recurrent candidates. A shared learned pseudo-query scores each candidate when it is produced, and a learned exponential decay discounts older candidates.
3. An exact constant-memory recurrence. Because the scorer is shared and decay is geometric, the aggregation can be maintained with a weighted-state vector N_t and a normalization scalar D_t per token, giving O(d) persistent memory and O(d) computation per iteration, and O(Td) total aggregation cost over T iterations. No approximation to the attention weights is introduced.
4. An interpretation as an adaptive residual update and an emergent self-verification behavior. The aggregation is shown to be equivalent to h_t = h_{t-1} + γ_t(z_t − h_{t-1}) with γ_t ∈ (0,1); empirically γ_t stays suppressed during unsolved search, surges at the update that completes a solution, and returns to baseline afterward.

Main Findings

• Depth degrades carry-last but not InfiLoop: On Sudoku-Extreme, TRM peaks at 62.2% exact accuracy at 468 layers and declines to 58.1% at 1,872 layers, while InfiLoop over the same depth reaches 90.9%. Each outer update executes 26 layers.
• Later loops overwrite correct cells: During updates 36 to 72, TRM changes 77.0 predictions per puzzle on initially empty cells and its number of incorrect cells rises by 0.21 per puzzle; InfiLoop makes 24.0 changes and removes 1.25 incorrect cells per puzzle.
• Correct predictions are lost: At update 36 both models predict 1.42 million initially empty cells correctly. During later updates TRM turns 20.9% of these into an incorrect value at least once; InfiLoop turns 2.4%.
• Solutions are retained or lost: TRM begins losing solutions after 9 updates and has lost 19.7% of them after 36 updates. InfiLoop's state change falls to 0.3% per update and every puzzle it solves remains solved 36 updates later.
• Full benchmark results (7M-parameter InfiLoop): 97.9% exact accuracy on Sudoku-Extreme (vs. FPRM 94.2%, reported TRM-Att 74.7%), 87.4% on Maze-Hard (vs. TRM-Att 85.3%, FPRM 87.0%), 49.1% pass@2 on ARC-AGI-1 (vs. FPRM 47.5%, TRM-Att 44.6%), and 13.6% pass@2 on ARC-AGI-2 (vs. TRM-Att 7.8%, FPRM 6.2%).
• Reproduced baselines: Retraining TRM on Sudoku-Extreme with the authors' own recipe gives 73.6%, close to the reported 74.7%; evaluating the released ARC-AGI-2 verification checkpoint under the paper's protocol yields 4.6% pass@2.
• Fewer layer evaluations for the same accuracy: InfiLoop reaches 90.9% at 1,872 layer evaluations, exceeding FPRM's 90.1% at 3,744 layer evaluations, halving the required computation. Matched at 3,744 layers, InfiLoop reaches 92.4% against FPRM's 90.1%.
• Test-time looping keeps helping: The accuracy InfiLoop reaches at 936 layer evaluations (86.5%) is never matched by TRM even at 24,960 layers (83.65%). At 24,960 layers, accuracy is 92.74% for InfiLoop, 91.63% for FPRM, and 83.65% for TRM. On Sudoku-Extreme, InfiLoop continues to improve beyond 20,000 effective test-time steps.
• Difficulty scaling: Evaluated on 32,768 test puzzles at matched 1,872 layers, InfiLoop and FPRM both reach 100.0% on puzzles with 45 to 51 empty cells. At 52 empty cells InfiLoop attains 100.0% versus 96.5% for FPRM and 95.8% for TRM; at 55 empty cells 93.4% versus 90.3% and 84.2%; at 58 empty cells 84.1% versus 77.9% (margin peaks at +6.2 points) with a +19.6-point margin over TRM; at 60 empty cells 59.5% versus 56.5% and 36.0% (+23.4 points over TRM).
• Learned step size reveals self-verification: With equal weights, γ_t would converge to 1 − β = 0.56. Observed, the mean step size stays near 0.40 during unsolved search, surges to 0.91 at the update that completes a puzzle, and returns to the 0.56 baseline within six updates. For unsolved puzzles it stays suppressed across all iterations.
• Content scores track correctness without supervision: The clipped score ln e_t = clip(s_t, −8, 8) has a median of 7.25 for candidates decoding to a fully correct board (upper clipping bound 8.0), falls to 2.75 at 20 wrong cells, and to −0.58 at 33 wrong cells. Removing the content score during training or at test time lowers accuracy by 6 to 7 percentage points.
• Robustness to perturbation: Injecting a relative perturbation ε ∈ [10⁻⁴, 0.3] into the state at update 36 and unrolling 36 more updates, InfiLoop retains 100% of solved puzzles at update 72 across all tested scales, while TRM loses 15–20% of solutions even without perturbation. For TRM even ε = 10⁻⁴ causes the trajectory to diverge to relative distance of order one within 12 updates; InfiLoop keeps displacement bounded.
• Both recurrent axes benefit (7,680-puzzle ablation): Applying the connection to both axes gives 82.1% at H3L6 versus 72.6% (inner alone) and 75.1% (outer alone), widening to 89.7% versus 79.4% and 80.0% at H32L6.
• Decay direction and initialization matter: With fixed β, reverse decay β = 1.82 solves only 2.0% at H3L6 (2.0% at H32L6), uniform weights β = 1 reach 38.4% (40.2%), forward decay β = 0.55 reaches 74.0% (77.9%). Learning β from β₀ = 0.55 gives the best results, 82.1% at H3L6 and 89.7% at H32L6. Initializations β₀ = 0.10 and β₀ = 0.90 reach 66.3% and 63.3% at H3L6.
• Constant persistent state: InfiLoop stores 582.8 KiB of persistent state per puzzle, independent of the loop budget (194.0 KiB for two BF16 carries plus 388.8 KiB for FP32 numerator and denominator per position on each axis). Carry-last memory of O(Td) is reduced to O(d), and total aggregation cost from O(T²d) to O(Td).

Methodology in Plain English

The authors start from the observation that in a looped Transformer, the only way information from an early iteration reaches a late iteration is through the single most recent state. If one update is noisy, that noise propagates and correct partial work is discarded.

Their fix is to change what gets passed between iterations. Instead of forwarding only the newest block output, they keep a running weighted average of all previous block outputs. Each output gets a score from a small learned query vector (the same query is used at every iteration), and older outputs get multiplied by a decay factor each step, so recent, high-scoring outputs dominate. Crucially, the same query and a geometric decay mean the running average can be updated with two small quantities — a vector sum and a scalar mass — rather than by storing every past output and attending over all of them. That keeps memory flat no matter how many iterations run.

They prove this running average is exactly equivalent to a residual update that nudges the old state toward the new candidate by a step size γ_t. They then plug the connection into a compact 7M-parameter nested recurrent reasoner (two Transformer layers, hidden dimension 512, eight heads, SwiGLU with expansion factor four, a depthwise 3×3 convolution, and a 513-parameter scorer–decay pair per recurrent axis), train it on Sudoku-Extreme, Maze-Hard, and ARC, and compare against carry-last TRM, FPRM, HRM, Attractor Models, EqR, URM, and frontier chain-of-thought LLMs. They also run diagnostic experiments: tracking state change after solving, perturbing the recurrent state, measuring accuracy against executed layers, and inspecting the learned step size and content scores.

Training used AdamW with peak learning rate 10⁻⁴, weight decay 1.0, and batch size 768 on Sudoku-Extreme, with H = 3 outer loops and L = 6 inner updates over 12 deep-supervision segments, each running 21 shared-block applications with gradients through the final seven. Maze-Hard and ARC used H = 3, L = 6 over up to 16 supervision segments with peak learning rate 3 × 10⁻⁴.

Why This Matters

The paper reframes a fundamental architecture choice: how a recurrent or looped model carries information across iterations. It shows that "keep only the latest state," inherited from standard deep networks, actively destroys correct intermediate reasoning as depth grows, and that replacing it with learned, decayed, content-weighted aggregation restores the benefit of added depth. This is relevant to the broader question of whether more test-time compute translates into better reasoning — the paper reports it does for InfiLoop on Sudoku-Extreme beyond 20,000 effective steps, whereas carry-last baselines plateau and degrade.

Real-world applications:

• Constraint-satisfaction and scheduling problems, where partial assignments must be preserved rather than overwritten during search.
• Logical and mathematical puzzle solving, including the Sudoku-like and shortest-path tasks the paper evaluates.
• Abstract visual reasoning and program induction, via the ARC-AGI benchmarks for inducing transformation rules from few demonstrations.
• Any deployed inference system that spends variable compute per query and needs stable, auditable intermediate reasoning states.

Industry relevance: The constant-memory property (582.8 KiB per puzzle, independent of loop count) means test-time depth can be increased without growing the persistent aggregation footprint, which matters for inference serving where memory per request is a hard constraint. The reported compute efficiency — matching FPRM's 90.1% at half the layer evaluations — is directly relevant to cost per query. The paper is careful to note its FLOP counts are analytical and do not compare measured runtime, and that FPRM's solver operations are outside the InfiLoop FLOP accounting.

Future Directions

• Whether the same learned aggregation helps at larger parameter scales. The paper's evaluation centers on a 7M-parameter model and lists larger recursive models (27M HRM, 27M Attractor Models, 14M URM) only as reported baselines.
• Extending the approach beyond the four benchmarks tested. The paper evaluates Sudoku-Extreme, Maze-Hard, ARC-AGI-1, and ARC-AGI-2; generalization to language-model pretraining and other modalities is not reported.
• Understanding the emergent self-verification more deeply — why the shared scorer assigns high weight to valid configurations without explicit correctness supervision, and whether that signal can be used or trained for directly.
• Whether additional loop-native connections (for example, iteration-specific queries rather than a shared one) can improve results while preserving the constant-memory recurrence, given that the no-decay case with a tied query recovers the Full Attention Residuals rule.

Target Audience

Researchers and engineers working on recurrent and looped Transformer architectures, latent reasoning, and test-time compute scaling. It is also relevant to practitioners building compact reasoning systems for constrained-satisfaction or abstract-reasoning tasks, and to readers interested in why deep recurrent models lose correct intermediate results. Readers need comfort with attention mechanisms, residual connections, and complexity notation to follow the derivation sections, though the experimental results and the adaptive-update interpretation (Equation 5 and Figure 4) are readable on their own.

Authors’ abstract

In this paper, we argue that looped Transformers need their own residual connections to prevent performance degradation as the number of iterations grows. We observe that increasing loop iterations can reduce reasoning accuracy: noisy state updates overwrite correct intermediate deductions and even undo completed solutions. This leaves subsequent iterations to recover lost information from an already degraded representation: once an error arises in an earlier loop, often as a result of long-range propagation through the recurrence, later loops find it difficult to correct. In this paper, we introduce InfiLoop, a loop-native residual connection that learns which past computations to retain and how much to accept from each new update. InfiLoop combines content-based weighting with learned temporal decay to maintain a running summary of recurrent states. An exact streaming recurrence keeps its persistent aggregation memory constant as the loop count grows. The resulting adaptive update suppresses unreliable proposals and preserves useful intermediate states. Across extensive reasoning tasks, a 7M-parameter InfiLoop model outperforms existing recursive architectures, reaching 97.9% exact accuracy on Sudoku-Extreme, and 13.6% pass@2 on ARC-AGI-2. Notably, on Sudoku-Extreme, InfiLoop continues to improve with test-time looping beyond 20,000 effective steps, showing that added depth translates directly into stronger reasoning. Our code is available at https://github.com/pixeli99/InfiLoop.