Looped language models can run more iterations at test time, but additional iterations need not help them solve longer problems. For sequential state tracking, final-answer supervision does not specify how far each iteration should advance, and repeated updates can change the scale of the recurrent state. We introduce Progress-Aligned Recurrence (PAR), which trains the prediction at iteration k to match the state after the first k operations and rescales recurrent states to their first-iteration norms. On swap-S5, a Qwen2.5-0.5B model trained on at most 32 moves achieves 95.5% mean full-state accuracy at 512 moves, or 16× the training length, compared with 55.7% reported for RLT and 0.8% for explicit chain of thought on the same backbone. Controlled S5 ablations show that removing norm anchoring reduces the extrapolation factor from 16× to below 2.2×, and replacing progress-aligned targets with final-answer supervision reduces it to at most 8×. A learned completion readout selects when to answer without a supplied stopping depth. When intermediate labels are unavailable, prefix self-distillation induces the correspondence from shorter problem prefixes using only final answers. The gains apply across sequential state-tracking tasks, establishing the task structure required for depth extrapolation. These findings support training intermediate predictions and state scale together to make additional recurrent depth useful for longer problems.
Figures & tables
Figure 1: (a) Depth–length correspondence : iteration k predicts intermediate task state zk , scaling compute with problem length. (b) Test accuracy on single-query S5 (trained on depth ≤16 , chance 0.2, dashed); PAR sustains near-perfect accuracy to 16× ( d=256 ), whereas prior supervision recipes collapse (§ 6.1 ).
Figure 2: Overview of the PAR framework. (a) Architecture (§ 3.1 , § 3.3 ): Retrofitted prelude P , weight-tied core C , and coda Q , stabilized across iterations via RMS norm anchoring. (b) Progress alignment (§ 3.2 ): Recurrent iteration k predicts intermediate state zk after k operations, establishing an explicit depth–length correspondence. (c) Prefix self-distillation (§ 3.5 ): An EMA teacher evaluated on prefix 1:k supervises iteration k , inducing intermediate states from final-answer labels alone. (d) Learned halting (§ 3.4 ): The coda jointly predicts task completion ( Yes / No ); inference terminates when a multi-step confirmation window is met.
S5 , train d≤16
A5 , train d≤16
Swap- S5 , train L≤32
Method
32 ( 2× )
64 ( 4× )
128 ( 8× )
G
G
256 ( 8× )
512 ( 16× )
Transformer
–
–
–
–
–
0.85±0.30
0.85±0.30
RLT
–
–
–
–
–
97.3±2.8
55.7±25.8
Explicit CoT
12
4
0
1.247
1.125
0.8
0.8
Fixed depth ( R=1 )
19
20
17
1.62
1.32
–
–
Untied 4× core (48 layers)
39
14
–
1.76
–
–
–
Table 1: Main results. Accuracy (%) with R=d iterations on S5 (one query, chance 20%) and full-state accuracy (%) on the swap- S5 protocol of Zhang et al. [41] ; column headers give the test length and its multiple of the training length. G : extrapolation factor. ∗ Best R per depth. PAR: mean over 3 seeds on S5 , 2 on A5 and 6 on swap- S5 (1,024 test sequences per length). –: not reported.
Method (swap- S5 , train L≤32 )
256 ( 8× )
512 ( 16× )
RLT
97.3±2.8
55.7±25.8
PAR, oracle R=L
98.7±3.1
93.2±16.3
PAR + learned halting
98.7±3.1
92.9±16.2
Table 4: Learned halting: full-state accuracy (%) without access to the problem length L (6 seeds, 1,024 sequences per length, step-6000 checkpoints). Grey: the same models run for exactly L iterations, for reference.
Figure 3: Halting at 16× the training length ( L=512 , all 5 queries). Solid purple: joint training (6 seeds, 1,024 sequences); dashed red: read-out added by fine-tuning (2 seeds, 1,024 sequences). The shaded region marks iterations before completion ( R<L ) and the dashed vertical line marks R=L . (a) The completion read-out stays at 0 until the iteration count reaches L . For jointly trained models it then stays near 1, while for the fine-tuned read-out it fades within two iterations. (b) Accuracy against the final answer peaks exactly when the read-out fires and decays a few iterations later, so the answer must be read where the “finished” signal starts. Before L the iterations hold intermediate states, which differ from the final answer (horizontal dashed line: chance).
Method (final-answer labels only)
d=256 ( 16× )
G
Fixed depth ( R=1 )
–
1.62
Untied 4× core (48 layers)
–
1.76
Random- R recurrence
–
1.95
Final answer at every loop
21 / 21
2.14 / 1.99
Final answer at loop d
15
1.48
Final answer at loop d , norm
26 / 19 / 36
< 4 / < 4 / 8
Table 5: Extrapolation without intermediate labels on S5 (train d≤16 , one query, chance 20%; accuracy in %, “/” separates seeds). The final-answer recipes are defined in Table 7 . –: not evaluated at d=256 ; these models are at chance by d=48 (Table 10 ).
Method
Rule chain
Graph reach.
( ≤8→64 )
( ≤5→24 )
Final answer at every loop
1.45
2.4
Final answer at loop d
–
3.3
Prefix targets, no norm
1.78
3.45
PAR-hold
≥8
fail ∗
PAR
6.0
2.6
Table 6: G on benchmark tasks from Shapiro [28] (rule chain) and Chen [4] (graph reachability). Methods are defined in Table 7 . –: not run. ∗ Predicts the majority class.
Stopping rule
256 ( 8× )
512 ( 16× )
Single iteration ( k=m=1 )
–
77.8±25.7
8 in a row
98.7±3.1
90.3±16.4
6 of 8 (PAR)
98.7±3.1
92.9±16.2
Table 8: Stopping rules for the jointly trained read-out (swap- S5 , full-state %, same 6 seeds and checkpoints as Table 4 ). –: not reported.
Teacher
d=64
d=256
G
EMA 0.99 + curriculum
–
–
2.6
Slow EMA + confidence filter
–
–
5.6
EMA 0.99, forced depth
99
57
≈16
EMA 0.99
97
61
≥16
Live student + depth curriculum
98
85
≥16
Table 9: Teacher configurations for prefix self-distillation on S5 (train d≤16 , single-query accuracy in % at R=d ). –: not evaluated. The last row is the configuration reported in Table 5 .
Model
24
32
48
64
128
256
384
512
Fixed depth
0.73
0.19
0.23
0.20
0.17
–
–
–
Random- R (best R per depth)
0.93
0.55
0.23
0.20
0.25
–
–
–
Untied 4 × core (48 layers)
0.88
0.39
0.20
0.14
–
–
–
–
Fixed + stretch FT
0.84
0.35
0.23
0.16
0.16
0.19
–
–
Random- R + stretch FT
0.93
0.60
0.18
0.18
0.17
–
–
–
PAR-hold w/o state-norm
0.98
0.23
0.18
0.17
0.20
–
–
–
Table 10: S5 accuracy by test depth (single-query proportions, chance 0.2). These earlier runs sample R≤32 in training; the c=1 , R=d rows use 16× the largest training loop count at d=512 . Positional stretch helps only together with progress-aligned targets.
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
256 ( 8× )
512 ( 16× )
Seed
R=L
6 of 8
8 in a row
R=L
6 of 8
8 in a row
0
100.0
100.0
100.0
99.6
99.6
99.6
1
99.9
99.9
99.9
99.9
98.3
99.9
2
92.4
92.4
92.4
59.9
59.9
59.9
3
100.0
100.0
100.0
99.9
99.9
99.9
4
99.9
99.9
99.9
99.7
99.7
82.4
Appendix
Table 11: Learned halting per seed (swap- S5 , full-state %, 1,024 sequences per length, step-6000 checkpoints). R=L : oracle iteration count. “6 of 8”: select the first answer iteration R with P(Yes)>0.5 at which at least 6 of the iterations R..R+7 exceed 0.5, then return the prediction from R after confirmation; “8 in a row”: all 8 must exceed 0.5.
Starting model
8×
16×
q-acc 16×
nk=2 , seed 0
73.8
0.4
22.1
+ positional stretch
76.2
26.7
74.9
nk=4 , seed 2
52.7
0.0
26.7
+ stretch
99.9
99.7
99.8
+ stretch, second FT seed
99.2
98.4
99.2
nk=4 , seed 4
72.3
5.5
39.4
Appendix
Table 12: Prefix self-distillation on swap- S5 (train L≤32 , full-state %; q-acc: per-query accuracy, chance 20%). Models before the stretch fine-tune use 256 sequences per length, all others 1,024; every stretch fine-tuned model scores at least 99.2% in range (32 moves). The seed-4 and seed-5 starting models are step-5000 checkpoints.
Repeating a small block of middle layers increases a language model's effective inference depth without adding parameters or generating extra tokens, and recent work shows that this latent recurrence improves reasoning. However, two design choices limit these gains. Each iteration sees only the previous output and cannot directly access earlier computations. Moreover, a fixed loop count wastes depth on easy inputs while leaving hard ones with too little computation. We introduce RecurTrace, which addresses both limitations using the loop's own trajectory. Specifically, Loop Memory Attention lets each looped layer attend to its own states from previous iterations along the loop-time axis, so the model can revisit earlier computations instead of relying on the latest state alone. A halting head then reads the loop state and predicts whether to continue, with supervision from an oracle that identifies when additional depth still reduces loss. In a controlled MathQA comparison on the same looped backbone, RecurTrace achieves 56.9% accuracy with an average of 2.0 loops, exceeding the best fixed loop depth by 2.2 points at matched compute. By comparison, ACT and PonderNet collapse to one loop, and CALM reaches only 54.1% with 5.6 loops, while the stronger LoopUS-Conf and TaH-Mismatch baselines reach 55.3% at 3.2 loops and 55.7% at 2.1 loops. Finally, RecurTrace improves generation accuracy over same-budget fine-tuned baselines at 0.6B, 1.7B, 4B, and 8B, with the gain growing with model size from 0.6 to 3.4 points.
Yuxiang Wang, Kunyu Feng, Yingda Shen +3
The Chinese University of Hong Kong, Shenzhen · The Chinese University of Hong Kong · Tianjin University +2
A dense, pretrained language model can be retrofitted with recurrent depth and learn an iterative latent transition that persists after outcome-only annealing. Qwen2.5-0.5B-Instruct is split into a Prelude, a weight-tied Recurrent Block, and a Coda, with an identity-preserving one-loop path and a re-entry bridge on later loops. At loop 1 the retrofit remains non-inferior to its base on a preregistered ARC battery. Three findings. First, the mechanism is a reusable procedure rather than terminal-answer lookup, and installs at two budgets: 6M trained parameters over frozen base weights and 180M full-block. With intermediate-step supervision, the model computes one task step per loop and persists when only final answers are graded. The adapter matched the full block overall (83.8% versus 84.0%), led through depth 11, and trailed beyond. Verbal fine-tuning reached 79-86% on controlled verbal renderings (zero-shot transfer was minimal), and adapter verbal training begun from the installed mechanism outpaced matched fresh training by 18.6 points, including on a held-out test set. Second, the operation extrapolates to roughly 1.5 times its supervised depth, holding 70% accuracy through depth 18. Third, a same-size scratchpad-trained model matched the recurrent model within its learned horizon but collapsed beyond it. The recurrent model won overall, 84% versus 72%, retained 53% versus 2.5% beyond depth 10, and answered 7.6 times faster. An iterative transformer can therefore perform deeper reasoning in latent space faster than comparable or larger models fine-tuned on the same task, in a system-level comparison. A second task, running the rule in reverse, exposed the limits: the inverse was learnable in isolation, but no continuation acquired it while preserving the installed mechanism and general capability, a catastrophic-interference boundary. Learned depth selection remains open.
Recent work on looped language models suggests that many reasoning problems benefit from greater computational depth rather than from additional independent parameters. Existing studies, however, focus almost exclusively on Transformer backbones, leaving open whether this principle also applies to state-space language models. We investigate Looped Mamba and Looped Hybrid Mamba-Transformer architectures, which repeatedly apply a shared Mamba (or hybrid) block to introduce explicit finite-depth recurrent computation. On two controlled reasoning tasks-Mano (modular-arithmetic manipulation) and p-hop induction-Looped Mamba consistently outperforms parameter-matched non-looped baselines and, in several settings, matches or exceeds non-looped models of equal effective depth. We then extend the study to language model pre-training under matched iso-parameter and iso-FLOPs protocols, which jointly disentangle the effects of parameter sharing and effective depth: looped models remain competitive on downstream benchmarks with substantially fewer distinct parameters, although deeper non-looped models retain an advantage in validation perplexity under strict iso-FLOPs comparisons. Finally, we adapt Ouro's two-stage exit gate to Looped Mamba for threshold-controlled selection among recurrent-step outputs. Executing such exits on a state-space backbone, however, leaves the recurrent state without its deeper updates, and validation perplexity then degrades severely. We therefore introduce a cache-hole adaptation that aligns continued training with skipped-state inference. At the scales studied, the adapted model keeps perplexity close to full computation and matches or exceeds full-compute exit-state selection on downstream benchmarks while executing roughly half of the recurrent steps, which translates into measured inference speedups once the prefill is compute-bound.