A Mathematical Framework for Transformer Circuits

interpretabilitymechanistic interpretabilitytransformersattention headsinduction headsresidual stream

Elhage, Nanda, Olsson et al. rewrite small attention-only transformers as sums of end-to-end paths, show that zero-layer models learn bigram statistics and one-layer models learn bigrams plus "skip-trigrams" readable from the weights, and find that two-layer models compose attention heads into induction heads, a general in-context copying mechanism.

Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan et al. (Anthropic) · 2021-12-22 · original

Can a transformer language model be read the way a programmer reads compiled code, by recovering the algorithm its weights implement? Earlier work of this kind, the Distill Circuits thread, had reverse engineered parts of an image model, InceptionV1, but there was no comparable project for language models. This paper, published on the Transformer Circuits Thread on 22 December 2021, is the earliest interpretability paper in this collection. It does not attempt a large model. It studies transformers with at most two layers and no MLP (feedforward) layers, and asks whether such toy models can be understood completely from their weights. Its main conceptual contribution is a way of rewriting a transformer, mathematically equivalent to the usual description, in which the model becomes a sum of separate paths from input tokens to output predictions. Its main empirical finding is the induction head, a pair of attention heads that lets a two-layer model continue patterns it has seen earlier in the same text. The follow-up paper, In-context Learning and Induction Heads, argues that this mechanism accounts for much of in-context learning in large models.

Terms used in this note

  • Token, logit. A language model reads text as a sequence of tokens (words or word pieces) and, at each position, outputs one score per vocabulary entry, a logit, for what the next token will be. A softmax turns the logits into probabilities.
  • Embedding and unembedding. The embedding matrix $W_E$ maps each token to a vector; the unembedding matrix $W_U$ maps a vector back to logits.
  • Residual stream. In a transformer, each layer reads from and adds its output to a running vector at every position. The paper calls this running sum the residual stream and treats it as a communication channel between layers, since it does no processing of its own (section "Virtual Weights and the Residual Stream as a Communication Channel"). See Assembling the Transformer Block for how residual connections are built.
  • Attention head. A unit that, at each position, chooses which earlier positions to look at (the attention pattern) and copies information from them. The usual description computes a query, key and value vector for each token; Building Self-Attention Step by Step builds one step by step.
  • Attention-only model. A transformer with the MLP layers removed, leaving only embedding, attention layers and unembedding.
  • Mechanistic interpretability. The paper's term for trying to reverse engineer the detailed computation a network performs, by analogy with turning a binary back into source code (Introduction).

What the paper does

The models

The models are zero-, one- and two-layer decoder-only transformers with attention but no MLP layers, typically with 12 heads of dimension 64 per layer, and one example with 32 heads of dimension 128; the context is 2,048 tokens (appendix "Model Details"). To keep the algebra clean the analysis ignores biases and layer normalization, which the authors argue can be folded into adjacent weights up to a scaling (section "Model Simplifications" and appendix "Handling Layer Normalization"). For the training data the appendix points to the dataset description in Anthropic's A General Language Assistant as a Laboratory for Alignment. The authors call dropping MLP layers "a very dramatic simplification" and say their lack of progress on MLP layers is "a major weakness of our work" (section "Model Simplifications").

Rewriting attention

The paper's first move is to rewrite standard operations in forms that are equivalent but easier to read (section "Transformer Overview"):

  • Heads are independent and additive. The usual "concatenate and multiply" formulation, which joins the heads' results and applies one output matrix, is mathematically the same as running each head separately and adding its output to the residual stream.
  • Two circuits per head. A head does two separable things. Which positions it attends to depends only on the product $W_{QK} = W_Q^T W_K$ of its query and key weights; what it copies, once it attends, depends only on the product $W_{OV} = W_O W_V$ of its output and value weights. Queries, keys and values are, in the paper's words, intermediate by-products of these two low-rank matrices. The paper calls the token-to-token versions of these matrices the QK circuit (how much a token at the current position wants to attend to each earlier token) and the OV circuit (how an attended token changes the output logits).
  • Freezing the attention pattern. If the attention pattern is held fixed, every head is a linear map, so the whole attention-only model becomes linear in its input tokens. The authors suggest imagining running the model twice, the first time to record attention patterns and the second time with those patterns frozen (section "OV and QK Independence").
  • Path expansion. Because the residual stream is a sum, the model's output can be multiplied out into a sum of terms, one per path from the embedding to the unembedding: a direct path that skips every layer, a path through each single head, and, in deeper models, paths through chains of heads.
  • Virtual weights. Multiplying the output weights of one layer by the input weights of a later layer gives "virtual weights" describing how much the later layer reads what the earlier one wrote, even across many intervening layers.

Composition between heads

In a two-layer model, a second-layer head can read what a first-layer head wrote in three ways (section "Three Kinds of Composition"): through its query (Q-composition), its key (K-composition) or its value (V-composition). Q- and K-composition change where the later head attends; V-composition chains what the two heads move, producing what the paper calls a virtual attention head. The paper measures composition from the weights, with a normalized product of the matrices involved, and uses an ablation procedure to estimate how much each order of path contributes to the loss (section "Term Importance Analysis").

What it finds

Zero layers: bigram statistics

A model with no layers maps each token directly to a prediction for the next token through $W_U W_E$. The best it can do is approximate bigram statistics, the log-probability of each token given only the one before it (section "Zero-Layer Transformers"; compare the hand-built The Bigram Language Model). The same direct path appears inside every deeper transformer, where the paper finds it tends to hold bigram facts not covered by more general rules, giving "Barack" followed by "Obama" as an example.

One layer: bigrams plus skip-trigrams

A one-layer attention-only model is, the paper shows, a sum of the bigram term and one term per head, and each head term can be read as a table of skip-trigrams: patterns of the form "A … B C", where the current token B attends back to an earlier token A and raises the probability that C comes next (section "One-Layer Attention-Only Transformers"). Both the QK and OV circuits can be expanded into vocabulary-by-vocabulary matrices and read directly. Examples the paper reports from the largest entries include:

  • Copying. Most heads devote much of their capacity to copying: an earlier token raises its own probability, and the QK circuit attends only to earlier tokens that could plausibly come next. The authors describe this as a very primitive form of in-context learning (section "Copying / Primitive In-Context Learning").
  • Code and markup. Python patterns such as "def … ( → self", LaTeX "\left … \ → right" and HTML table tags.
  • Bugs. Because each head stores its skip-trigrams in factored form, a head that raises "keep … in mind" and "keep … at bay" also raises "keep … in bay" and "keep … at mind". The authors present this as an early example of interpretability explaining a model's errors (section "Skip-Trigram 'Bugs'").

The paper also proposes a summary statistic for copying: a head whose OV circuit has mostly positive eigenvalues tends to make tokens raise their own probability. By this measure, 10 of the 12 heads in the one-layer model studied are "significantly copying", which the authors say agrees with inspection of the weights; they caution that positive eigenvalues are strong but not conclusive evidence of copying (section "Detecting Copying Behavior").

On whether one-layer models are now fully understood, the authors' answer is qualified. With a vocabulary of about 50,000 tokens, one expanded OV matrix has about 2.5 billion entries, and they describe the result as having "revealed the one-layer attention-only model to be a compressed Chinese room" (section "Interpretation as Skip-Trigrams"): no algorithmic mystery remains, but there is far too much to hold in mind without better ways of summarizing it (section "Do We 'Fully Understand' One-Layer Models?").

Two layers: induction heads

The two-layer model studied uses composition almost entirely for one purpose. One first-layer head attends mainly to the previous token. Several second-layer heads use K-composition with it, so that their keys encode the token before each earlier position. The result is an induction head: at the current token A, it searches the context for earlier occurrences of A, attends to the token that followed (B), and raises the probability of B (sections "Analyzing a Two-Layer Model" and "Induction Heads"). In the paper's notation the one-layer copying pattern is "[b] … [a] → [b]", while the induction pattern is "[a][b] … [a] → [b]". Because it relies on the structure of the context rather than learned statistics about which tokens follow which, it works even on sequences of random tokens repeated in the context, which the paper uses as its test.

The paper checks this account against the weights: the induction heads have both OV circuits with positive eigenvalues (copying) and a positive-eigenvalue K-composition term with the previous-token head (matching), and they sit together in an extreme corner of a plot of the two measures (section "Checking the Mechanistic Theory"). The ablation analysis finds that paths through two heads in sequence (the virtual attention heads of V-composition) contribute little to this model's loss, so the direct path and individual heads, mainly in the second layer, carry the behavior; the authors stress that this says nothing against Q- and K-composition, which act inside the attention computation (section "Term Importance Analysis").

The paper summarizes the difference between depths as a change in kind: with composition, two-layer models become "much more like a computer program running an algorithm" than the look-up tables of a one-layer model (section "Two-Layer Attention-Only Transformers").

A correction published later

After publication the authors found a bug in their low-rank matrix library that affected the composition diagram for the two-layer model (section "Correction: Attention Head Composition Diagram"). The faulty computation happened to emphasize exactly the composition an induction head needs. The corrected figure keeps the K-composition with the previous-token head but shows more composition than first reported, including two induction heads that also use a head attending to the last several tokens. The authors say this changes their interpretation of the two-layer section "in some ways" and fits their remark that real induction heads often use more than the minimal circuit.

Limitations and open questions

  • No MLP layers. The analysis covers attention-only models. The authors note that MLP layers hold about two thirds of a standard transformer's parameters, that many attention heads interact with them, and that progress there needs interpretable neurons, which they have had limited success finding (section "Where Does This Leave Us?"). In the appendix they report that MLP neurons have been much harder to form hypotheses about, apart from neurons about 5% of the way through the model.
  • Toy models. The paper deliberately does not apply its results to large models and calls its steps "initial, very preliminary" (Introduction). It anticipates that some attention-only circuits, induction heads in particular, will carry over.
  • Interpretation still needs summarizing. Reading billions of matrix entries is not understanding in a useful sense, and the paper says the right decomposition or dimensionality reduction for OV and QK matrices is still open (section "Summarizing OV/QK Matrices").
  • Not all the mathematics is new. The related-work section says most of the observations are not novel, citing for example Dong et al.'s path analysis of self-attention, and places the contribution in using them for interpretability.
  • Residual stream bandwidth and superposition. The paper argues that the residual stream has far fewer dimensions than the neurons that read and write it, so it must carry information in superposition, and it expects such "bottleneck activations" to be hard to interpret. In a follow-up comment dated February 2023, the authors list superposition as the main explanation they had suspected for the difficulty of MLP layers and point to Toy Models of Superposition and other work since (section "Summary of Follow-Up Research").

Research questions

  1. How much loss do skip-trigram bugs cost? The authors ask this themselves. Comparing a one-layer model's loss with the loss of a model whose factored skip-trigrams are replaced by unfactored ones on the affected token triples would put a number on it.
  2. Do the composition types matter differently at scale? Virtual attention heads contributed little in the two-layer model, but the number of possible virtual heads grows quadratically with depth. Repeating the term-importance ablation across models of increasing depth would show whether V-composition becomes important.
  3. Is the eigenvalue statistic a reliable test for copying? The paper notes that a matrix with all positive eigenvalues can still lower some tokens' own logits. Comparing the eigenvalue measure with direct measures (the diagonal, or how often a token most increases itself) across many heads would show where they disagree.
  4. How far does the path-expansion view extend once MLP layers are included? The appendix sketches the expansion for one MLP layer. Testing whether the same reading works for the final MLP layer of small standard transformers, where the effect on logits is linear, would show how much of the method survives.
  5. Which induction mechanism do models with positional embeddings use? The appendix reports one GPT-2 induction head built by "pointer arithmetic" on positional embeddings rather than K-composition. A survey of induction heads in models with different position schemes would show how often each mechanism appears.
  6. How much does the corrected composition picture change the account of the two-layer model? Reanalyzing the extra composition the correction revealed, including the heads attending to several recent tokens, would show whether the minimal induction circuit is the main mechanism or one of several.

How it connects

References

  • Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., DasSarma, N., Drain, D., Ganguli, D., Hatfield-Dodds, Z., Hernandez, D., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. (2021). A Mathematical Framework for Transformer Circuits. Transformer Circuits Thread, 22 December 2021. https://transformer-circuits.pub/2021/framework/index.html. Section names in this note refer to the web version, including its later correction and follow-up comments.
  • Anthropic. (2021, December 22). A Mathematical Framework for Transformer Circuits [research page]. https://www.anthropic.com/research/a-mathematical-framework-for-transformer-circuits
  • Companion exercises: https://transformer-circuits.pub/2021/exercises/index.html