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

AI summary

Overview

Research area: Natural language processing / machine learning — specifically looped (recurrent-depth) language models, length generalization, and latent reasoning on sequential state-tracking tasks.

Technical level: Intermediate. The paper assumes familiarity with transformer architecture, cross-entropy/KL training objectives, and the standard length-generalization framing, but its central idea — matching loop index to task progress — is explained in accessible terms.

Scope: The paper proposes Progress-Aligned Recurrence (PAR), a training recipe that makes each iteration of a retrofitted looped language model predict the state after a matching number of task operations, plus two supporting mechanisms (norm anchoring and learned halting), and evaluates it on group-word and rule-chain state-tracking tasks.

What This Paper Is About

Looped language models reuse a shared block of layers repeatedly, so in principle they can spend more computation on longer problems. In practice, extra iterations often do not help: the training signal (typically the final answer) never says how far each iteration should advance, and repeated application of the shared block can change the magnitude of the recurrent state. The paper's goal is to make additional iterations actually useful by defining what each iteration should compute, and to show that this yields extrapolation far beyond the training length.

Key Contributions

1. A reframing of length generalization as a depth–length correspondence. The authors identify two failure modes — state drift across iterations and unaligned loop targets — and propose PAR, which trains the prediction at iteration k to match the state after the first k operations, combined with rescaling each token's recurrent state to its first-iteration RMS norm.

2. A parameter-free learned halting rule. The existing coda read-out is trained to predict completion (Yes/No) using class-balanced cross-entropy at weight 0.5, so inference can stop without knowing the problem length; a multi-step confirmation window (τ = 0.5, 6 of 8 iterations) selects the answer.

3. Prefix self-distillation for label-free settings. When intermediate state labels are unavailable, prefixes that are themselves valid shorter problems are used as targets: a teacher (an EMA of the student at decay 0.99, or the live student) evaluated on the prompt truncated after c·k steps supervises iteration k via KL divergence, learning the correspondence from final answers alone.

4. Empirical demonstration of 16× extrapolation. On swap-S₅ with a Qwen2.5-0.5B backbone trained on at most 32 moves, PAR reaches 95.5% mean full-state accuracy at 512 moves, versus 55.7% reported for RLT and 0.8% for explicit chain of thought, with consistent benefits across Qwen2.5 backbones up to 7B and Qwen3 backbones up to 4B.

Main Findings

• Depth extrapolation on swap-S₅. A Qwen2.5-0.5B model trained on at most 32 moves reaches 95.5 ± 11.0% mean full-state accuracy at 512 moves (16× the training length), compared with 55.7 ± 25.8% for RLT and 0.8% for explicit CoT. RLT is stronger at 8× (97.3 ± 2.8% vs. PAR's 99.5 ± 1.2%). At 256 moves PAR scores 99.5 ± 1.2%.

• Single-query S₅ and A₅ results. With training depth ≤ 16 and R = d iterations, PAR achieves the maximum measurable extrapolation factor G = 16 on both S₅ and A₅ across all seeds, with 100 / 100 / 99.3% accuracy at depths 32 / 64 / 128 on S₅. Fixed-depth (R = 1), untied 4× core (48 layers), and random-R recurrence baselines all stay below 2× (G of 1.62, 1.76, and 1.95 respectively).

• Explicit CoT collapses. CoT is perfect within the training range but falls to 35.6% at 64 moves, 3.1% at 128, and 0.8% at 256 and 512 on swap-S₅; on single-query S₅ it drops from 100% at depth 16 to 12% at 32, 4% at 64, and 0% at 128.

• Both components are necessary. Removing norm anchoring drops the extrapolation factor from 16× to below 2.2× (prefix targets without norm stall at ~1.8× under both schedules). Replacing progress-aligned targets with final-answer supervision reduces it to at most 8× — supervising the final answer at every loop stalls near G ≈ 2, and supervising it only at iteration d reaches G = 1.48.

• Larger backbone scaling is real but seed-variable. On S₅ at d = 256 (16×), the best PAR seed hits 100% on Qwen2.5-1.5B, 3B, and 7B, while fixed-depth baselines stay at or below 26.5%; means are 71.7 ± 42.7, 57.8 ± 42.5, and 80.7 ± 33.1 respectively. On Qwen3, PAR reaches 36.2% (0.6B), 38.8% (1.7B), and 53.0% (4B) at d = 256 versus near-chance fixed depth. The authors explicitly state this does not establish a monotonic trend with model size.

• Learned halting works without supplied depth. With the "6 of 8" rule, PAR reaches 92.9 ± 16.2% at 512 moves (versus 93.2 ± 16.3% for the oracle R = L run and 55.7% for RLT), and 98.7 ± 3.1% at 256. A single-iteration stopping rule is brittle (77.8% at 512), and requiring 8 in a row is overly strict (90.3%).

• Halting must be trained jointly. Adding the read-out by fine-tuning an existing PAR checkpoint makes P(Yes) spike and fade within two iterations, so multi-step windows fail to trigger; training on extra iterations to create an absorbing state instead causes accuracy to collapse below 1% at 16×.

• Label-free learning works, with caveats. Prefix self-distillation is the only label-free recipe sustaining high accuracy at 16× (85% at d = 256, G ≥ 16). Label-supervised PAR remains stronger at the longest depths; positional-stretch fine-tuning lets label-free models reach up to 99.7% at 16× on swap-S₅.

• The inductive bias of the task matters. On the rule-table chain (an inherently sequential lookup task), PAR-hold reaches G ≥ 8 (94% at depth 64) and PAR reaches 6.0, while prefix targets without norm anchoring stall at 1.78. On graph reachability, which has no canonical step-wise intermediate state, all recipes stay below 3.5×.

• Format transfer. On a BIG-Bench-Hard-style natural-language object-tracking format, PAR-hold with positional stretch reaches 94–96% at 32 swaps and 79% at 64 swaps (chance 17%), versus at most 67% and 21% for random-R and fixed-depth baselines with the same stretch.

• Positional stretch alone does not create the correspondence. Applied to fixed-depth and random-R models it leaves both at chance by depth 48 (23% and 18%); applied to a PAR-hold model with step size c = 2 it keeps 99% accuracy at depth 512, i.e., 32× the training depth, in one fixed-R = 288 run.

Methodology in Plain English

The authors retrofit recurrence into a pretrained Qwen2.5-0.5B by splitting its 24 layers into a prelude (layers 1–8), a weight-tied core (layers 9–16) that is applied repeatedly, and a coda (layers 17–24) that reads out an answer from any iteration. An adapter initialized as [0 | I] ensures that with one iteration the model initially computes exactly what its parent did. Gradients flow through the last 8 iterations; training uses AdamW at learning rate 2×10⁻⁵ (10⁻⁴ for the adapter), 200 warm-up steps, and one epoch over approximately 400k examples spread uniformly over training depths.

The core idea is a bookkeeping rule for the training signal. If a problem takes d sequential moves, then iteration k is trained to predict the state after the first k moves (with step size c = 1) — not the final answer, which would reward shortcuts that don't scale. Two schedules satisfy this: "forced depth," where each batch has one depth and the model runs ⌈d/c⌉ iterations with no loss after completion, and "hold," where the iteration count is drawn from a heavy-tailed distribution (1 + Poisson(λ) with log-normal λ, mean 12, cap 48) and later iterations keep predicting the final state.

Two stabilizers keep the repeatedly-applied step map from drifting. State-norm anchoring rescales every token's state after each iteration to the RMS norm of its first-iteration state, with a stop-gradient on the reference. Positional stretch, used in a short fine-tuning stage, multiplies position indices by s ~ U[1, s_max] while ramping s_max from 1 to 8, then tests with ordinary positions.

For stopping, the coda's existing read-out is trained — with no added parameters — to put probability on "Yes" when c·k ≥ d and "No" otherwise. At test time the model picks the first iteration where P(Yes) > 0.5 and at least 6 of the next 8 iterations also exceed 0.5, then returns the answer from that starting iteration.

Experiments compare against explicit CoT, fixed-depth, a 48-layer untied model, random-R recurrence, and reported RLT/Transformer numbers, on group-word problems over five boxes (swap-S₅ with transpositions, A₅ with 3-cycles, chance 0.2), the rule-table chain of Shapiro (train ≤ 8, test to 64), and the graph reachability task of Chen (train ≤ 5 hops, test to 24). The extrapolation factor G = d₅₀/d_train measures the first depth where accuracy falls below the midpoint between chance and 100%.

Why This Matters

The paper argues that what an iteration means is a training target in its own right, alongside how many iterations to run — a design decision that current recurrent-depth work tends to leave implicit. It also shows that a learned completion signal can be extracted from the output head already present in a retrofitted model, without extra parameters, ponder losses, or knowledge of the problem length at test time. Since latent recurrence writes no intermediate text, keeping every iteration decodable into the state it holds also offers a route to inspecting latent computation that chain of thought provides by default.

Real-world applications:

• Multi-step tool-use and agent pipelines, where an agent must apply a sequence of operations whose length varies at inference time.
• Simulating or verifying discrete protocols — state machines, inventory or location updates, transaction sequences — where intermediate states are well defined.
• Efficient inference for long procedural inputs, since no intermediate tokens are generated at training or test time.
• Reasoning systems where the number of required steps is not known in advance, using the learned halting read-out instead of a supplied depth.

Industry relevance: The method is retrofitted onto an off-the-shelf pretrained model rather than trained from scratch, and the demonstrated gains appear on small backbones (0.5B, 0.6B) as well as larger ones — relevant for cost-sensitive deployment of sequential reasoning. The large cross-seed variance on larger models (standard deviations of 33–43 points at 16× on Qwen2.5-1.5B to 7B) is a caveat for production adoption.

Future Directions

• Open-ended natural-language reasoning. The authors leave application of the correspondence to open-ended reasoning to future work; the current evaluation is on algorithmic and templated state-tracking tasks where intermediate states are well defined.
• Tasks without sequential decomposition. Graph reachability, solved by expanding a frontier in parallel, falls outside the method's scope — understanding when a canonical step-wise state exists remains open.
• Seed variance and scale stability. The paper explicitly does not establish a monotonic trend with model size; characterizing the mechanisms behind seed variance on larger backbones and stabilizing recurrent scaling are named as open questions.
• Broader evaluation coverage. Swap-S₅, learned halting, and prefix self-distillation are evaluated only on Qwen2.5-0.5B, and the scaling study covers three seeds in the single-query S₅ setting only, leaving wider replication open.

Target Audience

Researchers and graduate students working on recurrent-depth or looped transformer architectures, length generalization, and latent reasoning; practitioners interested in retrofitting recurrence into pretrained models for sequential tasks; and anyone evaluating whether extra test-time computation reliably translates into solving longer problems. Readers without a background in transformer training objectives will find the architectural details dense, but the central progress-alignment idea is accessible.

Authors’ 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 \textbf{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.