Attention

"Mia took her doll to the park, to the shop, to her grandma's house… At night, Mia put the…" Put the what? Any child will tell you: the doll. But "doll" came up two sentences ago, and to remember it you have to be able to look back over the whole story, not just the last few words.

The model from the previous chapter can't do that. It sees exactly the last eight tokens, and everything further left simply doesn't exist for it. In this chapter we give Sprout the power to look back, through a mechanism called attention. Every modern language model stands on this invention, from GPT to Llama, and there is nothing mysterious about it: dot products, softmax and a weighted average. You already know all three from earlier chapters.

The window is too narrow

Recall how the network from chapter 6 worked. It took the last 8 tokens, swapped each one for its embedding of 48 numbers, glued them into a single vector of length $8 \times 48 = 384$ and passed that through a hidden layer. This design has three problems.

First, everything outside the window is lost. If the hero's name appears at the start of the story, a couple of sentences later it slides out of the window, and the model has no way to use it. No amount of training can fix this: there is simply no input through which that word could get in. See for yourself:

Second, the window is expensive to widen. In the first layer, every hidden neuron has a separate weight for each combination of "slot in the window × number in the embedding". With 8 slots, 48-number embeddings and 384 hidden neurons, that makes $8 \times 48 \times 384 = 147{,}456$ weights. Stretch the window to Sprout's 512 tokens and you get $512 \times 48 \times 384 \approx 9.4$ million weights in the first layer alone, most of which would hardly ever get a chance to learn anything.

Third, nothing is shared. The word "ball" in slot 3 of the window and the same word in slot 7 go through different weights. Whatever the network learns about balls in one slot, it has to learn all over again for every other slot.

So we want a mechanism that (1) can reach any earlier token, however far back; (2) treats every position with the same weights; and (3) decides for itself where to look, depending on what the text says. That mechanism is attention.

Where Sprout looks

Before any formulas, let's look at the finished product. The full Sprout has 8 layers, and each layer has 6 independent attention "heads". For every token, each head hands out weights to all the tokens before it, the token itself included, saying how much to take from each. The weights are non-negative and add up to one. So this is a probability distribution again, as in the first chapters, only now it is spread over earlier positions in the text rather than over the vocabulary.

Play with the switches. A few things are worth noticing:

• The future is closed. Whichever token you pick, everything to its right is hatched: the model cannot look ahead. We'll see why below, and it matters.
• Each layer looks in its own way. Compare layer 1 and layer 5 on the same token. Hunt for heads that almost always look at the neighbour on their left, and for heads that reach far back. In the example about "he", pick layer 6 and head 4: this head gives 82% of the attention of "he" to "Tom".
• A lot of attention goes to ◆. Heads often give the start mark a sizeable share, even though it carries no meaning. We'll explain this oddity once we have seen softmax at work.

Query, key and value

Picture a classroom. The word "she" stands up and asks: "Which of you is a girl?" That is its query. Every earlier word has a card on its desk: "I'm a girl's name", "I'm a colour", "I'm something you can throw". Those are the keys. The better a card answers the question, the more closely "she" listens to that word. And what it hears is the word's value: what the word has to say about itself. In the end "she" walks away with a blend of everything it heard, made mostly of the stories whose cards matched best.

Why keep the key and the value apart? Because "how I can be found" and "what I have to say" are different questions. You find a library book by the title on its spine, but what you read is what's inside.

Now the same thing in vectors. Every token has a vector $x$: its embedding, or whatever that embedding has become in earlier layers. Three learned matrices turn it into three new vectors:

$$q = x\,W_Q, \qquad k = x\,W_K, \qquad v = x\,W_V.$$

These are three ordinary linear layers, like the ones in chapter 5. After that come just three steps. Say a token with query $q$ looks at tokens $1, \dots, t$, whose keys are $k_j$ and values $v_j$.

1. Scores. How well the query matches each key is measured by a dot product: $s_j = q \cdot k_j$. It is large when the two vectors point the same way, near zero when they are perpendicular and negative when they point away from each other.
2. Weights. Softmax, our old friend from chapter 3, turns the scores into shares: $a_j = e^{s_j} / \sum_i e^{s_i}$. Every $a_j$ is positive, and together they add up to 1.
3. Blend. The output is the weighted average of the values: $o = \sum_j a_j\, v_j$.

That's all there is to attention. Try it on a flat plane: here queries and keys have two numbers each, so that we can draw them, and values have three, which we draw as a colour (a colour on a screen is exactly three numbers: red, green and blue).

Some experiments to try:

• A zero query. Every score is zero, so the weights come out equal: 25% each. A token with "nothing to ask" simply averages everyone.
• A long query. Same direction, longer vector: the scores grow and softmax gets sharper, until one key takes almost all the attention. The query's length works like an inverse temperature, the same stretching of the scores that sharpened softmax in chapter 3. The "Sharpness" slider does the same thing directly.
• Every key points away. Press "Nobody fits": the keys gather in one half of the plane and the query turns away from all of them, so every score is negative. The weights still add up to 1, and the least bad key wins. Attention cannot "look nowhere".

That last point solves the ◆ puzzle from the previous section. A head that has nothing to look for at this step still has to spend its 100% somewhere. The most convenient place is one fixed, harmless token whose value adds almost nothing to the blend. The start mark is perfect for the job: it is in every text, and always in the same place.

Here is the same arithmetic in Python, using the numbers from the widget's starting position:

import numpy as np # four earlier tokens: each has a key (how it can be found) # and a value (what it has to say); values are colours, i.e. 3 numbers tokens = ["girl", "ball", "red", "ran"] K = np.array([[ 1.6, 0.9], [-1.2, 1.3], [-1.5, -0.6], [ 0.4, -1.6]]) V = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [0.5, 0.5, 0.5]]) q = np.array([1.4, 0.7]) # the query of "she" scores = K @ q # one dot product per earlier token weights = np.exp(scores - scores.max()) weights /= weights.sum() # softmax: positive, adds up to 1 out = weights @ V # the weighted average of the values for tok, s, w in zip(tokens, scores, weights): print(f"{tok:5s} score {s:+.2f} weight {w:.3f}") print("output:", out.round(3))

A token has compared its query with four keys, and every score came out negative: −3, −5, −6, −9. What does softmax do?

Every weight becomes zero: no token matches The weights still add up to 1, and the token that scored −3 gets the biggest share The weights turn negative, and the token "pushes away" those values

Softmax only compares the scores with one another: $e^{-3}$ is far bigger than $e^{-9}$, and after dividing by the sum the shares are always positive and add up to 1. That is why heads need somewhere to park their attention, like the start mark.

Why divide by √d

In the widget the vectors are two-dimensional. In Sprout, each head's queries and keys have 64 numbers, so their dot product is a sum of 64 products. And the more terms a sum has, the wider it spreads.

Why is that bad? Remember chapter 3: with large logits softmax becomes almost "hard", with one weight close to 1 and the rest close to 0. And a hard softmax has almost no gradient: nudge the scores a little and the weights don't move. Attention like that can barely learn, right from the first step. The cure is simple: divide the scores by $\sqrt d$ so that their spread is about one again:

$$s_j = \frac{q \cdot k_j}{\sqrt{d}}.$$

No peeking at the future

How does the model learn? We give it a stretch of $T$ tokens and ask it to guess the next token at every position at once: $T$ predictions together, with their average as the loss (chapter 2). That's a bargain: one pass over the text yields $T$ training examples.

But there is a trap. If the token at position $t$ could look at position $t+1$, it would see the very answer it is supposed to guess. The model would quickly learn to copy it, and then be useless once it writes on its own, when the future doesn't exist yet and there is nothing to copy. So before softmax every score "from the future" is set to $-\infty$:

$$s_{tj} = -\infty \quad \text{for } j > t, \qquad e^{-\infty} = 0.$$

Those positions get a weight of exactly 0, and the remaining weights in the row are renormalised so that they add up to 1 again. This is the causal mask: every token sees itself and the past, never the future. Build it yourself:

The lower half of the widget shows what attention costs. Every token compares itself with every earlier token, so there are about $T^2/2$ comparisons: double the length of the text and the work quadruples. At Sprout's full context of 512 tokens, a single head computes 262,144 scores (half of them masked straight away), and all 48 heads together more than 12 million per pass. That is why long contexts are so expensive for large models.

All the tokens at once

So far we have looked at one token at a time. Going one by one is wasteful: the queries, keys and values of all $T$ tokens can be stacked into matrices $Q$, $K$ and $V$ (one row per token), and then all the scores come out of a single multiplication. Entry $(t, j)$ of $QK^\top$ is exactly $q_t \cdot k_j$. The whole chapter fits into one formula:

$$\operatorname{Attention}(Q, K, V) = \operatorname{softmax}\!\left(\frac{QK^\top}{\sqrt d} + M\right) V,$$

where $M$ is the mask (zeros on and below the diagonal, $-\infty$ above it) and softmax is taken along each row separately. As for sizes, $Q$, $K$ and $V$ are $T \times d$, the table of scores is $T \times T$, and the result is $T \times d$ again, one vector per token.

Why does this matter? As chapter 5 showed, GPUs (and CPUs too) are at their best multiplying big matrices. In matrix form, attention is two matrix multiplications and a softmax, exactly what the hardware does fastest. That is a large part of why transformers pushed aside recurrent networks, which read text strictly one word at a time.

Several heads

One query is one question. But a token usually needs to know several things at once: who the main character is, which word came just before it, where the sentence started, which object was mentioned. A single head, a single weighted average, cannot look in several places for different reasons.

The fix is multi-head attention. A token's vector in Sprout has 384 numbers. We cut the query, key and value into 6 pieces of 64 numbers each, and each piece works as a little head of its own: it computes its own scores, applies the mask and takes its own softmax. The six heads run in parallel, their outputs are glued back into a vector of 384 numbers, and a final matrix $W_O$ mixes them together:

$$\operatorname{MultiHead}(x) = \big[\,o^{(1)}, o^{(2)}, \dots, o^{(6)}\,\big]\, W_O.$$

Splitting into heads costs no extra parameters: $W_Q$, $W_K$ and $W_V$ are still $384 \times 384$; we just group their columns differently. Together with $W_O$ that makes $4 \times 384^2 = 589{,}824$ parameters per layer. Here are all 48 of Sprout's heads at once:

See if you can find three kinds of head:

• Previous token: a head that almost always looks at the token just before. On the map it is a bright stripe just below the diagonal. A head like this tells each token who stood in front of it.
• First token: a head that parks its attention on ◆. On the map it is a bright column on the left. Most likely it is one of the heads that often have nothing to look for.
• Induction: pick the example "A made-up name". The dragon is called Zorbin Plax, a name Sprout has never met, so the first time round there is no guessing it. But the second time, on reaching "Z", the model can look at what came after "Z" last time and copy it. An induction head lights up cells far to the left of the diagonal, where the name first appeared. At the bottom of the widget you can see how much better Sprout guesses pieces of text it has already seen.

Let's put it all together in numpy: matrices, the mask and heads. The function below is complete multi-head causal attention, just untrained; at the end we check that the matrix version agrees with a by-hand calculation for one token.

import numpy as np def softmax(x): x = x - x.max(axis=-1, keepdims=True) e = np.exp(x) return e / e.sum(axis=-1, keepdims=True) def attention(x, Wq, Wk, Wv, n_head): T, C = x.shape d = C // n_head q, k, v = x @ Wq, x @ Wk, x @ Wv # (T, C) each # cut the C columns into heads: (n_head, T, d) q, k, v = (a.reshape(T, n_head, d).transpose(1, 0, 2) for a in (q, k, v)) scores = q @ k.transpose(0, 2, 1) / np.sqrt(d) # (n_head, T, T) future = np.triu(np.ones((T, T), dtype=bool), k=1) # above the diagonal scores[:, future] = -np.inf # no peeking w = softmax(scores) # every row adds up to 1 y = w @ v # (n_head, T, d) return y.transpose(1, 0, 2).reshape(T, C), w # glue the heads back rng = np.random.default_rng(0) T, C, H = 5, 12, 3 x = rng.standard_normal((T, C)) Wq, Wk, Wv = (rng.standard_normal((C, C)) / np.sqrt(C) for _ in range(3)) y, w = attention(x, Wq, Wk, Wv, H) print("output:", y.shape) print("head 0, weights:\n", w[0].round(2)) # the same row by hand: token 3 in head 0 d = C // H q3 = (x[3] @ Wq)[:d] k0 = (x[:4] @ Wk)[:, :d] s = k0 @ q3 / np.sqrt(d) print("by hand matches:", np.allclose(np.exp(s) / np.exp(s).sum(), w[0, 3, :4]))

Look at the upper triangle of the printed matrix: all zeros. That is the mask at work. And every row adds up to 1.

Attention in model.py

And here is the same thing in the real Sprout: the class from the file the model was trained with, without a single change.

class Attention(nn.Module): def __init__(self, cfg): super().__init__() self.n_head = cfg.n_head self.head_dim = cfg.d_model // cfg.n_head self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False) self.proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False) self.q_norm = RMSNorm(self.head_dim) self.k_norm = RMSNorm(self.head_dim) def forward(self, x, cos, sin): B, T, C = x.shape q, k, v = self.qkv(x).split(C, dim=-1) # (B, T, C) -> (B, heads, T, head_dim) q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2) k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2) v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2) q, k = self.q_norm(q), self.k_norm(k) q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin) # softmax(q·k / sqrt(d)) · v, every token sees only itself and the past y = F.scaled_dot_product_attention(q, k, v, is_causal=True) y = y.transpose(1, 2).contiguous().view(B, T, C) return self.proj(y)

Line by line:

• self.qkv is $W_Q$, $W_K$ and $W_V$ glued into a single $384 \times 1152$ matrix, because one big multiplication is faster than three small ones. .split(C, dim=-1) cuts the result back into three parts of 384.
• B, T, C are the batch size (how many texts we process at once), the length of the text and the width of the vector, 384. Texts in a batch never see each other: a batch is just parallel work.
• .view(B, T, n_head, head_dim).transpose(1, 2) makes the same 6 pieces of 64 as the numpy version: the heads become one more "batch" dimension and from here on are computed in parallel.
• q_norm and k_norm are the QK-norm from the note above.
• apply_rope adds information about position to the queries and keys. That is the next chapter's subject, and in a moment you'll see why it is needed.
• F.scaled_dot_product_attention(..., is_causal=True) is the whole formula $\operatorname{softmax}(QK^\top/\sqrt d + M)V$ in one call. PyTorch divides by $\sqrt d$ and applies the mask itself, and on a GPU it picks a fast algorithm such as FlashAttention.
• transpose + view glue the heads back into 384 numbers, and self.proj is $W_O$.

This mechanism has one surprising blind spot. Look at the formula $o = \sum_j a_j v_j$ once more and try this question.

Take the attention from this chapter on its own, without apply_rope. Shuffle all the tokens before the current one, leaving the current token where it is. Does its output change?

Yes: word order changes the meaning No: the weights depend only on what the query and keys contain, and a sum doesn't care about the order of its terms Yes, because the mask now hides different positions

The scores $q\cdot k_j$ depend on what the vectors contain, and a weighted sum is the same in any order. The mask doesn't help either: the set of past tokens hasn't changed. To attention like this, "the dog bit the boy" and "the boy bit the dog" are the same bag of words. So a transformer needs a separate way to tell positions apart; in Sprout that is RoPE, the rotations of the next chapter.

Attention as a soft dictionary lookup

Python has dictionaries: d[key] finds the one entry with exactly that key and returns its value. Attention is a "soft" version of the same thing: instead of an exact match, a degree of similarity; instead of one entry, a blend of all of them, weighted by that similarity. If softmax were infinitely sharp, attention would turn into an ordinary lookup of the most similar key. The softness is there for the sake of learning: a hard choice has no gradient and a weighted blend does, so the model can adjust its queries and keys a little at a time.

Sprout right now

We have replaced the eight-token window with attention. This model is already a real little transformer: embeddings, one attention layer with four heads, and a small network after it (the next chapter explains what that does). It can see up to 256 tokens back, so a name from the start of a story no longer falls off the edge of a window. On held-out text its loss is 2.87 nats per token on average, against 3.31 for the window network from chapter 6. In other words, $e^{3.31} \approx 27$ while $e^{2.87} \approx 18$: where the old model hesitated between twenty-seven candidates for the next token, this one hesitates between eighteen.

But there is only one attention layer, so in a single pass information can hop from token to token only once. The model cannot first work out "who stood before whom" and then use it, the way an induction head does. We checked with the dragon Zorbin Plax: meeting "Z" for the second time, gpt-1 gives the continuation "or" a probability below one percent. It does already know about positions, by the way: that is the apply_rope line, whose workings we haven't looked at yet. In the next chapter we build the full transformer block, work out the rotating positions and stack blocks on top of each other.