Structure Development in List-Sorting Transformers

TL;DR

This paper investigates how a minimalist, single-layer attention-only transformer learns to sort lists and how its internal attention heads organize into interpretable structures. By decomposing attention into Output-Value (OV) and Query-Key (QK) circuits and using the Local Learning Coefficient (LLC) as a complexity proxy, the authors identify two principal end-states—vocabulary-splitting and copy-suppression—and show that the training data distribution over gaps between adjacent sorted numbers ($ar{ ho}$) dictates which specialization emerges. A developmental view reveals three recurring stages and shows that vocabulary-splitting corresponds to a simpler solution than overlapping ranges, surviving even without weight decay, while copy-suppression calibrates the copying head and can alter model confidence. The findings demonstrate that training data characteristics can shape internal transformer organization, offering a concrete, testable link between data distributions and emergent mechanistic structures with implications for interpretability of larger LLMs.

Abstract

We study how a one-layer attention-only transformer develops relevant structures while learning to sort lists of numbers. At the end of training, the model organizes its attention heads in two main modes that we refer to as vocabulary-splitting and copy-suppression. Both represent simpler modes than having multiple heads handle overlapping ranges of numbers. Interestingly, vocabulary-splitting is present regardless of whether we use weight decay, a common regularization technique thought to drive simplification, supporting the thesis that neural networks naturally prefer simpler solutions. We relate copy-suppression to a mechanism in GPT-2 and investigate its functional role in our model. Guided by insights from a developmental analysis of the model, we identify features in the training data that drive the model's final acquired solution. This provides a concrete example of how the training data shape the internal organization of transformers, paving the way for future studies that could help us better understand how LLMs develop their internal structures.

Figures (19)

• Figure 1: Illustration of the transformer architecture and an idealized version of the sorting circuits based on a similar figure in elhage2021mathematical. The Output-Value (OV) circuit copies forward the tokens in the context, whereas the Query-Key (QK) circuit attends more to smaller numbers larger than the current token.
• Figure 2: Baseline 2-head model trained on $D_{\overline{\delta}\approx4.7}$ undergoes three stages characterized by: rapid learning (left), heads copying partly overlapping vocabularies, as can be seen from the diagonal OV circuits (middle), and vocabulary-splitting head specialization with diagonal OV circuits covering contiguous regions (right). The loss on $D_{\overline{\delta}\approx 2.2}^d$ measures out-of-distribution loss on lists with closer elements.
• Figure 3: 2-head model trained on $D_{\overline{\delta}\approx2.2}^d$ (the baseline model is trained on $D_{\overline{\delta}\approx4.7}$). Initially, it evolves similar to the baseline model(Fig. \ref{['fig: baseline2']}), but develops copy-suppression at the end of training. This is a gradual process, where the vocabulary covered by head 1 decreases before switching to copy-suppression.
• Figure 4: Distribution of $\delta$ (separation between neighbouring list elements) in $D_{\overline{\delta}\approx 2.2}^d$ and $D_{\overline{\delta}\approx 4.7}$.
• Figure 5: Illustration of the active regions in the QK circuits at the end of training for the baseline 2-head model trained on $D_{\overline{\delta}\approx 4.7}$. The attention pattern in regions close to the diagonal diminishes from left to right for each row, establishing a gradient. We hypothesize: the stronger the mean normalized QK gradients $\hat{\nabla}_{\rm QK}$ of a model, the better the sorting of neighboring list elements with small separation.