Summary (Overview)

• This paper provides a mechanistic investigation of length generalization in Looped Transformers, comparing two configurations: Matched-Recurrence Looped Transformer (MR-Loop) and Decoupled-Recurrence Looped Transformer (DR-Loop).
• The authors evaluate three iterative task families—polynomial iteration, finite-state composition, and knowledge-graph traversal—using attention analysis, intermediate-state decoding (logit lens), and causal interventions (activation patching).
• Key finding 1: MR-Loop and DR-Loop achieve length generalization through distinct recurrent mechanisms: MR-Loop uses a fixed readout accumulator with advancing relation selection, while DR-Loop propagates a computational frontier through the sequence.
• Key finding 2: Extreme-length failure arises from degradation of computation-ready states and limited self-correction, allowing local errors to persist and compound.
• Key finding 3: Length generalization does not imply faithful step-by-step reasoning; models can exploit task compressibility shortcuts without representing every intermediate state.

Introduction and Theoretical Foundation

Background and Motivation

• Many reasoning tasks require executing the same computation for a variable, previously unseen number of steps (Schwarzschild et al., 2021; Thomm et al., 2024).
• Standard Transformers struggle with length generalization because their computational depth is fixed while the required number of sequential operations grows (Hahn, 2020).
• Looped Transformers repeatedly apply shared Transformer blocks, allowing computational depth to grow without adding parameters (Dehghani et al., 2019; Geiping et al., 2025).

Theoretical Framework

The paper defines the task horizon $H$ (number of ground-truth transitions) separately from the recurrent depth $K$ (number of times the shared Transformer stack is applied). These need not coincide because one recurrent pass updates all sequence positions simultaneously.

A looped Transformer is formalized as repeatedly applying a parameterized stack $B_\theta$:

$$h^{(t+1)} = B_\theta\left(\Phi(h^{(t)}, e)\right), \quad t = 0, \ldots, K-1$$

where $\Phi$ specifies how input embeddings enter the recurrent computation.

Two Loop Configurations

Sequential Composition Tasks

All tasks instantiate: input $[s_0, o_1, \ldots, o_H]$ with transitions $s_k = T_{o_k}(s_{k-1})$, trained only on the final state $s_H$ with no intermediate supervision.

Table 1: Sequential composition tasks

Methodology

Models

• All checkpoints use 3 shared Transformer blocks, 2 attention heads per block, width 256, no positional embeddings.
• Trained on inputs up to length $n=12$ (11 hops), evaluated on longer sequences.

Analytical Tools

1. Attention Analysis: Tracking how information routes across token positions across recurrent passes, identifying whether attention patterns remain stationary or advance.

2. Logit Lens Decoding: Applying the model's final normalization and output projection to intermediate residual states:

$$\ell^{(t)}_{b,j} = W_{\text{out}} \text{LN}_{\text{out}}\left(h^{(t)}_{b,j}\right), \quad p^{(t)}_{b,j} = \text{softmax}\left(\ell^{(t)}_{b,j}\right)$$

3. Activation Patching / State Swaps: Replacing internal residual states with donor activations:

$$\tilde{h}^{(t,s)}_{b,j} \leftarrow h^{(\tau,s)}_{b',j'}$$

Key Concepts Identified

• Live–Consumed direction (MR-Loop): A shared residual vector $v_s$ estimated from the difference between a relation token's "live" state (scheduled fetch loop) and "consumed" state (two loops later):

$$\Delta h_{s,i,k} = h^{\text{live}}_{s,i,k} - h^{\text{consumed}}_{s,i,k}, \quad v_s = \frac{\mathbb{E}_{i,k}[\|\Delta h_{s,i,k}\|_2]\mathbb{E}_{i,k}[\Delta h_{s,i,k}]}{\|\mathbb{E}_{i,k}[\Delta h_{s,i,k}]\|_2}$$

• Fresh–Aged direction (DR-Loop): A shared direction $v_z$ distinguishing newly computed entity states (usable for composition) from aged states (still decodable but unusable).

Empirical Validation / Results

RQ1: How Do Looped Transformers Achieve Length Generalization?

MR-Loop: Fixed Readout with Advancing Relation Selection

• Attention patterns: (1) Traveling previous-token pattern—query at $r_{t+2}$ attends to $r_{t+1}$; (2) Fixed-query tracking—[EQ] query attention shifts to $r_{t+1}$ at loop $t$. Both advance by one token per loop beyond training length.
• Intermediate value tracking: The [EQ] readout recovers intermediate entity $e_{t+1}$ with 100% accuracy beyond training length, while relation positions remain near zero.
• Causal tests: State-swap interventions confirm that relation-token states control selection and [EQ] supplies the entity operand. The update follows $e_{t+1} = r_{t+1}(e_t)$.
• Live–consumed cue transfer: At $k=14$ on length-20 inputs, consumed-state replacement reduces accuracy to 9.4%; adding $v_s$ restores it to 93.8%, matching native performance.

DR-Loop: Propagating State Frontier

• Attention propagation: Localized attention regions advance along the main diagonal and first subdiagonal by 1–2 positions per recurrent iteration, extending beyond training length.
• Intermediate-state propagation: Entities $e_k$ become decodable at relation position $k$; on 14-hop queries, $e_{12}$ and $e_{13}$ reach peak probabilities ~0.90–0.92 beyond training depth.
• Causal evidence: Patching position 5 makes counterfactual entities $e'_6$ and $e'_{14}$ decoded in 90.6% and 71.9% of examples; patching position 12 (beyond training) yields 56.25% counterfactual argmax.
• Fresh–aged distinction: An aged representation of $e_2$ remains decodable at 96.9% argmax accuracy, yet its next-hop score falls from 0.650 (fresh) to 0.038 (aged). Adding the fresh–aged direction restores it to 0.698.

RQ2: Why Does Length Generalization Fail at Extremely Long Horizons?

Two coupled causes identified:

1. Degradation of computation-ready states:

• MR-Loop: Transplanting a same-loop [EQ] state from a length-20 prefix into a length-30 run restores next-hop accuracy from 5.5% to 83.6%. The live–consumed cue projection declines from 28.0 (key 11) to 9.1 (key 23).
• DR-Loop: Mean per-example peak correct-entity probability falls from 0.982 (hop 5) to 0.713 (hop 16).

2. Limited self-correction:

• Survival probability formulation:

$$S_K(n) = \Pr\left(\bigcap_{k=1}^{K} C_k \mid n\right) = \prod_{k=1}^{K}(1 - h_k(n))$$

where $h_k(n) = \Pr(\neg C_k \mid C_1, \ldots, C_{k-1}, n)$ is the conditional first-error hazard. Even constant hazard yields geometric decay $S_K = (1-h)^K$.

• MR-Loop native analysis: Only 2/21 first-error trajectories at $n=20$ have a later correct intermediate; 0/21 sustained recovery; at $n=30$, no sustained recovery and 0/64 final-answer correctness after error.

RQ3: Does Successful Length Generalization Imply Faithful Iterative Computation?

Length generalization success is task- and model-dependent (Table 4):

Faithfulness gap: Successful models can produce correct final answers while ground-truth intermediates are only weakly recovered. This is explained by task compressibility: when $G_t(u) = G_t(v)$ for $u \neq v$ (the remaining composition merges states), distinguishing intermediate states is unnecessary. Compressive settings (composite moduli, non-bijective FSC) admit shortcut solutions; information-preserving settings (prime moduli, bijective FSC, permutation KG) require trajectory preservation.

Theoretical and Practical Implications

Common Representational Principle

Both architectures reveal that semantic content alone is insufficient for recurrent reuse. The paper identifies computational status as a distinct representational property:

• MR-Loop: Live vs. consumed relation states control selection.
• DR-Loop: Fresh vs. aged intermediate states control composition usability.

These transferable residual directions causally control whether represented information can participate in subsequent computation, including beyond the training horizon.

Distinguishing Behavioral from Mechanistic Length Generalization

The paper separates three distinct concepts:

1. Behavioral length generalization: Correct extrapolated answers.
2. Faithful step-by-step reasoning: Recovery of ground-truth intermediate states.
3. Reusable computation: Preservation of computation-ready states.

Implications for Recurrent Reasoning Design

• Reliable recurrent reasoning depends on preserving and refreshing internal states, not merely allocating more test-time computation.
• The survival analysis provides a framework for understanding how state degradation and error compounding interact.
• Task compressibility constrains the solution space, determining whether lossy shortcuts are available.

Conclusion

Main Takeaways

1. Distinct mechanisms underlie length generalization: MR-Loop uses a fixed accumulator with advancing operator selection; DR-Loop uses a propagating state frontier.
2. Computational status is critical: Representing correct information is not enough; it must remain in a computationally usable state for the next transition.
3. Extreme-length failure is systematic: Computation-ready states degrade, and learned dynamics provide limited correction, allowing errors to persist and compound.
4. Faithfulness is not guaranteed: Successful extrapolation does not imply faithful execution of the intended step-by-step procedure, especially when task structure permits lossy shortcuts.

Future Directions

• Architecture ablations: Systematic study of how recurrence scheduling, normalization, and readout design affect the identified mechanisms.
• Scaling: Test these principles at larger scales and in more naturalistic reasoning settings.
• Alternative configurations: Explore other looped Transformer designs beyond the matched/decoupled recurrence distinction.