PAR: Progress-Aligned Recurrence for Length Generalization in State Tracking

Pan Wang; Yihao Hu; Guangyuan Dong; Yitong Lin; Hang Wang; Zimi Liu; Zhihao Wen; Jianshe Li

cs.CLOct 20, 2026

PAR: Progress-Aligned Recurrence for Length Generalization in State Tracking

Authors: Pan Wang, Yihao Hu, Guangyuan Dong, Yitong Lin, Hang Wang, Zimi Liu, Zhihao Wen, Jianshe Li,

Organizations: Ant International · National University of Singapore

Abstract

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- $S_5$ , a Qwen2.5-0.5B model trained on at most 32 moves achieves 95.5% mean full-state accuracy at 512 moves, or $16\times$ the training length, compared with 55.7% reported for RLT and 0.8% for explicit chain of thought on the same backbone. Controlled $S_5$ ablations show that removing norm anchoring reduces the extrapolation factor from $16\times$ to below $2.2\times$ , and replacing progress-aligned targets with final-answer supervision reduces it to at most $8\times$ . 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 $z_{k}$ , scaling compute with problem length. (b) Test accuracy on single-query $S_{5}$ (trained on depth $\leq 16$ , chance 0.2, dashed); PAR sustains near-perfect accuracy to $16\times$ ( $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 $z_{k}$ 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.

	$S_{5}$ , train $d\leq 16$				$A_{5}$ , train $d\leq 16$	Swap- $S_{5}$ , train $L\leq 32$
Method	32 ( $2\times$ )	64 ( $4\times$ )	128 ( $8\times$ )	$G$	$G$	256 ( $8\times$ )	512 ( $16\times$ )
Transformer	–	–	–	–	–	$0.85\pm 0.30$	$0.85\pm 0.30$
RLT	–	–	–	–	–	$97.3\pm 2.8$	$55.7\pm 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\times$ core (48 layers)	39	14	–	1.76	–	–	–

Table 1: Main results. Accuracy (%) with

R=d

iterations on

S_{5}

(one query, chance 20%) and full-state accuracy (%) on the swap-

S_{5}

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

S_{5}

, 2 on

A_{5}

and 6 on swap-

S_{5}

(1,024 test sequences per length). –: not reported.

Method (swap- $S_{5}$ , train $L\leq 32$ )	256 ( $8\times$ )	512 ( $16\times$ )
RLT	$97.3\pm 2.8$	$55.7\pm 25.8$
PAR, oracle $R=L$	$98.7\pm 3.1$	$93.2\pm 16.3$
PAR + learned halting	$\mathbf{98.7\pm 3.1}$	$\mathbf{92.9\pm 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\times

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\times$ )	$G$
Fixed depth ( $R=1$ )	–	1.62
Untied $4\times$ 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

S_{5}

(train

d\leq 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.
	( $\leq 8\to 64$ )	( $\leq 5\to 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	$\mathbf{\geq 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\times$ )	512 ( $16\times$ )
Single iteration ( $k=m=1$ )	–	$77.8\pm 25.7$
8 in a row	$\mathbf{98.7\pm 3.1}$	$90.3\pm 16.4$
6 of 8 (PAR)	$\mathbf{98.7\pm 3.1}$	$\mathbf{92.9\pm 16.2}$

Table 8: Stopping rules for the jointly trained read-out (swap-

S_{5}

, 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	${\approx}16$
EMA 0.99	97	61	$\mathbf{\geq 16}$
Live student + depth curriculum	98	85	$\mathbf{\geq 16}$

Table 9: Teacher configurations for prefix self-distillation on

S_{5}

(train

d\leq 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 $\times$ 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:

S_{5}

accuracy by test depth (single-query proportions, chance 0.2). These earlier runs sample

R\leq 32

in training; the

c=1

R=d

rows use

16\times

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\times$ )			512 ( $16\times$ )
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-

S_{5}

, 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(\texttt{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\times$	$16\times$	q-acc $16\times$
$n_{k}{=}2$ , seed 0	73.8	0.4	22.1
+ positional stretch	76.2	26.7	74.9
$n_{k}{=}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
$n_{k}{=}4$ , seed 4	72.3	5.5	39.4

Appendix

Table 12: Prefix self-distillation on swap-

S_{5}

(train

L\leq 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.

Length Generalization

Length generalization studies whether sequence models trained on short inputs can reliably solve the same tasks on longer sequences.

See all 47 papers

RecurTrace: Adaptive Latent Reasoning with Loop-Time Memory
Sep 3, 2026Yuxiang Wang, Kunyu Feng, Yingda Shen +3Memory-Augmented Language Models Recurrent Transformers
The Chinese University of Hong Kong, Shenzhen · The Chinese University of Hong Kong · Tianjin University +2
Retrofitting Recurrent Depth into a Pretrained Language Model: Installation, Extrapolation, Transfer, and Retention at Two Parameter Budgets
Jul 31, 2026Mark ShapiroLLM Fine-Tuning Recurrent Transformers
CHASE: Cache-Hole-Adapted Skip Exit for Looped State-Space Language Models
Jul 11, 2026Zhenxuan Yu, Takeshi Kojima, Yutaka Matsuo +1Efficient Language Model Inference LLM Reasoning
Rikkyo University · The University of Tokyo

Sep 3, 2026J/K move · Enter open · S save

RecurTrace: Adaptive Latent Reasoning with Loop-Time Memory

Yuxiang Wang, Kunyu Feng, Yingda Shen, Haoning Xu+2

The Chinese University of Hong Kong, Shenzhen · The Chinese University of Hong Kong · Tianjin University+2

Memory-Augmented Language Models Recurrent Transformers LLM Inference Acceleration Adaptive Computation

Open paper AskPDF arXiv

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.

Abstract

Length Generalization

Explore similar work

RecurTrace: Adaptive Latent Reasoning with Loop-Time Memory

Retrofitting Recurrent Depth into a Pretrained Language Model: Installation, Extrapolation, Transfer, and Retention at Two Parameter Budgets

CHASE: Cache-Hole-Adapted Skip Exit for Looped State-Space Language Models