PaperlayerView exact revision on arXiv

Source-linked AI summary

Stability and Generalization in Looped Transformers

Asher Labovich

arXiv:2604.15259v2cs.LGcs.AI

TL;DR

It remains unclear which looped-transformer architectures extrapolate to harder problems when run for more iterations. This paper analyzes fixed points across three stability axes and finds that recall with outer normalization reliably supports all three, while empirical performance across tasks tracks the framework's predictions. The study also finds that internal recall becomes competitive with standard recall under outer normalization, and substantially better on sudoku.

  • Problem

    It remains unclear which architectural choices let looped transformers extrapolate to harder test problems rather than memorize training-depth solutions.

    Source: §1 · p. 1

  • Method

    The paper analyzes looped architectures through fixed-point reachability, input-dependence, and geometry, comparing recall existence and placement with outer normalization.

    Source: §1 · p. 2, §6 · p. 11

  • Results

    Recall with outer normalization reliably yields fixed points satisfying all three stability axes, and performance across chess, sudoku, and prefix-sums tracks the framework's predictions.

    Source: §1 · p. 2, §6 · p. 11

  • Takeaways & Limitations

    Internal recall becomes competitive with standard recall under outer normalization and substantially better on sudoku, though its relative ranking across tasks remains unresolved.

    Source: §6 · p. 11

    The framework predicts which architectures can reach a stable regime but not which will perform best by task, and its empirical analysis uses small, single-layer transformers.

    Source: §5.2 · p. 11, §6 · p. 11

Abstract

from arXiv · show

Looped transformers promise test-time compute scaling by spending more iterations on harder problems, but it remains unclear which architectural choices let them extrapolate to harder problems at test time rather than memorize training-specific solutions. We introduce a fixed-point based framework for analyzing looped architectures along three axes of stability -- reachability, input-dependence, and geometry -- and use it to characterize when fixed-point iteration yields meaningful predictions. Theoretically, we prove that looped networks without recall have countable fixed points and cannot achieve strong input-dependence at any spectral regime, while recall combined with outer normalization reliably produces a regime in which fixed points are simultaneously reachable, locally smooth in the input, and supported by stable backpropagation. Empirically, we train single-layer looped transformers on chess, sudoku, and prefix-sums and find that downstream performance tracks the framework's predictions across tasks and architectural configurations. We additionally introduce internal recall, a novel recall placement variant, and show that it becomes competitive with -- and on sudoku, substantially better than -- standard recall placement once outer normalization is applied.

1 Introduction

Looped transformers may scale reasoning by running a shared network for more iterations, but whether they extrapolate beyond training depth remains unclear. This paper uses fixed-point stability to explain how recall and outer normalization relate to that goal and evaluates the framework across tasks.

1 Introduction

  • Chain-of-thought uses autoregressive discrete tokens and has latency, energy, and trained-token-budget limitations, while looped transformers offer weight-tied iteration scaling.The paper frames additional test-time iterations as a possible route to solving problems harder than those seen during training.Source: §1 · p. 1
  • Whether looped transformers extrapolate to harder test problems remains unclear, despite recall and outer normalization appearing necessary for stable computation.Their theoretical roles and the consistency of generalization across tasks and scales had not been established.Source: §1 · p. 1
  • The paper identifies reachability, input-dependence, and geometric robustness as stability axes, and argues architectural choices affect these axes in ways that predict downstream performance.Its key theoretical result is that recall with outer normalization reliably yields a regime satisfying all three axes.Source: §1 · p. 2
  • The study evaluates single-layer looped transformers on chess, sudoku, and prefix-sums, and introduces internal recall as a placement variant.Internal recall is reported to become competitive with standard recall under outer normalization and substantially better on sudoku.Source: §1 · p. 2

2 Related Work

Prior work studies weight tying, extra loop iterations, equilibrium models, adaptive computation, and deep supervision, but leaves important questions about extrapolation beyond training depth unresolved. This paper situates its fixed-point stability analysis against those distinct aims.

Inductive Biases

  • Studies of weight tying report reasoning-related inductive biases, but generally evaluate fixed loop counts rather than whether extra test-time iterations solve harder problems.Saunshi et al. [2024, 2025] report a trade-off between memorization capacity and reasoning ability; Merrill and Sabharwal [2025] study parallelizable problems.Source: §2 · p. 2

Extrapolation in Looped Transformers

  • Prior extrapolation studies identify recall and outer normalization as important, yet report mixed generalization and do not theoretically justify recall's role.Bansal et al. [2022] study CNNs, Geiping et al. [2025] report performance plateaus beyond training depth, and Yang et al. [2024] report OOD failures.Source: §2 · p. 2

Deep Equilibrium Models

  • Deep Equilibrium Models use implicit root-finding to enforce convergence, whereas this paper analyzes when fixed points arise from the architecture itself.The fixed-point framing motivates examining whether additional iterations avoid overthinking.Source: §2 · p. 2

Adaptive Computation

  • Adaptive-computation methods use learned halting to reduce unnecessary inference computation rather than to increase iterations for harder problems.The Universal Transformer applies ACT to allocate fewer loops to easier tokens; Banino et al. [2021] add KL-divergence regularization.Source: §2 · pp. 2–3

Deep Supervision in Looped Transformers

  • Deep-supervision studies report gains on common reasoning tasks, but evaluate only through training depth and not on harder-than-training problems.That leaves the source of those gains—weight-tying inductive bias or extrapolation with extra iterations—unclear.Source: §2 · p. 3

3 Preliminaries

The preliminaries define looped computation and distinguish autonomous networks from recall networks, then frame successful iteration through reachability, input-dependence, and geometry. These axes capture different requirements for reaching stable, input-sensitive fixed points without fragile parameter regimes.

3.1 Definitions

  • A looped network iterates a parameterized function on hidden states, with iteration depth denoting the number of steps and initialization set by the first update.Autonomous networks depend only on the current state, whereas recall networks also condition on the initial input.Source: §3.1 · p. 3, Figure 1 · p. 3
  • Outer normalization applies a normalizing function to each iteration's output, while the function itself may also contain internal normalization.For RMSNorm or LayerNorm, this is termed post-norm in the paper; GRU outer normalization has no analogous standard name.Source: §3.1 · p. 4

3.2 Axes of Stability

  • The analysis assumes smooth models, and treats reaching a correct fixed point as the criterion for computation that can continue without overthinking or degradation.The paper separates easy training-distribution data from hard out-of-distribution data when assessing the three stability properties.Source: §3.1 · p. 4, §3.2 · p. 4
  • Reachability means iteration converges rather than diverges or cycles, while input-dependence requires fixed points to vary meaningfully with the input.A contraction can ensure reachability yet be useless as an algorithm if its output is independent of input.Source: §3.2 · p. 4
  • Geometry concerns the shape of stable parameter regions: highly anisotropic, narrow regions can make stable configurations fragile and hyperparameter-sensitive.The paper identifies anisotropy as directional bias and notes sharp geometric differences between internal and external recall without outer normalization.Source: §3.2 · pp. 4–5

4 Why Autonomous Networks Fail

The paper argues that autonomous networks have structurally limited input-dependence and unfavorable gradient dynamics at fixed points. Its theoretical analysis links these limits to empirical difficulty, including weaker hard-data performance and a particularly severe prefix-sums failure.

4 Why Autonomous Networks Fail

  • The analysis assumes smoothness and a mild transversality condition; it argues that autonomous fixed points cannot have Jacobian eigenvalue 1 for almost all parameter settings.The passage describes this condition as easier to fulfill as model flexibility and parameter count increase.Source: §4 · p. 5
  • Under the paper's assumptions, almost all autonomous networks have dimension-zero, hence countable, fixed-point sets, limiting input-dependent outputs to basin selection among predetermined points.The count is finite when the state is confined to a compact set, as under post-norm.Source: §4 · pp. 5–6
  • At a fixed point, no regime provides useful training signal: ρ(J) < 1 causes the input-gradient to vanish exponentially, ρ(J) > 1 fails reachability, and ρ(J) = 1 induces exploding parameter gradients.Source: §4 · p. 6
  • Autonomous models may avoid overthinking with stronger progressive-loss training, but still show validation-to-hard degradation and do not approach recall variants' hard accuracy.The paper distinguishes avoiding degradation under extra iterations from having reachable, input-dependent fixed points.Source: §4 · p. 7
  • Autonomous models achieve 0% accuracy on prefix-sums across configurations; its test sequences grow from 32 bits in training to 512 bits, unlike chess and sudoku inputs.The paper suggests basin-selection learned at training dimension largely cannot transfer when the state dimension changes.Source: §4 · p. 7
  • Recall adds the initial input at each step, but recall alone does not ensure fixed-point stability because placement and outer normalization matter.The paper next compares recall architectures along the stability axes.Source: §4 · p. 7

5 Stability and Generalization of Recall Networks

The section distinguishes external and internal recall placements and analyzes their fixed-point stability with and without outer normalization. Its results connect these stability properties to downstream performance while showing that the best architectural choices vary by task.

5 Stability and Generalization of Recall Networks

  • External recall replaces the residual stream with the recalled representation, whereas novel internal recall keeps the current state on the residual path and uses recall only in the update.The formal definitions use a two-sublayer model.Source: Figure 4 · p. 8, §5 · p. 8

5.1 Recall Without Outer Normalization

  • For external recall, x∗ is an attracting fixed point (and thus practically reachable) iff ρ(Mext) < 1; for internal recall, it is attracting iff ρ(Mint) < 1.Source: §5.1 · p. 9
  • Without normalization, external recall outperforms internal recall at low and medium learning rates, while expansive regimes produce task-dependent outcomes and degradation for both.Internal recall’s stable region is more anisotropic and axis-hugging, making strong configurations harder to reach even with contractive recall.Source: §5.1 · pp. 9–10
  • Recall models beat autonomous models on hard data across normalization and learning-rate pairs, but recall alone does not ensure favorable geometry, especially when ρ(Wx) > 15.Source: §5.1 · p. 10

5.2 Recall with Outer Normalization

  • Contractive outer-normalization Jacobians shrink the fixed-point spectral radius; when post-norm maps a compact convex set to itself, Brouwer’s theorem guarantees a fixed point.Source: §5.2 · p. 10
  • With outer normalization and ρ < 1, input-gradients converge to a finite, nonzero fixed-point gradient, and fixed points vary smoothly with x0.This contrasts with autonomous models, which select among a predetermined set of fixed points.Source: §5.2 · p. 10
  • Outer normalization also makes the fixed point locally independent of the initial recall iterate e.The passage connects this initialization-independence with prior evidence that path-independence improves reasoning.Source: §5.2 · p. 11
  • Normalized models outperform unnormalized and autonomous models across tasks; the best normalized models gain 12 points on hard sudoku and 2 on hard chess over the best unnormalized models.On prefix sums, normalized models are the only ones to achieve nonzero hard accuracy at the largest learning rate.Source: §5.2 · p. 11

6 Conclusion

The conclusion frames generalizable looped-transformer behavior through reachability, input-dependence, and geometry, and finds that architectural choices across these axes predict downstream performance. Recall with outer normalization is the most reliable route to all three properties, while internal recall becomes competitive once outer normalization is added.

6 Conclusion

  • The framework identifies reachability, input-dependence, and geometry as distinct stability axes affected by recall existence, recall placement, and outer normalization.Each architectural choice affects a distinct subset of the axes in ways that predict downstream performance.Source: §6 · p. 11
  • Recall with outer normalization is the most reliable route to fixed points that are reachable, input-dependent, and geometrically robust.Across architectures, outer-normalized recall models perform substantially better than recall models without outer normalization, which in turn outperform autonomous models.Source: §6 · p. 11
  • Internal recall becomes competitive with standard recall placement, and substantially better on sudoku, once outer normalization is added.Without outer normalization, internal recall has poor geometric stability.Source: §6 · p. 11
  • The framework predicts architectures that can reach a stable regime, but not which will perform best for a given task; the empirical study uses small, single-layer transformers.Whether the findings extend naturally to much larger models remains future work.Source: §6 · p. 11

A Proofs

The appendix establishes a transversality condition for length-L inputs and characterizes the dimension of autonomous networks’ fixed points, including after a final head layer. It also lays out derivative notation and begins the post-convergence gradient analysis.

A Proofs

  • The appendix distinguishes total derivatives, which include dependencies through the unrolled recurrence, from partial derivatives that hold other arguments fixed.This convention is intended to make chain-rule computations unambiguous.Source: A Proofs · p. 12

A.1 Transversality for Length-L Input

  • For L ≤ d and linearly independent columns of X, the derivative of f(X, A, b) with respect to A is full-rank regardless of A.The proof establishes surjectivity of the derivative map using the full column rank of X.Source: §A.1 · p. 12

A.2 Joint Fixed Points in Autonomous Networks

  • Under the stated transversality condition, autonomous networks’ joint fixed points form a dimension-zero set, finite when the state space is bounded.A final head layer does not change the dimension of the autonomous fixed-point set before applying the loss.Source: §A.2 · p. 12

A.3 Gradient Instability for Autonomous Looped Networks

  • The autonomous-network gradient proof uses the recurrence x_{t+1} = f_θ(x_t) and simplifies its post-convergence term once x_t = x*.The supplied proof passages do not state the resulting gradient bound.Source: §A.3 · pp. 12–13

A.4 Stability Cases for Autonomous Looped Networks

The appendix analyzes how spectral regimes and recall design affect stability, input sensitivity, and derivatives. It also gives explicit recall Jacobians and shows that stable recall can discard dependence on its initial iterate while retaining dependence on the original input.

A.4 Stability Cases for Autonomous Looped Networks

  • When an autonomous fixed point has spectral radius below one, the state becomes exponentially insensitive to its initial input; above one, almost every nearby start eventually leaves its neighborhood.As the spectral radius approaches one from below, the parameter gradient becomes arbitrarily unstable.Source: §A.4 · p. 13

A.5 Stability Regimes of Recall Models without Outer Normalization

  • Without outer normalization, external and internal recall fixed points are only locally stable when their respective Jacobian spectral radii are below one.The appendix states the conditions as ρ(Mext) < 1 for external recall and ρ(Mint) < 1 for internal recall.Source: §A.5 · p. 14

A.6 Stability of Recall Models with Outer Normalization

  • For recall models with outer normalization, the proof establishes invertibility needed for the implicit function theorem and a differentiable fixed point in a neighborhood of x0.The proof bounds the relevant series by a constant before establishing this local dependence.Source: §A.6 · pp. 14, 18

A.7 Explicit Jacobian Forms for External/Internal Recall

  • The explicit external and internal recall Jacobians include outer-normalization Jacobians, which can shrink the spectral radius when they are small.Source: §A.7 · p. 18

A.8 Local-Independence of e in Stable Regimes

  • In the stable regime, the fixed point becomes locally independent of the initial recall iterate e while remaining dependent on x0.Source: §A.8 · p. 20

B.1.1 Data

The three datasets probe different forms of generalization: sequence length in prefix-sums, solver-defined difficulty in sudoku, and puzzle rating in chess. Their train–test splits increase difficulty in different ways.

B.1.1 Data

  • Prefix-sums tests length generalization, training on 8,000 32-bit sequences and testing on 10,000 512-bit sequences.The passage identifies length as the problem’s direct determinant of difficulty.Source: §B.1.1 · p. 20
  • Sudoku training uses 40,000 puzzles below the 38-backtrack threshold, while testing uses 10,000 puzzles above it.The threshold is the dataset’s 80th percentile of backtracks on the Tdoku solver.Source: §B.1.1 · p. 20
  • Chess trains on 40,000 puzzles rated below 1600 and tests on 10,000 puzzles rated 1600–2000, with board size unchanged.Both train and test use 8 × 8 boards.Source: §B.1.1 · p. 20

B.1.2 Training Details

The experiments use single-layer looped transformers with several normalization and recall configurations, alongside fixed-depth comparison models. Training uses progressive losses and a staged learning-rate schedule; all models run on a single GPU.

B.1.2 Training Details

  • Models use a 256-dimensional single-layer transformer, eight attention heads, and a 1,024-dimensional GELU MLP, with task-specific attention choices.Chess and sudoku use full bidirectional attention; prefix-sums uses a width-5 sliding window because RoPE causes length-generalization failures.Source: §B.1.2 · p. 21
  • The sweep compares no outer normalization, pre-norm, post-norm, peri-norm, and GRU outer normalization, with linear recall in recall models.GRU outer normalization replaces the standard residual update with a GRU update.Source: §B.1.2 · p. 21
  • Recall adds 0.13M parameters, but fixed-depth models with roughly an order of magnitude more parameters are outperformed by recall variants on sudoku and prefix-sums.Fixed-depth models only match recall variants on chess, described as the least naturally iterative task.Source: §B.1.2 · p. 21
  • Training runs for 100 epochs with AdamW, a 10-epoch exponential warmup, constant learning rate through epoch 60, then a 10× cooldown.The progressive-loss scheme uses up to 30 iterations, with initial iterations without gradients and loss over only the final gradient-bearing iterations.Source: §B.1.2 · p. 21

B.1.3 Compute Usage

  • The final experiments used 120 NVIDIA B200 GPU-hours, with every model trained on a single GPU.Training used BF16 casting and PyTorch compile.Source: §B.1.3 · p. 22

B.2 Anisotropy Analysis

The appendix examines projected-point anisotropy, learning-rate effects on Jacobian spectral radius, and the relationship between external-recall spectral radius and accuracy. These analyses report greater anisotropy for internal recall and a performance boundary for external recall without outer normalization.

B.2 Anisotropy Analysis

  • The anisotropy analysis samples points from N(0, σ²I) and projects them onto internal and external stability regions for metric comparison.Balance approaches 1 for more isotropic points; the other metric approaches 0.Source: §B.2 · p. 22
  • Internal recall is more anisotropic across σ, with median log-range at least three times larger and median balance at least four times smaller than external recall.The simulation projects 10,000 sampled points per σ onto the stability regions.Source: §B.2 · p. 22, Table 3 · p. 22

B.3 Effect of learning rates on ρ

  • Across problems and recall choices, higher learning rates correspond to larger recall-Jacobian spectral radii.The comparison averages over pre-norm and peri-norm models.Source: §B.3 · p. 22

on External Recall Performance

  • External-recall models without outer normalization achieve no nonzero validation or hard accuracy when the recall-Jacobian spectral radius substantially exceeds 1.At large spectral radius, other eigenvalues must shift toward a complex unit ball centered around −1, making reachable fixed points difficult to create.Source: on External Recall Performance · p. 22

B.5 Seed-Averaged Results for All Models

The section reports seed-averaged accuracy across problems and summarizes learning-rate sweeps, including all four rates for sudoku. The supplied passages identify the reported comparisons and evaluation scope but do not provide table cell values.

B.5 Seed-Averaged Results for All Models

  • Validation and hard accuracy are reported across problems as total accuracy, averaged over three seeds.Total accuracy counts inputs whose complete answer is correct; the figure compares validation and hard accuracy across problems.Source: §B.5 · p. 22, Figure 7 · p. 23
  • Figure 6 compares non-outer-normalized internal- and external-recall Jacobian spectral radii across learning rates.The listed rates are 0.0001, 0.0003, and 0.001 for chess and prefix-sums, and 0.0003, 0.001, and 0.003 for sudoku.Source: Figure 6 · p. 23

B.6 Sudoku Results

  • Sudoku consistently required larger learning rates to reach equivalent accuracy, and its results are reported across all four rates.The appendix includes validation and hard accuracy across the full sudoku learning-rate sweep.Source: §B.6 · p. 23, Table 6 · p. 25

Related papers

Loading 2604.15259v2…