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

LoopMTP: A looped transformer guided by latent multi-token prediction
Aug 4, 2026Behzad Shomali, Markus Frey, David Berghaus +2Multi-Token Prediction Recurrent Transformers
Lamarr Institute · University of Bonn · Fraunhofer IAIS
Loop, Think, & Generalize: Implicit Reasoning in Recurrent-Depth Transformers
Apr 9, 2026Harsh Kohli, Srinivasan Parthasarathy, Huan Sun +1Position Extrapolation Implicit Reasoning in Language Models
The Ohio State University
Shared Weights, Selected Computations: How Looped Transformers Route What Each Loop Does
Sep 30, 2026Jiaju Wu, Yi Hu, Muhan ZhangTransformer Interpretability Transformer Attention
Institute for Artificial Intelligence, Peking University. · School of Mathematical Sciences, Peking University.

Aug 4, 2026J/K move · Enter open · S save

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

Behzad Shomali, Markus Frey, David Berghaus, Joachim Koehler+1

Lamarr Institute · University of Bonn · Fraunhofer IAIS

Multi-Token Prediction Recurrent Transformers LLM Reasoning

Open paper AskPDF arXiv

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.

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