Chapter 15 Self-Attention and Q/K/V

In Chapter 14, a decoder could consult several saved source states instead of relying on the last one. Its current query determined the mixture. But suppose there is only one sequence: red key .. Before predicting what follows, we want each position to build a representation that can use the other positions. Where should its query come from? What exactly should it compare, and what should it receive?

The simplest starting point is to use each position’s current vector for all three jobs. That works as a calculation, but it couples three choices: changing how a position matches another also changes what it contributes. We will give the jobs separate matrix multiplications, then follow every number. The matrix and Softmax conventions are those of Chapters 5–6. Reading the calculation needs no Python; the optional implementation uses Primer-PY.

15.1 One input, three different products

Example 15.1 — Three positions consult the same sequence. The positions, in order, are red, key, and .. Their two-coordinate representations are the rows of the following matrix. These are new, author-chosen teaching values, not the recurrent states from Chapter 14 and not measured word embeddings. Neither coordinate is inherently “redness” or “keyness.” All values are dimensionless.

$$X=\begin{bmatrix}1&0\\0&1\\1&1\end{bmatrix},\qquad W_Q=\begin{bmatrix}1&0\\0&1\end{bmatrix},\quad W_K=\begin{bmatrix}1&1\\0&1\end{bmatrix},\quad W_V=\begin{bmatrix}1&0\\1&2\end{bmatrix}.$$

Here X has shape [3,2]: three positions, two input coordinates. Each W has shape [2,2]: input-coordinate rows and output-coordinate columns. A W is a stored parameter matrix, shared across positions. Its values are fixed during this forward calculation. The three products also have shape [3,2]:

$$Q=XW_Q=\begin{bmatrix}1&0\\0&1\\1&1\end{bmatrix},\qquad K=XW_K=\begin{bmatrix}1&1\\0&1\\1&2\end{bmatrix},\qquad V=XW_V=\begin{bmatrix}1&0\\1&2\\2&2\end{bmatrix}.$$

For example, Q’s third-row first coordinate is $1\times1+1\times0=1$. K’s third-row second coordinate is $1\times1+1\times1=2$. V’s second-row second coordinate is $0\times0+1\times2=2$. Each product combines coordinates within one position. No position has read another yet. “Projection” here means this learned linear transformation; it need not be a perpendicular geometric projection or reduce dimension.

We can now name the jobs. A row of Q is a query: the vector used to compare one receiving position with possible contributors. A row of K is a key: the vector it compares against. The corresponding row of V is a value: the vector to mix into the result. These names suggest looking up a record, but the resemblance stops at matching and receiving. There is no exact-key database lookup, no guarantee of retrieving a fact, and usually no single selected record. Q, K, and V are computed arrays; the W matrices are parameters.

Because all three arrays come from this same sequence, this is self-attention. Chapter 14’s source and decoder were different sequences, giving cross-attention. The word “self” refers to the shared source sequence, not to a model having self-awareness.

15.2 Compare every receiving row with every contributing row

For the first query [1,0], compare against the three key rows [1,1], [0,1], [1,2]. The dot products are 1, 0, 1. For the second query [0,1], they are 1, 1, 2. The third query [1,1] gives 2, 1, 3. Gathering these results produces:

$$D=QK^{\mathsf T}=\begin{bmatrix}1&0&1\\1&1&2\\2&1&3\end{bmatrix}.$$

D has shape [3,3]. Rows identify receiving queries; columns identify contributing keys. Its third-row second-column element is $1\times0+1\times1=1$: query 3 compared with key 2. This explains the transpose. K originally stores one key per row; turning those rows into columns lets matrix multiplication perform the row-with-key comparison. [3,2] @ [2,3] contracts the two matching coordinates and retains both position axes. Using Q @ K would not even have compatible inner dimensions here.

Do not read D as a table of next-token probabilities. It has position columns, not vocabulary columns, and its rows do not sum to one. In particular, position 3 gives a score of 3 to itself and 2 to position 1; the calculation has not yet said how much of either value to use.

15.3 Make scores usable without making them unnecessarily sharp

We could immediately apply Chapter 6’s Softmax. A larger coordinate count, however, can make dot products vary more widely simply because more terms were added. To see the motivation, imagine independent coordinates whose positive and negative values balance around zero, each with variance 1. Variance measures average squared deviation. Under the additional independence assumptions, a sum of two query-key products has variance 2; a sum of 64 has variance 64. Their typical spreads are proportional to the square roots, about 1.414 and 8. Dividing by the square root of the coordinate count removes this particular growth in spread.

This motivates scaled dot-product attention. Our queries and keys each have two coordinates, so we divide every D element by $\sqrt{2}$. The divisor is not the three positions, the value width, or the batch size. Write $s=1/\sqrt{2}\approx0.707107$:

$$S=D/\sqrt{2}=\begin{bmatrix}s&0&s\\s&s&2s\\2s&s&3s\end{bmatrix}.$$

S still has shape [3,3] and the same row/column meanings. For instance, its last element is $3/\sqrt{2}\approx2.121320$. Scaling reduces score differences in this example, hence makes the distribution less concentrated. In high dimensions it helps avoid prematurely saturated Softmax probabilities and small gradients under the paper’s assumptions. It does not guarantee them: learned coordinates need not remain independent or have variance 1. Nor does division bound scores or replace the maximum-subtraction trick for safe exponentiation. Vaswani et al., Section 3.2.1.

Each receiving position needs its own mixture, so normalize each row over contributing positions. Let $r=e^s\approx2.028115$. The first row’s exponentials are [r,1,r], whose total is $2r+1\approx5.056230$. Thus its weights are those three numbers divided by that total. In the second row, [r,r,r²] can be divided throughout by r; in the third, [r²,r,r³] can also be divided throughout by r. The exact row expressions are:

$$A_1=\frac{[r,1,r]}{2r+1},\qquad A_2=\frac{[1,1,r]}{2+r},\qquad A_3=\frac{[r,1,r^2]}{r+1+r^2}.$$

Table 15.1 — Each query distributes a total weight of one. Six decimal places are displayed; calculations retain full precision.

Receiving position	Weight on value 1	Weight on value 2	Weight on value 3
1	0.401112	0.197776	0.401112
2	0.248255	0.248255	0.503490
3	0.283995	0.140029	0.575975

A has shape [3,3]. Column normalization would instead make different receivers compete for a given contributor; it would not generally give each receiver a unit-total mixture. Nothing here requires a contributor’s column to sum to one. Several queries can all use the same position heavily.

15.4 Receive values, not scores

The first output uses weight 0.401112 on [1,0], 0.197776 on [1,2], and 0.401112 on [2,2]. Multiply each entire value row by its weight, then add corresponding coordinates. Using the unrounded weights:

$$O_{1,1}=\frac{r\times1+1\times1+r\times2}{2r+1}\approx1.401112, \qquad O_{1,2}=\frac{r\times0+1\times2+r\times2}{2r+1}\approx1.197776.$$

The denominator is still the first query’s row normalizer. It does not become the number of value coordinates. All three output rows follow the same rule:

$$O=AV\approx\begin{bmatrix} 1.401112&1.197776\\ 1.503490&1.503490\\ 1.575975&1.432009 \end{bmatrix}.$$

The multiplication is [3,3] @ [3,2] → [3,2]. It sums over contributors, keeping the receiving position and value-coordinate axes. Values may be negative or greater than one; O is a representation, not a probability distribution. Each output coordinate lies between the smallest and largest corresponding V coordinates because its weights are nonnegative and sum to one. Subsequent projections or residual additions need not preserve that bound.

Separate three roles

1. A list of numbersX [3, 2]
2. Stored model numbersW_Q, W_K, W_V: each [2, 2]
3. A list of numbersQ, K, V: each [3, 2]

Projection sums over input coordinates; position rows stay separate.

Compare and normalize

1. An operationQ [3, 2] @ Kᵀ [2, 3] → D [3, 3]
2. A list of numbersRow 1: [1, 0, 1] / √2
3. An operationexp → [r, 1, r]; divide by 2r+1

r = exp(1/√2). Each receiving row has its own denominator.

Mix the values

1. A list of numbersA₁ = [0.401112, 0.197776, 0.401112]
2. A list of numbersV rows: [1, 0]; [1, 2]; [2, 2]
3. An operationA [3, 3] @ V [3, 2] → O [3, 2]
4. A list of numbersO₁ = [1.401112, 1.197776]

Mixing sums over source positions. Displayed decimals are rounded to six places.

Figure 15.1 — One query produces three weights and two output coordinates. Original computed teaching diagram, CC BY-SA 4.0. Inputs and parameters are authored; results follow the declared arithmetic.

Figure 15.1 separates the stored matrices from the arrays they produce. Follow the receiving first row through its three scores to its two output coordinates. The token position survives the entire route; the axis that is summed over changes from input coordinates during projection to source positions during mixing.

Having completed the arithmetic, we can compress it into the familiar formula:

$$\operatorname{Attention}(Q,K,V)=\operatorname{softmax}_{\text{keys}}\!\left(\frac{QK^{\mathsf T}}{\sqrt{d_k}}\right)V.$$

For one head, Q has shape [T,d_k], K [T,d_k], and V [T,d_v]. T counts positions; the positive integer $d_k$ counts comparison coordinates; $d_v$ counts value coordinates and may differ. The result has shape [T,d_v]. The subscript “keys” says exactly which axis Softmax normalizes. Substituting T=3 and both widths=2 recovers every matrix above. A batch adds a leading B axis and computes these products separately for each example; it does not mix examples.

15.5 The same heatmap can carry different information

Example 15.2 — Keep the weights; change what flows through them. Keep X, W_Q, and W_K fixed but replace W_V with a zero matrix. Q, K, D, S, and all of Table 15.1 remain unchanged. V becomes zero, so every output becomes [0,0]. A colorful attention table would look identical even though the values delivered have disappeared.

A less extreme change also has an exact consequence: if each V row receives the same added vector [2,-1], every O row receives [2,-1], because its weights sum to one. This is a controlled intervention on V, not a claim that editing raw text changes only V. Editing X would normally change all three projections. An attention weight is evidence about this mixture, not sufficient evidence for why a whole model answered a question or whether its answer is true.

15.6 Check the arithmetic with NumPy

After Primer-PY, the following complete head calculation can verify the prose. axis=-1 means the last axis; keepdims=True retains a one-column row total so division aligns each score row with its own total. .T is appropriate here because K has exactly two dimensions.

import numpy as np

x = np.array([[1., 0.], [0., 1.], [1., 1.]])
wq = np.eye(2)
wk = np.array([[1., 1.], [0., 1.]])
wv = np.array([[1., 0.], [1., 2.]])
q, k, v = x @ wq, x @ wk, x @ wv
scores = (q @ k.T) / np.sqrt(q.shape[-1])
e = np.exp(scores - scores.max(axis=-1, keepdims=True))
weights = e / e.sum(axis=-1, keepdims=True)
output = weights @ v

The shared program code/part-iv/run.py compares every Q/K/V, dot-product, weight, and output element against data/part-iv/attention.json. Those expected values were established by the integer expansions and symbolic r expressions above, independently of the vectorized implementation. In the recorded CPU NumPy 2.0.2 run, integer products matched exactly; the maximum absolute differences were 0.000000000000000111 for weights and 0.000000000000000222 for outputs, below the declared absolute tolerance of 0.000000000001. These are arithmetic agreement measurements, not model-quality measurements. The full trace is data/part-iv/numpy-run.json; setup and rerun commands are in code/part-iv/README.md.

15.7 What this one head has accomplished

Each position can now build a query, compare it with the sequence’s keys, and receive a query-dependent mixture of values. The W matrices determine the three transformations; the attention weights are temporary results that depend on this input. We have explained the calculation, not trained those matrices.

This is still one attention head, not a complete Transformer block. It has no residual path, normalization, position mechanism, or per-position nonlinear network. More immediately, position 1 has already read positions 2 and 3. If its task is to predict position 2 before that token exists, our calculation has seen the answer. Chapter 16 will make allowed information part of the specification.

15.8 Exercises

1. Replace the third query with [0,0] while keeping K and V fixed. Compute its three weights and two output coordinates. Explain why a zero query does not generally produce a zero output.
2. A program returns a [3,3] matrix whose columns each sum to one. Is that sufficient to reproduce this head? Give a concrete column-normalized matrix whose first row is not a valid mixture.
3. Keep Q and K fixed and add [2,-1] to every V row. Compute the new first output using Table 15.1’s unrounded result. Would the same argument apply if only the first V row changed?
4. A head has Q/K shapes [5,4] and V shape [5,3]. State the scale divisor, score shape, output shape, and the axis normalized by Softmax. Why is division by the number 5 a different operation?
5. A report points to a large weight on the word red and concludes that the model will reliably follow red-key instructions. Use Example 15.2 to identify the missing evidence and propose a more informative controlled comparison.

15.9 Reference answers

1. All dot products are zero, so each exponential is 1 and the normalizer is 3. The weights are [1/3,1/3,1/3]. Averaging [1,0], [1,2], and [2,2] gives [4/3,4/3]. Zero similarity differences mean equal weighting; they do not erase the values.
2. No. The matrix with rows [1,1,1], [0,0,0], [0,0,0] has unit column sums but first-row sum 3. Multiplying V would give a sum, not a unit-total mixture, at receiver 1. The contract requires one normalized row per receiving query.
3. The first output becomes [3.401112092679786,0.1977758146404283]. The extra contribution is the common added vector times the sum of the three weights, which is 1. If only V row 1 changes, the extra contribution is [2,-1] times 0.4011120926797859; it is no longer the entire added vector.
4. The divisor is $\sqrt{4}=2$; scores have shape [5,5]; output has shape [5,3]. Each query row is normalized across its five key positions. Four is the number of terms in each dot product; five is the number of contributors. The scaling motivation concerns the former, not the latter.
5. The value projection and later computation may discard or transform the information. The zero-V counterexample has the same weights and no delivered information. Compare outputs on a fixed set of paired red/blue prompts, keeping the task and other wording controlled, and inspect actual success and failure. Even that establishes behavior only on the tested conditions; a heatmap alone establishes neither reliability nor a full causal explanation.

15.10 References

• Vaswani et al. (2017), Attention Is All You Need. Sections 3.2.1–3.2.3 support the scaled product, projected roles, and distinction between self- and cross-attention. The paper’s learned networks are separate from our authored matrices.
• NumPy, matmul, version 2.0. Documents the matrix and batched-matrix behavior used to verify the arithmetic. The practice record pins the installed patch version separately.