# Learning to Extrapolate in Depth with Looped Transformers

> Looped transformers trained at depth K extrapolate to harder graph connectivity instances at test depths up to an exponential window, beyond which overthinking causes failures.

- **Source:** [arXiv](https://arxiv.org/abs/2610.14456)
- **Published:** 2026-11-14
- **Permalink:** https://picx.dev/p/QT4EMW
- **Whiteboard:** https://picx.dev/p/QT4EMW/image

## Summary

# Learning to Extrapolate in Depth with Looped Transformers

**Authors:** Huitao Yang, Jianhao Huang, Zixuan Wang, Baharan Mirzasoleiman (UCLA & Princeton)

## Summary (Overview)

- **Core question:** How far can a looped transformer block trained at depth $K$ be safely extrapolated to greater test-time depths $T > K$ on harder graph connectivity instances?
- **Key theoretical result:** Under exchangeability and symmetric positive initialization, population training at depth $K$ converges to an update rule governed by a single effective parameter $\gamma_*$ that balances graph propagation against unresolved training error.
- **Finite extrapolation window:** The trained block recovers connectivity labels only within a finite window of test depths—too few loops provide insufficient propagation, while too many create false connections.
- **Exponential extrapolation gain:** A test depth chosen from the training distribution can guarantee recovery on substantially harder graphs, with certified depth gain growing exponentially as training depth outpaces training hardness.
- **Empirical validation:** Experiments on Erdős–Rényi (ER) graphs and path graphs confirm predicted windows and show how compute budget and initialization affect optimal training depth.

## Introduction and Theoretical Foundation

### Background

Looped transformers reuse a single block across iterations, decoupling depth from parameter count. This enables scaling test-time compute by unrolling further at inference. However, performance can deteriorate under excessive test-time iteration ("overthinking"), and transformers often fail on inputs beyond their training distribution.

### The Problem Setting

**Graph connectivity task:** Given a graph $G \in \mathcal{G}_N$ on labeled vertices $[N]$, predict the binary connectivity matrix $Y_G \in \{0,1\}^{N \times N}$ where $(Y_G)_{ij} = 1$ iff vertices $i$ and $j$ are connected by a path (including zero-length paths when $i = j$).

**Input representation:** The input is a normalized graph operator:
$$H_0 = P_G := I - \frac{L_G}{2\max\{1, d_{\max}(G)\}}$$
where $L_G$ is the graph Laplacian and $d_{\max}(G)$ is the maximum degree.

**Supervision target:** $\Pi_G$ has the same support as $Y_G$ but is normalized by component size: $(\Pi_G)_{ij} = (Y_G)_{ij}/|C(i)|$ where $C(i)$ is the component containing vertex $i$.

**Looped self-attention model:** The hidden state evolves through one attention block with parameters $\theta = (W, V)$ shared across all loops:
$$H_{t+1} = \Phi_\theta(H_t) := \mathcal{N}\big((H_t W H_t^\top) H_t V\big)$$
where $\mathcal{N}(Z) = \text{diag}(Z\mathbf{1})^{-1}Z$ normalizes rows to sum to one.

**Readout:** At test time with frozen weights, a thresholded readout produces predictions:
$$\widehat{Y}_{ij}(T) = \mathbf{1}\left\{\frac{(H_T)_{ij}}{\sqrt{(H_T)_{ii}(H_T)_{jj}}} > \tau\right\}$$

### Related Work

The paper builds on prior work in:
- **Looped transformers:** Universal Transformers, ALBERT, programmable looped transformers
- **Expressivity results:** Logarithmic depth suffices for graph connectivity; looped transformers can simulate graph algorithms
- **Training dynamics:** Linear looped transformers learn multi-step preconditioned gradient descent; looped linear transformers with layer normalization learn the power method
- **Length generalization:** Transformers often fail on longer inputs; extrapolation succeeds when architecture encodes target nonlinearity

## Methodology

### Key Reduction: Single Effective Parameter

The central insight is that under symmetric initialization, the weight matrices maintain the structure $W = aF + bE$ and $V = cF + dE$ where $E = \mathbf{11}^\top/N$ and $F = I - E$. The normalized output depends only on the ratio:

$$\gamma = \frac{ac}{bd}$$

This yields the closed-form dynamics:
$$H_t = E + \gamma^{(3^t - 1)/2}\left(P_G^{3^t} - E\right)$$

Each loop cubes the graph operator, tripling the number of graph propagation steps.

### Training Dynamics (Theorem 1)

Define $\kappa = \mathbb{E}_G(c_G - 1)$ (expected number of components beyond one) and $\delta_K = \mathbb{E}_G\|P_G^{3^K} - \Pi_G\|_F^2$ (residual propagation error at training depth).

**Theorem 1 (Population training learns the effective ratio):** Under exchangeability and symmetric positive initialization, population matrix gradient flow converges to:

$$\gamma_* = \left(1 + \frac{\delta_K}{\kappa}\right)^{-2/(3^K - 1)}$$

with exponential convergence:
$$\mathcal{L}_K(r) - \mathcal{L}_* \leq e^{-2\mu r}\left(\mathcal{L}_K(0) - \mathcal{L}_*\right)$$
$$|\gamma(r) - \gamma_*| \leq C e^{-\mu r}$$

The minimum loss is $\mathcal{L}_* = \kappa\delta_K/[2(\kappa + \delta_K)]$.

**Interpretation:** Training balances preserving graph structure (component separation) against suppressing propagation error. As training depth $K$ grows, $\delta_K$ decreases and $\gamma_*$ approaches 1.

### Graph Hardness (Definition 1)

For a test graph $G$ with largest component size $s_{\max}$ and largest non-unit eigenvalue $\rho \in (0,1)$ of $P_G$:

$$\text{Hard}(G) := \log_3\left(\frac{\log[s_{\max}(1+\tau)/(1-\tau)]}{-\log\rho}\right)$$

Larger components or eigenvalues close to 1 (poorly connected components) increase hardness.

## Empirical Validation / Results

### Test-Depth Window (Proposition 1)

At the population optimum, the trained block recovers every connectivity label at every integer $T \geq 0$ satisfying:

$$\text{Hard}(G) < T < T_{\max}(G; K)$$

where:
$$T_{\max}(G; K) := \log_3\left[1 + \frac{3^K - 1}{\log(1 + \delta_K/\kappa)} \cdot \log\left(1 + \frac{\tau N}{(1-\tau)s_{\max}}\right)\right]$$

**Key findings from experiments (Figure 2):**
- Training on ER graphs with $N = 128$, $Np_{tr} = 0.8$ at depths $K = 2, 3, 4$
- Accuracy first rises with extra loops (propagation), then falls (false connections)
- More training loops widen the useful range of test depths
- Certified windows are conservative lower bounds on true successful ranges

### Generalization Guarantees (Theorem 2)

Under spectral regularity conditions, if training hardness $H \to \infty$ and $H \cdot 3^{K-H} \to \infty$, then there exists a single inference depth $T > K$ such that:

$$T - K = \Theta(H \cdot 3^{K-H})$$

and the population-optimal block solves every test graph with $\text{Hard}(G) < T$.

**Example 1 (Paths):** Training on paths of length $s$, the model extrapolates to paths of length $L = s^{q + o(1)}$ for $1 < q < 1 + 3^\Delta$ where $\Delta = \lim_{s\to\infty}(K - \text{Hard}_{tr})$.

**Example 2 (ER graphs):** Training on $G(n, p_{tr})$ with $np_{tr} < 1$ generalizes to denser ER graphs with:
$$p_{te} \leq \frac{1}{n}\alpha^{-1}\left(\frac{\log(n/\varepsilon)}{S_K}\right)$$
where $\alpha(c) = c - 1 - \log c$.

### Training Depth Selection (Section 5.3)

**Key empirical findings:**
- **Deeper training can be harder to optimize:** From $\gamma_0 \approx 0.96$, $K \leq 4$ converges while $K = 6, 8$ remain near initial plateau under matched compute
- **Best depth depends on compute:** At smaller budgets, $K = 4$ gives strongest generalization; with 10× more compute, $K = 6$ also works
- **Initialization matters:** Starting closer to $\gamma = 1$ allows deeper training to succeed

## Theoretical and Practical Implications

### Theoretical Contributions

1. **First guarantees for depth extrapolation** of trained looped transformers beyond training depth
2. **Complete characterization of training dynamics** through reduction to a single effective parameter
3. **Finite extrapolation window explanation:** The same learned update rule that propagates information farther also determines the limit of safe extrapolation
4. **Exponential gain result:** Certified depth gain $T - K = \Theta(H \cdot 3^{K-H})$ grows exponentially with the gap between training depth and hardness

### Practical Implications

1. **Test-time compute scaling:** More loops help only within a finite window; practitioners should validate on harder tasks rather than assume monotonic improvement
2. **Training depth selection:** Should be chosen jointly with initialization and compute budget
3. **Distribution shift:** Training on tractable instances with a few loops can enable solving substantially harder instances at test time

## Conclusion

### Main Takeaways

- Population training on graph connectivity learns an effective ratio $\gamma_*$ balancing component separation against propagation error
- This creates a finite window of test depths with exact recovery: too few loops under-propagate, too many create false connections
- A single training-selected depth can solve harder test graphs, with guaranteed extrapolation gain growing exponentially with $K - \text{Hard}_{tr}$
- Compute budget and initialization critically affect whether deeper training is beneficial

### Limitations and Future Directions

- **Architecture restrictions:** Analysis assumes linear attention, row normalization, population training, and symmetric initialization
- **Extension needed:** Softmax attention and finite-sample training remain open
- **Tightness:** The sufficient test-depth window could be tightened (experiments show success beyond certified depths)
- **Generality:** Whether similar behavior arises on other algorithmic tasks remains open

---

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