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

> Looped Transformers achieve length generalization via distinct recurrent mechanisms, but success does not guarantee faithful step-by-step reasoning when tasks permit lossy shortcuts.

- **Source:** [arXiv](https://arxiv.org/abs/2609.33144)
- **Published:** 2026-10-31
- **Permalink:** https://picx.dev/p/VgyMOs
- **Whiteboard:** https://picx.dev/p/VgyMOs/image

## Summary

## 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
| Feature | MR-Loop | DR-Loop |
|---|---|---|
| Recurrence-depth scheduling | Matched to input length ($K = T = n$) | Decoupled from input length |
| Recurrent update | $h^{(t)} = B(h^{(t-1)} + e)$ | $h^{(t)} = B(h^{(t-1)})$ |
| Answer readout | From [EQ] delimiter | From padded output position |
| Training length curriculum | Increases on fixed schedule | Increases when accuracy threshold met |
| Recurrent depth sampling | $T = n$ (input length) | $R \sim \text{Poisson}(\lambda=4)$, clipped to $\{2,\dots,8\}$ |

### 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**

| 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$ |

## 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):**

| Task | MR-Loop | DR-Loop |
|---|---|---|
| Polynomial (composite moduli) | ✓ | ✓ (some) |
| Polynomial (prime moduli) | × (except m=11) | × |
| FSC (all configs) | ✓ | × with 16 functions |
| KG (all 6 configs) | ✓ | ✓ (3 of 6) |

**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.

---

_Markdown view of https://picx.dev/p/VgyMOs, served by PicX — AI-generated visual whiteboard summaries of research papers._