Beyond the Training Horizon: Mechanisms and Limits of Length Generalization in Looped Transformers

Jia Liang; Xi Jin; Liangming Pan

cs.LGSep 27, 2026

Beyond the Training Horizon: Mechanisms and Limits of Length Generalization in Looped Transformers

Authors: Jia Liang, Xi Jin, Liangming Pan

Organizations: Institute for Computational and Mathematical Engineering, Stanford University · MOE Key Laboratory of Computational Linguistics, School of Computer Science, Peking University

Abstract

Looped Transformers can generalize to reasoning chains longer than those encountered during training, but the computations enabling this behavior and limiting its extent remain unclear. We mechanistically compare two looped-Transformer configurations, which we call the Matched-Recurrence Looped Transformer (MR-Loop) and Decoupled-Recurrence Looped Transformer (DR-Loop), reflecting their respective recurrence-training schemes. We evaluate polynomial iteration, finite-state composition, and knowledge-graph traversal using detailed mechanistic analysis. Attention analysis, intermediate-state decoding, and causal interventions reveal distinct mechanisms learned under final-answer supervision. MR-Loop updates an intermediate state at a fixed readout while advancing relation selection through adjacent-token interactions and a transferable progress cue. DR-Loop instead propagates intermediate states across relation positions, forming an advancing computational frontier. However, both mechanisms become unreliable at greater depths: MR-Loop exhibits degradation of its readout state and progress cues, while DR-Loop exhibits declining reliability of state propagation. Limited self-correction allows local errors to persist and compound. Across both models, we uncover a common representational principle: recurrent states encode not only task-relevant content but also its computational status, whether that content remains in a form that can support subsequent computation. Transferable live-consumed and fresh-aged residual directions causally control whether represented information can participate in subsequent computation, including beyond the training horizon. We further show that length generalization need not rely on faithful step-by-step reasoning, as Looped Transformers can exploit task structure without explicitly representing every intermediate state.

Figures & tables

Task	State	Operator / transition	Varied parameters
Polynomial	$s_{k}\in\mathbb{Z}_{M}$	$x_{k}:\;s_{k}=(s_{k-1}x_{k}+1)\bmod M$	$M$
FSC	$s_{k}\in[S]$	$o_{k}:\;s_{k}=f_{o_{k}}(s_{k-1})$	$S,F,\alpha$
KG	$e_{k}\in[E]$	$r_{k}:\;e_{k}=r_{k}(e_{k-1})$	$E,R,D$

Table 1: Sequential composition tasks. All tasks repeatedly apply operators to a running state but differ in their state spaces and transition families.

Figure 1: Distinct length-generalization mechanisms in looped Transformers . MR-Loop uses a fixed readout with advancing relation selection; DR-Loop advances a state-propagation frontier.

Figure 2: Computational status controls recurrent reuse. Live–consumed relation states govern selection in MR-Loop, while fresh–aged intermediate states govern composition in DR-Loop.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

	MR-Loop	DR-Loop
Recurrence depth vs. length	matched: loss after $T=n$ passes	decoupled: $R$ independent of $n$
Train curriculum	increasing on a fixed step schedule	increasing only after required accuracy is met
Train depth	length-dependent $T$	$R\sim\mathrm{Poisson}$ , clipped to $\{2,\ldots,8\}$
Eval depth	$T=n$ (default)	chosen $R$ , possibly $R\neq n$
Recurrent update	$h^{(t)}=B(h^{(t-1)}+e)$	$h^{(t)}=B(h^{(t-1)})$
Final LN in the loop	yes, inside $B$	no; $\mathrm{LN}_{f}$ at readout

Appendix

Table 2: MR-Loop versus DR-Loop in the checkpoints studied here.

Figure 3: Traveling previous-token attention in MR-Loop. Head L0H1 implements a recurrently advancing local attention pattern: at the recurrent pass (loop) $t$ , $r_{t+2}$ attends strongly to its predecessor $r_{t+1}$ . As recurrence proceeds, the attention peak shifts forward by one relation per loop and continues beyond the maximum training length $n=12$ , suggesting a shared local mechanism for propagating computational progress to unseen relation positions.

Figure 4: Fixed-query tracking attention in MR-Loop. Head L2H0 maintains its query at the [EQ] readout position while the attended relation advances across recurrent steps: at recurrent pass (loop) $t$ , [EQ] attends strongly to $r_{t+1}$ . This tracking pattern progresses by one relation per loop and extends beyond the maximum training length $n=12$ , providing a mechanism for sequentially routing relation information into the evolving readout state.

Figure 5: Intermediate-state tracking in MR-Loop. After each loop

t

, we decode the last-layer residual at

[\mathrm{EQ}]

, the just-used key

r_{t}

, and the next unread key

r_{t+1}

, and measure whether the argmax recovers the next walk entity

e_{t+1}

. On length-

15

walks, beyond the maximum training length

n=12

[\mathrm{EQ}]

tracks

e_{t+1}

with perfect accuracy, while predictions at

r_{t}

and

r_{t+1}

remain low. This shows that the evolving intermediate state is maintained as a look-ahead representation at the fixed

[\mathrm{EQ}]

readout position rather than on the relation tokens.

Figure 6: Traveling self-attention head in DR-Loop. A localized block of elevated attention advances toward later sequence positions along the main diagonal across recurrent iterations. The highlighted region spans approximately two to three neighboring positions and moves forward by roughly one to two positions per loop, indicating that the head sequentially tracks the advancing computational frontier. The pattern continues through relation positions beyond the training horizon until reaching the end of the sequence, after which the localized late-position attention disappears in subsequent iterations.

Figure 7: Traveling previous-position attention head in DR-Loop. A localized block of elevated attention advances toward later sequence positions along the first subdiagonal across recurrent iterations, corresponding to queries attending strongly to the immediately preceding position. The highlighted region spans approximately two to three positions and moves forward by roughly one to two positions per loop, qualitatively tracking the advancing intermediate-state frontier. This traveling pattern extends to relation positions beyond those encountered during training and weakens after reaching the end of the sequence.

Figure 8: Intermediate-state propagation in DR-Loop. Logit-lens probability assigned to each ground-truth intermediate entity $e_{k}$ across residual positions at selected recurrent iterations. The dashed diagonal marks the expected alignment between intermediate entity $e_{k}$ and relation position $k$ . As recurrence proceeds, a high-probability band advances toward later positions, showing that successive intermediate entities become decodable at their corresponding relation sites. This frontier continues beyond the maximum training depth of 11 hops: $e_{12}$ and $e_{13}$ become strongly decodable at unseen relation positions. The progression of this representational frontier qualitatively tracks the traveling attention patterns in Figure 7 , providing evidence that DR-Loop propagates intermediate states through the sequence during length generalization.

$F$	Aged	Fresh	Adjusted	Donor target	Continuation
5	0.038	0.650	0.698	0.099	0.617
8	0.015	0.614	0.636	0.040	0.484
11	0.005	0.551	0.666	0.071	0.444
12	0.016	0.336	0.555	0.056	0.354
14	0.009	0.111	0.217	0.021	0.084
16	0.029	0.057	0.115	0.024	0.048

Appendix

Table 3: Reuse across hop indices with the same source-

2

donor. Aged, Fresh, and Adjusted report the recipient-relation target at

F+1

. The final two columns report the donor target at

F+1

and the counterfactual continuation at

F+2

, both under the adjusted intervention. All entries are separately maximized batch-mean probabilities over loops. Rows below the divider evaluate the first counterfactual hop beyond the training depth of 11.

Polynomial
Modulus	11	19	25	29	30	35	40	50
MR-Loop	$\checkmark$	$\times$	$\checkmark$	$\times$	$\checkmark$	$\checkmark$	$\checkmark$	$\checkmark$
DR-Loop	$\times$	$\times$	$\checkmark$	$\times$	$\checkmark$	$\checkmark$	$\checkmark$	$\times$
Finite-State Composition (FSC)
$\alpha$	0.5				1.0
$N_{\text{state}}$	16		32		16		32

Appendix

Table 4: Length-generalization success across task configurations.

\checkmark

indicates successful generalization and

\times

indicates failure.

Figure 9: Accuracy of ground-truth intermediate values at the corresponding recurrent steps across tasks for MR-Loop . Easier settings and tasks that do not admit shortcut solutions consistently recover the correct intermediate states throughout the computation. In contrast, more difficult or compressible settings, such as polynomial iteration with

m=50

and FSC with

S=16,F=16

, fail to faithfully represent the intermediate trajectory and recover the correct states only near the final steps.

Figure 10: DR-Loop can recover the final answer without faithfully tracking all intermediate states. Batch-mean softmax probability assigned to the ground-truth intermediate state at each write site for finite-state composition (FSC) with $S=16$ , $F=8$ , and bijective-function fraction $\alpha=0.5$ , evaluated at task lengths $n\in\{8,15,20\}$ . Despite weak recovery of many intermediate states, particularly at early recurrent steps, the probability assigned to the final target becomes high near the final write site. This dissociation shows that successful final-answer prediction does not necessarily imply faithful step-by-step intermediate-state formation.

Task $(n,t,k)$	Metric	Native	Cons.	Cons. $+v$	Live $-v$	Cons. $+r$
KG $(20,13,14)$	Accuracy (%)	93.8	7.8	93.8	7.8	7.8
KG $(20,13,14)$	L2H1 mass	.831	.001	.869	.013	.000
Polynomial $(20,12,12)$	Accuracy (%)	92.2	17.2	90.6	17.2	17.2
Polynomial $(20,12,12)$	L1H1 mass	.824	.004	.766	.006	.009
FSC $(20,12,13)$	suffix mass	.194	.025	.177	.027	.040

Appendix

Table 5: Representative live–consumed interventions across separately trained MR-Loop checkpoints. “Cons.” replaces the target token with its consumed residual state; “Cons.

+v

” adds the separately estimated live–consumed direction; “Live

-v

” subtracts the direction from a live state; and

r

denotes the tested norm-matched random direction. The attention metrics use different aggregation domains across tasks and should be compared only within each task.

Checkpoint	Fresh	Aged	Aged $+v_{2}$	Fresh $-v_{2}$	Aged $P(e_{2})$	Aged argmax
KG-8	.612	.050	.935	.032	.999	100%
KG-12	.434	.008	.553	.008	.233	65.6%

Appendix

Table 6: Fresh–aged interventions in two additional DR-Loop KG checkpoints, evaluated on 64 examples with input length

n=12

(11 relation applications) and

R=12

recurrent passes. Source-position-2 residuals are transplanted to frontier

F=5

at zero-indexed loop

t^{\star}=3

. The four intervention columns report the maximum, over loops

t=3,\ldots,11

, of the batch-mean logit-lens probability of the counterfactual next-hop target at position 6. This maximum is computed separately for each intervention condition; the scores are not final-answer accuracies. Directions are estimated separately within each checkpoint.

Length Generalization

Length generalization studies whether sequence models trained on short inputs can reliably solve the same tasks on longer sequences.

See all 44 papers

Aug 4, 2026

LoopMTP: A looped transformer guided by latent multi-token prediction

Looped transformers have emerged as a parameter-efficient alternative to scaling depth for strong reasoning. By reusing one stack of layers across

T

iterations, they attain the effective depth and reasoning capabilities of larger models at a fixed parameter count. Yet existing approaches suffer from latent overthinking and undifferentiated computation, largely because intermediate representations receive no guidance across loops. Multi-token prediction (MTP) supplies exactly the dense, forward-looking supervision the loop is missing. We propose LoopMTP, which links the two through a structural correspondence in latent space: a model that loops

T

times can anticipate

T

future tokens. LoopMTP realizes this by softly aligning the hidden state of loop

t

with the embedding of the token

t

steps ahead, while a lightweight gate preserves useful information across iterations. LoopMTP improves average accuracy by up to 8.1% (relative) over the non-looped baseline, with training remaining stable for up to 15 loops.

Behzad Shomali, Markus Frey, David Berghaus +2

Lamarr Institute · University of Bonn · Fraunhofer IAIS

Multi-Token Prediction Recurrent Transformers LLM Reasoning

Ask PDF

Apr 9, 2026

Loop, Think, & Generalize: Implicit Reasoning in Recurrent-Depth Transformers

We study implicit reasoning, i.e. the ability to combine knowledge or rules within a single forward pass. While transformer-based large language models store substantial factual knowledge and rules, they often fail to compose this knowledge for implicit multi-hop reasoning, suggesting a lack of compositional generalization over their parametric knowledge. To address this limitation, we study recurrent-depth transformers, which enables iterative computation over the same transformer layers. We investigate two compositional generalization challenges under the implicit reasoning scenario: systematic generalization, i.e. combining knowledge that is never used for compositions during training, and depth extrapolation, i.e. generalizing from limited reasoning depth (e.g. training on up to 5-hop) to deeper compositions (e.g. 10-hop). Through controlled studies with models trained from scratch, we show that while vanilla transformers struggle with both generalization challenges, recurrent-depth transformers can effectively make such generalization. For systematic generalization, we find that this ability emerges through a three-stage grokking process, transitioning from memorization to in-distribution generalization and finally to systematic generalization, supported by mechanistic analysis. For depth extrapolation, we show that generalization beyond training depth can be unlocked by scaling inference-time recurrence, with more iterations enabling deeper reasoning. We further study how training strategies affect extrapolation, providing guidance on training recurrent-depth transformers, and identify a key limitation, overthinking, where excessive recurrence degrades predictions and limits generalization to very deep compositions.

Harsh Kohli, Srinivasan Parthasarathy, Huan Sun +1

The Ohio State University

Position Extrapolation Implicit Reasoning in Language Models Compositional Generalization Recurrent Transformers

Ask PDF

Sep 30, 2026

Shared Weights, Selected Computations: How Looped Transformers Route What Each Loop Does

Looped Transformers repeatedly apply the same set of Transformer layers, giving them a recurrent architecture for latent computation. Their strong performance on iterative reasoning and length-generalization tasks suggests an appealing explanation: recurrence may provide an inductive bias that lets the model reuse a learned algorithm across loops. However, weight sharing alone does not imply that every loop performs the same operation. This raises a basic question: is each loop actually repeating the same computation, and if not, what routes the shared parameters to different operations? We study this question using graph walks as a test case. In the model's native trajectories, decoded predictions can advance by different numbers of graph steps or remain at a reached target, showing that recurrent progress need not follow a fixed one-loop-one-step pattern. We then show that a frozen loop can be steered toward different transitions by modifying its entering hidden state: a learned linear layer

J

selects the desired transition without changing the shared Transformer layers. To test how this steering works, we use activation patching and find that attention patterns can recover its effects and switch the selected transition. Across five matched pairs of graph models, changing intermediate supervision during backbone training changes which transitions

J

can induce. This suggests that

J

selects computations learned by the backbone rather than creating new algorithms. Together, these results show that the hidden state can control shared computation, with attention routing as a causal pathway.

Jiaju Wu, Yi Hu, Muhan Zhang

Institute for Artificial Intelligence, Peking University. · School of Mathematical Sciences, Peking University.

Transformer Interpretability Transformer Attention Recurrent Transformers Conditional Computation

Ask PDF

Abstract

Length Generalization

Explore similar work

LoopMTP: A looped transformer guided by latent multi-token prediction

Loop, Think, & Generalize: Implicit Reasoning in Recurrent-Depth Transformers

Shared Weights, Selected Computations: How Looped Transformers Route What Each Loop Does