PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning to Extrapolate in Depth with Looped Transformers

Huitao Yang, Jianhao Huang, Zixuan Wang, Baharan Mirzasoleiman

arXiv:2610.14456v1cs.LG

TL;DR

The paper asks how far a looped transformer trained at one depth can safely extrapolate on graph connectivity. It analyzes population training of a shared block and derives the finite test-depth window and hardness guarantees. Experiments validate the windows and show that compute and initialization affect training-depth choice.

  • Problem

    The paper asks how far a looped transformer trained at fixed depth can be safely unrolled to solve harder graph-connectivity instances.

    Source: §1 · pp. 1–2

  • Method

    Under exchangeable training distributions and symmetric positive initialization, the analysis reduces population training to a single effective parameter governing a shared block.

    Source: §4 · p. 5

  • Results

    The learned block has a finite test-depth window for exact recovery, and a training-distribution-chosen depth can guarantee recovery on harder graphs with gain growing exponentially as training depth exceeds training hardness.

    Source: §5.2 · p. 8, §6 · p. 10

  • Takeaways & Limitations

    Experiments show that the useful training depth depends on compute and initialization, and more compute does not automatically favor the deepest training run.

    Source: §5.3 · p. 9

    The theory assumes linear attention, row normalization, population training, and symmetric initialization; extension to softmax attention and finite-sample training remains open.

    Source: §6 · p. 10

Abstract

from arXiv · show

Looped transformers reuse a single block across iterations, raising a natural question: how far can a block trained at one depth be safely extrapolated to greater depths? We study this question for all-pairs graph connectivity, where harder graphs require longer-range information propagation. We show that under exchangeability and symmetric positive initialization, population training at depth K converges to an update rule governed by a single effective parameter that balances graph propagation against fitting structure unresolved at training depth. This trade-off creates a finite extrapolation window: too few test-time loops provide insufficient propagation, while too many eventually create false connections. We derive explicit endpoints for this window and bound the graph hardness solvable beyond training as a function of training depth and hardness. In particular, a test depth chosen solely from the training distribution can guarantee recovery on substantially harder graphs, with a certified depth gain that grows exponentially as training depth outpaces training hardness. Experiments validate the predicted windows and show how compute budget and initialization affect the optimal training depth.

1 Introduction

The paper asks how far a looped transformer trained at fixed depth can safely extrapolate, using graph connectivity to study the question. It characterizes the finite window in which extra loops help and quantifies how much harder a graph the trained block can solve.

1 Introduction

  • Graph connectivity offers a controlled test because harder instances require information to propagate over longer distances.This connects instance hardness to the number of iterations needed for successful computation.Source: §1 · p. 1
  • A fixed-depth trained block balances graph propagation against structure unresolved at training depth, so additional loops can first help and later create false connections.This makes successful extrapolation a finite test-depth window rather than a monotonic benefit of more computation.Source: §1 · p. 2
  • Prior results establish that looped transformers can represent iterative algorithms and sometimes learn them, but do not characterize a trained block unrolled beyond its training depth.At fixed depth, training need not recover an ideal iterative rule; it balances propagation against residual unresolved structure.Source: §1 · p. 2
  • The paper derives extrapolation guarantees and studies shifts from short to longer paths and from sparse to denser Erdős–Rényi graphs.It also examines how compute and initialization affect training-depth choice.Source: §1 · p. 2
  • The paper characterizes when frozen blocks extrapolate beyond training depth and when further test-time iterations become harmful.The authors describe these as the first guarantees quantifying both the distance of safe extrapolation and its harmful endpoint.Source: §1 · p. 2

2 Related Work

Prior work demonstrates the promise of shared-depth computation, algorithmic expressivity, and learned iterative procedures. This paper distinguishes its focus: the behavior of a trained block when frozen and run beyond its training depth.

Looped transformers and test-time depth

  • Looped transformers can scale test-time computation by reusing a block, but excessive iteration can degrade performance.Prior work also reports improved language-model performance with further unrolling and length generalization using input-adaptive step counts.Source: §2 · pp. 2–3

Expressivity for algorithmic and graph reasoning

  • Expressivity results show that looped transformers can simulate graph algorithms and achieve logarithmic-depth connectivity, but do not characterize post-training unrolling.The paper’s architecture has a 3L connectivity range after L loops.Source: §2 · p. 3

Training dynamics of algorithm-implementing transformers

  • Prior training-dynamics work studies learned optimization algorithms and graph connectivity; here the object is a trained block run beyond its training depth on harder inputs.The cited graph result finds the matrix-power algorithm only when most training graphs lie within model capacity.Source: §2 · p. 3

Length and size generalization

  • Length-generalization work studies longer inputs and graph-size shifts, whereas this setting keeps the architecture fixed and extrapolates only through extra loops.The paper certifies how much harder a test graph can be in component size and spectral gap as a function of training depth and distribution.Source: §2 · p. 3

3 Problem Setting

The study trains one shared attention block for graph connectivity and evaluates its frozen weights with more loops on harder graphs. Exact success means recovering every pairwise connectivity label.

3 Problem Setting

  • The task predicts whether every pair of vertices in a simple undirected graph is connected, including each vertex’s zero-length self-connection.Pairs in different components receive label zero.Source: §3 · p. 3
  • Training uses a fixed input map with no learned input embedding and a component-normalized target that preserves the connectivity labels’ support.The input encodes each vertex and its neighbors; the training input is PG and target is ΠG.Source: §3 · p. 3
  • One linear-attention block with shared parameters is trained for K loops using terminal population loss, without separate supervision of intermediate states.The block uses row normalization, a residual connection, and no softmax or feed-forward sublayer.Source: §3 · p. 4
  • At inference, frozen weights run for T loops and a fixed thresholded readout succeeds only on exact recovery of all pair labels.Depth extrapolation uses T > K without additional parameters or test-time optimization; Figure 1 shows K = 3 training and exact recovery on longer paths at T = 6.Source: §3 · p. 4, Figure 1 · p. 4

4 Training Dynamics

Under exchangeable training data and symmetric positive initialization, population training preserves a symmetric weight structure and learns a single effective ratio. That ratio balances component separation against residual propagation error, shaping both the asymptotic solution and the optimization cost.

4 Training Dynamics

  • Population training reduces the symmetric weight matrices to a single effective ratio whose optimum balances component separation against propagation error remaining after 3K steps.κ measures the expected number of components beyond one, while δK measures expected residual propagation error.Source: §4 · p. 5
  • At training depth K, the optimal graph-dependent coefficient is γeK* = κ/(κ+δK), reflecting the trade-off between preserving component structure and suppressing propagation error.The minimum loss in this positive parameter family is L* = κδK/[2(κ + δK)].Source: §4 · p. 5
  • As training depth grows, propagation error decreases and γ* approaches one, so the learned block attenuates the graph-dependent component more slowly on later test loops.When γ = 1 the block computes P_G^(3t); when γ < 1, looping also moves the output toward E.Source: §4 · p. 5
  • Relabeling invariance and symmetric positive initialization keep gradient flow in span{E, F}, where the loss is a convex quadratic with exponential convergence.Gradient descent reaches the same limit with a sufficiently small constant learning rate.Source: §4 · p. 5
  • Deeper training improves the converged solution but can slow optimization because the gradient becomes suppressed as K grows when the effective ratio is below one.The time required to reach the solution depends on initialization.Source: §4 · p. 5

5 Extrapolation

The section characterizes when a population-optimal looped block recovers connectivity beyond its training depth, and how far that recovery can extend to harder graphs. It also shows that practical gains from deeper training depend on whether the optimization can reach the population optimum under the available compute and initialization.

5.1 Extra loops and harder test graphs

  • For a fixed graph, exact recovery is guaranteed when Hard(G) < T < Tmax(G; K), with the upper restriction absent for connected graphs or when δK = 0.The lower endpoint reflects propagation needs; the upper endpoint reflects training-induced bias that can create false-positive connections.Source: §5.1 · pp. 6–7
  • Exact accuracy rises and then falls as test loops increase; more training loops widen the useful depth range, and certified windows conservatively lower-bound successful depths.The certified fraction is a lower bound on exact accuracy for the same test samples.Source: §5.1 · pp. 6–7, Figure 2 · p. 6

5.2 Generalization from shallow to deep tasks

  • Under spectral regularity, guaranteed additional inference loops scale as Θ(H·3^(K−H)); when K−H = O(1), they scale as Θ(H).The gain is controlled by the gap between training depth and training hardness.Source: §5.2 · p. 7
  • For paths, a constant gap K−Hardtr → ∆ permits a training-selected test depth to solve lengths growing as s^(q+o(1)) for q < 1 + 3∆.The guarantee applies to feasible test sizes under the path-family setup.Source: §5.2 · p. 7
  • For denser ER test graphs, a training-selected depth and component-size budget yield exact-match accuracy of at least 1−ε when the stated depth and size conditions hold.The conditions include TK > K and SK > log(n/ε)/α(nptr).Source: §5.2 · p. 8

5.3 How many training loops?

  • At matched compute from γ0 ≈ 0.96, K ≤ 4 reaches the population optimum, while K = 6 and 8 stay near initial loss despite separately tuned learning rates.The learning-rate sweep also found that larger rates could destabilize the deeper runs.Source: §5.3 · p. 9, Figure 4 · p. 9
  • At the smaller budget, K = 4 generalizes best; with ten times more compute, K = 4 and 6 solve the full catalog, while K = 8 does not improve over initialization.More compute expands the useful range of training depths but does not make the deepest run preferable automatically.Source: §5.3 · p. 9
  • Near-identity initialization lets K = 4, 6, and 8 solve the full catalog at the smaller budget, but also provides a stronger untrained baseline.The paper recommends choosing depth jointly with initialization and compute, and comparing candidates at matched compute using harder-task validation.Source: §5.3 · p. 9

6 Conclusion & Limitations

Population training creates a finite window of test depths with exact recovery, and a training-selected depth can solve harder graphs. The analysis assumes linear attention, row normalization, population training, and symmetric initialization.

6 Conclusion & Limitations

  • A single depth chosen from the training distribution can solve harder test graphs, with a guaranteed extrapolation gain that grows exponentially with the gap between training depth and graph hardness.Experiments validate the predicted windows and show how compute and initialization affect training-depth choice.Source: §6 · p. 10
  • The analysis assumes linear attention, row normalization, population training, and symmetric initialization.Extending the theory to softmax attention and finite-sample training remains an important direction.Source: §6 · p. 10

A Problem Setting and Training: Supporting Results

The supporting analysis reduces population training under exchangeability and symmetric positive initialization to an invariant parameter family and a one-dimensional effective objective. It establishes the objective's optimum and convergence of population gradient flow to that optimum.

A Problem Setting and Training: Supporting Results

A.1 Graph spectrum and connectivity

  • The graph operator P_G is symmetric, positive semidefinite, and row stochastic; its eigenvalue-one space consists of component-constant vectors, and its powers converge to Π_G.The remaining eigenvalues are strictly below one.Source: §A.1 · p. 13

A.2 Vertex relabeling and parameter invariance

  • Exchangeability and permutation symmetry preserve the weight forms W = aF + bE and V = cF + dE under population gradient flow and gradient descent.The invariant subspace consists of matrices fixed by every permutation conjugation.Source: §A.2 · pp. 13–14

A.3 Closed-form loop recursion

  • For positive weights in the invariant family, the forward-map derivation uses commutation and updates its exponent by e_{t+1} = 3e_t + 1.The proof establishes the recursion by induction.Source: §A.3 · p. 14

A.4 Effective loss and stationary points

  • The effective objective has a unique minimizing ratio γ*, which is also the ratio of every positive restricted stationary pair and yields a stationary point of the full matrix objective.The objective is a strictly convex quadratic in q = γe_K, with q* = κ/(κ + δ_K).Source: §A.4 · pp. 14–15

A.5 Population gradient-flow convergence

  • Population gradient flow remains in a compact subset of the positive orthant and converges to a positive pair with ratio γ*, with an exponentially decaying gradient bound.The trajectory exists for all optimization times.Source: §A.5 · p. 15

B Discrete Population Gradient Descent

Under the assumptions of Theorem 1, population gradient descent with a sufficiently small constant step size preserves positivity and converges to a positive stationary pair with the optimal effective ratio.

B Discrete Population Gradient Descent

  • A sufficiently small constant step size keeps the iterates positive and yields geometric convergence to a positive stationary pair with effective ratio γ*.The parameter matrices converge to the same effective ratio as population gradient flow.Source: B Discrete Population Gradient Descent · p. 16, §B.1 · p. 17

B.1 An explicit sufficient step size

  • The sufficient step-size argument keeps the ratio between its initial and optimal values, maintains positive coordinates, and prevents overshooting the optimum.The proof establishes descent and geometric convergence within an invariant region.Source: §B.1 · pp. 16–17

C.1 Proof of the useful test-depth window

The certified depth window requires enough loops to propagate connectivity within components but not so many that disconnected pairs become false positives. Explicit conditions characterize when this window contains an integer depth, including one strictly beyond training depth.

C.1 Proof of the useful test-depth window

  • Exact recovery is certified when Hard(G) < T and, if δK > 0, T < Tmax(G; K), separating sufficient propagation from false connections.The upper restriction comes from cross-component scores eventually exceeding the threshold; it is absent when δK = 0 or G is connected.Source: §C.1 · p. 18
  • A sufficient failure condition is distinct from the upper endpoint of the success window.The endpoint bounds successful depths; the separate failure condition certifies a different regime.Source: §C.1 · p. 18

C.2 Explicit certified region

  • For extrapolation beyond training, an integer depth exists exactly when max{K, ⌊Hard(G)⌋} + 1 < Tmax(G; K).Both endpoints refer to the same graph, and the upper endpoint depends on its largest component size.Source: §C.2 · p. 18
  • The continuous certified window narrows as the largest component grows, while degenerate spectra remove the lower restriction.When δK = 0, the upper restriction is also absent.Source: §C.2 · pp. 18–19, §C.1 · p. 18

C.3 A test depth uniform over test graphs

This section studies training-selected loop depths that guarantee connectivity recovery across test graphs, deriving how those guarantees scale with training depth, hardness, and graph families. It also covers recovery beyond training hardness, including finite-size certificates, path generalization, ER cutoffs, and convergence across positive invariant initializations and learning rates.

C.3 A test depth uniform over test graphs

  • A single training-selected integer depth recovers every connectivity label for every test graph whose hardness is below its certified cutoff.The guarantee is simultaneous across graphs and uses a depth chosen without inspecting the test graph.Source: §C.3 · p. 19

C.4 Monotonicity with training depth

  • Increasing training depth makes the residual nonincreasing and preserves every test depth already certified by Proposition 1.The common certified depths are also nondecreasing with training depth at the population optimum.Source: §C.4 · p. 19

C.5 Proof of Theorem 2

  • Under spectral regularity, the common test-depth guarantee extends beyond training depth by order R, where R = H3K−H.The normalized training residual is exponentially small in R, and the common depth bound depends logarithmically on its inverse.Source: §C.5 · p. 20

C.6 High-probability shallow-to-deep extrapolation

  • At the training-selected depth, test-law mass of at least 1−ζ below RK yields exact-match accuracy of at least 1−ζ, while the training tail probability controls the population residual.Source: §C.6 · p. 21

C.7 Short to Long Path Generalization

  • For packed-path training with Ks−Hs → ∆, one common depth recovers paths of lengths up to sq for any fixed 1 < q < 1 + 3∆.The guarantee holds for every feasible integer length through min{Ns, ⌊sq⌋} and for every relabeling.Source: §C.7 · p. 23

C.8 ER-to-ER generalization

  • For ER training, the component-size tail yields a high-probability training cutoff, which bounds training-graph hardness and feeds the common-depth guarantee.The supplied passages give Pr(smax(Gtr) > Str) ≤ ε and Hard(G) ≤ Htr for nondegenerate training graphs with smax ≤ Str.Source: §C.8 · pp. 24–25
  • Asymptotically, a training-selected depth T > K with T − K = Θ(Mthm) solves all test graphs with Hard(G) < T.This guarantee applies when the theorem’s conditions hold and does not assert T > K for arbitrary fixed n, K, and ε.Source: §C.8 · p. 25
  • At finite graph size, an explicit component-size budget yields one depth TK that exactly recovers every graph with smax ≤ SK.Strict extrapolation beyond training additionally requires TK > K; retaining explicit spectral decay gives a sharper certificate than the direct theorem invocation.Source: §C.8 · pp. 25–26
  • The finite-size certificate proves recovery for sampled subcritical test graphs, also covers pte = 0, and permits denser-than-training graphs when the stated interval condition and TK > K hold.The denser-than-training interval is nonempty precisely under the condition stated in the passage.Source: §C.8 · p. 26

D.1 Population convergence

  • Four positive invariant initializations approach the same optimal ratio on uniform relabelings of C10 ⊔ C10 with N = 20 and K = 3.The experiment uses a constant learning rate of 0.002.Source: §D.1 · p. 26

D.2 Learning rates and finite-budget optimization

  • At K = 6 and 8, none of the tested constant learning rates converges within the budget, and increasing the rate can cause divergence.Several rates converge for K ≤ 4, so the useful learning-rate range depends on training depth.Source: §D.2 · p. 26

D.3 Initialization and useful training depth

The experiments examine initialization and optimization under different training depths, alongside population-training and ER test-loop protocols. They show that untrained performance matters when judging the benefit of additional training loops.

D.3 Initialization and useful training depth

  • Near-identity initialization already solves harder graphs before training, so trained curves must be compared with their respective untrained baselines.Figure 8 reports mean solved hardness and three-run ranges at matched FLOPs, with arrows marking success across the test catalog.Source: §D.3 · p. 26, Figure 8 · p. 27
  • At K = 3, both excess loss and error in γ approach zero across the four-initialization population-GD comparison.The figure caption identifies the two quantities and the training depth; the passage does not specify their plotted trajectories.Source: Figure 6 · p. 27

E.1 Population training and exact recovery

  • With exchangeable population training and symmetric initialization, full-matrix GD stays in the E–F subspace, making four-coordinate GD an exact reduction in exact arithmetic.This is a discrete-time identity, and the experiments use no minibatch sampling that could break the symmetry.Source: §E.1 · pp. 27–28

E.2 ER test-loop curves

  • The ER test-loop experiment trains on N = 128 graphs at K = 2, 3, 4, then evaluates independent test graphs across densities and loop depths.It also reports the fraction of test graphs certified by Proposition 1.Source: §E.2 · p. 28

E.3 Path and ER transfer

The transfer experiments use path and ER graph families to examine common test-depth choices across graph sizes or densities. The reported ER accuracies use a single test pool and test-outcome-based depth selection.

E.3 Path and ER transfer

  • The path study trains on relabeled collections of paths at several hardness levels and tests packed paths from L = 2 to 4096.It chooses one common depth maximizing the consecutive solved prefix, using test outcomes and allowing T ≤ K.Source: §E.3 · p. 28
  • The path table covers solved lengths at N = 8192 and includes a final comparison holding training depth fixed.The caption gives the table scope but supplies no cell values.Source: Table 1 · p. 28
  • ER tests span training and test densities, with 128 independent test graphs per density and directly evaluated depths T = 0 to 18.All graphs in a cell share a depth selected using test outcomes, and reported accuracies use one sampled test pool.Source: §E.3 · p. 29

E.4 Optimization at matched compute

The section describes matched-compute optimization and test-catalog evaluation, then numerically illustrates certified depth guarantees for path families and training distributions with rare long paths.

E.4 Optimization at matched compute

  • Population GD compares K = 1, 2, 3, 4, 6, 8 using learning-rate sweeps and selects the rate reaching convergence in the fewest mean updates when all runs converge.The protocol uses three paired runs and declares convergence below excess loss 10^-8.Source: §E.4 · p. 29
  • Full-matrix Adam trains all 32,768 entries of W and V, with pilot-selected learning rates based on held-out endpoint loss for each depth, initialization family, and budget.Table 2 lists the selected rates; the budgets are 2.21 × 10^11 and 10B GEMM FLOPs.Source: §E.4 · p. 29, Table 2 · p. 29
  • The test catalog evaluates 13 path, seven broom, and four cycle templates at a common validation-selected depth, reporting the longest prefix with at least 95% exact accuracy on every template.Each template uses 64 validation relabelings and 128 independent test relabelings; bands show the range over three runs.Source: §E.4 · p. 29

E.5 Numerical illustrations of the guarantees

  • Certified depth exceeds test-graph hardness for paths with L ≃s1.5 and L ≃s3, while rarer long training paths permit more certified test loops until the bound saturates.Figure 9 illustrates these guarantees for growing path families and a training distribution with rare long paths.Source: Figure 9 · p. 30, §E.5 · p. 30
  • The guarantee illustrations use packed paths for two path-family settings and a mixture of G128,4 and G128,64 for the rare-long-path setting.The mixture assigns probability ε to G128,64, with ε ranging from 10^-50 to 10^-1.Source: §E.5 · p. 30

Related papers

This paper cites

Loading 2610.14456v1…