PAR: Progress-Aligned Recurrence for Length Generalization in State Tracking
Abstract
Progress-Aligned Recurrence trains looped language models to match intermediate states and stabilize recurrent norms, enabling strong length extrapolation on sequential tracking tasks.
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 16times 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 16times to below 2.2times, and replacing progress-aligned targets with final-answer supervision reduces it to at most 8times. 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.
Community
🚀 Introducing PAR, our new work on length extrapolation in looped language models.
📦 The task is simple: five boxes, one object in each. Follow a sequence of swaps, then answer what’s inside. We train the model on short sequences and test whether it can keep up as the sequences get much longer.
Simply running more loops is not enough: performance breaks down at just 2× the training length. Scaling up the model alone doesn’t solve the problem either. On single-query S5, fixed-depth Qwen2.5 models from 0.5B to 7B achieve only around 20% accuracy at 16× the training length—roughly random guessing.
We identify two key obstacles: supervising only the final answer leaves the model without clear guidance on how far to advance in each loop; repeatedly applying the same layers hundreds of times causes the hidden-state scale to drift.
💡 PAR makes two straightforward changes: at loop k, predict the state after the first k steps; after each loop, rescale the hidden states to match their scale in the first loop. This aligns the loops with task progress: one loop, one step.
📈 On swap-S5, a 0.5B model trained on sequences of at most 32 steps achieves 95.5% full-state accuracy at 512 steps—a 16× length extrapolation. For comparison, the previous best method, RLT, achieves 55.7%, while the same model using chain-of-thought (CoT) reaches just 0.8%.
On single-query S5, PAR substantially outperforms fixed-depth models at every tested size, spanning Qwen2.5 0.5B, 1.5B, 3B, and 7B, and Qwen3 models from 0.6B to 4B. Across all four Qwen2.5 sizes, the best training runs achieve 100% accuracy at 16× the training length.
The model can also learn when to stop. We further show that PAR can be trained with final-answer supervision alone, without intermediate-state annotations.
Get this paper in your agent:
hf papers read 2610.27628 Don't have the latest CLI?
curl -LsSf https://hf.co/cli/install.sh | bash Models citing this paper 0
No model linking this paper
Datasets citing this paper 0
No dataset linking this paper
Spaces citing this paper 0
No Space linking this paper