PAR: Length generalization in state tracking with Looped language models that align iterations with problem progress

Problem

Looped language models, which repeat the same Transformer block, can increase iterations at test time, but that does not mean they solve longer problems. Training only on the final answer leaves undefined how much progress each iteration should make. The magnitude of the Hidden state also changes over repeated updates.

Method

The 24 layers of Qwen2.5-0.5B were split into a prelude (layers 1–8), a weight-shared core (layers 9–16) and a coda (layers 17–24) to turn it into a Looped model. The output of iteration k is trained to match the intermediate state after applying the first k operations (progress-aligned target), and after every iteration the state is rescaled to the RMS norm of the first iteration (norm anchoring). The same output layer also predicts whether the computation is complete, so the model can stop without knowing the problem length (learned halting). When intermediate labels are unavailable, prefix self-distillation is applied: the front part (prefix) of a problem is treated as a shorter problem, and its prediction is used as the Teacher signal.

Results

• After training on up to 32 moves in swap-S5, the model reached 95.5% full-state accuracy at 512 moves (16×). Under the same conditions, the cited RLT scored 55.7% and explicit CoT 0.8%.
• In controlled S5 experiments, removing norm anchoring dropped the Extrapolation factor from 16× to under 2.2×, and switching to final-answer supervision dropped it to at most 8×.
• Even with learned halting, which does not know the problem length, the model scored 92.9% at 16× length, nearly the same as when given the correct number of iterations (93.2%).
• Using only prefix self-distillation without intermediate labels gave 85% on S5 d=256 (16×), while other final-answer training methods stayed near chance (around 20%).
• On Qwen3-4B, 16×-length accuracy was also 53% versus 17.7% for fixed depth, but larger models showed greater variance across seeds.

Limitations

The method was validated only on synthetic sequential tasks such as permutation state tracking, rule-table chaining and graph reachability. On 1.5B–7B backbones, performance was sometimes near 0 depending on the seed, so variance is large (e.g., 3B mean 57.8±42.5). The RLT figures are cited from values reported in another paper. Iterating past the completion point quickly breaks the answer, so the method depends on the halting rule.

Why it matters

The experiments show that for 'recurrent depth,' which increases computation at test time, to actually lead to solving longer problems, one must define what each iteration should compute and control the state's magnitude. This gives design guidance for research on reasoning in latent space without text CoT.

Key terms

Looped language model

A model that applies the same Transformer block repeatedly to increase computational depth without adding parameters.

Length generalization

The ability of a model trained on short problems to also solve longer problems that follow the same rules.

Norm anchoring

A technique that resets the magnitude of the Hidden state to the first iteration's RMS norm at every iteration, so the state does not drift even over hundreds of iterations.

Prefix self-distillation

A method that trains without intermediate labels by using the model's (EMA Teacher's) prediction on a shorter problem, keeping only the first k steps, as the target for the k-th iteration on the full problem.

Original

• arXiv abstract
• PDF
• HTML full text
• Hugging Face

Authors: Pan Wang, Yihao Hu, Guangyuan Dong, Yitong Lin, Hang Wang, Zimi Liu, Zhihao Wen, Jianshe Li, et al. · Posted 2026-10-20 · License arXiv 비독점 배포 라이선스 · Based on the full text

This review is an AI-written summary of the paper. Please treat the original paper as the reference.