Looped transformers stop forgetting solutions: a 7M model hits 97.9% on Sudoku-Extreme

Looped transformers are supposed to get smarter the longer they run, but in practice they often solve a puzzle and then overwrite the answer. InfiLoop fixes that with one change: instead of passing only the latest iteration forward, it passes a learned, exponentially decaying average of every iteration so far. A 7M-parameter model built this way reaches 97.9% exact accuracy on Sudoku-Extreme and 13.6% pass@2 on ARC-AGI-2, and its accuracy keeps climbing as you add loops at test time.

• 97.9% on Sudoku-Extreme vs. 94.2% for the next-best 7M recursive model (FPRM) and 74.7% for TRM
• 13.6% pass@2 on ARC-AGI-2, about 1.7x TRM’s reported 7.8%, from a model with 7M parameters
• 20.9% vs. 2.4%: the share of correct cells that TRM vs. InfiLoop later flips to a wrong value at least once
• 86.5% at 936 layers beats TRM run 26x deeper (83.65% at 24,960 layers)
• 582.8 KiB of state per puzzle, constant regardless of how many loops you run

Carry-last recurrence throws away good work

Models like TRM (Tiny Recursive Model) and HRM reuse one small transformer block many times. Each pass reads the previous pass’s output and writes a new one. That “carry-last” rule means a single noisy update can overwrite a correct intermediate deduction, and nothing remembers what the model knew five steps ago.

The authors measure this directly. On Sudoku-Extreme, TRM peaks at 62.2% accuracy at 468 layers of depth, then declines to 58.1% at 1,872 layers. More compute makes it worse. Of 1.42M cells TRM had filled in correctly at update 36, it later flips 20.9% to a wrong value at least once. Among puzzles it had fully solved, it loses 19.7% within 36 more updates, because the state keeps churning at about 10% per update even after the answer is found.

InfiLoop, trained the same way, reaches 90.9% at that same 1,872-layer depth, flips only 2.4% of correct cells, and keeps every solved puzzle.

How it works: a softmax over your own history

The fix is a weighted average of all past iteration outputs, where each weight combines two things:

• A content score: a single learned query vector dot-producted with each normalized output. Higher-scoring states get more weight.
• A temporal decay: a learned factor β that downweights older iterations exponentially.

Naively this needs every past state in memory. The trick is that softmax-with-exponential-decay can be computed as a running sum. Each token keeps one numerator vector and one scalar normalizer, updated as N = βN + e·z and D = βD + e. The carried state is just N / D. Memory drops from O(Td) to O(d), and total aggregation cost from O(T²d) to O(Td). For nested TRM-style loops, each axis (inner and outer) gets its own pair, adding 513 parameters per axis.

Rewritten, the update is h = h + γ(z − h): a residual with a learned, input-dependent step size γ between 0 and 1. That step size turns out to be the interesting part.

It learned to verify its own answers

Nobody gave the scorer correctness labels, but it learned to track correctness anyway. The median content score is 7.25 for fully correct boards, 2.75 at 20 wrong cells, and −0.58 at 33 wrong cells. The step size behaves like a commit signal: it sits around 0.40 while the model is searching, spikes to 0.91 on the exact update that completes the puzzle, then settles. In effect, the model holds new proposals at arm’s length until one looks right, then locks it in.

That makes the solved state stable. When the authors inject random perturbations at update 36 (relative size from 10⁻⁴ to 0.3), InfiLoop retains 100% of solved puzzles at update 72 at every scale. TRM loses 15 to 20% even with no perturbation, and a 10⁻⁴ nudge diverges to an order-one difference within 12 updates.

The stability is what lets test-time depth pay off. InfiLoop keeps improving all the way to 24,960 layers:

The gap is largest on the hardest puzzles. At 60 empty cells, InfiLoop solves 59.5% vs. 56.5% for FPRM and 36.0% for TRM. Across the four main benchmarks the headline numbers are 97.9% Sudoku-Extreme, 87.4% Maze-Hard, 49.1% ARC-AGI-1, and 13.6% ARC-AGI-2. For reference, o3-mini-high and DeepSeek R1 score 0% on the first two and 3.0% and 1.3% on ARC-AGI-2.

Caveats

• Both pieces matter, and the decay is fragile. Applying the aggregation to only the inner or outer loop drops accuracy from 82.1% to 72.6 to 75.1%. A uniform average (no decay) gets 38.4%. Learned decay initialized at 0.10 or 0.90 instead of 0.55 lands at 66.3% and 63.3%. Setting β to 0 at test time, which recovers carry-last, collapses accuracy to 1.5%.
• It is not just damping, but the content score is a smaller share of the gain. Zeroing the content query during training costs 6.8 points (90.9% to 84.1%). The temporal decay does most of the stabilizing.
• One small architecture, puzzle-style tasks. Everything is a 7M-parameter block on the TRM codebase. Whether this carries over to looped language models is untested.
• Baselines are a mix of reported and reproduced numbers. The authors’ TRM reproduction gets 4.6% on ARC-AGI-2, not the reported 7.8%, so the ARC-AGI-2 gain is 1.7x or 2.2x depending on which you use. URM’s ARC figures are pass@1, so not directly comparable. ARC-AGI-2 evaluation is 120 tasks.
• No wall-clock numbers. Compute is reported in FLOPs (1.27 TFLOPs per puzzle at 1,872 layers, 17.0 at 24,960), not runtime.
• The “self-verification” reading is interpretive. The step-size spike and content-score correlation are suggestive, but not separately tested as a mechanism.

The takeaway for anyone building iterative inference: if more iterations make your model worse, the problem may not be the model but the recurrence. A running softmax over your own history costs one vector and one scalar of memory, and it turned test-time depth from a liability into a dial that keeps paying out.