How a token reads its context: queries, keys, values, and a weighted sum.
Chapter 2 gave each token a vector and a way to transform it, but every token was processed alone. The embedding for "it" is the same row of W_E in every sentence ever written, and the MLP transforms it the same way every time. Nothing so far lets one token's vector depend on the other tokens. This chapter adds that, and it is the single mechanism the whole transformer is named after. Chapter 4 wraps it into a block; Chapter 5 shows how its matrices are learned.
The question this chapter answers: how does a token gather information from the tokens before it, and how is "which tokens" decided?
A conference reception. Everyone wears a name badge, and on the badge, under the name, is a line about what they can help with: "compilers", "tax law", "knows where the good coffee is". Everyone also carries something to give away: a business card, a tip, a story.
You walk in with a question in your head. Not a sentence, a need: "I need someone who knows compilers". You do not read every badge with equal care. You scan them, and the ones that match your need light up. You spend most of your evening with the compiler people, some with the person whose badge said "linkers", almost none with tax law. By the end you have absorbed what they gave you, in proportion to how much time you spent with each. You leave changed by the room, and the change is a weighted blend of what the people you attended to had to offer.
Three things about this scene map exactly onto the mechanism. Your need and their badge lines are different kinds of thing, even though they are about the same topic; matching happens between the two. What a person gives you is a third thing, different from their badge. And the blend is weighted, not a vote: you can take 60% from one person and 5% from another.
| At the reception | In the machine | The word we will use |
|---|---|---|
| Your need, as you walk in | The token's vector projected by a learned matrix into "what I am looking for" | query q = x W_Q |
| The line on each badge | Each token's vector projected by another matrix into "what I offer" | key k = x W_K |
| What each person hands you | Each token's vector projected by a third matrix into "what I pass on" | value v = x W_V |
| How well a badge matches your need | The dot product of your query with their key | score q·k |
| How you divide your evening | Softmax over the scores: positive weights summing to 1 | attention weights |
| What you absorbed | The weighted sum of the values | output Σ wⱼ vⱼ |
| You can only talk to people who arrived before you | Positions after the query are excluded | causal mask |
| Everyone is simultaneously a guest and a badge-wearer | Every token is a query and a key and a value, all at once | self-attention |
Two sentences:
The cat sat on the mat because it was tired.
The cat sat on the mat because it was soft.
The token "it" refers to the cat in the first and to the mat in the second. Yet after the embedding lookup of Chapter 2, both occurrences hold the identical vector, row W_E["it"]. The MLP will transform that vector identically too, because it only ever sees one token at a time. To predict "tired" versus "soft", or to do anything that depends on what "it" refers to, the model needs a way to move information from "cat" or "mat" into the vector at position "it". And which one to move from cannot be fixed in advance; it depends on the sentence.
So the requirement is precise. We need an operation that, for each token, (1) decides which earlier tokens are relevant, based on the vectors themselves, and (2) copies something from those tokens into this one, in proportion to relevance. Attention is that operation, and the elegant part is that both "which" and "something" are computed with the tools of Chapter 2: dot products and matrices.
Start with the token vectors stacked as rows: X of shape [T, C]. Attention multiplies X by three learned matrices to get three new sets of vectors:
Q = X · W_Q W_Q : [C, d] queries "what am I looking for?" K = X · W_K W_K : [C, d] keys "what do I offer to those looking?" V = X · W_V W_V : [C, d] values "what do I hand over if chosen?" each of Q, K, V has shape [T, d]; row t belongs to token t
Why three matrices and not one? Because the three jobs are different, and forcing them through one representation would make them fight. What "it" looks for (a recent noun) is not what "it" offers to later tokens (a pronoun, probably the subject) and neither is what it should hand over if a later token attends to it. Three matrices let training tune each role independently. Each is a change of viewpoint in the sense of Chapter 2: the same 4096-dimensional token, re-expressed for a purpose.
Note the width d. It is smaller than C: Llama 3 uses d = 128 against C = 4096 or 16384 public. A single attention head works in a narrow slice of the representation. Chapter 4 explains why that is fine: there are many heads, each with its own slice.
Now the reception scan. Each query is dotted with every key. For token i and token j, the score is qᵢ · kⱼ: large when what i seeks aligns with what j offers. Doing this for all pairs at once is a single matrix multiplication:
S = Q · Kᵀ [T, d] × [d, T] → [T, T] S[i, j] = qᵢ · kⱼ = how much token i wants to look at token j row i = token i's scores over every key
Read the shape. T × T: one number for every ordered pair of tokens. For a 4-token sentence that is 16 numbers. For a 128k-token context it is 16 billion per head per layer, and that single fact drives much of Part 2's engineering. For now, four tokens.
Raw scores are not weights yet. Three steps turn a row of scores into a row of weights that sum to one and only look backwards.
Scores are dot products of length-d vectors. If the entries of q and k have spread about 1, the dot product has spread about √d: it is a sum of d terms, and sums of random terms grow like the square root of their count. With d = 128 the raw scores would routinely differ by 10 or 20, and the softmax of Chapter 1 would turn that into a near-certain choice of one key before any learning has happened. Dividing every score by √d restores the spread to about 1, so the softmax starts soft and the model can learn to sharpen it where sharpness is warranted. This is the "scaled" in scaled dot-product attention, and it is a fix for a numerical fact, not a modelling idea.
Each row of the scaled score matrix goes through the softmax from Chapter 1: exponentiate, divide by the row's sum. Now each row is a set of positive weights summing to 1: token i's attention distribution over the keys. Exactly the same function that turns logits into next-token probabilities, used here to turn "how relevant" into "what fraction of my attention". The properties carry over: order is preserved, big gaps are amplified, nothing is ever exactly zero.
One more rule, and it comes from the game in Chapter 1. When the model predicts the token after position i, it must not have seen positions i+1, i+2, …. Otherwise the training is a cheat: "predict the next token" becomes "read the next token". So before the softmax, every score where the key is after the query is replaced by −∞. Exponentiate −∞ and you get exactly 0: those positions receive no weight and contribute nothing. Everything above the diagonal of the [T, T] matrix is switched off.
The mask has a consequence that is easy to miss and central to Chapter 7: with it in place, every position's prediction can be computed at the same time from one pass over the text, because position i's output provably used nothing after i. Training runs all T next-token games of a sequence in parallel. Without the mask that would be impossible.
The last step is the one that actually moves information. Each token's output is the sum of all value vectors, each multiplied by that token's weight on it:
out_i = Σ_j W[i, j] · v_j for all rows at once: Out = W · V [T, T] × [T, d] → [T, d]
For "it", with weights [0.13, 0.56, 0.14, 0.18] over [The, cat, sat, it]: the output is 13% of The's value, 56% of cat's, 14% of sat's, 18% of its own. It is now a vector that is mostly what "cat" hands over. That is the information movement the chapter set out to build: the vector at position "it" now carries content that came from position "cat", chosen because their query and key matched, in a proportion set by softmax.
Geometrically, in the language of Chapter 2: the output is a point inside the shape spanned by the value vectors, pulled toward the ones with large weight. Attention cannot invent a value that no token offered; it can only blend what is there. New content is the MLP's job. Attention's job is routing.
Four tokens, C = 3, one head of width d = 2. The token vectors are deliberately simple so the projections are readable. The numbers below are the same ones the widget starts from and the same ones code/ch03/one_head_by_hand.py prints.
X (T=4, C=3) W_Q (3×2) W_K (3×2) W_V (3×2)
The [1.0, 0.0, 0.0] ┌ 0 1 ┐ ┌ 0.5 0 ┐ ┌ 1 0 ┐
cat [0.0, 1.0, 0.0] │ 2 0 │ │ 4 0 │ │ 0 1 │
sat [0.0, 0.0, 1.0] └ 0 1 ┘ └ 0 2 ┘ └ 0.5 0.5 ┘
it [0.2, 0.3, 0.0]
step 1 project Q = X·W_Q K = X·W_K V = X·W_V
The [0.0, 1.0] [0.5, 0.0] [1.0, 0.0]
cat [2.0, 0.0] [4.0, 0.0] [0.0, 1.0]
sat [0.0, 1.0] [0.0, 2.0] [0.5, 0.5]
it [0.6, 0.2] [1.3, 0.0] [0.2, 0.3]
step 2 scores for "it" q_it · k_j = [0.30, 2.40, 0.40, 0.78]
step 3 scale by √2 ÷ 1.414 = [0.21, 1.70, 0.28, 0.55]
step 4 mask nothing after "it", so unchanged
step 5 softmax exp → [1.24, 5.46, 1.33, 1.74], sum 9.76
weights = [0.13, 0.56, 0.14, 0.18]
step 6 blend values 0.13·[1,0] + 0.56·[0,1] + 0.14·[0.5,0.5] + 0.18·[0.2,0.3]
= [0.13 + 0 + 0.07 + 0.04, 0 + 0.56 + 0.07 + 0.05]
= [0.23, 0.68]
Look at what W_Q and W_K did. "it"'s vector has a little weight on the second coordinate (the "noun-like" one, 0.3). W_Q amplifies that coordinate ×2 into the first query dimension, and W_K amplifies it ×4 into the first key dimension. So a token that is noun-like seeks along dimension 1, and a token that is a noun advertises along dimension 1. The match between "it" and "cat" is manufactured by those two matrices. In a trained model nobody designs this; the matrices are pushed there by the loss, because sentences where "it" fetches its referent are easier to continue.
$ python code/ch03/one_head_by_hand.pyweights (softmax per row) = [[1. 0. 0. 0. ] [0.007 0.993 0. 0. ] [0.164 0.164 0.673 0. ] [0.127 0.559 0.136 0.178]] output = weights · V = [[1. 0. ] [0.007 0.993] [0.5 0.5 ] [0.23 0.681]]
Now do it live. Change what "it" is like and watch where its attention goes; pick other query tokens; switch the mask off and see the future leak in.
Here is one head, end to end, with every shape. This is the picture to memorise; everything in Chapter 4 is this picture repeated and wrapped.
Nothing in the mechanism says what a head should attend to. W_Q and W_K are learned, so the matching rule is learned; W_V is learned, so what gets moved is learned. Training discovers whatever routing lowers the next-token loss. When researchers look inside trained models, the heads they find are surprisingly legible public:
i−1. They give every token a copy of its predecessor, which is the raw material for detecting pairs and phrases.The reception analogy has a limit worth naming here. Real heads are not people with intentions. A head is two small matrices and a softmax, and "previous-token head" is a description of what its weights happen to do, discovered after the fact. Chapter 26 shows how such descriptions are found and how far they can be trusted.
| Quantity | Llama 3 8B | Llama 3.1 405B | Evidence |
|---|---|---|---|
Head width d | 128 | 128 | public |
| Query heads per layer | 32 | 128 | public (so heads × d = C; Chapter 4) |
| Key/value heads per layer | 8 | 8 | public (fewer than query heads: grouped-query attention, Chapter 12) |
| Context length at release | 8k → 128k | 128k | public |
| Closed frontier models | Head counts and widths not disclosed; contexts of 200k to 1M tokens advertised | public for contexts, unknown for internals | |
How big is the score matrix? One head, one layer, in a long context.
context T 128,000 tokens score matrix entries T² = 1.6 × 10¹⁰ per head, per layer bytes at 2 bytes/entry 3.3 × 10¹⁰ ≈ 33 GB per head, per layer × 32 heads ≈ 1 TB per layer, one forward pass compare: a top GPU holds 80–192 GB of memory in total
Materialising the full weight matrix for a long context does not fit on any GPU, by a factor of ten per layer. Yet 128k-token contexts are served every day. The resolution is that nobody materialises it: the scores are computed in tiles and consumed by the softmax and the value blend before the next tile is produced, so only a tile ever exists in memory. That technique is FlashAttention, and it is Chapter 10's centrepiece. The arithmetic does not go away (it is still 2·T²·d operations per head), but the memory does.
And the cost per token during generation? When the model produces token T+1, its one new query is scored against T keys: T·d multiplies per head, then a blend of T values. Linear in T, per new token. That is why long conversations get slower and more expensive as they grow, and why the keys and values of past tokens are worth keeping around rather than recomputing (the KV cache, Chapter 7).
Remove the scaling. With d = 128, raw scores have spread about 11. Softmax of scores that differ by 11 is one-hot: each token attends to a single key, chosen almost at random at initialisation. Chapter 5's gradients through a saturated softmax are near zero, so the head barely learns. Training is slower and less stable; some heads never recover. The fix costs one division.
Use the same matrix for queries and keys. Then a token's score against itself, q·q, is its squared length: always the largest possible match for a token of that length, so every token attends mostly to itself. The "what I seek" and "what I offer" roles collapse, and the asymmetric relations language needs (a verb seeking its subject is not a subject seeking its verb) cannot be expressed. Models with tied Q and K exist as experiments and are worse inferred from published ablations.
Drop the softmax and use raw scores as weights. Weights can now be negative and unbounded; the output is no longer a blend inside the span of the values but an arbitrary combination whose scale grows with T. Longer contexts produce larger outputs for no reason, and the downstream layers see inputs whose size depends on position. Softmax makes attention scale-free in T: whether there are 4 keys or 100k, the output is an average.
Remove the causal mask. During training, the token at position i can attend to position i+1, which is exactly the token it is supposed to predict. Loss falls to nearly zero in a few steps, and the model has learned to read one position ahead. At inference there is no position ahead, and the model is useless. The mask is the only thing that keeps the training game honest.
Replace the weighted sum with "take the top-1 value". A hard choice: no blending, and no gradient signal about the keys that were nearly chosen. The model cannot represent "60% cat, 30% mat" and cannot learn smoothly. Soft weights are what make the routing trainable.
Say it back. The embedding gives every token the same vector regardless of context, and the MLP treats each token alone; something must move information between positions, and which positions must depend on content. Attention does this with three learned viewpoints on each token vector: a query (what I seek), a key (what I offer), a value (what I hand over). Every query is dotted with every key to get a T × T score matrix; the scores are divided by √d so their spread does not depend on head width; scores against future positions are set to −∞; each row is softmaxed into weights that sum to one. Each token's output is the weighted sum of the values, a blend of what the tokens it attended to handed over, pulled toward the closest matches. Attention routes; it does not create. Which routes exist is learned through W_Q and W_K, and trained heads turn out to implement legible jobs like "copy the previous token" and "find what followed this token last time". The T² in the score matrix is the cost of letting every token consider every other, and it shapes the engineering of long contexts.
x_it = [0.0, 0.0, 1.0] (make "it" look exactly like "sat"). Recompute its query, its four scores, the scaled scores, the softmax weights, and the output. Check with the widget. Which token does it now attend to, and why does the answer follow from W_Q and W_K rather than from W_V?Q·Kᵀ step cost across the whole model for one forward pass? How does that compare with the projections X·W_Q, X·W_K, X·W_V at C = 16384? At what context length do the two costs cross?one_head_by_hand.py to remove the /√d scaling and multiply all of X by 4. Print the weight matrix. Then restore the scaling and print again. Explain in two sentences why the first matrix is nearly one-hot and what that would do to learning.