Transformers Learn Two Competing Sorting Algorithms, and Only One Length-Generalizes

Khashayar Gatmiry; Nathan W. Henry; Chloe H. Su; Emmanuel Abbe; Samy Bengio; Melanie Weber; Jonathan A. Kelner; Peter L. Bartlett

cs.LGOct 18, 2026

Transformers Learn Two Competing Sorting Algorithms, and Only One Length-Generalizes

Authors: Khashayar Gatmiry, Nathan W. Henry, Chloe H. Su, Emmanuel Abbe, Samy Bengio, Melanie Weber, Jonathan A. Kelner, Peter L. Bartlett

Organizations: UC Berkeley & Apple · UC Berkeley · Kempner Institute, Harvard University · EPFL, Apple ·

Abstract

We show that two-layer transformers trained to near-zero loss on a sorting task converge to one of two mechanistically distinct solution classes depending on their initialization. One solution class consists of solutions that do not use the second attention layer to sort, but also fail to length-generalize. The other, which uses both attention layers, length-generalizes perfectly, despite only being trained to sort sequences of a fixed length. By carefully choosing the weight initialization, we demonstrate that we can steer the model toward learning one solution or the other. We identify input types that are rare at the training length but increasingly common at longer sequence lengths as a source of this disparity in length-generalization, and find that layer-wise specialization of the second class of solutions allows them to handle these rare sequences more accurately. In the second class, we discover a mechanism by which the second layer attention can use the contextual information gathered by the first layer attention to reduce its search space. On the related task of finding the minimum of a sequence, we find a similar two solution-type phenomenon, but with the opposite generalization pattern, where a different form of layer-wise specialization causes the two-layer solution to length-generalize worse than the single-attention solution.

Figures & tables

Figure 1: Two-layer transformers trained from different initializations learn distinct solution types on three tasks that one layer could solve, and which type length-generalizes depends on the task. Each curve is one seed’s token accuracy, evaluated beyond its training length. Curves are colored by removability, defined as the model’s sequence accuracy when the second attention layer is ablated. The side histograms show how removability is distributed across seeds. Single-stage models are unaffected by the ablation, while two-stage models have poor accuracy without the second attention layer. (a) Sorting: removability is either 0 or 1, and two-stage solutions length-generalize almost perfectly while single-stage solutions do not. (b) Minimum-finding: The pattern reverses where single-stage solutions generalize better than two-stage solutions. (c) Greedy coin denomination: the same pattern as sorting occurs but with a broader spread across seeds and a weaker correlation between solution type and generalization.

Figure 2: Attention probability heatmaps for a two-stage seed (top) and a single-stage seed (bottom), with tokens ordered by numerical value. Both layers of the two-stage model more reliably attend to the correct next token, whereas the single-stage model’s first layer distributes attention across a broader set of nearby candidate tokens.

Panel of Figure 3 — Figure 3: (a) Fraction of the first attention’s mass on input-sequence tokens as a function of gap size, averaged over two-stage models. As the gap increases, the first attention increasingly attends to the separator token or sorted-output tokens rather than unsorted input tokens, indicating that the first attention may be playing a less significant role in the computation. (b) Token accuracy after removing the second attention, broken down by gap size, across two-stage seeds. For many two-stage seeds, even though the accuracy on random inputs goes to 0 after removing the 2nd attention, for tokens of gap size 1, the accuracy remains at $100\%$ after the ablation.

Figure 4: Hijack success rates averaged over all values of $i$ and 23 two-stage seeds (sequence length $32$ , vocabulary size $512$ ). Each panel shows a different gap $g$ . Red: redirecting the first attention’s direct circuit (bypassing the second attention). Blue: redirecting the second attention’s focus. Purple (dashed): fraction of instances where both interventions independently succeed on the same input. At small gaps, both circuits contribute, to varying extents depending on the seed; as the gap grows, the second attention dominates entirely.

Panel of Figure 5 — Figure 5: (a) Most-attended key of the second attention as a function of the query-side token $x$ that the first attention is forced to attend to. The second attention’s focus typically lies between the current token and the first attention’s target, increases approximately monotonically, and saturates for large first attention targets. For small interventions, resembling when the next token is very close in value to the current token, occurring at large input lengths, this property does not hold. (b) Average number of keys with attention weight ${>}0.04$ for varying sequence length for a two-stage and single-stage model.

Figure 6: Top: Fraction of seeds converging to the two-stage solution as a function of the relative initialization scale difference

\Delta=\sigma_{2}-\sigma_{1}

(sequence length 16, vocabulary size 128, 100 seeds). A small positive

\Delta

drives all seeds to the two-stage solution and a small negative

\Delta

drives all seeds to the single-stage solution. Bottom: Length-generalization curves for all models in each condition, colored by solution type. Two-stage models generalize beyond the training length, while single-stage models do not.

Appendix figures & tables24 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: Fraction of seeds converging to the two-stage solution as a function of absolute initialization scale (100 seeds per condition, 99% Wilson confidence intervals). The regime in which the two-stage solution dominates is narrow; outside it, nearly all seeds converge to single-stage solutions.

Figure 8: Best vs. worst length-generalizing steered seeds, where the second layer is initialized to be larger than the first (

\sigma_{1}=0.01,\sigma_{2}=0.013

). Of the models steered with asymmetric initialization and a larger layer 2 initialization scale than layer 1, we analyze the seed that length-generalized the best (blue) and the seed that length-generalized the worst (red). We find that the seed that length-generalized the best had similar specialization patterns across token gaps to two-stage models with symmetric initializations. However, there is indication that the model that generalized the worst, while being recorded as a two-stage model, relied more on the 2nd layer than the first.

Figure 9: Solution type vs. vocabulary size

N

. Proportion of seeds converging to each solution type as a function of vocabulary size

N

with sequence length

16

, reporting the probability of convergence to a two-stage model with 99% confidence intervals, estimated with 100 seeds for each

N

. Larger vocabularies increasingly favor the two-stage solution.

Figure 10: Solution type vs. embedding dimension

E

. Proportion of seeds converging to each solution type as a function of embedding dimension

E

, with sequence length

16

and vocabulary size

N=256

held fixed, reporting the probability of convergence to a two-stage model with 99% Wilson confidence intervals. Every bar is estimated from a common subsample of

70

trained seeds—the largest number available at every width—so that all widths share a denominator and their intervals are directly comparable. Seeds that never learned the task are excluded, since ablating the second attention layer is only informative for models that solve the task. Wider models increasingly favor the single-stage solution, with the even-odds crossing between

E=60

and

E=62

. This is the opposite direction to vocabulary size (Figure 9 ) and consistent with the same explanation: the two-stage solution is recruited when the task is hard relative to the capacity available to solve it. Seeds are judged at their

60{,}000

-iteration checkpoint, matching the training length used elsewhere in the paper, although the cosine schedule was set for

100{,}000

iterations so the learning rate has not fully annealed at that point. Repeating the classification at

100{,}000

iterations changes no seed’s assignment between the single-stage and two-stage classes; instead

16

775

seeds stop solving the task altogether,

14

of them at

E=32

, which moves each reported share by at most one point.

Figure 11: Solution type against training length. Proportion of trained models converging to the two-stage and single-stage solutions at each training length

K

, with the vocabulary fixed at

N=256

, using

50

seeds per choice of

K

. The lighter bands are

99\%

Wilson confidence intervals on the probability of converging to a two-stage model.

Figure 12: Length generalization of two-stage models, colored by first-layer attention on the input at gap $1$ , as a proxy for layer-wise specialization, studied in Section 4.1 . Two-stage curves are colored by the attention mass that the first layer places on the input sequence for queries with gap size $1$ to their next value, measured at the training length. Single-stage curves are dotted gray. The value $r$ is the within-length Pearson correlation between color and token accuracy at length $256$ .

Figure 13: Length generalization of two-stage models, colored by the second attention’s candidate set size. Color gives the average number of keys receiving second-attention weight above $0.04$ at training length. Single-stage curves are dotted gray. The value $r$ is the within-length Pearson correlation between color and token accuracy at length $256$ .

Figure 14: Small-gap adversarial examples across training lengths. Accuracy on runs of

n

consecutive values

Figure 15: Large-gap adversarial examples across training lengths. We report accuracy on the prediction that crosses a gap of size

g

. Panel titles give the training length

K

and the mean gap size seen during training for reference. Lines show the median over models and bands show the range. Gaps below 16 are handled perfectly and are omitted.

Figure 16: The multiplicity of solution classes, and their steerability, persists at depth four. Each column removes a deeper suffix of attention layers simultaneously (diagram at bottom) and re-evaluates. Plotted points are individual models, colored by the fewest leading layers the model can run on. Under equal initialization scales, no model reduces to a single attention layer, even though one would suffice. Lowering the scale of layers

2

4

steers

14

of the

50

models to a one-layer solution.

Figure 17: Mirroring our results for transformers of depth 2, when transformers of depth 4 are trained to sort sequences of numbers, only the solutions that utilize at least 2 layers are able to length-generalize.

Figure 18: Whether a second layer is used varies by seed, but when we train models with multiple heads, all heads are always used when the layer is being used. We report sequence accuracy after removing one unit from each model in each row. For heads, we remove the head whose removal costs the model the least. On the right side of the plot, we display the number of models that keep at least

95\%

of their accuracy. Top. Removing the second attention layer leaves some seeds unaffected and breaks others, both with one head per layer and with two. Bottom. Removing a single head breaks every model, whether it is one of two heads in a one-layer model (where each head alone would be enough to perfectly solve the task) or a head of either layer in the two-head, two-layer models.

Figure 19: A second head prevents the length-generalization failure. We evaluate token accuracy across input lengths for models with embedding dimension

64

that are trained on length-16 inputs. (a) Two-layer models with one head per layer. Single-stage models (red) fail at long lengths, while two-stage models (blue) do not. (b) Two-layer models with two heads per layer. This set-up is chosen so that every model is single-stage; however, we find that all single-stage models length-generalize when trained with 2 heads per layer. (c) We train both one-layer models with one head and one-layer models with two heads and find that only those with two heads length-generalize.

Figure 20: Second attention QK score heatmaps for query token $z=250$ . Left: The score is nearly invariant to the key-side first attention target $y$ , showing uniform horizontal bands. Right: The score is strongly modulated by the query-side target $x$ : as $x$ increases, a high-score band emerges in $[z,x]$ .

Figure 21: Mean $\ell_{1}$ distance between normal and MLP1-only second attention distributions for 13 two-stage checkpoints (outliers excluded).

Figure 22: Full-sequence sorting accuracy under normal versus MLP1-only second attention input. The accuracy drop is zero for 12 of 13 checkpoints and $0.05\%$ for the remaining one.

Figure 23: First attention QK score (black, left axis) versus second attention QK score slices (colored, right axis) as a function of key token value. The second attention’s scores are smoother within their respective windows, reflecting the benefit of the narrower search range.

Figure 24: Second attention QK score heatmaps for $n\in\{1,2,3,4\}$ equally attended tokens ( $\delta=5$ , $z=250$ ). The band structure is preserved across all cases.

Figure 25: Second attention QK score slices with $n=2$ equally attended tokens for $\delta\in\{1,5,20\}$ .

Figure 26: Second attention QK score slices with $n=3$ equally attended tokens for $\delta\in\{1,5,20\}$ .

Figure 27: Second attention QK score slices with $n=4$ equally attended tokens for $\delta\in\{1,5,20\}$ .

Figure 28: Argmax bias analysis ( $z=250$ ). Left: Argmax $t^{*}$ of the second attention vs $x_{\min}$ for $n\in\{1,2,3,4\}$ at $\delta=5$ . Middle: Same for $n=3$ at varying $\delta$ . Right: Scatter of $(t^{*}-x_{\min})$ vs $(t^{*}-x_{\mathrm{mean}})$ across all configurations; diamonds mark per- $n$ means. This shows that the argmax mostly varies with the minimum of the first attention concentration set, $x_{min}$ , and has almost no dependence on $n,\delta$ , or $x_{mean}$ . In particular, the argmax consistently tracks $x_{\min}$ , not $x_{\mathrm{mean}}$ .

Figure 29: Residual stream component removal across 17 two-stage checkpoints. Left: Per-checkpoint argmax agreement rates for Claim 1 (remove $\mathbf{a}_{1}$ , teal) and Claim 2 (remove $\mathbf{m}_{1}$ at readout, coral). The argmax agreement rate is the fraction of times removing the redundant circuit does not change the argmax of the output logits at a given token position. Hatched bars mark known outliers. Right: Scatter of Claim 1 vs Claim 2 agreement. The shaded region ( ${\geq}90\%$ on both axes) contains the 12 standard two-stage models where both direct paths are redundant.

Figure 30: Per-token accuracy without the second attention as a function of current sorted value $i$ , for gaps $g\in\{1,10,50\}$ (left to right). Grey dots show per- $i$ accuracy; the solid curve is a rolling average (window 20); shaded bands show 95% bootstrap confidence intervals. For small gaps the first-layer pathway alone achieves high accuracy once $i$ is sufficiently large; for large gaps the second attention is indispensable.

Length Generalization

Length generalization studies whether sequence models trained on short inputs can reliably solve the same tasks on longer sequences.

See all 46 papers

Oct 6, 2026

Exact-Solution Volume and Length Generalization in Transformers

Research on transformer expressivity shows whether a transformer is capable of solving a given task, but gives little indication of whether the solution, if learned, is generalizable to longer input lengths. We study this question through normalized exact-solution volume (NESV): the fraction of a bounded parameter region that achieves an exact solution on every input of length

n

. For fixed-width, single-layer transformers with

\log n

-scaled attention, we establish asymptotic bounds on NESV for four tasks: FIRST (

Θ(1)

), MAJORITY (

Θ(1/(n\log n))

), INDEX (

Θ(1/n^3)

), and PARITY (

0

). These results are consistent with previous empirical results: the faster the exact-solution volume decays with input length, the harder it is to length-generalize on that task. Looking deeper into INDEX, our volume analysis reveals two error sources that grow with

n

. Consequently, we study a transformer model that would structurally eliminate one of the terms, theoretically improving the NESV bound to

Θ(n^{-1})

, and empirically achieving 85% accuracy when tested at

10\times

the training length, compared with the 60% accuracy of the original model. We conclude that volume analysis may be a useful approach to identify concrete sources of length sensitivity and thus provide insights into task-specific model refinements.

Yijia Jessica Zhu, David Chiang

University of Notre Dame

Transformer Expressivity Transformer Length Generalization

Ask PDF

Oct 14, 2026

Training Dynamics of Transformers on Unbounded Inputs: Attention Learning Can Lag Output Layer Learning

The success of transformers is attributed to several factors, including their capacity to express formal languages over variable-length inputs, and ability of the attention mechanism to focus on relevant input features. However, there is limited understanding of the training dynamics of shallow transformers, specifically how to decouple the learning of the attention and the output (feedforward) layer. A natural extension of the theory of one-step gradient feature learning in feedforward networks suggests that one could first learn the attention weights with a single gradient step, and then train the output layer. We show that this fails for a shallow transformer trained to recognize a star-free language from unbounded-length input. We prove that the transformer can be trained to recognize the language with a three-phase gradient descent: the output layer is updated first, followed by an attention update, and the output layer is fine-tuned in the third phase. In contrast, the language cannot be learned by training only the attention or output layer, or by a single attention-layer update followed by finetuning the output layer. We empirically show that ``attention learning lagging output layer learning" is beneficial in other language recognition, regression, and in-context learning tasks as well.

Yaomengxi Han, Debarghya Ghoshdastidar

Technical University of Munich

Transformer Training Dynamics Neural Network Training Dynamics Transformer Attention Transformer FFNs

Ask PDF

Oct 3, 2026

Length Generalization Needs Proper Regularization

Length generalization is the ability of sequential models to perform well on context lengths unseen during training. In this work, we show that the challenge of achieving length generalization is related not only to architectural choices such as positional encoding and the attention mechanism but also to the training procedure itself. We study how regularization affects length generalization and find that weight decay can hinder extrapolation. In contrast, dropout improves extrapolation when its placement within the architecture is reconsidered. In that regard, we show that the standard placement of dropout before layer normalization introduces a systematic distributional mismatch, and that applying dropout just before the linear projection resolves this issue. For example, a modified SmolLM3 with sliding window attention, continually pre-trained with dropout, can extrapolate perfectly to 64

\times

on Needle-in-a-Haystack and far beyond the pre-training context size on RULER and HELMET. Mamba2 also benefits from dropout, suggesting an architecture-agnostic nature of the problem. We further propose Variance-Preserving Affine Dropout (VPAD), a new dropout strategy that substantially reduces the resulting pre-activation variance mismatch, leading to further extrapolation improvements in transformers.

Pavlo Vasylenko, Matthias Lindemann, André F. T. Martins +1

Instituto Superior Técnico, Universidade de Lisboa · Instituto de Telecomunicações · ELLIS Unit Lisbon +1

Length Generalization Language Model Generalization Neural Network Generalization LLM Training Long-Context Language Modeling

Ask PDF

Abstract

Length Generalization

Explore similar work

Exact-Solution Volume and Length Generalization in Transformers

Training Dynamics of Transformers on Unbounded Inputs: Attention Learning Can Lag Output Layer Learning

Length Generalization Needs Proper Regularization