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