PaperlayerView exact revision on arXiv

Source-linked AI summary

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, Jiang-Ming Yang, Wei Wu, Yongqi Tong

arXiv:2610.27628v1cs.CL

TL;DR

Sequential state-tracking models can run more iterations without necessarily solving longer problems, and final-answer supervision does not specify each iteration’s progress. PAR aligns iteration k with the state after k operations and anchors recurrent-state norms, with learned halting and prefix self-distillation as extensions. On swap-S5, it achieves 95.5% full-state accuracy at 16× the training length, exceeding reported RLT and explicit-CoT results.

  • Problem

    Length generalization asks whether a model trained on short problems can solve longer instances of the same task, while recurrent iterations alone do not ensure extrapolation.

    Source: §1 · p. 1

  • Method

    PAR trains iteration k to predict the state after the first k operations and rescales recurrent states to their first-iteration norms, with learned halting and prefix self-distillation extensions.

    Source: §1 · p. 2, §3.5 · p. 3

  • Results

    95.5% full-state accuracy at 16× on swap-S5 contrasts with 55.7% for RLT and 0.8% for explicit CoT.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Defining what each iteration computes is a training target alongside choosing how many iterations to run.

    Source: §7 · p. 8

    Evaluation is limited to algorithmic and templated state-tracking tasks with well-defined intermediate states; open-ended reasoning is left to future work.

    Source: Limitations · p. 8

Abstract

from arXiv · show

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.

1 Introduction

Length generalization asks whether models trained on short sequential problems can solve longer instances, but more recurrent computation alone does not ensure extrapolation. PAR aligns each iteration with task progress and anchors recurrent-state scale, alongside learned halting and prefix self-distillation.

1 Introduction

  • Length generalization tests whether models trained on short problems can solve longer instances of the same task, including sequential state-tracking problems with unchanged rules.The paper defines problem length as the number of task operations.Source: §1 · p. 1
  • More recurrent iterations or greater model depth alone yield less than 2× extrapolation on S5, while larger fixed-depth Qwen2.5 models remain near chance at 16×.These comparisons motivate studying how to sustain accurate execution over longer operation sequences.Source: §1 · p. 1
  • PAR trains iteration k to predict the state after the first k operations and rescales each token’s recurrent state to its first-iteration RMS norm.The targets define intended progress while norm anchoring fixes state magnitude as the shared block is reused.Source: §1 · p. 2
  • A learned completion readout selects when to answer without a supplied stopping depth, and prefix self-distillation derives intermediate targets from final-answer labels.The prefix method uses valid shorter problems as targets for corresponding iterations on the full input.Source: §1 · p. 2
  • 95.5% full-state accuracy at 16× on swap-S5 contrasts with 55.7% for RLT and 0.8% for explicit CoT.The reported comparison is for the Qwen2.5-0.5B backbone.Source: §1 · p. 2

2 Related Work

Prior work studies recurrent depth, length extrapolation, halting, and latent reasoning, but this paper focuses on defining what each iteration computes. It combines depth–length correspondence with existing output-head completion detection and prefix-based self-supervision.

2 Related Work

  • Looped architectures have been used to extend computation, but depth is often underused; this work addresses that issue by defining each iteration’s computation.Prior work also connects recurrent depth with implicit models and chain-of-thought simulation.Source: §2 · p. 2
  • Prior length-extrapolation work ties loop count to input length [7] or reasoning steps [8], while RLT [41] uses token-level feedback for swap-S5.Per-loop supervision has shown modest extrapolation [28], and step-wise final-answer supervision can impair compositional generalization [4].Source: §2 · p. 2
  • Unlike adaptive computation time [10] and PonderNet [2], PAR uses the existing output head to detect completion without extra parameters or ponder losses.For supervision, it uses intermediate labels or prefix self-distillation inspired by mean-teacher consistency [31].Source: §2 · p. 2

3 Depth–Length Correspondence

PAR establishes a depth–length correspondence by combining a retrofitted recurrent model, progress-aligned targets, state-scale control, learned halting, and prefix self-distillation. The design assigns each iteration a task stage and supplies ways to continue, stop, or learn that correspondence from final answers alone.

3 Depth–Length Correspondence

  • PAR combines retrofitted recurrence and RMS norm anchoring with progress-aligned targets, learned halting, and prefix self-distillation.These components are presented as ways to realize a depth–length correspondence.Source: §3 · p. 2

3.1 A retrofitted looped language model

  • The retrofitted Qwen2.5-0.5B model splits 24 layers into a prelude, weight-tied core, and coda, then reads an answer from any iteration.An adapter combines static prelude features with the previous recurrent state before each core update.Source: §3.1 · p. 2

3.2 Defining an iteration: progress-aligned targets

  • Iteration k is trained to hold the state after the first k steps, so the number of required iterations grows with problem length.Forced depth runs for the needed iterations; under hold, iterations after completion keep the final state.Source: §3.2 · p. 3, Figure 1 · p. 3

3.3 Keeping the step map stationary

  • State-norm anchoring rescales each token’s recurrent state to its first-iteration RMS norm, while positional stretch fine-tunes position indices for longer prompts.Label-supervised PAR reaches 16× without positional stretch; the label-free variant needed stretch to reach 16×.Source: §3.3 · p. 3

3.4 Reading completion: learned halting

  • The coda predicts Yes when execution is complete and No otherwise, and inference selects an answer after a multi-step confirmation window.The completion readout adds no parameters and does not require the problem length at test time.Source: §3.4 · p. 3

3.5 Correspondence without step labels: prefix self-distillation

  • Prefix self-distillation trains early iterations against a teacher’s answer distribution on the matching shorter prefix, using final answers rather than intermediate labels.The teacher may be an exponential moving average of the student or the live student.Source: §3.5 · p. 3

3.6 Training

  • Training uses AdamW over approximately 400k examples, with examples spread uniformly over training depths and no intermediate tokens generated.The stated learning rate is 2×10^-5, or 10^-4 for the adapter.Source: §3.6 · pp. 3–4

4 Experimental Setup

The experiments evaluate sequential state-tracking tasks at increasing lengths and compare PAR with recurrent, fixed-depth, untied, and explicit-CoT baselines. The swap-S5 protocol uses a full-state criterion and includes published Transformer and RLT comparisons.

4 Experimental Setup

  • Swap-S5 prompts contain moves on five boxes, and a sequence is correct only when all five boxes are answered correctly.The protocol trains on at most 32 moves and tests 1,024 sequences per length.Source: §4 · p. 4
  • The S5 and A5 extrapolation factor G is d50/dtrain, where d50 is the first test depth whose accuracy falls below the midpoint between chance and 100%.With test depths up to 256 and training depth 16, G=16 is the largest measurable value.Source: §4 · p. 4
  • Baselines include explicit CoT, fixed-depth and recurrent models, a four-times-repeated untied core, and published Transformer and RLT results.The same-backbone comparisons use retrofitted Qwen2.5-0.5B; published swap-S5 values use Zhang et al.’s best split, 4 + 4.Source: §4 · p. 4

5 Main Results

PAR sustains long-horizon state-tracking accuracy across swap-S5 and other backbones, while learned halting and prefix self-distillation extend the approach to autonomous stopping and settings without intermediate labels. Results also show that its gains are strongest on sequential tasks with a canonical progression of intermediate states.

5.1 Depth extrapolation

  • Over 95% full-state accuracy at 16× on swap-S5, while RLT reaches 55.7%; PAR also reaches the maximum measured extrapolation factor, G=16, on S5 and A5.Fixed-depth, recurrent, and 48-layer untied baselines stay below 2×; explicit CoT falls to chance by 8× on swap-S5.Source: §5.1 · p. 4, Table 1 · p. 5

5.2 Larger and different backbones

  • At 16× on S5, the best PAR seed maintains 100% accuracy on Qwen2.5 1.5B, 3B, and 7B, while fixed-depth baselines remain at or below 26.5%.On Qwen3, PAR remains near ceiling at 2× and has a wide margin through 8×; at 16× it reaches up to 53% mean accuracy on 4B as fixed-depth models collapse near chance.Source: §5.2 · pp. 4–5, Table 2 · p. 5, Table 3 · p. 5

5.3 Learning when to stop

  • 92.9% full-state accuracy at 16× with learned halting, without access to problem length L, closely matching the oracle iteration count across seeds.Forced-depth models can answer correctly at R=L but fall to near-zero accuracy within a few extra iterations; jointly trained models enter a persistent finished state.Source: §5.3 · p. 5, Table 4 · p. 5

5.4 Learning without step labels

  • Prefix self-distillation is the only evaluated label-free approach sustaining high accuracy at 16×, reaching 85% at d=256; positional-stretch fine-tuning reaches up to 99.7% on swap-S5.Label-supervised PAR remains stronger at the longest depths, while prefix self-distillation learns the correspondence from final answers alone.Source: §5.4 · pp. 5–6, Table 5 · p. 5

5.5 Prior benchmarks and scope

  • On the rule-table chain, PAR-hold reaches G≥8 with 94% accuracy at depth 64, while all recipes stay below 3.5× on graph reachability, which lacks a canonical intermediate state for each iteration.On natural-language object tracking, PAR-hold with positional stretch reaches 94–96% at 32 swaps and 79% at 64 swaps, versus at most 67% and 21% for random-R and fixed-depth baselines, respectively.Source: §5.5 · p. 6, Table 6 · p. 6

6 Ablation Studies

The ablations find that long-range extrapolation requires both progress-aligned targets and norm anchoring, while stopping and prefix self-distillation depend on distinct design choices. Positional stretch improves label-free extrapolation when combined with the self-distillation objective.

6.1 Both conditions are necessary

  • Final-answer supervision stalls at G≈2 when applied at every iteration and reaches only G=1.48 when applied only at iteration d.Adding norm anchoring to iteration-d supervision gives one seed G=8, but mean accuracy at depth 256 is 26.8% across seeds, far below PAR.Source: §6.1 · p. 6, Table 7 · p. 7
  • Prefix targets without norm anchoring stall at approximately 1.8× under either schedule, consistent with recurrent state representations drifting during repeated updates.Ablations vary one condition at a time on S5 with training depth ≤16; all rows share backbone, data, optimizer, and step budget.Source: §6.1 · p. 6, Table 7 · p. 7
  • PAR reaches G=16 across all seeds with near-ceiling accuracy at depth 256, while PAR-hold extrapolates beyond 12× with lower seed consistency.Positional stretch alone leaves fixed-depth and random-R models at chance by depth 48; applied to an existing PAR-hold model, it keeps 99% accuracy at depth 512 in one run.Source: §6.1 · pp. 6–7

6.2 Stopping rule and alternatives to halting

  • The 6-of-8 completion rule matches oracle stopping within 1.6 points across evaluated sequences, balancing transient probability spikes against delayed stopping.Single-iteration stopping is brittle at long horizons, while requiring eight consecutive confirmations can delay stopping until predictions drift.Source: §6.2 · p. 7, Table 8 · p. 7
  • Training extra iterations to preserve the final answer collapses accuracy below 1% at 16×, whereas an independent completion read-out avoids competing fixed-point and state-advancement objectives.Fine-tuning a read-out on an existing PAR checkpoint makes its completion probability decay within two iterations; joint training has minor accuracy degradation on one seed at 16× versus fine-tuning.Source: §6.2 · p. 7

6.3 Teacher and positional stretch in prefix self-distillation

  • Positional stretch fine-tuning restores label-free swap-S5 extrapolation across starting checkpoints, raising 16× accuracy from near chance up to 99.7%.Initial label-free models collapse at 16×, while label-supervised PAR reaches its results without stretch.Source: §6.3 · pp. 7–8
  • Three of five teacher configurations reach G≈16 or higher, and the live-student teacher with a depth curriculum has the highest accuracy at depth 256.Table 9 reports S5 single-query accuracy at R=d for training depth ≤16.Source: §6.3 · p. 7, Table 9 · p. 7

7 Conclusion

The conclusion argues that depth extrapolation depends on defining the job of each loop iteration and keeping the step map stationary. This correspondence supports longer swap-S5 sequences, autonomous stopping, and learning intermediate states from final answers.

7 Conclusion

  • A retrofitted 0.5B model surpasses the best published latent-recurrence model on swap-S5 at 8× and 16× the training length.The conclusion frames the result as depth extrapolation from assigning each iteration a defined job and keeping the step map stationary.Source: §7 · p. 8
  • The same depth–length correspondence lets the model learn when to stop and learn intermediate states from final answers alone.Source: §7 · p. 8
  • The paper treats defining an iteration’s meaning as a training target alongside choosing how many iterations to run.Source: §7 · p. 8

Limitations

The evaluation focuses on algorithmic and templated state-tracking tasks with well-defined intermediate states. Open-ended natural-language reasoning and broader conclusions about model scaling remain outside the demonstrated scope.

Limitations

  • The evaluation covers algorithmic and templated state-tracking tasks where intermediate states are well defined, leaving open-ended natural-language reasoning for future work.Source: Limitations · p. 8
  • Swap-S5, learned halting, and prefix self-distillation are evaluated only on Qwen2.5-0.5B, while the scaling study uses the single-query S5 setting.Source: Limitations · p. 8
  • Larger-backbone results vary substantially across seeds, Qwen2.5 accuracy does not rise monotonically with size, and no Qwen3 backbone approaches Qwen2.5-0.5B at 16×.The authors say these results do not establish a monotonic model-size trend and leave variance mechanisms and recurrent-scaling stability open.Source: Limitations · p. 8

Ethics Statement

The work studies training methods on synthetic reasoning tasks using publicly available models. It uses no human subjects or personal data, and discusses the relative transparency of latent reasoning and chain of thought.

Ethics Statement

  • The study examines training methods on synthetic reasoning tasks with publicly available models.Source: Ethics Statement · p. 8
  • The work involves no human subjects and uses no personal data.Source: Ethics Statement · p. 8
  • Latent reasoning is less transparent than chain of thought, but the correspondence allows each iteration to be decoded into its intermediate answer.Source: Ethics Statement · p. 8

A Additional Results

The appendix reports per-depth S5 results alongside explicit-CoT, learned-halting, and prefix-self-distillation evaluations.

A Additional Results

  • 84.5 ± 29.4% at 8× and 76.9 ± 38.0% at 16×: PAR-hold with positional stretch trails forced-depth PAR in mean accuracy and cross-seed consistency.These runs evaluate 1,024 sequences per length across six seeds.Source: A Additional Results · p. 10
  • Explicit CoT full-state accuracy falls from 100% at 32 moves to 0.8% at both 256 and 512 moves.At 64 moves it is 35.6%, and at 128 moves it is 3.1%.Source: A Additional Results · p. 10
  • Positional stretch helps only together with progress-aligned targets.Source: A Additional Results · p. 11
  • The learned-halting table compares oracle iteration counts with rules requiring six of eight or eight consecutive iterations to pass the completion threshold.The six-of-eight rule returns the first qualifying iteration’s prediction after confirmation.Source: Table 11 · p. 11
  • Every stretch-fine-tuned prefix-self-distillation model scores at least 99.2% in range at 32 moves.The table reports full-state accuracy for models trained on lengths up to 32.Source: Table 12 · p. 11

Related papers

Loading 2610.27628v1…