Notes on A Mathematical Framework for Transformer Circuits

Notes on A Mathematical Framework for Transformer Circuits (Elhage et al., 2021), the first paper in Anthropic’s Transformer Circuits thread.

The paper gives a way to rewrite small, attention-only transformers as sums of paths from input tokens to output logits. Each path is a chain of matrices with a fairly specific job: choose which position to read from, move information between positions, or turn that information into a change in the output distribution. The authors use this view to reverse engineer zero-, one-, and two-layer models.

The residual stream is the shared state

A transformer is usually drawn as a sequence of blocks, which makes it look as though each layer replaces the representation produced by the one before it. The paper uses a different view. Token embeddings initialize a residual stream, and each attention head or MLP reads from that stream and adds an update to it. The unembedding reads the final stream and converts it to logits.

The residual-stream diagram makes the change in viewpoint concrete. The long horizontal bar is the shared state; the heads and MLP sit beside it, each with one path for reading and another for adding its result back. There is no direct arrow from one component box to the next.

Figure 1: The transformer as a residual stream. Each attention head is an independent operation, and every component — heads, then the MLP, layer after layer — reads the current stream with a linear projection and adds its result back in. The logits are the unembedding of the accumulated sum. Redrawn from the paper.

No component owns the stream. The residual connection preserves earlier writes unless a later component changes them. A head can add new information alongside an earlier write or cancel it by writing the negative of what is already there. Two components interact only when the first writes into a direction that the second reads. The subspace diagram (Fig. 1) shows both cases: overlapping reads and writes create a connection, while reading a subspace and writing its negative deletes the information carried there.

The residual stream divided into subspaces; layers interact where their read and write subspaces overlap, and a layer can delete information by writing the negative of what it read — **Figure 2:** Components use subspaces of the residual stream. Disjoint write and read subspaces do not interact; aligned ones create an implicit connection. Figure from the paper.

Suppose one component writes with a matrix $W_A$ and a later component reads with $W_B$ . Their effective connection is the product $W_BW_A$ . That product is a virtual weight: there is no parameter with that value stored in the model, but it describes how the two components communicate through the residual stream. Multiplying weights through the stream exposes an implicit connection between every compatible pair of components. The left side of the next figure (Fig. 2) shows the literal read-and-write architecture; the right side has multiplied through the stream and made those pairwise connections explicit.

Layers reading from and writing to the residual stream with linear projections; multiplying the weights through reveals virtual weights connecting each pair of layers — **Figure 3:** A virtual weight is the product of one component's write matrix and another component's read matrix. Figure from the paper.

An attention head has two jobs

A single attention head computes

\text{attention}(Q,K,V) = \text{softmax}\big(\frac{QK^T}{\sqrt{d_K}}\big)V.

(1)

The queries, keys, and values in Equation 1 are linear projections of the residual stream ( $Q = W_Qx$ , and so on). The softmax term is the attention pattern $A$ . At each position, it gives a distribution over earlier positions to read from. The $\sqrt{d_K}$ scaling keeps the scores in a range where the softmax has usable gradients.

The paper combines the four projection matrices into two. Queries and keys meet only through $W_Q^TW_K$ , while values reach the residual stream only through $W_OW_V$ :

W_{QK}=W_Q^T W_K, \qquad W_{OV}=W_O W_V.

The QK circuit determines where the head attends. The OV circuit determines what the head does with the token it finds. Both matrices are low rank because they factor through the smaller head dimension. In the diagram (Fig. 3), the pink QK path touches both the destination and source tokens to produce a score. The gold OV path starts at the selected source and ends at the output logits.

Diagram tracing the QK circuit and OV circuit of an attention head as separate weight paths from source and destination tokens — **Figure 4:** The QK circuit scores source and destination token pairs. The OV circuit maps the selected source token to an update in the residual stream. Figure from the paper.

In a one-layer model, multiplying the QK circuit through the token embeddings gives $W_E^T W_{QK}W_E$ , a matrix that scores pairs of tokens. Multiplying the OV circuit through the embedding and unembedding gives $W_UW_{OV}W_E$ , which maps a source token to its direct effects on the output logits.

For a pattern such as [keep] … [in] → [mind], the QK circuit gives the token in a reason to attend to the earlier keep. Once attention lands there, the OV circuit makes keep vote for mind. The split lets us inspect two questions separately: which source tokens does the head select, and what do those source tokens contribute?

Expanding the model into paths

The attention pattern $A$ depends on the input because it comes from a softmax over query-key scores. If $A$ is held fixed for a particular input, an attention head is linear. An attention-only model can then be expanded into a sum of paths through the model.

Zero layers. Embed, then immediately unembed. The logits are $W_UW_E$ applied to the current token. A model with no attention can only learn bigram statistics: which tokens tend to follow the current one.

One layer. With attention patterns fixed, the model is a sum of paths:

T = \underbrace{\mathrm{Id} \otimes W_U W_E}_{\text{direct path}} + \sum_h \underbrace{A^h \otimes W_U W_{OV}^h W_E}_{\text{one term per head}}.

The tensor product $\otimes$ keeps two operations separate. The left factor describes movement between token positions; the right factor describes a transformation within the residual-stream space. The direct path is the zero-layer bigram model. Each head adds another path in which $A^h$ selects a source position and $W_UW_{OV}^hW_E$ turns the token at that position into a logit update.

This makes a one-layer attention-only model a collection of “skip-trigram” rules. Each contribution can depend on the current token and one earlier token, but the earlier token has not yet received information from anywhere else. Its key and value are functions of its own embedding.

Why a second layer is different

In a two-layer model, the second layer reads a residual stream that already contains the first layer’s outputs. An earlier head can therefore affect the values, queries, or keys used by a later head.

With V-composition, the later head moves information that the earlier head already moved. The path contains an attention pattern $A^{h_2}A^{h_1}$ and a transform

W_U W_{OV}^{h_2}W_{OV}^{h_1}W_E.

This behaves like a virtual attention head spanning both layers.

Q- and K-composition affect the later head’s attention pattern itself. An earlier head writes information that becomes part of a later query or key, so the later head can select a token based on context gathered from another position. The induction-head circuit uses K-composition: a previous-token head writes the identity of the preceding token into the current position, and a later head uses that write as a key. The later head can then find a token according to what came before it.

That operation is unavailable to a one-layer model. One layer can compare token identities; two layers can compare a token with information that another head has moved beside it.

What is exact here

The additive decomposition of attention heads and the path expansion are algebraic facts about the architecture. Given fixed attention patterns, the paths sum exactly to the model’s output. The decomposition does not guarantee that an individual path will have a simple interpretation. It gives the analyst a smaller object to inspect.

The paper demonstrates the method on attention-only models with at most two layers. In that setting, the authors can read behaviors from the QK and OV matrices, identify composition between heads, and recover the induction-head circuit. These models are small enough for fairly complete analysis, but they leave out MLPs and the complications that come with scale.

Freezing attention also sets aside part of the computation. The resulting paths explain what information flows through a head under the observed attention pattern. They do not, by themselves, explain why the head attended to those positions. QK analysis handles that question in the small models studied here; doing the same in a large model is harder.

What carried forward

TransformerLens (opens in a new tab) exposes models using this residual-stream and per-head decomposition, and the ARENA curriculum (opens in a new tab) uses it as its introduction to mechanistic interpretability.

Later work separates the two gaps left by the paper. Sparse autoencoders and transcoders try to find workable units for MLP computation. Attribution graphs return to the path-expansion idea: they freeze attention patterns for a prompt and trace linear effects through residual connections and OV circuits. A later attention-circuits paper (opens in a new tab) adds QK interactions back into those graphs. The July 2025 circuits update (opens in a new tab) also revisits the original framework using learned features instead of matrix eigenvectors as the objects moving through the residual stream.

Further material

The paper itself (opens in a new tab) — the exercises and videos linked in its appendix are good, and the two-layer section rewards a slow read.
Neel Nanda’s mechanistic interpretability quickstart (opens in a new tab) — orientation for the field this paper started.
TransformerLens (opens in a new tab) — the standard library for poking at models in this framework.

References

These are notes on the work marked ★ — cite the original rather than this page.

Elhage, N., Nanda, N., Olsson, C., & others. (2021). A Mathematical Framework for Transformer Circuits. Transformer Circuits Thread. https://transformer-circuits.pub/2021/framework/index.html ★

Table of Contents

The residual stream is the shared state
An attention head has two jobs
Expanding the model into paths
Why a second layer is different
What is exact here
What carried forward
Further material
References