cs.LGOct 10, 2026

Learning to Extrapolate in Depth with Looped Transformers

Authors: Huitao Yang, Jianhao Huang, Zixuan Wang, Baharan Mirzasoleiman

Organizations: University of California, Los Angeles · Princeton University

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

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.

Figures & tables

Figure 1: Depth extrapolation on graph connectivity. Top: Training and inference with a shared block. Bottom: Learned model weights and test-time predictions across iterations. We train a shared block ( K=3K=3 loops) on eight disjoint 4-vertex paths and evaluate extrapolation performance on two disjoint 16-vertex paths. T=6T=6 test loops perfectly recovers all test ground-truth labels, whereas insufficient or excess test iterations lead to under- or over-prediction of pairwise connectivity.
Figure 2: Test-depth window at population optimum. We train on 8192 ER graphs with N=128N=128 and Nptr=0.8Np_{\rm tr}=0.8 at depths K=2,3,4K=2,3,4 , and evaluate at varying test depths and densities. Top: exact graph accuracy averaged over test data. Bottom: the fraction of test graphs whose certified test-depth window in Proposition 1 contains each TT . The certified fraction lower-bounds exact accuracy on the same test samples.
Figure 3: Transfer to longer paths and denser ER graphs. We train at the population optimum on disjoint short paths or ER graphs, and evaluate transfer to longer paths or denser ER graphs. (a) Solved path length at fixed N=8192N=8192 , with training loops KK tracking training hardness HH . Both axes use log scales. (b) Maximum exact accuracy on 128 ER test graphs per density at N=64N=64 , using a common T∈{0,…,18}T\in\{0,\ldots,18\} in each cell.
Figure 4
Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Population GD from four initializations. Both excess loss and error in γ\gamma approach zero at training depth K=3K=3 .
Figure 7: Outcomes of the constant-learning-rate sweep at the largest budget. Diamonds indicate outcomes that vary across runs; triangles mark the rates used in Figure 5 .
Figure 8: Full-matrix Adam training at matched FLOPs. Curves and bands show mean solved hardness and ranges over three runs. Dash-dotted lines show untrained performance; the dotted line marks training hardness. Arrows indicate success on the entire test catalog.
sshhK=h−2K=h-2K=h−1K=h-1K=hK=hK=h+1K=h+1K=5K=5
135185610004096∗4096^{*}1000
216329615494096∗4096^{*}96
347562454096∗4096^{*}4096∗4096^{*}56
568986184096∗4096^{*}4096∗4096^{*}31
93917310244096∗4096^{*}4096∗4096^{*}32
Appendix
Table 1: Solved path lengths at N=8192N=8192 . An asterisk marks the catalog limit L=4096L=4096 ; the last column holds training depth fixed.
 γ0≈0.96\gamma_{0}\approx 0.96γ0≈0.999\gamma_{0}\approx 0.999
KKBB10B10BBB10B10B
13×10−53\times 10^{-5}10−610^{-6}3×10−53\times 10^{-5}10−610^{-6}
210−610^{-6}10−710^{-7}3×10−63\times 10^{-6}3×10−73\times 10^{-7}
33×10−63\times 10^{-6}3×10−73\times 10^{-7}3×10−73\times 10^{-7}3×10−83\times 10^{-8}
43×10−63\times 10^{-6}3×10−73\times 10^{-7}10−710^{-7}10−810^{-8}
610−510^{-5}10−510^{-5}3×10−73\times 10^{-7}3×10−63\times 10^{-6}
Appendix
Table 2: Adam learning rates selected by an independent pilot run.
Figure 9: Numerical illustrations of the extrapolation guarantees. (a–b) The certified depth exceeds test-graph hardness for paths with L≃s1.5L\simeq s^{1.5} and L≃s3L\simeq s^{3} . (c) Rarer long training paths permit more certified test loops, until the bound saturates. The left legend applies to (a–b); the right legend applies to (c).
More in this topic

Position Extrapolation

See all 8 papers

Explore similar work

CardsList

  1. Inside the Loop: A Mechanistic Study of Weight-Tied Transformers on Depth-Bound Algorithmic Tasks

    Oct 16, 2026Magnus Sesodia, Juil Sock, Philip Torr +1Recurrent TransformersAlgorithmic Reasoning

    University of Oxford · BBC R&D

  2. DeepLoop: Depth Scaling for Looped Transformers

    Jul 15, 2026Shuzhen Li, Yifan Zhang, Jiacheng Guo +2Residual LearningRecurrent Transformers

    Princeton University · University of California, Los Angeles

  3. Loop-μμP: Residual Scaling and Hyperparameter Transfer in Looped Transformers

    Jun 16, 2026Shaowen Wang, Bingrui Li, Tansheng Zhu +4Residual LearningHyperparameter Transfer

    Tsinghua University · ByteDance Seed · M-A-P