Progress-Aligned Recurrence for Length Generalization in State Tracking

Summary (Overview)

• Core contribution: The paper introduces Progress-Aligned Recurrence (PAR), a training method for looped language models that establishes a depth–length correspondence by training iteration $k$ to predict the state after the first $k$ operations, combined with RMS norm anchoring to stabilize recurrent state scale.

• Key result: On swap-$S_5$ state tracking, a Qwen2.5-0.5B model trained on at most 32 moves achieves 95.5% mean full-state accuracy at 512 moves (16× the training length), compared to 55.7% for RLT and 0.8% for explicit chain-of-thought on the same backbone.

• Two essential components: Controlled 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×.

• Additional mechanisms: PAR includes a learned halting rule (no parameters added) for autonomous stopping without a supplied depth, and prefix self-distillation that induces the correspondence from final-answer labels alone when intermediate labels are unavailable.

• Scope: Gains generalize across sequential state-tracking tasks and multiple backbones (Qwen2.5 up to 7B, Qwen3 up to 4B), but not to tasks without sequential decomposition (e.g., graph reachability).

---

Introduction and Theoretical Foundation

Background and Motivation

Length generalization asks whether a model trained on short problems can solve longer instances of the same task. Sequential state tracking provides a controlled setting: each operation updates a state used by later operations (e.g., rule-table chains, object tracking). The rules remain unchanged as the sequence grows; the model must apply them correctly over more steps.

Looped language models offer a way to make computational depth grow with length by repeatedly applying a shared transformer block to hidden states. However, on longer problems, neither additional parameters nor additional iterations help by itself:

• An untied model with 4× core depth extrapolates less than 2×
• A recurrent model with more iterations extrapolates less than 2×
• Fixed-depth Qwen2.5 models from 1.5B to 7B stay near chance at 16×

Theoretical Foundation: Depth–Length Correspondence

The paper frames length generalization in looped models as establishing a depth–length correspondence: iteration $k$ should hold the intermediate task state $z_k$ (the answer after the first $k$ operations). This addresses two twin failure modes:

1. Unaligned loop targets: Final-answer supervision at every iteration encourages early predictions and weakens extrapolation (consistent with Chen et al. [4] on graph reachability).
2. State drift: Recurrent states change scale as iterations accumulate, breaking the stationarity of the step map when the shared block is reused.

The correspondence is formalized by the training loss:

$$\mathcal{L}_{\mathrm{PAR}} = \sum_{k \in \mathcal{K}} \mathrm{CE}\left(Q(s_k), z_{\min(ck, d)}\right)$$

where $z_j$ is the answer after the first $j$ steps, $c$ is the step size (typically $c=1$), $d$ is the total number of operations, and $\mathcal{K}$ contains iterations inside the gradient window.

---

Methodology

Architecture: Retrofitted Looped Language Model

Following McLeish et al. [20], the 24 layers of Qwen2.5-0.5B are split into:

• Prelude $P$ (layers 1–8): computes static input features $e = P(x)$
• Weight-tied core $C$ (layers 9–16): iterated recurrently
• Coda $Q$ (layers 17–24): reads answers

The recurrence is:

$$s_k = C\left(W[e; s_{k-1}]\right), \qquad k = 1, \ldots, R$$

where $W \in \mathbb{R}^{h \times 2h}$ is an adapter initialized to $[0 \,|\, I]$, so with $R=1$ the retrofitted model initially computes the same function as its parent.

Key Components

1. Progress-Aligned Targets: Each iteration $k$ is trained to predict $z_{\min(ck, d)}$ — the state after exactly $ck$ operations. Two depth schedules satisfy this:

• Forced depth (default): each batch has a single depth $d$, model runs $R = \lceil d/c \rceil$ iterations, no loss after completion
• Hold: $R$ drawn from heavy-tailed distribution (1 + Poisson(λ)), iterations past $d/c$ keep $z_d$

2. RMS Norm Anchoring: Rescales each token's state to its first-iteration RMS norm after each iteration:

$$s_k \gets s_k \cdot \frac{\mathrm{rms}(s_1)}{\mathrm{rms}(s_k)}$$

with a stop-gradient on the reference. This removes drift in state magnitude as the shared block is reused.

3. Learned Halting: The coda read-out is jointly trained to predict completion (Yes when $ck \geq d$, No otherwise) using class-balanced cross-entropy with weight 0.5. At test time, a "6 of 8" confirmation window selects the answer: first iteration $R$ with $P(\text{Yes}) > \tau$ such that at least 6 of iterations $R, \ldots, R+7$ also exceed $\tau$ ($\tau = 0.5$).

4. Prefix Self-Distillation (for label-free settings): Uses the fact that the first $j$ steps of a problem form a shorter problem with answer $z_j$. Iterations $k \geq \lceil d/c \rceil$ receive cross-entropy on $z_d$; earlier iterations are trained via KL divergence to match a teacher (EMA of student, or live student) evaluated on the prompt truncated after $ck$ steps.

5. Positional Stretch: During a short fine-tuning stage, position indices are multiplied by $s \sim U[1, s_{\max}]$ with $s_{\max}$ ramped from 1 to 8, addressing long-prompt drift.

---

Empirical Validation / Results

Main Results: Depth Extrapolation

Table 1 (key results on Qwen2.5-0.5B):

G = extrapolation factor ($d_{50}/d_{train}$), where $d_{50}$ is the first test depth at which accuracy falls below the midpoint between chance and 100%.

Larger Backbones

Table 2 (accuracy at $d=256$, 16×, on $S_5$):

Learned Halting Results

Table 4 (swap-$S_5$, full-state accuracy without knowing problem length):

Label-Free Learning (Prefix Self-Distillation)

Table 5 (on $S_5$, train $d \leq 16$, final-answer labels only):

Ablation: Both Conditions Are Necessary

Table 7 (on $S_5$, train $d \leq 16$):

Prior Benchmarks

Table 6 (G values on prior benchmark tasks):

Graph reachability has no canonical intermediate state (solved by parallel frontier expansion), so all recipes stay below 3.5×.

---

Theoretical and Practical Implications

Theoretical Implications

1. Defining iteration semantics matters: The paper demonstrates that "what an iteration computes" is a training target in its own right, alongside choosing how many iterations to run. Simply having more iterations or parameters does not confer extrapolation ability.

2. Stationarity of the step map: RMS norm anchoring is not an optional regularizer but a necessary condition for reusing a shared block hundreds of times. Without it, state representations drift and the step map changes what it computes.

3. Task structure determines feasibility: The correspondence targets sequential multi-step computation — the regime where fixed-depth transformers fail. Tasks without a sequential decomposition (like graph reachability's parallel frontier expansion) fall outside its scope.

4. Conflicting objectives break recurrence: Training the core to simultaneously advance intermediate states and act as a fixed point (absorbing state) introduces conflicting objectives that severely harm extrapolation. An independent completion read-out avoids this competition.

Practical Implications

1. Autonomous stopping: The learned halting rule works without knowing problem length at test time, matching oracle stopping within 1.6 accuracy points — critical for real-world deployment where input lengths are unknown.

2. Label efficiency: Prefix self-distillation enables learning the correspondence from final-answer labels alone, making the approach applicable to real data where intermediate states are not annotated.

3. Backbone compatibility: PAR works across Qwen2.5 (0.5B–7B) and Qwen3 (0.6B–4B) families, though results vary across seeds on larger models (standard deviations of 33–43 points at 16×), suggesting open questions about stabilizing recurrent scaling.

4. Comparison to alternatives: PAR surpasses RLT (a state-of-the-art latent recurrence method) and explicit chain-of-thought on the same backbone, demonstrating that retrofitted recurrence with proper supervision can outperform both token-level feedback and explicit reasoning traces.

---

Conclusion

Main Takeaways

• A looped language model extrapolates in depth when each iteration has a defined job: aligning iteration $k$ with the first $k$ steps of the problem, and keeping the step map stationary via norm anchoring, turns extra iterations into extra reasoning.

• A retrofitted 0.5B model trained with PAR surpasses the best published latent-recurrence model (RLT) on swap-$S_5$ at both 8× and 16× the training length.

• The same correspondence enables the model to learn when to stop (via a parameter-free learned halting rule) and to learn its intermediate states from final answers alone (via prefix self-distillation).

Limitations and Future Directions

• Scope: Evaluated only on algorithmic and templated state-tracking tasks with well-defined intermediate states; open-ended natural-language reasoning is left to future work.

• Task constraints: Tasks without sequential decomposition (e.g., graph reachability) fall outside the method's scope.

• Scaling concerns: Results do not establish a monotonic trend with model size; Qwen3 backbones do not approach Qwen2.5-0.5B at 16×. Characterizing seed variance on larger backbones and stabilizing recurrent scaling remain open questions.

• Future work: Applying the correspondence to open-ended natural-language reasoning, understanding the mechanisms of seed variance on larger backbones, and stabilizing recurrent scaling across model families.