A Mathematical Framework for Transformer Circuits (2021)

Additional Resources

This article has a companion problem set of problems we found helpful in building our intuition for attention heads.

We've released one of the libraries we made for our research, PySvelte (which makes it convenient to use interactive visualizations from Python), and written an article describing another one called Garcon (our tooling for probing the internals of models, especially large ones).

Acknowledgments

In writing this paper, our thinking was greatly clarified and encouraged by correspondence with Martin Wattenberg, Vladimir Mikulik, Jeff Wu, Evan Hubinger, and Peter Hase; their generous and detailed feedback exceeded anything one might reasonably hope for.

We're also deeply grateful to Daniela Amodei, Jia Yuan Loke, Liane Lovitt, Timothy Telleen-Lawton, Kate Rudolph, Matt O’Brien, Jeffrey Ladish, Samuel Bowman, Geoffrey Irving, Tom McGrath, Michela Paganini, Paul Christiano, Allan Dafoe, Gabriel Goh, Nick Cammarata, Chelsea Voss, Katherine Lee, Beth Barnes, Jan Leike, Tristan Hume, Nate Thomas, Buck Shlegeris, Alex Tamkin, Quinn Tucker, and Rob Harries, for their support, for comments on this work, and for conversations that contributed to the background thinking on interpretability and safety this work is based on.

Neel Nanda would like to extend particular thanks to Jemima Jones for providing significant motivation, accountability and support as he navigated a difficult period while contributing to this paper.

We're grateful to several readers for reporting errata. Ken Kahn found several typos in the paper. Tuomas Oikarinen caught several typos in the accompanying exercises.

Matthew Rahtz reported several broken links.

Author Contributions

Theoretical Framework: The framework described in this paper developed as part of ongoing conversations and empirical investigation between Nelson Elhage, Catherine Olsson, Neel Nanda, and Chris Olah over the last year as they've tried to understand language models. It's impossible to disentangle all the contributions that contributed to its development, but some examples of the kind of contributions involved follow. Nelson Elhage discovered evidence of neurons performing memory management in the residual stream, informing the "residual stream as a communication channel" framing. Catherine Olsson carefully characterized dozens of neurons and attention heads (with other authors also helping out), greatly informing our overall thinking on how to think about different parts of the model. Neel Nanda greatly improved our thinking on how attention heads compose. Chris Olah developed the foundational ideas of the framework, discovered induction heads, and guided theoretical progress. These examples are only a small slice of the contributions each mentioned author made.

In addition to these core research contributors, other members of Anthropic contributed important insights to the framework. Jared Kaplan, Sam McCandlish, Andy Jones, and Dario Amodei contributed especially important insights.

Analysis of Toy Models: The analysis of toy models presented in this paper was done by Chris Olah and Neel Nanda. However, they were heavily based on previous experiments, including many by Nelson Elhage and Catherine Olsson. The special attention-only toy models used for this article were created by Nelson Elhage.

Writing: This article was drafted by Chris Olah. Neel Nanda made major pedagogical improvements to the draft. Dario Amodei contributed heavily to the high-level framing. Nelson Elhage did careful editing to improve clarity. Many members of Anthropic, including Nicholas Joseph, Zac Hatfield-Dodds and Jared Kaplan provided valuable feedback throughout the writing process. Others, including Deep Ganguli, Jack Clark, Ben Mann, Danny Hernandez, Liane Lovitt, and Tom Conerly significantly improved exposition by testing and giving feedback on approaches to explaining these ideas.

Infrastructure for Analyzing Models: The crucial piece of software infrastructure which made our analysis of models possible is a library called Garcon (described in another article) created by Nelson Elhage with help from Tom Brown, Sam McCandlish, and Chris Olah. Garcon makes it easy to access model activations and parameters, regardless of scale. Nelson and Catherine Olsson also built an internal website, based on top of Garcon, which allows users to see dataset examples as well as interactively edit text and see model activations and attention patterns. Although not presented here, this website was very important to the evolution of our thinking.

Visualizations: Having interactive visualizations to explore model activations was important in enabling us to explore large amounts of data. Chris Olah created PySvelte (see github) and many of our visualizations, with contributions from Nelson Elhage and Catherine Olsson. Conceptual illustrations of ideas in this article are by Chris Olah. Ben Mann set up infrastructure to allow convenient secure sharing of visualizations.

Model Training: All our interpretability research is enabled by having access to models to study, including large models. (Although the models used as examples in this particular paper are small, our framework was informed by experiments on large models.) Led by Tom Brown, Sam McCandlish, and Jared Kaplan, the majority of Anthropic's technical staff contributed to the development of our efficient distributed training infrastructure and the underlying machine learning. Core contributors include Nicholas Joseph, Tom Henighan, and Ben Mann. Nelson Elhage, Kamal Ndousse, Andy Jones, Zac Hatfield-Dodds, and Danny Hernandez also contributed to this infrastructure.

Cluster: Tom Henighan and Nova DasSarma, advised by Tom Brown and Sam McCandlish with contributions from many others at Anthropic, managed the research cluster our research depended on and maintained its stability. Nova provided important support when interpretability required unorthodox jobs to run on our cluster.

Other Contributions: 

All staff at Anthropic provided valuable feedback throughout. Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish all provided essential high-level feedback.

Chris Olah led the project.

Citation Information

Please cite as:

Elhage, et al., "A Mathematical Framework for Transformer Circuits", Transformer Circuits Thread, 2021.

BibTeX Citation:

@article{elhage2021mathematical,
   title={A Mathematical Framework for Transformer Circuits},
   author={Elhage, Nelson and Nanda, Neel and Olsson, Catherine and Henighan, Tom and Joseph, Nicholas and Mann, Ben and Askell, Amanda and Bai, Yuntao and Chen, Anna and Conerly, Tom and DasSarma, Nova and Drain, Dawn and Ganguli, Deep and Hatfield-Dodds, Zac and Hernandez, Danny and Jones, Andy and Kernion, Jackson and Lovitt, Liane and Ndousse, Kamal and Amodei, Dario and Brown, Tom and Clark, Jack and Kaplan, Jared and McCandlish, Sam and Olah, Chris},
   year={2021},
   journal={Transformer Circuits Thread},
   note={https://transformer-circuits.pub/2021/framework/index.html}
}

Additional Intuition and Observations

MLP Layers

This article has focused on attention-only transformers, without MLP layers. How can we extend this approach to understanding MLP models with MLP layers?

Recall that an MLP layer m computes its activations a^m from the residual stream x_i by performing a matrix multiply and applying its activation function. The activation vector is then projected back down and added into the residual stream:

The presence of the GeLU activation means we can’t linearize through MLP layers as we did attention layers. Instead, we likely need to take the approach of the Circuits project in reverse engineering vision models: understanding what the neurons represent, how they’re computed, and how they’re used.

In theory, there’s a lot of reason to be optimistic about understanding these neurons. They have an activation function which should encourage features to align with the basis dimensions. They’re four times larger than the residual stream, and information doesn’t need to flow through them, which are both factors one might expect to reduce polysemanticity. Unfortunately, things are much more challenging in practice. We’ve found the neurons much harder to develop hypotheses for, with the exception of neurons at ~5% depth through the model, which often respond to clusters of short phrases with similar meanings. We’ve focused this paper on understanding attention heads because we got a lot more early traction on attention heads.

Despite this, it’s worth noting that there’s a fairly clean story for how to mechanistically reason about neurons, which we can use when we find interpretable neurons.

Path Expansion of MLP Layer in One-Layer Model: For simplicity, let’s consider a standard one layer transformer (ignoring layer norm and biases). The pre-activation values are a linear function of the residual stream. By applying our equation for attention heads, we get essentially the same equation we got for the one-layer transformer logits:

a^m_\text{pre} ~=~ W^m_I \cdot \bigg(Id+\sum_{h\in H_1} A^h \otimes W_{OV}^h\bigg)\cdot~W_E

a^m_\text{pre} ~=~ W^m_I W_E ~+~ \sum_{h\in H_1} A^h \otimes \left(W^m_I W_{OV}^h W_E\right)

This means we can study it with the same methods. W^m_I W_E tells us how much different tokens encourage or inhibit this neuron through the residual stream, while W^m_I W_{OV}^h W_E does the same for tokens if they’re attended to by a given attention head. One might think of this as being similar to a neuron in a convolutional neural network, except where neuron weights in a conv net are indexed by relative position, these neuron weights are indexed by attention heads. The following section, Virtual Weights and Convolutional-Like Structure, will explore this connection in more detail.

What about the downstream effect of these neurons? W_U W^m_O tells us how the activation of each neuron affects the logits. It’s a simple linear function! (This is always true for the last MLP layer of a network, and gives one a lot of extra traction in understanding that layer.)

This approach becomes more difficult as a network gets deeper. One begins to need to reason about how one MLP layer affects another downstream of it. But the general circuits approach seems tractable, were it not for neurons being particularly tricky to understand, and there being enormous numbers of neurons.

Virtual Weights and Convolution-like Structure

When we apply path expansion to various terms in a transformer, we generally get virtual weights of the form:

y = (\text{Id} \otimes W_\text{Id} + \sum_h A^h \otimes W_h) x + …

(In the general case, h might include virtual attention heads.)

The main example we've seen in this paper is having y be logits for output tokens and x the one-hot encoded input tokens. But this can also arise in lots of other cases. For example, y might be MLP pre-activations and x might be the activations of a previous MLP layer.

Multiplying by (\text{Id} \otimes W_\text{Id} + \sum_h A^h \otimes W_h) can be seen as a generalization of a convolution, with the weights [W_\text{Id}, ~ W_{h_0}, ~ W_{h_1}...~] and attention heads taking the place of relative position.

Exact Equivalence to Standard Convolution in Limited Cases: We claim that all convolutions can be expressed as one of the tensor products above, and thus (1) attention is a generalization of convolution; (2) specific configurations of attention correspond exactly to convolution.

Consider the convolution W \ast x. We claim that it can be rewritten as:

W \ast x ~=~ \sum_v A^v \otimes W_v

where v is the convolution offset, A^v is an attention pattern always attending to the relative position v (such as a previous token head), and W_v is the weight entry corresponding to that offset.

This equivalence is explored in depth by Cordonnier et al. , who also empirically find that vision models often have many 2D relative position heads, similar to the previous token head we've observed.

Analogy to Convolution Generalizations: The correspondence to standard convolution breaks down for "dynamic" attention heads which attend based on some pattern other than fixed relative positions. However, it seems like there are many cases where they still are spiritually very convolution-like. For example, a "previous verb" head isn't a literal relative position (at least if we parameterize our input based on token index), but it seems morally like a type of convolution.

One way to resolve the above might be to consider graph convolutions. If we treat each attention pattern as a different set of weights on a graph, we could treat the tensor product as a sum of graph convolutions. (For an informal discussion of analogies between transformers and graph neural networks, readers may wish to look at an article by Joshi .) However, needing to consider multiple different weight graphs seems inelegant. There are other exotic kinds of convolutions based on abstract algebra (see tutorial ) and perhaps the right version could be found to find an exact correspondence.

But the fundamental thing is that when you see tensor products like the ones above, they can be thought of as something like a convolution with dynamic, soft, possibly semantic positions.

Why This is Useful: In reverse engineering convolutional neural networks, the Distill Circuits Thread benefited greatly from the fact that transformer weights are organized into convolutions. For example, one can look at the weights between two neurons and see how different spatial positions affect things.

In the above example, we see the weights between two curve detectors; along the tangent of the curve, the other curve excites it. This would be much harder to understand if one had to look at 25 different weights; we benefit from organization.

We expect that similar organization may be helpful in understanding transformer weights, especially as we begin considering MLP layers.

Activation Properties

We often find it helpful to think differently about various activations in transformers based on whether they have the following properties:

Privileged Basis vs Basis Free: A privileged basis occurs when some aspect of a model’s architecture encourages neural network features to align with basis dimensions, for example because of a sparse activation function such as ReLU. In a transformer, the only vectors with privileged bases are tokens, attention patterns and MLP activations. See the appendix table of variables for a full list of which activations have or don’t have privileged activations.

Some types of interpretability only make sense for activations with a privileged basis. For example, it doesn't really make sense to look at the "neurons" (ie. basis dimensions) of activations like the residual stream, keys, queries or values, which don't have a privileged basis. This isn't to say that there aren't ways to study them; a lot of interesting work is done on word embeddings without assuming a privileged basis (e.g. ). But having a privileged basis does open up helpful approaches.

Bottleneck Activations: We say that an activation is a bottleneck activation if it is a lower-dimensional intermediate between two higher dimensional activations. For example, the residual stream is a bottleneck activation because it is the only way to pass information between MLP activations, which are typically four times larger than it. (Additionally, in addition to it being a bottleneck for the communication between adjacent MLP layers, it’s also the only pathway by which arbitrary early MLP layers can communicate with arbitrary late MLP layers, so the stream may simultaneously be conveying different pieces of information between many different pairs of MLP layers, a much more extreme bottleneck than just 4x!) Similarly, a value vector is a bottleneck activation because it’s much lower dimensional than the residual stream, and it’s the only way to move information from the residual stream for one token position in the context to another (without the model dedicating an additional attention head to the same pattern, which it sometimes will). See above table for a list of which activations are or are not bottleneck activations.

Pointer Arithmetic with Positional Embeddings

Our models use a slightly unusual positional mechanism (similar to ) which doesn't put positional information into the residual stream. The popular rotary attention approach to position has this same property. But it's worth noting that when models do have positional embeddings (as in ), several possibilities open up.

At their most basic level, positional embeddings are kind of like token addresses. Attention heads can use them to prefer to attend to tokens at certain relative positions, but they can also do much more:

• Transformers can do "pointer arithmetic" type operations on positional embeddings. We observed at least one case in GPT-2 where an induction head was implemented using this approach, instead of the one described above. First, an attention head attended to a similar token, returning its positional embedding. Next, a query vector for another attention head was constructed with q-composition, rotating the positional embedding forward one token. This results in an induction head.
• We speculate that having an identifier for each token may be useful in some cases. For example, suppose you want all tokens in the same sentence to share an identifier. In a positional model, an attention head can attend to the previous period and copy its positional embedding into some subspace.

Identity Attention Heads?

It's worth noting that one could simplify all the math in this paper by introducing an "identity attention head," h_\text{Id}, such that A^{h_\text{Id}} = \text{Id} and W_{OV}^{h_\text{Id}} = \text{Id}. This identity head would correspond to the residual stream, removing extra terms in the equations. In this paper, we chose not to do that and instead keep the residual stream as an explicit, differentiated term. The main reason for this is that it really does have different properties than attention heads. But in other cases, it can be useful to think about things the other way.

Notation

Variable Definitions

Main Model Activations and Parameters
Variable	Shape / Type	Description
T(t)	[n_\text{context},~ n_\text{vocab}]	Transformer logits, for tokens t [activation, privileged basis]
t	[n_\text{context},~ n_\text{vocab}]	One-hot encoded tokens [activation, privileged basis]
x^n	[n_\text{context},~ d_\text{model}]	“Residual stream” or “embedding” vectors of model at layer n (one vector per context token) [activation, not privileged basis]
W_E	[d_\text{model}, n_\text{vocab}]	Token embedding [parameter]
W_P	[d_\text{model}, n_\text{context}]	Positional embedding [parameter]
W_U	[n_\text{vocab}, d_\text{model}]	Unembedding / softmax weights  [parameter]
Attention Heads Activations and Parameters
Variable	Shape / Type	Description
H_n	Set	Set of attention heads at layer n
h(x)	[n_\text{context},~ d_\text{model}]	Output of attention head h [activation, not privileged basis]
A^h	[n_\text{context},~ n_\text{context}]	Attention pattern of attention head h [activation, privileged basis]
q^h, k^h, v^h, r^h	[n_\text{context},~ d_\text{head}]	Query, key, value and result vectors of attention head h (one vector per context token) [activation, not privileged basis]
W^h_Q, W^h_K, W^h_V	[d_\text{head},~ d_\text{model}]	Query, key, and value weights of attention head h [parameter]
W^h_O	[d_\text{model},~ d_\text{head}]	Output weights of attention head h [parameter]
W^h_{OV}	[d_\text{model},~ d_\text{model}]	W^h_{OV} = W^h_{O}W^h_{V} [parameter, low-rank]
W^h_{QK}	[d_\text{model},~ d_\text{model}]	W^h_{QK} = W^h_{Q^T}W^h_{K} [parameter, low-rank]
MLP Layer Activations and Parameters
Variable	Shape / Type	Description
m(x)	[n_\text{context},~ d_\text{model}]	Output of MLP layer m [activation, not privileged basis]
a^m	[n_\text{context},~ d_\text{mlp}]	Activations of MLP layer m [activation, privileged basis]
W^m_I	[d_\text{mlp},~ d_\text{model}]	Input weights for MLP layer m [parameter]
W^m_O	[d_\text{model},~ d_\text{mlp}]	Output weights for MLP layer m [parameter]
Functions
Variable	Shape / Type	Description
\text{GeLU}()	Function	Gaussian Error Linear Units
\text{softmax}^*()	Function	Softmax with autoregressive mask for attention distribution creation.

Tensor Product / Kronecker Product Notation

In machine learning, we often deal with matrices and tensors where we want to multiply "one side". In transformers specifically, our activations are often 2D arrays representing vectors at different context indices, and we often want to multiply "per position" or "across positions". Frequently, we want to do both! Tensor (or equivalently, Kronecker) products are a really clean way to denote this. We denote them with the \otimes symbol.

Very pragmatically:

• A product like \text{Id} \otimes W (with identity on the left) represents multiplying each position in our context by a matrix.
• A product like A \otimes \text{Id} (with identity on the right) represents multiplying across positions.
• A product like A \otimes W multiplies the vector at each position by W and across positions with A. It doesn't matter which order you do this in.
• The products obey the mixed-product property (A \otimes B) \cdot (C \otimes D) = (AC) \otimes (BD).

There are several completely equivalent ways to interpret these products. If the symbol is unfamiliar, you can pick whichever you feel most comfortable with:

• Left-right multiplying: Multiplying x by a tensor product A\otimes W is equivalent to simultaneously left and right multiplying: (A\otimes W) x = A x W^T. When we add them, it is equivalent to adding the results of this multiplication: (A_1\otimes W_1 + A_2\otimes W_2) x = A_1 x W_1^T + A_2 x W_2^T.
• Kronecker product: The operations we want to perform are linear transformations on a flattened ("vectorized") version of the activation matrix x. But flattening gives us a huge vector, and we need to map our matrices to a much larger block matrix that performs operations like "multiply the elements which previously corresponded to a vector by this matrix". The correct operation to do this is the Kronecker product. So we can interpret \otimes as a Kronecker product acting on the vectorization of x, and everything works out equivalently.
• Tensor Product:  A\otimes W can be interpreted as a tensor product turning the matrices A and W into a 4D tensor. In NumPy notation, it is equivalent to A[:,:,None,None] * W[None, None, :, :] (although one wouldn’t computationally represent them in that form). More formally, A and W are “type (1,1)” tensors (matrices mapping vectors to vectors), and A\otimes W is a “type (2,2)” tensor (which can map matrices to matrices).

Technical Details

Model Details

The models used as examples in this paper are zero, one, and two layer decoder-only, attention-only transformers . For all models, d_\text{model} = n_\text{heads} * d_\text{head}, typically with n_\text{heads}=12 and d_\text{head}=64, but also with one explicitly noted example where n_\text{heads}=32 and d_\text{head}=128.

Models have a context size of 2048 tokens and use dense attention. (Dense attention was preferred over sparse attention for simplicity, but made a smaller context preferable.) We use a positional mechanism similar to Press et al. , adding sinusoidal embeddings immediately before multiplying by W_Q and W_K to produce queries and keys. (This excludes pointer-arithmetic based algorithms without the distorted QK matrices like rotary.)

The training dataset is as described in Kaplan et al. .

Handling Layer Normalization

Our transformer models apply layer normalization every time they read off from the residual stream. In practice, this is done to encourage healthy activation and gradient scales, making transformers easier to train. However, it also means that the "read" operation of each layer is more complicated than the idealized version we've described above. In this section, we describe how we work around this issue.

Before considering layer normalization, let's consider the simpler case of batch normalization. Transformers typically use layer normalization over batch normalization due to subtleties around the use of batch normalization in autoregressive models, but it's an easy case to work out and would make our theory exactly align with the model. Batch normalization tracks a moving estimate of the mean and variance of its inputs, and uses that to independently normalize them, before rescaling with learned parameters. At inference time, a fixed estimate of the mean and variance is used. This means that it's simply a fixed linear transformation and bias, which can be multiplied into the adjacent learned linear transformation and bias. As such, batch normalization in models can typically be absorbed into other operations and ignored for the purposes of interpretability on the inference time model.

Layer normalization is a little more subtle. For each residual stream vector, we subtract off the average activation, normalize by variance, and then multiply by a learned set of diagonal weights and add a learned bias vector. It turns out that subtracting off the average activation is a fixed linear transformation, since it just zeros out a single dimension in the vector space. This means that everything except for normalizing by variance, layer normalization applies a fixed affine transformation. Normalizing by variance multiplies the vector by a scalar, and multiplying by a scalar commutes with all the other operations in a path. As a result, we can fold everything but normalization into adjacent parameters, and then think of the normalization scaling as a variable reweighting of the set of path terms going through that layer normalization. (The main downside of this is that it means that paths going through different sets of layers aren't easily comparable; for example, we can't trivially compare the importance of virtual attention heads composed of heads from different layers.)

A final observation is that layer normalization in the first attention layer of the model can alternatively be applied to every vector of the embedding matrix and then ignored. In some cases, this is more convenient (for example, we do this for skip-trigrams in our one-layer models).

In the future, it might be worth trying to train models with batch normalization to avoid this as a pain point, despite the complexities it presents.

Working with Low-Rank Matrices

In this paper, we find ourselves dealing with very large, but extremely low-rank matrices. Picking the right algorithm is often the difference between painfully slow operations on the GPU, and everything being instantaneous on your CPU.

First, we recommend keeping matrices in factored form whenever possible. When multiplying a chain of matrices, identify the "smallest" bottleneck point and represent the matrix as a product AB split at that bottleneck.

Eigenvalues: Exploit the fact that \lambda_i(AB) = \lambda_i(BA) . Calculating eigenvalues of a 64×64 matrix is greatly preferable to a 50,000×50,000 matrix.

SVD: The SVD can also be calculated efficiently, and a variant can be used to efficiently compute PCA:

• Compute the SVDs of A and B: U_AS_AV_A = A and  U_BS_BV_B = B.
• Define C=S_AV_AU_BS_B
• SVD C:  U_CS_CV_C = C
• SVD of AB is U_{AB} = U_AU_C, S_{AB} = S_C, V_{AB} = V_CV_B.