cs.LGOct 18, 2026
SaveTransformers Learn Two Competing Sorting Algorithms, and Only One Length-Generalizes
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.
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 after the ablation.
Figure 4: Hijack success rates averaged over all values of and 23 two-stage seeds (sequence length , vocabulary size ). Each panel shows a different gap . 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.
Figure 5: (a) Most-attended key of the second attention as a function of the query-side token 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 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 (sequence length 16, vocabulary size 128, 100 seeds). A small positive drives all seeds to the two-stage solution and a small negative 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 ( ). 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 . Proportion of seeds converging to each solution type as a function of vocabulary size with sequence length , reporting the probability of convergence to a two-stage model with 99% confidence intervals, estimated with 100 seeds for each . Larger vocabularies increasingly favor the two-stage solution.
Figure 10: Solution type vs. embedding dimension . Proportion of seeds converging to each solution type as a function of embedding dimension , with sequence length and vocabulary size 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 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 and . 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 -iteration checkpoint, matching the training length used elsewhere in the paper, although the cosine schedule was set for iterations so the learning rate has not fully annealed at that point. Repeating the classification at iterations changes no seed’s assignment between the single-stage and two-stage classes; instead of seeds stop solving the task altogether, of them at , 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 , with the vocabulary fixed at , using seeds per choice of . The lighter bands are 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 , 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 to their next value, measured at the training length. Single-stage curves are dotted gray. The value is the within-length Pearson correlation between color and token accuracy at length .
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 at training length. Single-stage curves are dotted gray. The value is the within-length Pearson correlation between color and token accuracy at length .
Figure 14: Small-gap adversarial examples across training lengths. Accuracy on runs of consecutive values
Figure 15: Large-gap adversarial examples across training lengths. We report accuracy on the prediction that crosses a gap of size . Panel titles give the training length 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 to steers of the 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 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 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 . Left: The score is nearly invariant to the key-side first attention target , showing uniform horizontal bands. Right: The score is strongly modulated by the query-side target : as increases, a high-score band emerges in .
Figure 21: Mean 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 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 equally attended tokens ( , ). The band structure is preserved across all cases.
Figure 25: Second attention QK score slices with equally attended tokens for .
Figure 26: Second attention QK score slices with equally attended tokens for .
Figure 27: Second attention QK score slices with equally attended tokens for .
Figure 28: Argmax bias analysis ( ). Left: Argmax of the second attention vs for at . Middle: Same for at varying . Right: Scatter of vs across all configurations; diamonds mark per- means. This shows that the argmax mostly varies with the minimum of the first attention concentration set, , and has almost no dependence on , or . In particular, the argmax consistently tracks , not .
Figure 29: Residual stream component removal across 17 two-stage checkpoints. Left: Per-checkpoint argmax agreement rates for Claim 1 (remove , teal) and Claim 2 (remove 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 ( 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 , for gaps (left to right). Grey dots show per- 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 is sufficiently large; for large gaps the second attention is indispensable.
More in this topic
Length Generalization
Length generalization studies whether sequence models trained on short inputs can reliably solve the same tasks on longer sequences.
See all 47 papers