Attention in tensor notation

How to read this page

An attention head answers three questions about each position in the text: where should I look, what should I bring back from there, and what does that do to the next-token prediction. The poster asks the same three questions of a real model. This page shows the arithmetic behind them, twice: first on a toy you can multiply out on paper, then in the general notation the papers use.

Four words carry most of the weight, so here they are in plain terms:

residual stream

a short list of numbers that every position carries with it, like a notepad; each layer reads it and writes a little back

query

what a position is looking for, computed from its own notepad

key

what a position advertises about itself, so that queries can find it

value

what a position hands over once it has been chosen

Query against key decides where (section 1, step 1, and section 4). Value decides what (section 1, step 2, and section 5). The unembedding turns what arrived into a change in the logits (section 1, step 3, and section 6). Sections 7 and 8 stack two heads.

1. One head by hand: five tokens

Take the five tokens result = 0 ans return, the end of the poster's prompt shrunk to its bones. Give each token a notepad of four numbers, and give the four slots meanings so that every weight below can be read out loud:

Written as a matrix with one row per position, that table is the input \(X\), five rows of four numbers. A trained model does not hand out slots this cleanly; that is what superposition and the learned dictionaries on the timeline are about. The arithmetic is the same either way.

Step 1. Where to look: query meets key

In words. The head we build has one job: when it stands on return, find the variable names that came before, and prefer result. So the query should be non-zero only at return, and the keys should be non-zero only at names.

Each position turns its notepad into a query by one small matrix, \(W_Q\); the query is what that position is looking for. We want only return to look, so \(W_Q\) should read slot 4 and nothing else. Give it two coordinates, with weights 2 and 1; the 2 : 1 split is what will later make result beat ans. Each position also turns its notepad into a key by another small matrix, \(W_K\); the key is what that position advertises about itself. We want names to be findable, so \(W_K\) should read slots 1 and 2, one coordinate each. The smallest matrices that do exactly that are

Two rows because the head is two numbers wide, four columns because the notepad has four slots. \(X\) holds one row per position, so the queries of all five positions at once are \(Q = X W_Q^{\top}\): a \(5 \times 4\) times a \(4 \times 2\) gives a \(5 \times 2\), one query per row. For a single position that is \(q_p = W_Q\, x_p\). The keys are the same product with \(W_K\): \(K = X W_K^{\top}\).

Only row 5 of \(Q\) is non-zero, because only return has slot 4 set and \(W_Q\) reads nothing else.

Only rows 1 and 4 of \(K\) are non-zero, because only the two names have slot 1 or slot 2 set.

A score says how well a query matches a key: multiply coordinate by coordinate and add, \(q \cdot k\). All twenty-five scores at once are \(S = Q K^{\top}\), a \(5 \times 5\) matrix with one row per destination. From return, the row is

\(2 \cdot 1 + 1 \cdot 1 = 3\) for result, \(2 \cdot 1 + 1 \cdot (-1) = 1\) for ans, zero for the rest. The other four rows of \(S\) are zero, because the other four rows of \(Q\) are.

Here is the first thing the notation is for. \(W_Q\) and \(W_K\) never act alone; only their product matters. Multiply them into one \(4 \times 4\) matrix, \(W_{QK} = W_Q^{\top} W_K\), and it has exactly two non-zero entries:

That is the head's where rule in one line: "a position carrying return scores each earlier position by \(2 \times\) (is a name) \(+\, 1 \times\) (which name)". This matrix is the QK circuit. The split into \(W_Q\) and \(W_K\) is not unique: \(W_Q = \big[\begin{smallmatrix} 0 & 0 & 0 & 1 \\ 0 & 0 & 0 & 1 \end{smallmatrix}\big]\) with \(W_K = \big[\begin{smallmatrix} 2 & 0 & 0 & 0 \\ 0 & 1 & 0 & 0 \end{smallmatrix}\big]\) gives the same \(W_{QK}\) and the same scores, which is why the pair is read as one matrix. Section 4 gives its general form.

Scores are not yet weights. A softmax turns the five scores into five numbers that are positive and add up to one, so that the head can take a weighted average in the next step: divide each score by \(\sqrt{2}\) (the head's width; this keeps scores from growing with the head size), take \(e^{x}\), and divide by the sum.

From return, 62% of the attention goes to result, 15% to ans, and the rest is spread thinly. That row is what one row of the poster's heatmap is: one head, one query position, a weight for every earlier token. The two matrices here were chosen by hand for the toy; in a trained model \(W_Q\) and \(W_K\) are learned parameters of each head, and the poster reads their effect straight off the attention row.

Step 2. What to move: value and output

In words. Attention weights only say where to look. The head still has to decide what to carry back from the places it looked at, and where on the notepad to write it. Two more small matrices do that. \(W_V\) turns a source's notepad into its value, what it hands over if chosen; \(W_O\) decides which slots of the destination's notepad the result is written to. We want the name bits to travel, so \(W_V\) should read slots 2 and 1 and \(W_O\) should write them back into slots 2 and 1. The smallest matrices that do that are

\(W_V\) is two rows by four slots, like \(W_Q\): it reads a notepad (four slots in) and produces a value (two numbers out). \(W_O\) goes the other way, four slots by two, because it writes: two numbers from the head go in, four slots come out, so the result has the shape of a notepad and can be added to one. The values of all five positions at once are \(V = X W_V^{\top}\), the same product as for \(Q\) and \(K\):

Only the two names offer anything. The head's read at return, call it \(r_5\), is the attention row from step 1 times \(V\): a \(1 \times 5\) times a \(5 \times 2\), which is the weighted average of the five values.

\(W_O\) writes those two numbers into slots 2 and 1 of the notepad at return. The result, \(h_5\), is what the head adds to return's notepad; in the row convention it is \(r_5 W_O^{\top}\), a \(1 \times 2\) times a \(2 \times 4\):

So return's notepad gains 0.78 of "a name is here" and \(+0.47\) of "and it leans toward result". As before, \(W_V\) and \(W_O\) only ever act as a pair. Their product \(W_{OV} = W_O W_V\) is one \(4 \times 4\) matrix, the OV circuit, and here it simply copies slots 1 and 2 and ignores slots 3 and 4:

The chain only closes because \(W_O\) is transposed relative to \(W_V\): the head's width, 2, is what gets multiplied away in the middle, and the output comes back with four slots.

Step 3. What it does to the prediction

In words. The model's guess for the next token is a score per vocabulary entry, the logits. They are read off the last notepad by the unembedding \(W_U\), a fixed table with one row per vocabulary entry: the logit for a token is that token's row dotted with the notepad. Let the row for result be (name \(+\) which-name), the row for ans be (name \(-\) which-name), and the other three rows read slots 3 and 4. The head added \(h_5\) to the notepad, so it changes each logit by \(W_U\) applied to \(h_5\); call that \(\Delta L_5\):

from \(0.78 + 0.47 = 1.25\) and \(0.78 - 0.47 = 0.30\). Both names go up, result four times as much. Notice that the two increments are exactly twice the two attention weights (\(2 \times 0.62\) and \(2 \times 0.15\)). That is because, end to end, this head is a copying head: attending to result adds 2 to the result logit, attending to ans adds 2 to the ans logit, and attending to anything else adds nothing. Section 6 gives that end-to-end matrix a name, and section 8 needs exactly this shape for the induction head.

The whole example on one card

import torch
X   = torch.tensor([[1,1,0,0],[0,0,1,0],[0,0,1,0],[1,-1,0,0],[0,0,0,1.]])   # result = 0 ans return
W_Q = torch.tensor([[0,0,0,2],[0,0,0,1.]]);  W_K = torch.tensor([[1,0,0,0],[0,1,0,0.]])
W_V = torch.tensor([[0,1,0,0],[1,0,0,0.]]);  W_O = torch.tensor([[0,1],[1,0],[0,0],[0,0.]])
W_U = torch.tensor([[1,1,0,0],[0,0,1,0],[0,0,1,0],[1,-1,0,0],[0,0,0,1.]])

W_QK = W_Q.T @ W_K                     # [4, 4], two non-zero entries
W_OV = W_O @ W_V                       # [4, 4], identity on the name slots
S = torch.einsum("di, ij, sj -> ds", X, W_QK, X)
mask = torch.tril(torch.ones(5, 5))    # a position may only look backwards
A = torch.softmax(S.masked_fill(mask == 0, -1e9) / 2**0.5, dim=-1)
H = torch.einsum("ds, si, oi -> do", A, X, W_OV)
print(A[4])            # tensor([0.624, 0.075, 0.075, 0.152, 0.075])
print(H[4] @ W_U.T)    # tensor([1.248, 0.000, 0.000, 0.303, 0.000])

2. Shapes: token ids, one-hot columns, and the embedding table

In words. In the example each token arrived with its notepad already filled in. In a real model the notepad comes from a lookup: the token's id picks one row out of a big table. The maths writes that lookup as a matrix product, so that the whole network becomes one chain of matrices. This section is only about that bookkeeping.

Three sizes appear everywhere: the context length (five here, \(n_{\text{context}} = 5\)), the notepad width \(d_{\text{model}}\) (four in the example, a few thousand in a real model), and the vocabulary size \(n_{\text{vocab}}\) (a few hundred thousand).

What the code holds: five integers

The tokenizer turns the five tokens into five ids, say [42, 105, 7, 900, 13]: an array of shape (5,), or (1, 5) with a batch axis. nn.Embedding uses each integer directly to pick a row of its table. No \(n_{\text{vocab}}\)-wide matrix is ever built in the program.

What the maths writes: a one-hot matrix \(T\)

Each token becomes a column of length \(n_{\text{vocab}}\) with a single 1 at its id. Five columns side by side make a very tall, very sparse matrix:

The table \(W_E\) does not depend on the sentence

\(W_E\) is the model's fixed lookup table, one column per vocabulary entry, each column \(d_{\text{model}}\) long. Multiplying it by a one-hot column with its 1 at row 42 returns column 42 of the table. So the input notepads are

which is exactly "pull these five columns out of the table and put them side by side". The program does it with five integer lookups; the maths writes it as a product so that it can be expanded and differentiated like everything downstream.

Run it on section 1's numbers

Suppose the vocabulary is only the five tokens of the example, in that order. Then each token's one-hot column has its 1 at its own position, so \(T\) in row form is the \(5 \times 5\) identity, and \(W_E\) is the example's five notepads placed as columns, one per vocabulary entry:

In the row convention of section 1 the lookup is \(X = T W_E^{\top}\), here \(I\, W_E^{\top}\), which is exactly the \(X\) of section 1: the five columns pulled out and laid down as rows. Column 4 of \(W_E\) is ans's notepad \((1, -1, 0, 0)\), so row 4 of \(X\) is that notepad. If a token appeared twice, the same column would be pulled twice.

3. Einstein summation: name the axes once

In words. Every operation from here on is "multiply some entries and add them up over one axis". Rather than writing sums, give each axis a letter and agree on one rule: a letter that appears twice in a product is summed over, and the summation sign is dropped. torch.einsum uses the same letters, so the formula and the code read the same.

\(p\)

position in the sequence, 1 to 5

\(v\)

vocabulary entry, 1 to \(n_{\text{vocab}}\)

\(d\)

slot on the notepad, 1 to \(d_{\text{model}}\)

The embedding lookup of section 2, with the input as \(T_{p,v}\) and the table as \((W_E)_{d,v}\):

# T: [5, n_vocab] one-hot rows, W_E: [d_model, n_vocab]
# X: [5, d_model]
X = torch.einsum("pv, dv -> pd", T, W_E)

Read the string left to right: the inputs have axes pv and dv; the letter v appears in both and not in the output, so it is summed away; what is left, pd, is the output's shape. Every formula below is one more einsum with a couple of extra letters.

Run it on section 1's numbers

Take one entry of the lookup, \(X_{p,d} = T_{p,v}\,(W_E)_{d,v}\), at \(p = 4\) (ans) and \(d = 2\). The sum runs over \(v\); \(T_{4,v}\) is 1 only at \(v = 4\), so the sum collapses to \((W_E)_{2,4} = -1\), and \(X_{4,2} = -1\): the which-name bit of ans.

The query of section 1 in the same notation is \(Q_{p,h} = X_{p,i}\,(W_Q)_{h,i}\), summed over the slot index \(i\). At \(p = 5\): \(Q_{5,1} = X_{5,i}(W_Q)_{1,i} = X_{5,4} \cdot 2 = 2\) and \(Q_{5,2} = X_{5,4} \cdot 1 = 1\), so \(q_5 = (2, 1)\). In code that is torch.einsum("pi, hi -> ph", X, W_Q), which is the same thing as X @ W_Q.T: the letters just say which axis is shared.

4. The QK circuit in general: where to look

In words. Step 1 of the example, with the sizes left open. A query is the destination's notepad pushed through \(W_Q\); a key is a source's notepad pushed through \(W_K\); a score is their dot product; the two weight matrices collapse into one, the QK circuit.

Three more letters:

\(d\) and \(s\)

destination position (the one making the prediction) and source position (the one being read); both run over the 5 positions

\(i, j, o\)

notepad slots, \(d_{\text{model}}\) wide (input side \(i, j\); output side \(o\))

\(h\)

the head's internal coordinate, \(d_{\text{head}}\) wide (two in the example)

Queries come from the destination, keys from the source: \(Q_{d,h} = X_{d,i}\,(W_Q)_{h,i}\) and \(K_{s,h} = X_{s,j}\,(W_K)_{h,j}\). The score sums their product over \(h\):

\(W_Q\) and \(W_K\) only ever appear together, so the framework paper folds them into one matrix, \((W_{QK})_{i,j} = (W_Q)_{h,i}(W_K)_{h,j}\), of size \(d_{\text{model}} \times d_{\text{model}}\) but of rank at most \(d_{\text{head}}\). In the example it had two entries; in general it is dense, and the score is a bilinear form in the two notepads:

# S: [n_ctx, n_ctx] scores; W_QK: [d_model, d_model]
S = torch.einsum("di, ij, sj -> ds", X, W_QK, X)

# causal mask (look backwards only), then softmax over the source axis
A = torch.softmax(S.masked_fill(mask == 0, -1e9) / d_head**0.5, dim=-1)

Each row of \(A\) is a probability distribution over sources; the kernel-estimator page shows it is a kernel weight. The poster's heatmap is one such row per head, \(d = \) the return position, for every head in the model.

Run it on section 1's numbers

With the example's \(W_{QK}\), only the entries \((4,1) = 2\) and \((4,2) = 1\) are non-zero, so in \(S_{d,s} = X_{d,i}\,(W_{QK})_{i,j}\,X_{s,j}\) only \(i = 4\) survives:

\(X_{d,4}\) is 1 only for \(d = 5\), so rows 1 to 4 of \(S\) are zero and row 5 is \((3, 0, 0, 1, 0)\), as before. Now the two things the general form adds. First the causal mask: it hides nothing from row 5, but for rows 1 to 4 an all-zero row becomes uniform over the sources the mask allows. Second the rank: \(W_{QK}\) is \(4 \times 4\) but has one non-zero row, so its rank is 1, within the bound \(d_{\text{head}} = 2\). The whole attention matrix is

one row per destination, each row adding to one, zeros above the diagonal because a position cannot look ahead. Section 1 only used the last row.

5. The OV circuit in general: what to move

In words. Step 2 of the example, with the sizes left open. Each source offers a value; the destination takes the attention-weighted average of the values and writes it back through \(W_O\); the two weight matrices collapse into the OV circuit.

With \((W_{OV})_{o,i} = (W_O)_{o,h}(W_V)_{h,i}\) the head coordinate \(h\) disappears and what is left is

# A: [n_ctx, n_ctx], X: [n_ctx, d_model], W_OV: [d_model, d_model]
# H: [n_ctx, d_model], added to the residual stream
H = torch.einsum("ds, si, oi -> do", A, X, W_OV)

Run it on section 1's numbers

\(H_{d,o} = A_{d,s}\,X_{s,i}\,(W_{OV})_{o,i}\) can be done in two products. The inner one, \(X W_{OV}^{\top}\), keeps slots 1 and 2 of every notepad and zeroes the rest: rows \((1, 1, 0, 0)\), \((0,0,0,0)\), \((0,0,0,0)\), \((1, -1, 0, 0)\), \((0,0,0,0)\). The outer one mixes those rows with each row of \(A\):

Row 5 is \(h_5\) from section 1. Row 4, the ans position, is \((0.5, 0, 0, 0)\): it attends a quarter each to result and ans, and their which-name bits cancel. The einsum reached the same numbers without ever forming \(r\) or \(W_O\) on its own; that is what folding \(W_V\) and \(W_O\) into \(W_{OV}\) buys.

6. End to end: from tokens to logits

In words. Step 3 of the example, with the sizes left open. Feed the head's write-back through the unembedding and substitute the one-hot input for \(X\). What comes out is a matrix from source token to output token, the end-to-end OV circuit, and in the example it was the copying matrix with a 2 on each name.

Writing \((W_U)_{v_{\text{out}},o}\) for the unembedding, \(v_{\text{src}}\) for the source token's vocabulary index and \(v_{\text{out}}\) for the predicted token's:

# T: [n_ctx, n_vocab] one-hot input
# M_OV: [n_vocab, n_vocab] = W_U @ W_OV @ W_E   (the end-to-end OV circuit)
delta_logits = torch.einsum("ds, sv, ov -> do", A, T, M_OV)

A one-layer attention-only model is now two matrices that can be read separately:

Together they implement the paper's skip-trigram, [source] ... [destination] [out]: a lookup table with a gap in it. In the example the table said "[result] ... [return] → result".

Run it on section 1's numbers

With \(T = I\) and \(W_E = X^{\top}\) from section 2, both end-to-end matrices are \(5 \times 5\) tables over the five tokens. The OV one, \(M_{OV} = W_U W_{OV} W_E\), has a 2 at (result, result) and a 2 at (ans, ans) and zeros elsewhere: attend to a name, raise that name. The QK one, \(W_E^{\top} W_{QK} W_E\), has a 3 at (return, result) and a 1 at (return, ans): it is the score table of section 1, and it coincides with \(S\) only because here every token sits at one position. Then \(\Delta L = A\,T\,M_{OV}\):

Row 5 is section 1's answer. Row 4 is the same head standing on ans: it pushes both names equally, because its attention row split evenly. The skip-trigrams this head encodes are [result] ... [return] → result, with 3 in the QK table and 2 in the OV table, and [ans] ... [return] → ans, with 1 and 2.

7. Two layers: composition and virtual heads

In words. A head in layer 2 does not see the raw tokens; it sees notepads that layer-1 heads have already written on. So a layer-2 head can build its query, its key, or its value out of what a layer-1 head left behind. That is composition, and it is what lets two layers do things one layer cannot.

Expand every path from input to logits. Besides the direct path \(W_U W_E\) (bigram statistics) and each head's own path, there are compositional paths where a layer-2 head \(h_2\) reads what a layer-1 head \(h_1\) wrote. It can read it through any of its three input matrices; the three names below are from the framework paper's section Three kinds of composition (Elhage et al., 2021):

V-composition is the tidy one, because moving information twice is still just moving information. Compose the two tensor-product operators and use the mixed-product rule \((A \otimes B)(C \otimes D) = (AC) \otimes (BD)\):

This has exactly the form of a single head, hence the name virtual attention head. \(A^{h_2} A^{h_1}\) is a product of row-stochastic matrices, so it is again a valid attention pattern: if \(h_1\) moves position \(k\) to \(m\) and \(h_2\) moves \(m\) to \(d\), the product moves \(k\) to \(d\) in one hop. Inside the composite OV matrix the middle factor \(W_V^{h_2} W_O^{h_1}\) is the paper's virtual weight: how much of \(h_1\)'s output \(h_2\) actually reads.

# composite pattern: [n_ctx, n_ctx]
A_virtual = torch.einsum("dm, ms -> ds", A_h2, A_h1)

# composite OV: [d_model, d_model]
W_OV_virtual = torch.einsum("oj, ji -> oi", W_OV_h2, W_OV_h1)

# the virtual head in one step
H_virtual = torch.einsum("ds, si, oi -> do", A_virtual, X, W_OV_virtual)

Run it on section 1's numbers

Give the example a layer-1 partner \(h_1\): a previous-token head, whose pattern \(P\) puts all its weight one position back (\(P_{p,\,p-1} = 1\), and \(P_{1,1} = 1\) because position 1 has nothing before it), with the same copy matrix for its OV circuit. Let \(h_2\) be section 1's head. The composite pattern is \(A^{h_2} P\); its row 5 is \((A P)_{5,s} = A_{5,m}\,P_{m,s}\), summed over \(m\):

Every weight moved one position earlier: the 0.15 that sat on ans now sits on 0, the token before it, and the 0.62 on result and the 0.07 on = both land on position 1. The composite OV matrix is \(W_{OV} W_{OV}\), which is the same copy matrix, because copying twice is copying. So the virtual head reads the name bits of the token before whatever \(h_2\) attended to, in one hop. In code: A_virtual = torch.einsum("dm, ms -> ds", A, P).

8. The induction head, step by step

In words. Q- and K-composition change the attention pattern itself, and that is what lets a two-layer model do something a lookup table cannot: notice that a token has appeared before and predict what followed it last time. The sequence is

and at position \(d\) the model should predict \([B]\). Two heads do it together.

Step 1. A layer-1 previous-token head tags \([B]\) with "the token before me is \([A]\)"

Its QK circuit attends one position back, so at \(i+1\) it looks at \(i\). Its OV circuit copies \([A]\)'s features into the notepad at \(i+1\). Position \(i+1\) now carries both "I am \([B]\)" and, in a separate set of slots, "the token before me is \([A]\)".

Step 2. The layer-2 induction head builds its key from that tag (K-composition)

\(W_K^{h_2}\) reads the slots the previous-token head wrote, not the raw embedding. So the key at \(i+1\) means \(k_{i+1} \approx \text{Key}(\text{"the token before me is } [A]\text{"})\): position \(i+1\) is broadcasting "if the token you are looking at is \([A]\), look at me".

Step 3. The QK circuit locks on to \([B]\)

At the destination \(d\) the query is built from the current token, \(q_d \approx \text{Query}(\text{"my token is } [A]\text{"})\). The score \(S_{d,i+1} = q_d^{\top} k_{i+1}\) is large exactly where the tag matches, and after the softmax \(A_{d,i+1}\) carries almost all the weight: the row has a single sharp peak on the position after the previous \([A]\).

Step 4. The OV circuit copies \([B]\) into the prediction

\(W_V^{h_2}\) reads position \(i+1\)'s own token features, \([B]\). The end-to-end OV matrix \(W_U W_{OV}^{h_2} W_E\) of an induction head is a copying matrix, like the one in the example with its 2 on each name: whatever token it attends to, it raises that token's own logit. So \(\Delta L_d \approx W_U W_{OV}^{h_2} x_{i+1} \propto \text{logit}([B])\).

The pair never needed to see this particular \([A][B]\) in training; one occurrence in the context is enough. That is the difference between a skip-trigram table and an algorithm. The four steps are the ones the framework paper describes; for a version drawn out head by head, see Induction heads, illustrated (McDougall, 2023).

Run it on section 1's numbers

Section 1's head already has half of an induction head. Its OV side is the copying matrix step 4 needs: \(M_{OV}\) is 2 on each name. What it lacks is step 3's sharp peak; its row puts 0.62 on result, not almost all. Sharpen the QK circuit by tripling \(W_Q\) and the same formulas give

and the logit increments, still twice the weights, become result \(+1.96\) and ans \(+0.03\). That single spike is what a trained induction head's row looks like. What the toy cannot do is find the spike by K-composition, because it has one layer: the tag "the token before me is \([A]\)" needs the layer-1 partner of section 7.

Sources

• A Mathematical Framework for Transformer Circuits

  Elhage, Nanda, Olsson et al., 2021. The QK / OV split, path expansion, skip-trigrams, Q/K/V-composition, virtual heads, and the two-layer induction head. The index conventions above follow its appendix on tensor notation.

• In-context Learning and Induction Heads

  Olsson, Elhage, Nanda et al., 2022. The induction score used for the numbers quoted on the poster.

• Induction heads, illustrated

  Callum McDougall, 2023 (LessWrong). The two-head mechanism of section 8 as diagrams: the previous-token head, K-composition, and the copying OV circuit, one picture per step.

• Attention as a kernel estimator

  The same split from the Nadaraya-Watson side: the kernel is QK, the values are OV.