Lost in Latent Space

Memo: On the Specialization of Attention Heads

Background

Mechanistic interpretability has produced a catalog of named attention heads: induction heads, previous-token heads, name-mover heads, duplicate-token heads, and so on. The catalog implies a tidy picture in which each head learns one job. The empirical picture is messier: most heads in a trained transformer do not admit a single clean description, and the well-named heads are a minority of a model’s compute. This memo collects what is and isn’t known about why specialization happens at all.

What is observed

• Specialized heads exist and are stable across seeds. Induction heads emerge reliably in 2-layer attention-only models trained on natural language. Replicating across seeds, they appear in the same layer positions with similar circuit structure (Olsson et al., 2022). Stability across seeds is the strongest evidence that the architecture biases certain solutions over others.
• Most heads are polysemantic. When mech interp work scales beyond the toy circuits, the median head looks like a smear of many partial functions, none individually clean. Quantitatively, sparse-coding decompositions of attention outputs routinely find more “features” than there are heads.
• Head pruning is highly uneven. Many heads can be pruned with little performance loss; a small number cannot be pruned without significant degradation. The distribution of “head importance” is heavy-tailed.

What plausibly causes specialization

We believe — with varying confidence — that the following forces shape the head landscape:

1. Architectural bottlenecks. A single attention head has a low-rank structure (rank ≤ head_dim). A function that needs more rank than the head provides cannot live in one head. This pushes the model toward either splitting the function across heads or compressing many partial functions into one. Evidence: induction heads need two heads in series (previous-token + induction), because no rank-d head can do both jobs alone.
2. Gradient interference. When two functions both want a head, SGD tends to either pick one and saturate it (specialization) or smear both into a polysemantic mix. Which outcome obtains seems to depend on how decorrelated the two functions’ input distributions are. We do not have a good predictive theory here.
3. Width vs depth pressure. Adding heads (width) increases capacity for parallel specialization. Adding layers (depth) increases capacity for serial composition. Empirically, width has diminishing returns past a point — likely because there is a finite number of “natural primitive operations” that benefit from being on their own head.

What is not known

• Why specific heads in specific layer positions appear at specific training-step counts. The phase-transition behavior of induction-head formation is observed but not derived.
• Whether polysemantic heads are doing genuine superposition (multiple features encoded in a non-orthogonal basis) or are just noisy averages of partial circuits. The sparse-autoencoder literature is closing on this question but has not resolved it.
• The functional form of the tradeoff between number of heads and per-head dimensionality. Empirically, very narrow heads (head_dim < 32) underperform; very wide heads waste capacity. The bottom of the curve is poorly characterized.

Tradeoffs for downstream work

• Treating named heads as load-bearing for safety-relevant behavior is premature. The named heads we have are a biased sample of the easy-to-name ones.
• Studying families of heads (all heads that attend to the previous token with > some threshold) is likely more robust than studying individual heads.

Open questions

• Is there a width regime where polysemanticity vanishes?
• Do MoE-style architectures change the picture by replacing one polysemantic head with many specialized experts gated by routing?
• How much of “specialization” is just an artifact of human pattern recognition on whatever heads happen to be easiest to label?

What’s next

The strongest direction is the sparse-autoencoder line: decompose head outputs into a learned dictionary of features and ask whether the dictionary is sparser, more stable, and more semantically interpretable than the head basis itself. If yes, “heads” stop being the right unit of analysis. That conclusion, if it holds, will quietly invalidate a lot of the current mech-interp vocabulary — which is what makes it the most interesting bet on the table.