A Mathematical Framework for Transformer Circuits

The residual stream is the sum of every layer's output, so heads act independently and add. Every layer reads from the stream with a linear map and adds its output back, so the stream at any depth is the embedding plus everything written so far. Any two layers are therefore joined by virtual weights $W_I^2W_O^1$, the later layer's read times the earlier layer's write, however many layers sit between them. The concatenate-then-project form of multi-head attention is likewise a sum, $$W_O^H\begin{pmatrix}r^{h_1}\\ r^{h_2}\\ \vdots\end{pmatrix}=\sum_iW_O^{h_i}r^{h_i},$$ so each head writes to the stream on its own and can be studied alone. Rotating the stream together with every matrix that touches it leaves the model unchanged, so the paper works with weight products, never stream coordinates.

A head is an attention pattern tensored with one low-rank matrix. Write a head's three stages (a value vector per position, mixing across positions, an output projection per position) as tensor products and collapse them by the mixed-product property: $$h(x)=(\mathrm{Id}\otimes W_O)(A\otimes\mathrm{Id})(\mathrm{Id}\otimes W_V)\,x=(A\otimes W_{OV})\,x.$$ So a head does two separate things. $A$, computed from the bilinear form $W_{QK}$, says between which positions information moves, and $W_{OV}$ says what is read at the source and written at the destination. Queries, keys and values are only one factorization of these two matrices. The softmax is the only nonlinearity, so running the model once to record every pattern and again with the patterns frozen makes the logits linear in the tokens. Products of heads are again heads, $(A^{h_2}\otimes W_{OV}^{h_2})(A^{h_1}\otimes W_{OV}^{h_1})=(A^{h_2}A^{h_1})\otimes W_{OV}^{h_2}W_{OV}^{h_1}$, which is where virtual heads come from.

A one-layer model is a bigram table plus skip-trigram tables. Expand $(\mathrm{Id}\otimes W_U)\big(\mathrm{Id}+\sum_hA^h\otimes W_{OV}^h\big)(\mathrm{Id}\otimes W_E)$ into the sum over paths in (ii). The direct path cannot move information between positions, so it is a bigram table (Barack is followed by Obama). Each head adds two vocabulary-by-vocabulary tables, the QK circuit $W_E^\top W_{QK}^hW_E$, which scores how much a destination token attends to a source token, and the OV circuit $W_UW_{OV}^hW_E$, which gives the logit change from attending to that source. Together they encode skip-trigrams "A … B $\to$ C": from B attend back to A and raise C. The largest entries are mostly copying: attend back to tokens that could plausibly come next and raise them. Because the form is factored, the logit effect depends only on the attended source, so a head that raises "keep … in mind" and "keep … at bay" also raises "keep … in bay". Copying is summarized by the OV circuit's eigenvalues (cheap to compute, since $\lambda(AB)=\lambda(BA)$): a copying matrix maps a combination of tokens to the same combination of logits, so its eigenvalues are positive. The authors call this suggestive only, since non-orthogonal eigenvectors allow an all-positive spectrum while some tokens lower their own logits.

Two layers add composition, and only Q- and K-composition adds expressiveness. A second-layer head reads $x_1=x_0+\sum_{h\in H_1}h(x_0)$ instead of the embeddings. In the logits this only adds the virtual heads of Step 2 (V-composition), with pattern $A^{h_2}A^{h_1}$, which still only moves information between positions. What is new appears in the attention scores. Substituting $x_1$ on both sides of $x_1^\top W_{QK}^hx_1$ gives terms $A_q\otimes A_k\otimes W$, where $A_q,A_k$ are $\mathrm{Id}$ or a first-layer pattern according to whether a head first moved the query-side or key-side information and $W$ is the bilinear form, for example $$\mathrm{Id}\otimes A^{h_1}\otimes W_E^\top W_{QK}^hW_{OV}^{h_1}W_E\qquad\text{(K-composition)}.$$ A second-layer head can thus attend by what an earlier head moved into a position, which no table over token pairs expresses. The paper scores each pair of heads by $\|W_{QK}^{h_2}W_{OV}^{h_1}\|_F/(\|W_{QK}^{h_2}\|_F\|W_{OV}^{h_1}\|_F)$ against random matrices. Most heads do not compose, and the significant composition is K-composition from one first-layer head that attends mainly to the previous token.

An induction head is K-composition with a previous-token head plus a copying OV circuit. An induction head completes [a][b] … [a] $\to$ [b]. Take the K-composition term $\mathrm{Id}\otimes A^{h_{-1}}\otimes W$ with a previous-token pattern $A^{h_{-1}}$: the key at position $j$ carries the token at $j-1$, so a same-token bilinear form $W$ raises the score of exactly the positions whose predecessor matches the current token, and a copying OV circuit raises the token found there. Unlike one-layer copying heads, it uses how the token was used earlier, whatever the corpus statistics. Two checks support this. The heads behave the same way on uniformly random tokens repeated three times, where corpus statistics cannot help. And every induction head has strongly positive eigenvalues in both its OV circuit and its K-composition QK term; the authors note this is not circular, since the heads were selected for large composition, which says nothing about its sign.

Path ablations show that the virtual heads barely matter here. Freeze every attention pattern at its recorded value and rerun the model several times. The first frozen run replaces every head's output with zero, leaving the direct path, and each later run adds to the stream, in place of each head's current output, the output that head produced on the run before. The $n$-th run then contains exactly the OV paths through at most $n-1$ heads, and successive loss differences give the marginal value of each order. In the two-layer model the second-order (V-composition) terms add little, and among first-order terms the second-layer heads carry most of the gain. Since the patterns are frozen, this says nothing about Q- and K-composition, which act only through the patterns.