One-sentence Summary

Researchers from Ant International and National University of Singapore propose Progress-Aligned Recurrence (PAR), which trains looped models to align iteration predictions with partial states and renormalize recurrent states, enabling a Qwen2.5-0.5B model trained on $32$32 moves to reach $95.5\%$95.5% full-state accuracy at $512$512 moves, a $16\times$16× extrapolation over prior methods.

Key Contributions

• The paper introduces Progress-Aligned Recurrence (PAR), which aligns each loop iteration with the state after the corresponding number of problem operations and applies RMS norm anchoring to stabilize recurrent state scale across iterations.
• PAR is extended with a learned completion readout that selects when to answer without a supplied stopping depth and with prefix self-distillation that learns intermediate targets from shorter problem prefixes using only final answers.
• Experiments show that a Qwen2.5-0.5B model trained on at most 32 moves reaches 95.5% mean full-state accuracy at 512 moves on swap-??_5, compared with 55.7% for RLT and 0.8% for explicit chain of thought, and ablations isolate the necessity of both norm anchoring and progress-aligned targets.

Introduction

Length generalization asks whether models trained on short problems can solve longer instances of the same task, and sequential state tracking tasks such as rule-table chains and object tracking provide a controlled setting because the rules stay fixed while the number of operations grows. Looped language models can increase computational depth with sequence length by repeatedly applying a shared transformer block to hidden states, but prior work shows that neither adding parameters nor running more iterations is sufficient on its own, and final-answer supervision at every loop can encourage early predictions and weaken extrapolation. Intermediate-state targets help define progress, yet they still lose accuracy beyond short extrapolation distances when recurrent state scale is uncontrolled. The authors introduce Progress-Aligned Recurrence (PAR), which matches each iteration to the state after the first k operations and uses RMS norm anchoring to control the scale of states passed between loops. They further add a learned halting rule and prefix self-distillation to obtain intermediate targets from final-answer labels alone, demonstrating 16x extrapolation on a state-tracking task.

Method

The authors introduce the PAR framework to establish a depth-length correspondence in language models, enabling them to scale computation dynamically with problem complexity. Refer to the framework diagram for a visual overview of the architecture and training pipeline.

To realize this architecture, the authors retrofit a standard Transformer by splitting its layers into three distinct components: a prelude, a weight-tied core, and a coda. The prelude processes the initial prompt to generate static input features. These features, along with the previous latent state, are concatenated and projected into the shared core via an adapter. The core iterates over the state according to the equation:

and the coda reads the answer from any given iteration. To ensure that the initial retrofitted model computes the exact same function as its parent network when run for a single iteration, the adapter is initialized to project only the static input features.

Defining an iteration requires aligning the model's recurrent steps with the actual progress of the computation. For a problem requiring a specific number of sequential steps to solve, the authors enforce a strict correspondence where each recurrent iteration must advance the computation by exactly a fixed number of steps. The training loss is computed by comparing the model's output at each iteration to the intermediate answer corresponding to that specific progress point:

This prevents the model from taking shortcuts and ensures that the number of iterations scales linearly with the problem depth.

Because the model applies this step map hundreds of times during extrapolation, any drift in the state representation could alter the computed function. To maintain a stationary step map, the authors implement state-norm anchoring. After each iteration, every token's state is rescaled to match the RMS norm of its state from the first iteration, preventing norm drift without affecting single-iteration behavior. Additionally, to handle long prompts where the query is far from early moves, a positional stretch technique is applied during a short fine-tuning stage to multiply position indices, ensuring robust extrapolation.

For autonomous inference, the framework incorporates a learned halting mechanism. The coda is trained to output a completion probability, predicting whether the computation has finished at the current iteration. During inference, the model selects the first iteration where the completion probability exceeds a threshold, provided that a confirmation window of subsequent iterations also agrees. This allows the model to terminate computation dynamically without requiring external depth signals.

Finally, to extend this correspondence to settings lacking intermediate step labels, the authors employ prefix self-distillation. In this approach, an exponential moving average teacher model evaluates truncated prompts to generate soft targets for the student model's earlier iterations. The student is trained to match these distributions using KL divergence, effectively inducing intermediate state representations using only final-answer supervision.

The model is trained using AdamW with specific learning rates for the adapter and the rest of the network, utilizing a warm-up schedule and a single epoch over a diverse set of training depths. No intermediate tokens are generated during either training or testing, ensuring efficient computation.

Experiment

The experiments evaluate PAR on permutation group word problems and prior sequential benchmarks, comparing it with fixed-depth, recurrent, untied, and explicit chain-of-thought baselines. The main results show that PAR's learned recurrent step map generalizes to much longer sequences, maintaining near-ceiling accuracy up to 16x training depth across model sizes, while fixed-depth and chain-of-thought models fail to extrapolate. Ablations confirm that correct iteration supervision and stationary state dynamics are both necessary, and that a learned halting signal or positional-stretch fine-tuning helps when step labels or long-distance position encodings are issues. On structured tasks the method supports sequential multi-step reasoning but does not confer the same benefit for parallel graph reachability.

PAR substantially outperforms the baselines on depth extrapolation, maintaining near-perfect full-state accuracy on the swap protocol out to the longest tested multiple and reaching the maximum measurable extrapolation factor on the permutation tasks. Fixed-depth, recurrent, untied, and explicit chain-of-thought models degrade instead, with explicit chain-of-thought falling to chance before the longest extrapolation. Ablations show that both supervised iteration targets and norm anchoring are necessary for reliable extrapolation. PAR keeps full-state accuracy above 95% at 16x extrapolation on the swap protocol, while RLT drops to 55.7%. On the permutation tasks, PAR reaches the maximum measurable extrapolation factor across seeds, whereas fixed-depth, recurrent, and untied baselines stay below 2x. Explicit chain-of-thought performs well within the training range but collapses to chance by 8x on the swap task. Removing norm anchoring or altering final-answer supervision sharply reduces long-depth accuracy compared with PAR.

PAR substantially outperforms fixed-depth baselines on the S5 16x extrapolation setting across all evaluated Qwen2.5 backbones. The best PAR seed retains perfect accuracy at horizon 256 on every backbone, while fixed-depth models stay near chance. PAR's cross-seed mean gains are large but variable, and the benefit does not scale monotonically with backbone size. Best PAR runs reach 100% accuracy at the 16x horizon on all tested Qwen2.5 backbones. Fixed-depth baselines remain near chance, with best accuracies under 27% across every backbone. PAR provides large mean accuracy gains over fixed depth at every model size, though seed-to-seed variance is substantial. Mean PAR accuracy is highest for the 7B model and lowest for the 3B model, showing no simple monotonic scaling trend.

Across the Qwen3 backbones tested, PAR consistently outperforms fixed-depth models at every evaluated test depth. PAR accuracy starts near ceiling at the shortest extrapolation depth, while fixed-depth performance falls quickly and remains near chance at longer horizons. The advantage persists through sixteen times the training depth, with the largest backbone showing the strongest long-horizon result. At the shortest test depth, PAR reaches near-ceiling accuracy on all three Qwen3 backbones, whereas fixed-depth models score only around half correct. Fixed-depth accuracy drops to near chance by the second test depth and remains low through the longest horizon on every backbone. PAR maintains a wide margin over fixed depth through eight times the training depth, and at sixteen times the 4B PAR model achieves the highest long-horizon mean accuracy.

A jointly trained learned halting mechanism tracks an oracle stopping policy and maintains high full-state accuracy without access to problem length. At the longest evaluated extrapolation, it avoids much of the degradation seen by the recurrent baseline and stays close to oracle performance. This indicates that explicit stopping is important for reliable length extrapolation. Learned halting matches the oracle stopping baseline closely at both evaluated extrapolated lengths. At the longest extrapolation, learned halting preserves high accuracy while the RLT baseline degrades substantially and becomes more variable across seeds.

Among final-answer-only training recipes evaluated on the S5 task at 16 times training depth, prefix self-distillation is the only method that sustains high accuracy, reaching 85%. The alternatives degrade substantially, with most showing limited extrapolation and some falling to chance before the longest depth. This indicates that final-answer labels alone can support strong length extrapolation when paired with prefix self-distillation. Prefix self-distillation is the only final-answer-only approach evaluated that keeps high accuracy at 16 times training depth, reaching 85%. Other final-answer supervision recipes show weak extrapolation, with reported accuracies mostly in the teens or twenties at 16 times training depth and much lower extrapolation gaps.

Across swap and permutation depth-extrapolation tasks, PAR maintains near-perfect full-state accuracy well beyond training depth and up to 16 times the training horizon, while fixed-depth, recurrent, untied, and explicit chain-of-thought baselines degrade or fall to chance. Evaluations on multiple Qwen2.5 and Qwen3 backbones show PAR consistently outperforming fixed-depth models at every tested depth, though gains vary by seed and do not scale monotonically with model size. Ablations confirm that supervised iteration targets and norm anchoring are necessary for reliable extrapolation, learned halting closely tracks oracle stopping, and prefix self-distillation is the only final-answer-only training variant that preserves high accuracy at long horizons.