Chapter 3 — Attention: A Deep Dive with Intuition

By the end of this chapter you can compute attention by hand, derive why it divides by \sqrt{d_k}, mask it correctly, backpropagate through it, count its FLOPs and memory for a real model, size a KV cache, and write it in numpy and PyTorch without the bugs that bite in production. LLM Ch 2 turned text into tokens and tokens into vectors, and left one problem open: the vector for bank is the same in "the river bank was steep" and "the bank raised rates". This chapter is the fix. One idea carries all of it: each token's new vector is a weighted average of the other tokens' values, with weights from a softmax over how well its query matches their keys. Attention is a learned, soft dictionary lookup, and everything else in this chapter keeps that average trainable (\sqrt{d_k}), legal (masks), expressive (heads) and affordable (the KV cache).

How to read it. The chapter is a chain of 15 questions in three acts. Build it (Sections 1–6): where attention came from, what it fixes, and how to compute it by hand. Understand it (Sections 7–11): what it computes, how it learns, what it cannot see, why it has heads, and how it fails. Pay for it (Sections 12–15): FLOPs, memory, the KV cache, code, and what modern models run. Each section opens with its question and a 30-second version, and ends with the question it leaves open.

Two passes. On the first pass, read only the question, the 30-second version and the figures of each section. You'll follow the whole arc in under an hour. On the second pass, read everything, doing the Remember box at the end of each section from memory before you look. Practice lives in its own tabs: Derivations (what to prove on a whiteboard) and Interview Bank (every question by level, with the intuition behind each answer).

The running example. One five-token sentence, small enough to check every number by hand: "the river bank was steep". Each token has a 4-number embedding with named dimensions, [nature, money, function word, describes], so you can watch meaning move. For anything that needs scale we use the published shapes of GPT-2 small, Llama 2 7B and 70B, and Llama 3 8B.

1. How We Got Here: From a Bottleneck to a Lookup

The question: why does attention exist at all? What was broken before it, and why has every large language model since 2017 been built on it?

The 30-second version: in 2014, neural translators squeezed a whole sentence into one vector and translated from that vector alone, and they got worse as sentences got longer. Bahdanau, Cho and Bengio (2014) let the decoder look back at every input word and take a weighted average, with weights chosen by content. Luong et al. (2015) simplified the scoring to a dot product. Vaswani et al. (2017) threw out the recurrent network and built the whole model from attention. The price was a T \times T table of scores and a cache that grows with every token, and the years since have been spent paying it down: shared keys and values (MQA, 2019; GQA, 2023) and exact attention without the table (FlashAttention, 2022). Hook: Squeeze, Look, Only look, Share, Stream. First read: all of it.

The problem: one vector for a whole sentence

In 2014 the best neural translators were sequence-to-sequence models (Sutskever, Vinyals and Le, 2014; Cho et al., 2014). An encoder, a recurrent neural network (RNN), read the English one word at a time and updated a running summary vector. Its final summary was handed to a decoder RNN, which wrote the French one word at a time.

Picture an interpreter who must hear a whole speech, write a one-line summary on an index card, throw the speech away, and translate from the card alone. For "Hello, how are you?" the card is enough. For a forty-word sentence with a subordinate clause, it isn't. That is the bottleneck: a fixed-size vector must hold a sentence of any length. Cho et al. (2014, "On the Properties of Neural Machine Translation") measured the consequence: quality fell as sentences got longer.

The fatal flaw: the decoder can only read the card. It can never go back to the speech.

Step 1 · Bahdanau, Cho and Bengio (2014): let the decoder look back

The idea. The paper went on arXiv in 2014 and appeared at ICLR 2015. Keep the encoder's state after every input word. At each output step, score every input position against what the decoder is about to write, turn the scores into weights with a softmax (positive, summing to 1; Section 4 teaches it from zero), and read a weighted average of the encoder states: a fresh context vector for every output word.

e_{tj}

the score: how relevant input word j is when writing output word t

t

which output (French) word is being written

j

which input (English) word is being scored

\mathbf s_{t-1}

the decoder's state just before writing word t (bold = vector)

\mathbf h_j

the encoder's state at input word j: that word, read in context

W_a

a learned matrix mapping the decoder state into a shared comparison space

U_a

a learned matrix mapping each encoder state into the same space

\mathbf v_a^\top

a learned vector, transposed (rows become columns), that turns the comparison into one number

\alpha_{tj}

the attention weight on input word j; the weights over j sum to 1

\sum_{k}

add up over every input word k, to normalize

e_{tk}

the score of input word k

k

an index over all input words

\mathbf c_t

the context vector: a weighted average of all encoder states, rebuilt for each output word

\sum_{j}

add up over every input word j

The intuition. The interpreter keeps the full transcript and, before saying each French word, glances back at the English words that matter for it. The paper called the weights an alignment; plotted, they showed the soft word-to-word correspondence a human translator would draw. Because the two states are added inside a small network, this is called additive attention.

The fatal flaw. Attention was an add-on to an RNN. Step t still waited for step t - 1, so the time dimension stayed serial on hardware that wants to do thousands of things at once.

Step 2 · Luong, Pham and Manning (2015): simpler scores

The idea. Luong et al. compared scoring functions and found cheap multiplicative scores work well: the plain dot product of decoder and encoder states, or a "general" form with a learned matrix in the middle. Factor that matrix into two pieces and you have \mathbf q \cdot \mathbf k with separate query and key maps: exactly Section 3. They also split global attention (look at every source word) from local (a window around a predicted position), an early version of the "which pairs are allowed" question that masks answer in Section 6.

The fatal flaw. Still an RNN underneath. Better scores did not make it any more parallel.

Step 3 · Vaswani et al. (2017): attention is the whole model

The idea. "Attention Is All You Need" removed recurrence entirely. Every token makes a query, a key and a value; every token attends to every token in one matrix multiply; scores are divided by \sqrt{d_k}; several heads run in parallel; and order is added as a separate signal, because attention alone cannot see it (Section 9). Stack that with a small per-token network and you have the Transformer:

Q

the queries: one row per token, saying what that token is looking for

K

the keys: one row per token, saying what that token offers to be found by

K^\top

K transposed (rows become columns), so QK^\top scores every query against every key

V

the values: one row per token, the content it hands over when attended to

d_k

the length of each query and key vector

The intuition. Stop passing a note down a line of people. Put everyone in one room and let each person ask the whole room at once. Any two words are one step apart, and all words are computed simultaneously.

The fatal flaw. "Everyone asks everyone" is T^2 pairs: compute and memory grow with the square of the context (Section 12). And generation, one token at a time, must re-read every earlier token's keys and values, so a KV cache grows with every token and layer (Section 13).

Step 4 · Shazeer (2019): Multi-Query Attention

The idea. "Fast Transformer Decoding: One Write-Head is All You Need" observed that generation is limited by memory traffic, not arithmetic: each new token reloads the whole cache. So let all query heads share one key head and one value head. The cache shrinks by the number of heads. Picture thirty-two reporters sharing one filing cabinet of interviews, each still asking their own question.

The fatal flaw. One cabinet is a squeeze: Ainslie et al. (2023) report MQA can lose quality and destabilize training.

Step 5 · Dao et al. (2022): FlashAttention

The idea. The T \times T table never has to exist. Process keys in blocks, keep a running maximum and a running sum per query, and correct the partial answer whenever the maximum grows. The result is exactly softmax attention with memory linear in T. The math is the online softmax (Milakov and Gimelshein, 2018; Rabe and Staats, 2021); Dao et al. made it fast by keeping blocks in the GPU's small on-chip memory, followed by FlashAttention-2 (Dao, 2023) and -3 (Shah et al., 2024). To average a river you don't store the river: keep a running total, and fix it up if you change units mid-way.

The fatal flaw. It saves memory and memory traffic, not arithmetic: FLOPs are still O(T^2).

Step 6 · Ainslie et al. (2023): Grouped-Query Attention

The idea. The middle ground between one KV head (MQA) and one per query head (multi-head attention, MHA): split the query heads into groups, one key/value head per group. Ainslie et al. also showed how to convert an MHA model: mean-pool its key and value heads within each group, then continue training ("uptraining") with about 5% of the original pre-training compute. Llama 2 70B (Touvron et al., 2023) uses 8 KV heads for 64 query heads: eight cabinets, one per team of eight reporters.

What it couldn't do. The cache still grows linearly with context and layers; at 128K tokens it is measured in tens of GiB (Section 13).

Where it stands today

The core hasn't changed since 2017: softmax of scaled dot products, then a weighted average of values. Everything around it has. Sharing K and V is standard: Llama 3 (2024) uses 8 KV heads at every size, Mistral 7B (Jiang et al., 2023) adds a 4,096-token sliding window, and DeepSeek-V2 (DeepSeek-AI, 2024) compresses keys and values into a small latent vector (multi-head latent attention). Exact fused kernels are the default: PyTorch's F.scaled_dot_product_attention dispatches to FlashAttention-style kernels when it can. Bounded logits keep training stable: QK-norm and Gemma 2's soft-capping (Section 11). Linear-time alternatives (linear attention, state-space models) exist and are LLM Ch 15's subject.

The through-line: every step since 2014 keeps the same move, look things up by content and average what you find, and spends its effort making that lookup cheaper to train, store and serve.

Remember: Squeeze, Look, Only look, Share, Stream. Seq2seq squeezed a sentence into one vector; Bahdanau looked back; Vaswani made looking the whole model; MQA/GQA share keys and values; FlashAttention streams through them without the T^2 table.

Next question → Attention is a lookup. But what exactly is wrong with a token's vector before it looks around? (Section 2)

2. Why "Bank" Needs to Look Around

The question: in "the river bank was steep", bank enters the model as [1, 1, 0, 0]: half "nature", half "money", the same vector it has in "the bank raised rates". How can its vector come to mean riverbank?

The 30-second version: an embedding table gives each token one vector, whatever its context. To pick a sense, a token has to read its neighbours. A fixed window sees only a few words; a recurrent network squeezes the prefix into one vector and can't run in parallel over time. Attention lets every token build its new vector as a weighted average of other tokens' vectors, with weights chosen by content. Hook: any distance, all at once, by content. First read: all of it.

The problem

Recap of LLM Ch 2 in one line: a tokenizer splits text into tokens, and an embedding table maps each token to a fixed vector. Our five tokens, with named dimensions (hand-chosen for teaching; a real model's dimensions have no labels):

The word that settles bank, river, is right next to it. In "I walked along the bank of the river" it comes later and further away. The mechanism must reach any position and choose which by what is there, not where it is.

Three ways to read context

(T = tokens, d = vector width.) Attention buys the short path and the parallelism with a T^2 term. Section 12 sends the bill.

The idea: a soft dictionary lookup

A hard lookup is a library catalogue: type an exact call number, get exactly one book. A soft lookup is a librarian: describe what you want, they judge how well every book matches, and hand you a blend, mostly the best match with a little of the next best.

\mathbf q

the query: a description of what this token is looking for (bold = vector)

\mathbf k_j

the key of token j: the label it is found by, like a book's spine

\mathbf v_j

the value of token j: the content it hands over, like the book's text

\text{sim}(\mathbf q, \mathbf k_j)

how well the query matches key j (Section 3 makes it a scaled dot product)

\mathbf k_l

the key of token l, used in the normalizer

a_j

the weight on token j: positive, and all the weights sum to 1

\sum_{j=1}^{T}

add up over every visible token j, from 1 to T

\sum_{l=1}^{T}

add up over every token again, to normalize

T

the number of tokens in the context

j

which token is being looked at

l

a second index over tokens, used only in the normalizer

A hard lookup puts weight 1 on one key and 0 elsewhere. The soft version is differentiable, so gradient descent can learn what to ask for and how to be found.

Self- vs cross-attention, bidirectional vs causal

In self-attention, queries, keys and values all come from the same sequence (every layer of a GPT-style LLM). In cross-attention, queries come from one sequence and keys and values from another: Bahdanau's decoder reading the English, or a vision-language model's text reading image features. An encoder (BERT-style) lets every token see every other token; a decoder (GPT-style) must not let a token see the ones after it. That difference is a mask (Section 6).

Remember: attention swaps "remember everything in one vector" for "look up what you need: any distance, all at once, by content".

Next question → How does a token decide whom to look at, and what does it take from them? (Section 3)

3. Queries, Keys and Values: Three Roles From One Vector

The question: from bank's single vector [1, 1, 0, 0], how do we get a question to ask (what disambiguates me?), a label to be found by, and a message to send?

The 30-second version: three learned linear maps act on the same vector: \mathbf q = \mathbf x W_Q is what I'm looking for, \mathbf k = \mathbf x W_K is what I advertise, and \mathbf v = \mathbf x W_V is what I deliver. Hook: Ask, Advertise, Deliver. In the toy, bank asks \mathbf q = [2, -1], and the score from bank to river is 2 while the score from river to bank is 1: separate maps make relevance asymmetric and learned. Only the product W_Q W_K^\top affects the scores. First read: the three maps, the toy table, the shapes table.

The three maps

Back to the library. The query is the request slip you write ("something about rivers"). The key is the label on each book's spine. The value is what's inside the book. Each token plays all three roles, through three learned matrices:

\mathbf x_i

token i's input vector, a row of length d_{\text{model}} ([nature, money, function word, describes])

i

which token: 1 to T (the, river, bank, was, steep)

W_Q

the learned query matrix, d_{\text{model}} \times d_k: turns a token into the question it asks

W_K

the learned key matrix, d_{\text{model}} \times d_k: turns a token into the label it is found by

W_V

the learned value matrix, d_{\text{model}} \times d_v: turns a token into the message it delivers

\mathbf q_i

token i's query (Ask)

\mathbf k_i

token i's key (Advertise)

\mathbf v_i

token i's value (Deliver)

Stack the T rows into X (T \times d_{\text{model}}) and one multiply does every token: Q = XW_Q, K = XW_K, V = XW_V.

The toy uses d_k = d_v = 2. Its matrices are hand-chosen, not trained, so every number can be checked; no trained head is claimed to look like this. W_Q has rows [1, 0], [1, −1], [0, 0], [1, 0]; W_K has rows [1, 0], [0, 1], [0, 0], [0, 0]; W_V has rows [1, 0], [−1, 1], [0, 0], [0, 0]. Values live in a 2-dimensional "sense space" [geography, money]: \mathbf v = [\text{nature} - \text{money},\ \text{money}].

Bank's query [2, −1] says "I want things high on the first key axis (nature) and low on the second (money)". Its key [1, 1] advertises both. Its value [0, 1] says that, alone, bank delivers pure money. The has query [0, 0]: a token with no question (Sections 4 and 7 come back to it).

d_k = 2)" loading="lazy" />

Why not just use \mathbf x_i \cdot \mathbf x_j?

1. It is symmetric. "How much should bank read river" would always equal "how much should river read bank". Language isn't like that: a verb looks for its subject, not the reverse. With separate maps, \mathbf q_{\text{bank}} \cdot \mathbf k_{\text{river}} = 2 but \mathbf q_{\text{river}} \cdot \mathbf k_{\text{bank}} = 1.
2. Every token prefers itself. When vectors have similar lengths, \mathbf x_i \cdot \mathbf x_i is the largest dot product in the row (Cauchy–Schwarz), so each token mostly reads itself.
3. Relevance isn't learned. The embedding was trained to identify tokens, not to say which ones help each other.

Shapes at every step

Only W_Q W_K^\top matters for the scores

s_{ij}

the raw score: how well token i's question matches token j's label

\mathbf q_i

token i's query, a row vector

\mathbf k_j^\top

token j's key, transposed into a column so the product is one number

\mathbf x_i

token i's input vector

W_Q W_K^\top

one d_{\text{model}} \times d_{\text{model}} matrix of rank at most d_k: all the scores depend on

\mathbf x_j^\top

token j's input vector, as a column

i

the token doing the looking

j

the token being looked at

So W_Q and W_K are not separately identifiable: replace them with W_Q A and W_K A^{-\top} for any invertible d_k \times d_k matrix A and every score is unchanged. Interpretability work calls W_Q W_K^\top the QK circuit (where to look) and W_V W_O the OV circuit (what gets written back; W_O arrives in Section 10) (Elhage et al., 2021).

Remember: Ask, Advertise, Deliver: \mathbf q = \mathbf x W_Q, \mathbf k = \mathbf x W_K, \mathbf v = \mathbf x W_V. Separate maps make relevance asymmetric and learned. Scores depend only on W_Q W_K^\top.

Next question → How exactly do queries and keys turn into weights, and bank's new vector? (Section 4)

4. Attention by Hand: Score, Scale, Mask, Softmax, Sum

The question: compute bank's new vector by hand, using only the, river and bank: what it may see while the model predicts the next word.

The 30-second version: Score: \mathbf q_{\text{bank}} \cdot \mathbf k for each visible key gives 0, 2, 1. Scale: divide by \sqrt{d_k} = 1.414: 0, 1.414, 0.707. Mask: was and steep get -\infty (Section 6). Softmax: exponentiate (1, 4.113, 2.028; sum 7.141) and normalize: 0.140, 0.576, 0.284. Sum: 0.140 \cdot [0, 0] + 0.576 \cdot [1, 0] + 0.284 \cdot [0, 1] = [0.576, 0.284]. Bank's own value was pure money [0, 1]; after attention it reads 0.576 geography against 0.284 money. Hook: Score, Scale, Mask, Softmax, Sum (four S's and an M). First read: the softmax box, the bank table, the matrix formula.

Softmax from zero

Softmax turns any list of real numbers into weights that are positive and sum to 1:

\text{softmax}(\mathbf u)_j

the weight given to item j

\mathbf u

the list of scores, one per item (bold = vector)

u_j

the score of item j

j

which item (here: which key)

\tau

the temperature: 1 normally; small makes it sharper, large makes it flatter

\sum_{l}

add up over every item l

u_l

the score of item l

l

an index over all items, for the normalizer

e

the base of the natural log, about 2.718

Properties used later:

• Keeps order. A bigger score always gets a bigger weight.
• Shift-invariant. Add the same number to every score and nothing changes (e^{u + c} = e^c e^u; the e^c cancels). Add 10 to bank's scores: still 0.140 / 0.576 / 0.284. That's why we can subtract the row max for numerical safety (Sections 11–12).
• Scale-sensitive. Multiplying scores is a temperature change. As \tau \to 0 softmax becomes argmax; as \tau \to \infty, uniform. It's a soft version of "pick the max".
• Exponential. Gaps matter, not ratios: a gap of 2.3 already means a 10× weight ratio. Section 5 is about this.

The bank row, step by step

River wins because bank's query points along river's key. Bank's own key gets half credit because the query's −1 on the money axis cancels part of it. The gets 1 / 7.141 = 0.140 though it matches nothing: a zero score is not a zero weight.

All rows at once

\text{Attn}(Q, K, V)

every token's new vector, one row per token (T \times d_v)

Q

all queries stacked, T \times d_k

K

all keys stacked, T \times d_k

K^\top

K transposed (rows become columns), so QK^\top is the T \times T table of scores

d_k

the query/key width (2 in the toy, 128 in Llama 2)

M

the mask: 0 where looking is allowed, -\infty where forbidden (Section 6)

V

all values stacked, T \times d_v

The softmax runs along each row. Shapes: (T \times d_k)(d_k \times T) \to T \times T, then (T \times T)(T \times d_v) \to T \times d_v. The full causal table (row = the token asking; · = masked, weight exactly 0):

Without the mask, bank also sees steep: weights 0.082 / 0.336 / 0.165 / 0.082 / 0.336, output [0.671, 0.165], more geographic still. And the, with its zero query, scores 0 against every key: uniform weights of 0.2 and an output equal to the mean of all five values, [0.4, 0.2]. A head with no question returns the average of everything.

What the output is

Each output row is a convex combination of the value rows (weights non-negative, summing to 1). So it is linear in V, nonlinear in Q and K: double every value and the output doubles; double every query and the weights change shape. And it creates no new values: bank's [0.576, 0.284] lies on the segment from river's [1, 0] to bank's [0, 1], pulled toward the origin by the.

How spread out is a row? Entropy

H(\mathbf a_i)

the entropy of row i in nats: 0 if all weight is on one key, \ln T if uniform

\mathbf a_i

token i's row of attention weights

a_{ij}

the weight token i puts on token j

\sum_{j}

add up over the keys j that token i may see

i

the token doing the looking

j

the token being looked at

The causal entropies are 0, 0.634, 0.951, 1.327, 1.446, against maxima of 0, 0.693, 1.099, 1.386, 1.609. The effective number of keys e^H is the size of a uniform row that is equally spread. It's the single most useful number to monitor per head (Section 11).

Two traps. A score is not a weight: a dashboard showing "attention −3.2" is showing a score (any real number), not a weight (in [0, 1]). And weights are not explanations: a large weight says what was averaged, not what changed the prediction. Jain and Wallace (2019) showed attention maps can often be altered without changing outputs; Wiegreffe and Pinter (2019) replied that it depends on what you call an explanation, and neither paper licenses reading a weight map as the reason for a prediction. For "what mattered", use gradients or ablations.

Remember: four S's and an M: Score, Scale, Mask, Softmax, Sum. Each output is a convex combination of the values. Bank: 0.140 / 0.576 / 0.284 → [0.576, 0.284].

Next question → At d_k = 2 the scaling barely mattered (unscaled, bank's weights are 0.090 / 0.665 / 0.245). Llama 2's heads have d_k = 128. Why divide by \sqrt{d_k}? (Section 5)

5. Why Divide by \sqrt{d_k}

The question: Llama 2 7B's heads have d_k = 128 and see up to 4,096 keys. What happens to a row of attention if we skip the 1/\sqrt{d_k}?

The 30-second version: with independent, unit-variance components, \mathbf q \cdot \mathbf k is a sum of d_k terms of variance 1. Variances add, so its variance is d_k (simulated: 128.4 at d_k = 128) and its standard deviation 11.3. Softmax over logits that spread is nearly one-hot: over 4,096 keys the median top weight is 0.868 and the entropy 0.49 nats (about 1.6 effective keys). Scaled: top weight 0.005, entropy 7.83 (about 2,510 effective keys), the same at every d_k from 4 to 512. Dividing by d_k overshoots: entropy 8.314 against the uniform 8.318, so nothing is selected. Hook: variance adds, so divide by the standard deviation. Two caveats: it holds at initialization only, and the gradient story is about where gradient goes, not just its size. First read: the variance derivation, the d_k table, Figure 3.3.

The variance of a dot product

Flip d_k coins worth ±1 and add them. One flip lands within ±1; 128 flips typically land around ±11. The same happens to \mathbf q \cdot \mathbf k:

\text{Var}(\mathbf q \cdot \mathbf k)

how much the raw score varies from one random (query, key) pair to the next

\mathbf q \cdot \mathbf k

the dot product of a query and a key: the unscaled score

\sum_{m=1}^{d_k}

add up over the d_k coordinates m

q_m

coordinate m of the query: mean 0, variance 1, independent of everything else

k_m

coordinate m of the key, same assumptions

\text{Var}(q_m k_m)

the variance of one coordinate's product

\mathbb{E}[q_m^2]

the average of q_m^2 over random draws: 1 for unit variance

\mathbb{E}[k_m^2]

the same for the key: 1

d_k

the query/key width

m

which coordinate

Step 2 uses independence (uncorrelated terms: variances add). Step 3 uses mean 0: \text{Var}(q_m k_m) = \mathbb{E}[q_m^2]\mathbb{E}[k_m^2] - (\mathbb{E}[q_m]\mathbb{E}[k_m])^2, and the last term is 0. Dividing the score by \sqrt{d_k} divides its variance by d_k, back to 1 (Card 3). Simulated with 100,000 pairs: unscaled variance 0.993, 8.02, 32.0, 128.4, 509 at d_k = 1, 8, 32, 128, 512; after scaling, between 0.994 and 1.003 at every width.

What softmax does with big logits

Feed those scores to a softmax over T = 4,096 keys (Llama 2's context), 400 random queries per width. Medians:

The uniform maximum is \ln 4096 = 8.318 nats. No scaling: rows go one-hot as d_k grows. 1/\sqrt{d_k}: the row's shape is the same at every width, so you can change head width without retuning. 1/d_k: rows drift to uniform and the head just averages (Section 7). Hook: too big → one-hot, too small → uniform.

The temperature view

If a row's logits have standard deviation s, then for moderate s the expected entropy is about \ln T - s^2/2 nats. At T = 4,096: s = 1 gives 7.818 (simulated 7.818); s = 2 gives 6.32 (simulated 6.41). It breaks once s nears \sqrt{2 \ln T} = 4.08: at s = 4 it predicts 0.32, but the simulation gives 3.17. Unscaled d_k = 128 logits have s = 11.3, deep in the one-hot regime.

Two keys, exactly

Two keys, d_k = 128: the top weight exceeds 0.99 when the score gap exceeds \ln 99 = 4.595. Given the query, the gap \mathbf q \cdot (\mathbf k_1 - \mathbf k_2) is normal with variance 2\lVert \mathbf q \rVert^2, and \lVert \mathbf q \rVert^2 itself varies (chi-squared with d_k degrees of freedom). Averaging over it exactly: P(\text{top weight} > 0.99) = 0.773 unscaled vs 0.00135 scaled. The usual shortcut (gap normal with variance 2d_k) gives 0.774 and 0.00116; a million-draw simulation gives 0.773 and 0.00142.

What it does to gradients

The usual line is "unscaled → gradients vanish". The truth is more specific, and Vaswani et al. (2017) themselves wrote that they "suspect" this effect. A softmax row's Jacobian is \text{diag}(\mathbf a) - \mathbf a \mathbf a^\top (Section 8), exactly zero for a one-hot row. In simulation, its size (Frobenius norm) is not monotone in d_k without scaling: 0.020 at d_k = 1, a peak of 0.30 at d_k = 32, then 0.181 at 128 and 0.041 at 512. At 128 it's even larger than the scaled one (0.0253). But it is concentrated on about 1.6 effective keys instead of 2,511, so the other keys can't learn to compete, and 22% of rows have top weight above 0.99 (46% at d_k = 512), where it is near zero. The correct statement: unscaled logits make softmax nearly one-hot, so only the winner gets gradient, and that gradient goes to zero as the row saturates.

When the assumption breaks

The derivation assumed \mathbf q and \mathbf k independent, which holds at initialization. Training breaks it two ways:

1. Alignment. If a head learns \mathbf k \approx \mathbf q, then \mathbf q \cdot \mathbf k \approx \lVert \mathbf q \rVert^2, which grows like d_k. Simulated with \mathbf k = \mathbf q at d_k = 128: the mean of \mathbf q \cdot \mathbf k / \sqrt{d_k} is 11.3; of \mathbf q \cdot \mathbf k / d_k, 1.00. This is why μP (Yang et al., 2022) scales attention by 1/d_k, for hyperparameters that transfer across widths (Card 16).
2. Logit growth. Scale W_Q and W_K each by c and every logit scales by c^2: entropy over 4,096 keys goes 7.82 (c = 1) → 3.24 (c = 2) → 0.339 (c = 4). That's entropy collapse (Section 11).

Two modern fixes bound the logits directly. QK-norm normalizes each query and key to unit length (times a scale) before the dot product, so no logit can exceed that scale (Henry et al., 2020; Dehghani et al., 2023, used it to train a 22B-parameter vision transformer). Soft-capping squashes each logit smoothly into (-c, c):

\text{softcap}(u)

the capped logit: almost u when small, never beyond \pm c

u

the logit before capping

c

the cap; Gemma 2 uses 50 on attention logits (and 30 on the final output logits)

With c = 50: 10 → 9.87, 50 → 38.08, 200 → 49.97 (Gemma Team, 2024).

The trap at d_k = 2. The toy can't show saturation: there the 1/\sqrt 2 makes bank's row softer (0.090 / 0.665 / 0.245 → 0.140 / 0.576 / 0.284). Scaling matters at d_k = 128, where it divides by 11.3, not at d_k = 2, where it divides by 1.41. A unit test at small d_k passes with the scale deleted (Q4).

Remember: \text{Var}(\mathbf q \cdot \mathbf k) = d_k: variance adds, so divide by the standard deviation. Too big → one-hot, too small → uniform. In training, alignment and logit growth break the assumption; QK-norm and soft-capping bound the logits.

Next question → While training on "the river bank was steep", bank must not see was and steep, and in a padded batch river must not read padding. How do we forbid it? (Section 6)

6. Masks: No Peeking at the Future, No Attending to Padding

The question: we train the model to predict every next token of "the river bank was steep" in one pass. What stops bank from reading was? And when the sentence shares a batch with "the river" padded to length 5, what stops river from reading the padding?

The 30-second version: add -\infty to every forbidden score before the softmax, since e^{-\infty} = 0. The causal mask allows 15 of the 25 pairs. Test it: change steep, and the first four outputs don't move at all (difference exactly 0); without the mask, bank's output moves by 0.217. A padding mask blocks fake keys: unmasked, river spends 49.8% of its attention on three PAD tokens; masked, it gets the same 0.330 / 0.670 as the unpadded sentence. Rule: mask the scores, not the weights. Mask types: Future, Fake, Far (causal, padding, sliding window). First read: the rule, the test, the two traps.

The additive mask

M_{ij}

the mask entry for the pair (token i asks, token j is looked at)

i

the token doing the looking (the row)

j

the token being looked at (the column)

a_{ij}

the attention weight after masking

s_{ij}

the scaled score \mathbf q_i \cdot \mathbf k_j / \sqrt{d_k}

\sum_{l}

add up over every key l in the row

s_{il}

the scaled score of key l in row i

M_{il}

the mask entry for key l in row i

l

an index over all keys in the row

e

the base of the natural log; e^{-\infty} = 0, which is what makes the mask work

A masked key contributes nothing to the numerator or the denominator, so the remaining weights renormalize among the allowed keys and still sum to 1. Why not mask after the softmax? Zeroing weights afterwards leaves rows that don't sum to 1, and the forbidden keys already shaped the denominator, so the future leaks through the normalizer. Mask the scores, not the weights.

The causal mask, and one pass for T predictions

The causal mask allows j \le i: the lower triangle including the diagonal (Section 4's table, where · marks the blocked pairs). That's 15 of 25 pairs; in general T(T+1)/2 of T^2, a fraction (T+1)/2T: 0.6 at T = 5, 0.500122 at T = 4,096. A kernel that skips masked blocks does half the work; one that computes the full table and then masks does not (Section 12).

With the causal mask, row i depends only on tokens 1 to i, so it can predict token i + 1. One forward pass over "the river bank was steep" gives five predictions at once, each from its own prefix. Feeding the true previous tokens during training is teacher forcing; it makes training parallel over T. Generation can't do this: each new token depends on the one just sampled (Section 13).

The correctness test: change the future, the past must not move

A mask bug doesn't crash; it just makes training loss too good. So test it:

import numpy as np
X = np.array([[0,0,1,0],[1,0,0,0],[1,1,0,0],[0,0,1,1],[1,0,0,1]], float)   # the river bank was steep
WQ, WK = np.array([[1,0],[1,-1],[0,0],[1,0]], float), np.array([[1,0],[0,1],[0,0],[0,0]], float)
WV = np.array([[1,0],[-1,1],[0,0],[0,0]], float)

def attend(X, blocked=None, fill=-np.inf):
    S = (X @ WQ) @ (X @ WK).T / np.sqrt(2)
    if blocked is not None:
        S = np.where(blocked, fill, S)               # mask the SCORES, before the softmax
    E = np.exp(S - S.max(axis=1, keepdims=True))
    return (E / E.sum(axis=1, keepdims=True)) @ (X @ WV)

causal = np.triu(np.ones((5, 5), bool), k=1)         # block j > i
X2 = X.copy(); X2[4] = [0, 3, 0, 1]                  # change the future: a different 'steep'
print("causal, rows 0-3 moved by:", np.abs(attend(X, causal)[:4] - attend(X2, causal)[:4]).max())
print("bidirectional, bank moved by:", np.abs(attend(X)[2] - attend(X2)[2]).max().round(3))
diag_too = np.triu(np.ones((5, 5), bool), k=0)       # the off-by-one: also blocks the diagonal
with np.errstate(invalid="ignore"):
    print("triu(k=0), row 0:", attend(X, diag_too)[0])                       # [nan nan]
print("same bug, -1e9 fill, row 0:", attend(X, diag_too, fill=-1e9)[0].round(3))  # silent: the mean

With the mask, the first four outputs move by exactly 0.0; without it, bank moves by 0.217. The same fact makes the KV cache valid: if the past never depends on the future, the past's keys and values never need recomputing (Card 6; Section 13).

Padding: masking fake keys

To batch sequences of different lengths, short ones are filled with PAD tokens. Put "the river" in a batch of length 5 with three PADs (zero vectors here), bidirectional, no padding mask: river's weights are 0.166 / 0.336 / 0.166 / 0.166 / 0.166, so 0.498 of its attention goes to padding. With a key-padding mask (-\infty on PAD columns) its weights are 0.330 / 0.670 / 0 / 0 / 0, identical to running "the river" alone. Mask PAD columns (keys); PAD rows (queries) produce junk nobody reads, provided the loss ignores them. Combine masks with OR: blocked = causal OR padding.

Two traps

1. The off-by-one. triu(ones, k=1) blocks the future; k=0 also blocks the diagonal, so token 0 has no allowed keys: 0/0 = NaN, which spreads through the next layer. Fill with −1e9 instead of -\infty and the bug turns silent: row 0 becomes a uniform 0.2 over keys it should never see (its output is the mean value, [0.4, 0.2]). The change-the-future test finds it; staring doesn't (Q17).
2. Left vs right padding in decoders. Batched generation wants the last real token at the same position in every row, so it usually pads on the left and shifts positions to match. Right padding is fine for training (the loss ignores PAD positions).

Patterns, one line each

Bidirectional (BERT-style encoder): no causal mask. Causal (GPT-style decoder): look back only; every LLM in this course. Prefix-LM: bidirectional inside a prompt prefix, causal after it. In general a mask is any pattern of allowed pairs: a sliding window allows only the last W tokens (Mistral 7B: 4,096; LLM Ch 4). Mask types to recall: Future, Fake, Far.

Remember: mask the scores, not the weights, with -\infty before the softmax. Future, Fake, Far. Test it: change the future, and the past must not move.

Next question → We can compute attention correctly. What is that output, mathematically, and what can it never do? (Section 7)

7. What Attention Computes: A Learned Kernel Smoother

The question: bank's output [0.576, 0.284] is a weighted average of values. Weighted by what kind of similarity, and what can such an average never produce?

The 30-second version: estimating a curve at x by averaging the y_j of nearby points, with weights K(x, x_j) / \sum K, is Nadaraya–Watson kernel regression (Nadaraya, 1964; Watson, 1964). Attention is the same thing with a learned exponential kernel \exp(\mathbf q \cdot \mathbf k / \sqrt{d_k}). A Gaussian-kernel smoother is dot-product attention once each key carries a bias -\lVert \mathbf k \rVert^2 / 2h^2 (simulation agrees to 5.6 \times 10^{-16}). Consequences: outputs stay inside the hull of the values; a zero query gives the mean; uniform weights over T values shrink their variance by 1/T; and the scores have rank at most d_k. Hook: linear in V, nonlinear in Q and K. First read: the Nadaraya–Watson box and the three limits.

From the librarian to kernel smoothing

Section 2's librarian blends books by how well each matches. Statisticians had that idea in 1964, for smoothing noisy data: to estimate a curve at x, average the observed y_j, weighting each by how close x_j is.

\hat f(x)

the smoothed estimate of the curve at point x (hat = estimated from data)

x

where we want an estimate (plays the query)

x_j

the location of observation j (plays a key)

y_j

the value observed at x_j (plays a value)

K(x, x_j)

the kernel: a similarity, large when x_j is close to x

\sum_{j}

add up over all observations j

\sum_{l}

add up over all observations again, to normalize

K(x, x_l)

the kernel weight of observation l

j

which observation

l

an index over observations for the normalizer

Query ↔ x, keys ↔ x_j, values ↔ y_j, kernel ↔ \exp(\mathbf q \cdot \mathbf k / \sqrt{d_k}). Attention is Nadaraya–Watson with a learned kernel and learned values (Tsai et al., 2019, develop this view). In one dimension, 40 noisy samples of \sin x smoothed with a Gaussian kernel of width 0.5 give 0.913 at x = \pi/2, where the truth is 1: a little bias for much less noise.

The Gaussian kernel is dot-product attention with a key bias

\lVert \mathbf q - \mathbf k_j \rVert^2

squared distance between the query and key j

\mathbf q

the query

\mathbf k_j

key j

h

the kernel bandwidth: how far "nearby" reaches

\mathbf q \cdot \mathbf k_j

the dot-product score

\lVert \mathbf q \rVert^2

the query's squared length: the same for every key in the row

\lVert \mathbf k_j \rVert^2

key j's squared length: a per-key bias

The \lVert \mathbf q \rVert^2 term is the same for every key, so softmax's shift invariance cancels it. What's left is a dot product at temperature h^2 plus a per-key bias: Gaussian kernel smoothing = dot-product attention with a bias -\lVert \mathbf k_j \rVert^2 / 2h^2 on each key. Simulated: they agree to 5.6 \times 10^{-16}; drop the bias and they differ by up to 1.09. When all keys have equal length (as with QK-norm), the bias is constant and plain dot-product attention matches exactly (6.9 \times 10^{-17}) (Card 8).

Three things an average can never do

1. Leave the hull. Each output is a convex combination of the values. Attention can choose and blend; it can't extrapolate. New directions come from W_O, the feed-forward network and depth (Section 15).
2. Avoid averaging when it has no question. A zero query gives equal weights and the plain mean. For T independent values of variance \sigma^2, the output's variance is \sigma^2 / T: at T = 4,096, theory 2.44 \times 10^{-4}, simulation 2.45 \times 10^{-4}. A uniform head over a long context washes its input out to a near-constant.
3. Exceed rank d_k in the scores. QK^\top is (T \times d_k)(d_k \times T), so its rank is at most d_k. Simulated at T = 64, d_k = 8: scores have rank 8, while the softmax weights, a nonlinear function of them, have full rank 64 (Bhojanapalli et al., 2020, call the limit a low-rank bottleneck). The extreme is d_k = 1: the score is q_i k_j, so a row's top key is always the largest k (if q_i > 0) or the smallest (if q_i < 0). In 2,000 random instances, every row's top choice was one of those two. A d_k = 2 head with keys on a circle can single out each of 12 keys (Card 10).

Mass must go somewhere. Bank gives 14% of its attention to the, whose key and value are both zero. Softmax rows sum to 1, so a head with nothing useful to read must still park its mass somewhere. Real models learn to park it on a harmless, always-visible token: an attention sink (Section 11).

Remember: attention is kernel smoothing with a learned kernel: linear in V, nonlinear in Q and K. Each row is a convex combination, so it can't extrapolate; a zero query gives the mean; the scores have rank at most d_k.

Next question → The kernel is learned. How does a gradient reach W_Q, W_K and W_V? (Section 8)

8. Gradients Through Attention

The question: the loss falls if bank reads more like geography, so the upstream gradient is \mathbf g = \partial L / \partial \mathbf o_{\text{bank}} = [-1, 0]. When gradient descent takes a step (it moves every parameter against its gradient), which way do bank's query, the keys and the values move?

The 30-second version: three paths. Values: \partial L / \partial \mathbf v_j = a_j \mathbf g, so the step pushes each value toward geography in proportion to its weight (river 0.576). Scores, through the softmax Jacobian: \partial L / \partial s_j = a_j(\mathbf g \cdot \mathbf v_j - \mathbf g \cdot \mathbf o). River's gradient is −0.244, the's +0.081, bank's +0.164, summing to 0: the descent step raises river's score, because its value beats the current average, and lowers the other two. Queries and keys: \partial L / \partial \mathbf q_{\text{bank}} = [-0.057, 0.116]; the step goes the opposite way, along [0.057, -0.116], and turns bank's query toward river's key. A finite-difference check agrees to 2.4 \times 10^{-12}. When a row saturates, the middle path closes. Hook: push values by weight; pull scores by surprise, a(v - o). First read: the three paths and the toy numbers.

The softmax Jacobian

How does each weight move when you nudge each score? For one row \mathbf a = \text{softmax}(\mathbf s):

\frac{\partial a_i}{\partial s_j}

how fast weight i changes as you nudge score j, holding the others fixed: a slope, not a fraction to divide

a_i

the weight on key i

a_j

the weight on key j

s_j

the score of key j

\delta_{ij}

1 if i = j, else 0

i

one key in the row

j

another key in the row

J

the Jacobian: the full table of those slopes, one row and column per key

\text{diag}(\mathbf a)

a matrix with the weights on its diagonal and zeros elsewhere

\mathbf a \mathbf a^\top

the outer product of the weight vector with itself

\mathbf a

the row of weights as a vector

Hook: one-hot minus p. Raising your own score raises your weight by a_i(1 - a_i); raising someone else's lowers yours by a_i a_j. On bank's row (diagonal 0.120, 0.244, 0.203): rows sum to 0 (shift invariance); J is symmetric, positive semi-definite, rank T - 1, with eigenvalues 0.393, 0.175 and 0; its largest eigenvalue is never above 1/2 (two keys at 50/50; simulation max 0.500); and J = 0 when \mathbf a is one-hot: Section 5's gradient story (Card 4).

The backward pass in matrix form

Forward: S = QK^\top / \sqrt{d_k} (plus mask), P = \text{softmax}(S) by rows, O = PV. Given dO = \partial L / \partial O:

dO

the gradient of the loss with respect to the outputs (from the layer above), T \times d_v

dV

the gradient for the values: each value collects the output gradients of whoever read it, by weight

P

the attention weights after softmax, T \times T

P^\top

the weights transposed (rows become columns): "who read me, and how much"

dP

the gradient with respect to the weights

V^\top

the values, transposed

V

the values, T \times d_v

dS

the gradient with respect to the scores, after the softmax Jacobian

\odot

element-wise product

\text{rowsum}(P \odot dP)

for each row, \sum_j P_{ij} dP_{ij}: subtracted so each row of dS sums to 0

dQ

the gradient for the queries

K

the keys, T \times d_k

dK

the gradient for the keys

dS^\top

dS transposed

Q

the queries, T \times d_k

d_k

the query/key width

Then dW_Q = X^\top dQ, and likewise for W_K and W_V (Card 5).

The toy numbers

Per row, with upstream gradient \mathbf g, the score gradient reads cleanly:

\frac{\partial L}{\partial s_j}

how fast the loss changes as you nudge key j's scaled score

L

the loss; here simply minus the geography coordinate of bank's output, so it falls as bank reads more geography

\mathbf g

the upstream gradient \partial L / \partial \mathbf o for this row: [-1, 0]; the loss falls as geography rises

a_j

key j's current weight

\mathbf v_j

key j's value

\mathbf o

the row's current output, [0.576, 0.284]

\mathbf g \cdot \mathbf v_j - \mathbf g \cdot \mathbf o

the surprise: how much better key j's value is than the current average, along \mathbf g

With \mathbf g \cdot \mathbf o = -0.576: river 0.576 \times (-1 + 0.576) = -0.244; the 0.140 \times (0 + 0.576) = +0.081; bank 0.284 \times (0 + 0.576) = +0.164. They sum to 0: attention is a zero-sum budget. Descent moves each score against its gradient, so river's score rises and the other two fall. The query gradient is \sum_j (\partial L / \partial s_j)\, \mathbf k_j / \sqrt 2 = [-0.057, 0.116]. The descent step points the other way, along [0.057, -0.116]: per unit of step it raises bank's score on river's key [1, 0] by 0.040 and lowers the score on its own key [1, 1] by 0.041, turning the query toward river. Each value moves by -a_j \mathbf g, toward geography. Read the sign before the story: the gradient points uphill, and descent goes the other way. Hook: push values by weight; pull scores by surprise, a(v - o).

import numpy as np

def forward(Q, K, V, mask):
    S = Q @ K.T / np.sqrt(Q.shape[1])
    S = np.where(mask, -np.inf, S)
    P = np.exp(S - S.max(1, keepdims=True)); P /= P.sum(1, keepdims=True)
    return P @ V, P

def backward(Q, K, V, P, dO):
    d_k = Q.shape[1]
    dV = P.T @ dO                                        # push values by weight
    dP = dO @ V.T
    dS = P * (dP - (P * dP).sum(1, keepdims=True))       # pull scores by surprise: a(g.v - g.o)
    return dS @ K / np.sqrt(d_k), dS.T @ Q / np.sqrt(d_k), dV, dS

X = np.array([[0,0,1,0],[1,0,0,0],[1,1,0,0],[0,0,1,1],[1,0,0,1]], float)
WQ = np.array([[1,0],[1,-1],[0,0],[1,0]], float)
WK = np.array([[1,0],[0,1],[0,0],[0,0]], float)
WV = np.array([[1,0],[-1,1],[0,0],[0,0]], float)
Q, K, V = X @ WQ, X @ WK, X @ WV
mask = np.triu(np.ones((5, 5), bool), k=1)

O, P = forward(Q, K, V, mask)
dO = np.zeros_like(O); dO[2] = [-1.0, 0.0]              # loss L = minus bank's geography coordinate
dQ, dK, dV, dS = backward(Q, K, V, P, dO)
print("dL/ds (bank row):", dS[2, :3].round(3))           # [ 0.081 -0.244  0.164], sums to 0
print("dL/dq_bank:", dQ[2].round(3))                     # [-0.057  0.116]
print("descent step on q_bank:", (-dQ[2]).round(3))      # [ 0.057 -0.116]: toward river's key [1, 0]

eps, fd = 1e-6, np.zeros(2)                              # finite-difference check on q_bank
for j in range(2):
    Qp, Qm = Q.copy(), Q.copy(); Qp[2, j] += eps; Qm[2, j] -= eps
    fd[j] = -(forward(Qp, K, V, mask)[0][2, 0] - forward(Qm, K, V, mask)[0][2, 0]) / (2 * eps)
print("finite-difference error:", np.abs(fd - dQ[2]).max())

Two things the backward pass needs

• Unsaturated rows. A one-hot row has dS = 0: W_Q and W_K get no signal from it however wrong the output is. The value path still works, but the head can no longer change where it looks.
• The weights P. The formulas use P, T \times T per head. Storing it for every layer is Section 12's memory wall. FlashAttention-style kernels store one number per row (the log of the softmax denominator) and recompute P in the backward pass.

Remember: dS = P \odot (dP - \text{rowsum}(P \odot dP)). Push values by weight; pull scores by surprise, a(v - o). The Jacobian is one-hot minus p: at most 1/2, and zero when saturated.

Next question → Nothing so far knows where a token sits. Does that matter? (Section 9)

9. Attention Is Order-Blind: Permutation Equivariance and Why Positions Are Needed

The question: reverse the sentence to "steep was bank river the". With no causal mask, what happens to bank's output?

The 30-second version: nothing: every output simply moves with its token (max error 1.1 \times 10^{-16}). Self-attention without positions is permutation-equivariant: it treats the input as a set. Pool over tokens and it becomes invariant. "Dog bites man" and "man bites dog" get the same vectors, rearranged. The causal mask breaks the symmetry (reversing changes outputs by up to 0.426), so a decoder with no positional encoding can infer position weakly: with a start token and equal scores, the start token's share at position t is 1/t. Real models add order explicitly. Hook: a set, not a sequence. Trap: "permutation-invariant" is the wrong word. First read: the experiment and the table.

The experiment

Run the toy bidirectionally on "the river bank was steep", then on the reversed order. Bank sits in the middle both times and its output is [0.671, 0.165] both times; every other token's output moves with its token. Maximum difference after undoing the reordering: 1.1 \times 10^{-16}, rounding error. Nothing in \mathbf q_i \cdot \mathbf k_j mentions i or j, only what the tokens are.

Equivariant, not invariant

\text{Attn}(\Pi X)

self-attention (no mask, no positions) applied to the shuffled tokens

\Pi

a permutation matrix: it reorders the rows (tokens)

X

the input tokens, one row each

\text{Attn}(X)

self-attention applied to the original order

\frac{1}{T}\mathbf 1^\top

average over the T rows (mean pooling)

T

the number of tokens

Equivariant: shuffle the input and the output is shuffled the same way. Invariant: shuffle the input and the output doesn't change. Attention is the first; mean-pooling its output gives the second (simulated difference 1.1 \times 10^{-16}). The old version of this course called attention "permutation-invariant"; that's wrong, and interviewers notice (Card 7).

The causal mask leaks position

With a causal mask, token i sees exactly i tokens, so order matters: reversing our sentence changes the outputs by up to 0.426. That gives a decoder a weak positional signal with no positional encoding at all. Give a start token (BOS) value 1 and every other (identical) token value 0, with equal scores: at position t the output is BOS's share, 1/t (1, 0.5, 0.333, 0.25, …). The model can read its position off how diluted the start token is; Haviv et al. (2022) and Kazemnejad et al. (2023) found decoders without positional encodings do learn position. A bidirectional model gets no such signal: every position sees the same 1/T.

Why real models still add positions

The 1/t signal is weak and blurs as t grows. Many heads need sharp relative position: a previous-token head attends to i - 1, and induction heads (Olsson et al., 2022) combine one with a content match to copy patterns ("Harry Potter … Harry → Potter"). Content alone can't say "the token just before me", so production models inject order:

This chapter stops at why. LLM Ch 4 covers the move from learned positions to RoPE; LLM Ch 5 derives RoPE and context extension.

Remember: without positions, attention sees a set, not a sequence: permutation-equivariant (pooling makes it invariant). A causal mask leaks a little order (1/t); real models add positions (LLM Ch 4–5).

Next question → One head blends geography against money. What if a token needs to ask two questions at once? (Section 10)

10. Multi-Head Attention: Several Questions at Once

The question: was needs to find its subject (bank), while bank needs to find its sense (river). Can one 2-dimensional head do both, and what do two heads cost?

The 30-second version: run h heads in parallel. Each has its own W_Q, W_K, W_V of width d_k = d_{\text{model}} / h; concatenate their outputs and mix them with W_O. In the toy, head 1 is a "sense" head (bank → river 0.576) and head 2 a "subject" head (was → bank 0.449, steep → bank 0.476). The parameter count is 4 d_{\text{model}}^2 (64 in the toy, 67,108,864 per layer in Llama 2 7B) whatever h is. Heads are free in parameters and projection FLOPs, not free in rank (each score matrix has rank at most d_k), and not free in KV cache once they stop sharing keys and values (Section 13). Hook: Split, Attend, Concat, Project. First read: the formula, the toy heads, the parameter count.

Split, Attend, Concat, Project

One head computes one set of weights per token: one question. A token often needs several (who is my subject? which sense am I? what came just before me?). Think of a newsroom: several reporters interview the same crowd, each with a different question, and an editor merges their notes into one story.

\text{head}_r

the output of head r: one d_v-wide vector per token

r

which head: 1 to h

X

the input tokens, T \times d_{\text{model}}

W_Q^{(r)}

head r's query matrix, d_{\text{model}} \times d_k

W_K^{(r)}

head r's key matrix, d_{\text{model}} \times d_k

W_V^{(r)}

head r's value matrix, d_{\text{model}} \times d_v

\text{MHA}(X)

multi-head attention's output, T \times d_{\text{model}}

\text{Concat}(\text{head}_1, \ldots, \text{head}_h)

the heads' outputs side by side, T \times h d_v

h

the number of heads

W_O

the output matrix, h d_v \times d_{\text{model}}: the editor that mixes the heads back together

In code nobody loops over heads: project once with a d_{\text{model}} \times d_{\text{model}} matrix, then reshape (B, T, d_{\text{model}}) \to (B, T, h, d_k) \to (B, h, T, d_k) (Split), attend per head (Attend, scores (B, h, T, T)), reshape back to (B, T, h\,d_k) (Concat) and multiply by W_O (Project). For Llama 2 7B that is (B, 4096, 4096) \to (B, 32, 4096, 128) and back.

The toy's two heads

Head 1 is the one we've used all chapter. Head 2 (also hand-chosen) has \mathbf q = [\text{describes}, \text{function}], \mathbf k = [\text{nature} + \text{money} - \text{describes}, \text{function}] and \mathbf v = [\text{money}, \text{nature}]: predicates go looking for a noun. Its causal weights put was → bank 0.449 and steep → bank 0.476 (Figure 3.5 shows both heads' full tables).

Head 1 answers "which sense is bank?" (bank → river 0.576); head 2 answers "what is being described?" (was and steep → bank). Bank has a zero query in head 2, so it averages its three visible values to [0.333, 0.667]. With W_O = I (identity, chosen only so you can check by hand; trained W_O is a full learned matrix), bank's multi-head output is the concatenation [0.576, 0.284, 0.333, 0.667]. The block adds this back to the input (the residual connection of Section 15): [1, 1, 0, 0] + [0.576, 0.284, 0.333, 0.667] = [1.576, 1.284, 0.333, 0.667].

The parameter count does not depend on h

h

the number of heads

d_{\text{model}}

the model width

d_k

the query/key width per head, d_{\text{model}} / h

d_v

the value width per head, d_{\text{model}} / h

4\, d_{\text{model}}^2

four square matrices' worth of weights, whatever h is

W_Q

the query matrices of all heads together

W_K

the key matrices of all heads together

W_V

the value matrices of all heads together

W_O

the output matrix that mixes the heads

Splitting into heads just partitions the same matrices. Projection FLOPs don't change either. (Card 9.)

So what do heads cost? Rank: each head's scores have rank at most d_k = d_{\text{model}} / h, so more heads means more patterns per layer but each one coarser (Section 7; Bhojanapalli et al., 2020). Cache: each KV head stores its own keys and values during generation (Section 13). And very small d_k uses hardware poorly; most LLMs settle on d_k = 128 (Llama 2 and 3) or 64 (GPT-2).

W_O is not decoration

Split W_O into h horizontal blocks W_O^{(r)}, one per head. Then concatenating and multiplying is the same as summing:

W_O^{(r)}

the block of rows of W_O that multiplies head r's output (d_v \times d_{\text{model}})

\sum_{r=1}^{h}

add up over all heads r

Each head writes its own contribution into the token's vector, and the contributions add. That is the "residual stream" view used in interpretability (Elhage et al., 2021), and it's why heads can be studied, or pruned, one at a time.

The implementation trap

The correct split is X @ Wq → .view(B, T, h, d_k) → .transpose(1, 2). Writing .view(B, h, T, d_k) directly gives the right shape and the wrong data: it slices the flat memory into chunks that mix tokens and head dimensions. Nothing crashes; a model loaded with pretrained weights just gets much worse (Q20).

Heads in practice. Trained models grow specialized heads (previous-token, syntactic, delimiter heads; Clark et al., 2019; Voita et al., 2019), and many redundant ones: Michel, Levy and Neubig (2019) pruned many heads at test time with little loss.

Remember: Split, Attend, Concat, Project. 4 d_{\text{model}}^2 parameters whatever h is. Heads cost rank and cache, not parameters. W_O sums the heads' contributions.

Next question → What goes wrong when attention trains at scale? (Section 11)

11. How Attention Fails: S-C-A-L-E

The question: a training run's loss spikes. The dashboard shows attention entropy in one layer falling toward 0, and several heads putting most of their mass on the first token. What is happening, which of these is a bug, and which is a feature?

The 30-second version: five failure modes, hook S-C-A-L-E. Saturation: logits too large, rows one-hot, gradient reaches only the winner (Section 5). Collapse: entropy collapse, where logits grow during training (scaling W_Q and W_K by 4 takes entropy from 7.82 to 0.339), and rank collapse, where a stack of pure attention layers makes every token identical doubly exponentially fast (0.78 → 0.27 → 0.013 → 2 \times 10^{-6}); residual connections slow it greatly. Attention sinks: a row must sum to 1, so a head with nothing to do parks its mass on an always-visible token, usually the first. That's a feature, but it breaks naive cache eviction. Leakage: mask bugs. Exploding numerics: fp16 overflows above e^{11.09}, and a naive fp16 softmax over 4,096 unscaled keys produces NaN in every row. First read: the SCALE table.

The SCALE table

Collapse, part 1: entropy collapse

Section 5 showed what happens when logits grow: scale W_Q and W_K by c and every logit scales by c^2. Over 4,096 keys the median entropy goes 7.82 (c = 1) → 5.95 (1.5) → 3.24 (2) → 0.339 (4), and the top weight from 0.005 to 0.919. Nothing in the loss stops this: sharper attention can lower training loss right up until the rows saturate and gradients through them die. Zhai et al. (2023) linked low attention entropy to training instability and proposed σReparam, which reparameterizes each weight matrix by its spectral norm. The other fixes bound the logits directly: QK-norm and soft-capping (Section 5).

Collapse, part 2: rank collapse

A head averages; a stack of heads averages averages. With pure attention layers (no residual, no feed-forward network), each layer pulls every token toward a common mean. Measure each token's distance from the mean token, relative to the size of the input:

\text{res}(X)

the relative residual: 0 means every token's vector is identical

X

the token vectors after a given layer, one row per token

\mathbf 1

a column of T ones

\bar{\mathbf x}^\top

the mean token vector, as a row (bar = mean)

\lVert X \rVert

the size (Frobenius norm) of the matrix

\lVert X - \mathbf 1 \bar{\mathbf x}^\top \rVert

how far the tokens are, in total, from all being the mean

Simulated (32 tokens, width 64, random weights, 20 repeats): pure attention goes 0.99 → 0.78 → 0.27 → 0.013 → 2.2 \times 10^{-6} over layers 0–4, and machine zero by layer 5. That's doubly exponential, the rate Dong, Cordonnier and Loukas (2021) proved in "Attention is not all you need: pure attention loses rank doubly exponentially with depth". With residual connections the same stack falls only slowly: 0.97 after layer 1, 0.75 after layer 5, 0.10 after layer 12. Skip connections mitigate collapse; they do not abolish it at random initialization. That is one reason the Transformer block (Section 15) is attention plus a residual plus a feed-forward network. (Card 17.)

Attention sinks: a feature that bites

A softmax row must sum to 1. A head that has nothing useful to read in a given context can't output "nothing"; it has to put the mass somewhere. Trained models learn to park it on a token that is always visible and carries a near-useless value, usually the first token. BERT parks it on [SEP] and punctuation (Clark et al., 2019). How loud must the sink be? Give it a logit advantage a over T - 1 equal keys:

w

the share of the row's attention the sink takes

a

the sink's logit advantage over each of the other keys

T

the number of keys in the row

e

the base of the natural log

For a 90% share: a = 7.04 at T = 128, 9.13 at T = 1,024, 10.51 at T = 4,096. An advantage of 5 at T = 1,024 buys only 12.7%. So a sink needs a large logit, which grows with context: one reason large logits aren't always a bug.

Where it bites: a serving system that keeps only the most recent tokens in the cache evicts the first token, the heads lose their parking spot, and quality collapses. Xiao et al. (2023) found that keeping the first 4 tokens plus a recent window fixes it (StreamingLLM).

Leakage

Two tests catch almost every mask bug: change the future (outputs before the change must not move) and batch vs single (a sequence alone and inside a padded batch must match to float precision).

Exploding numerics

Half precision (fp16) has a largest finite value of 65,504, so e^x overflows to infinity for x > 11.09. Unscaled d_k = 128 logits have standard deviation 11.3, so 16.4% of them exceed that (simulated: 16.1%), and with 4,096 keys every row has at least one. A naive fp16 softmax (\exp, then divide) turned every row into NaN in simulation; the stable version, which subtracts the row maximum first, produced none (the largest scaled logit was 4.87). Rules: subtract the max before exponentiating (shift invariance makes it free); accumulate the softmax sum in fp32 even with fp16 or bf16 inputs; and fill masks with the dtype's minimum (torch.finfo(dtype).min) rather than -\infty when a row can be fully masked, or handle empty rows explicitly, because a fully masked row gives 0/0 = NaN in a hand-written softmax (frameworks differ: Section 14).

What to monitor, per head and layer: entropy (or e^H), the maximum logit, and the fraction of mass on token 0. Entropy falling while the max logit rises is the early warning for entropy collapse.

Remember: S-C-A-L-E: Saturation, Collapse, Attention sinks, Leakage, Exploding numerics. Most fixes keep logits bounded and masks right. Sinks are a feature; just don't evict them.

Next question → Correct and stable. Now what does it cost? (Section 12)

12. The T^2 Bill: FLOPs, Memory, and Attention Without the T \times T Matrix

The question: for one layer of Llama 2 7B at 4,096 tokens, how many FLOPs does attention cost, how much memory would the T \times T weights take, and can we avoid storing them?

The 30-second version: the Q, K, V and O projections cost 8Td^2 FLOPs (linear in T); scores and the weighted sum cost 4T^2 d (quadratic). At T = 4,096 that's 550 GFLOP for projections, 275 GFLOP for the quadratic part and 1,108 GFLOP for the FFN, so the quadratic part is 14.2% of the layer (7.7% with causal skipping). It overtakes everything else only at T^* = 2d + 1.5\, d_{\text{ff}} = 24,704. Memory is the harder wall: the fp16 weights are 32 MiB per head, 1 GiB per layer, 32 GiB for 32 layers if kept for the backward pass. The fix is exact: process keys in blocks, keep a running max and a running sum, and rescale the past whenever the max grows. That's softmax attention to 7.8 \times 10^{-16} with O(T) extra memory. Hooks: d^2 for the weights, T^2 for the pairs and keep a running max, rescale the past. First read: the FLOP table, the memory table, the online-softmax trace.

Counting FLOPs

One rule does all the counting: multiplying an m \times k matrix by a k \times n matrix is mn dot products of length k, so it costs 2mnk FLOPs (one multiply and one add per term). Applied to one layer (forward pass, one sequence, multi-head attention with h d_k = d):

F_{\text{proj}}

FLOPs for the four projections W_Q, W_K, W_V, W_O

T

tokens in the sequence

d

the model width d_{\text{model}} (4,096 for Llama 2 7B)

F_{\text{quad}}

FLOPs for the scores QK^\top and the weighted sum PV, summed over heads

F_{\text{FFN}}

FLOPs for the feed-forward network (Llama's SwiGLU has three matrices)

d_{\text{ff}}

the FFN's hidden width (11,008 for Llama 2 7B)

QK^\top

the scores: every query against every key (T \times T per head)

PV

the weighted sum: weights times values

Softmax's exponentials are also O(T^2) but aren't matrix multiplies, so they're left out. Llama 2 7B, per layer:

At Llama 2's own context length, the "quadratic" part is a minority of the layer. The T^2 term matters for long contexts. Hook: d^2 for the weights, T^2 for the pairs. (Card 11.)

The crossover

When does the quadratic part cost as much as everything else? Set them equal and divide by 4Td:

T^*

the context length where attention's pairwise part equals the projections plus FFN

For Llama 2 7B: 2 \times 4,096 + 1.5 \times 11,008 = 24,704 tokens. With a kernel that skips masked blocks the pairs cost 2Td(T + 1) and the crossover is 49,407 exactly; "double it" (2 \times 24,704 = 49,408) is the large-T shortcut (Card 11). Per token at T = 4,096: projections 134M FLOPs and FFN 271M, against 67M for the pairs.

Memory: the real wall

FLOPs are a minority at 4K. Memory is not. The naive implementation materializes P, a T \times T matrix per head, and keeps it for the backward pass (Section 8):

At 128K tokens a single head's weight matrix wouldn't fit on an 80 GB GPU. The last column is what an exact algorithm needs instead: one number per (head, query row). (Card 12.)

Online softmax: exact attention, one block at a time

Remember the river: to average a stream you don't store the stream. The obstacle is the softmax denominator, which needs every score before any weight is known, and the max we subtract for stability, which isn't known until the end. The fix is to keep a running version of each and correct it as you go. For one query row, process the keys in blocks n = 1, 2, \ldots:

m_n

the running max: the largest score seen after block n

n

which block of keys we've just processed

m_{n-1}

the running max before this block (start at -\infty)

\max_{j \in n}

the largest score among the keys in block n

s_j

the scaled score of key j

j \in n

key j belongs to block n

l_n

the running denominator, measured relative to the current max

e^{m_{n-1} - m_n}

the rescale factor: 1 if the max didn't change, smaller if it grew

l_{n-1}

the running denominator before this block (start at 0)

\sum_{j \in n}

add up over the keys in block n

\mathbf o_n

the running unnormalized output (a weighted sum of values)

\mathbf o_{n-1}

the running output before this block (start at 0)

\mathbf v_j

key j's value

e

the base of the natural log

At the end, the output is \mathbf o / l, and m + \ln l is the row's logsumexp, the one number kept for the backward pass. By induction, after every block l = e^{-m} \sum e^{s_j} and \mathbf o = e^{-m} \sum e^{s_j} \mathbf v_j over the keys seen so far, so \mathbf o / l is exactly the softmax-weighted average. Hook: keep a running max, rescale the past. (Card 15.)

By hand, on bank's row (one key per block):

Exactly Section 4's answer. When river's score raised the max, everything accumulated so far was multiplied by 0.243 to put it on the new scale.

On a real row (4,096 keys, d_v = 128, 32 blocks of 128): the max changed only 3 times (counting the first block); the logsumexp is 12.6882 both ways; the maximum error against direct softmax is 7.8 \times 10^{-16} in fp64 and 1.7 \times 10^{-7} in fp32. Block size doesn't matter (blocks of 1, 16, 128, 1,024 and 4,096 all agree to about 10^{-15}).

Merging partial results

The same formula merges two partial results computed on disjoint sets of keys, for example two GPUs that each saw half the context, or several thread blocks during decoding:

m_1

the max score seen by part 1

m_2

the max score seen by part 2

m

the combined max

l_1

part 1's denominator, relative to m_1

l_2

part 2's denominator, relative to m_2

l

the combined denominator

\mathbf o_1

part 1's unnormalized output

\mathbf o_2

part 2's unnormalized output

\mathbf o

the combined unnormalized output; the answer is \mathbf o / l

Simulated merge error: 1.0 \times 10^{-15}. This is exact, unlike low-rank or sparse approximations of attention.

In code it is ten lines: a loop over blocks that updates m, l and o exactly as above, then returns o / l and m + log(l). Derivations Card 15 proves it; Interview Bank Worked Problem 2 traces it on 4,096 keys.

Backward by recomputation. Keep only the row's logsumexp \text{lse}_i = m_i + \ln l_i; in the backward pass recompute each block of scores and get P_{ij} = e^{s_{ij} - \text{lse}_i} exactly. This recurrence (Milakov and Gimelshein, 2018; Rabe and Staats, 2021) is the math inside FlashAttention (Dao et al., 2022). FLOPs don't drop; the speedup comes from moving far fewer bytes between GPU memory levels (LLM Ch 5).

Remember: d^2 for the weights, T^2 for the pairs: 8Td^2 + 4T^2 d, crossover T^* = 2d + 1.5\, d_{\text{ff}}. Keep a running max, rescale the past: exact attention in O(T) memory, with one logsumexp per row for the backward.

Next question → Training sees all T tokens at once; generation adds them one at a time. What do we recompute, and what can we keep? (Section 13)

13. The KV Cache, and Sharing Keys and Values (MQA, GQA)

The question: generating token 4,097 of a chat with Llama 2 7B: what must we keep from the first 4,096 tokens, how big is it, and how does Llama 2 70B make it smaller?

The 30-second version: at decode step t only the new token's \mathbf q, \mathbf k and \mathbf v are new. Under a causal mask, old keys and values never change, so we cache them. Per token the cache is 2 \cdot L \cdot H_{kv} \cdot d_h \cdot b = 2 \cdot 32 \cdot 32 \cdot 128 \cdot 2 = 512 KiB for Llama 2 7B, or 2 GiB per 4,096-token sequence. After 13.5 GB of fp16 weights on an 80 GB GPU, that leaves room for at most 30 such sequences. Decoding is memory-bound: each token reads every weight (6.6 ms on an A100, so at most 151 tokens/s at batch 1) plus the cache. Sharing K and V across query heads shrinks the cache: MQA uses one KV head, GQA uses groups. Llama 2 70B has 8 KV heads for 64 query heads: 320 KiB per token instead of 2.5 MiB, which means 16 vs 2 concurrent 4K sequences on 2×80 GB. Hook: share K and V, keep Q. First read: the formula, the model table, the budget.

Prefill vs decode

Serving a chat has two phases with opposite bottlenecks:

• Prefill: the whole prompt goes through at once, like training. Big matrix multiplies, lots of arithmetic per byte loaded. At T = 4,096 the projections do about 1,365 FLOPs per byte of weights. An A100 does 312 TFLOP/s and reads 2,039 GB/s, so it needs about 153 FLOPs per byte (its ridge point) to be compute-bound. Prefill is far above it: compute-bound.
• Decode: one new token per step. Every weight is loaded to do one multiply-add for one token: about 1 FLOP per byte at batch 1. Memory-bound: the GPU waits on memory, not arithmetic.

Why K and V can be cached, and Q need not be

At step t the new token needs its own query against all keys and values so far. The old tokens' outputs are already computed, and the causal mask guarantees they never change (Section 6's test), so their keys and values never need recomputing and their queries are never needed again. In the library: you write a new request slip each time, but the spine labels and book contents stay on the shelf. Hook: share K and V, keep Q: Q is used once and thrown away.

The cache formula

2

one key and one value

L

the number of layers (each layer has its own cache)

H_{kv}

the number of key/value heads (= h for MHA, 1 for MQA, a group count for GQA)

d_h

the width of each head (d_k = d_v), usually 128

b

bytes per number: 2 for fp16/bf16, 1 for int8

Hook: "K and V, every layer, every KV head, every dimension, every byte." Multiply by tokens T and batch B for the total. (Card 13.)

An int8 cache halves every row (Llama 2 7B: 256 KiB per token). Note the units: cache sizes are in GiB (2^{30} bytes); weights and GPU capacity below are in GB (10^9), as vendors state them.

The serving budget

These are upper bounds: they ignore activations, memory fragmentation and framework overhead. Still, GQA is the difference between a product and a demo.

The decode speed limit. At batch 1, every token must stream all 13.5 GB of Llama 2 7B's weights from memory: 13.5 \text{ GB} / 2,039 \text{ GB/s} = 6.61 ms, at most 151 tokens/s, however fast the arithmetic is. A full 4K cache adds 1.05 ms per token (130 tokens/s). Batching many users amortizes the weight reads, which is why serving systems batch; but each user's cache is their own, so the cache reads don't amortize.

MQA and GQA: what sharing buys

Group the h query heads into H_{kv} groups; every query head in a group uses the group's one key head and one value head.

d_{\text{model}}

the model width (8,192 for Llama 2 70B)

H_{kv}

the number of key/value heads

d_h

the width per head (128)

W_Q

the query projection, d_{\text{model}} \times d_{\text{model}}

W_O

the output projection, d_{\text{model}} \times d_{\text{model}}

W_K

the key projection, d_{\text{model}} \times H_{kv} d_h: shrinks with fewer KV heads

W_V

the value projection, same shape as W_K

The last column is the subtle win. During decode, each cached key and value is read once from memory and used by every query head that shares it, so attention does h / H_{kv} FLOPs per byte: 1 for Llama 2 7B, 8 for Llama 2 70B, 4 for Llama 3 8B. Still far below the A100's ridge of 153, so decode attention stays memory-bound, but 8× less so. (Card 14.)

• MQA (Shazeer, 2019): H_{kv} = 1. Smallest cache; Ainslie et al. (2023) report quality loss and training instability. PaLM (Chowdhery et al., 2022) used it.
• GQA (Ainslie et al., 2023): 1 < H_{kv} < h. Close to MHA quality with a cache close to MQA's. An existing MHA checkpoint can be converted by mean-pooling its K and V heads within each group, then uptraining with about 5% of the original pre-training compute.
• GQA with H_{kv} = h is MHA; with H_{kv} = 1 it is MQA.

What comes next: paged KV caches (vLLM; Kwon et al., 2023; LLM Ch 12), quantized caches (int8, int4; LLM Ch 11), latent caches (DeepSeek-V2's multi-head latent attention, 2024), and sliding windows that cap the cache (Mistral 7B: 4,096; LLM Ch 4).

Remember: KV per token = 2 \cdot L \cdot H_{kv} \cdot d_h \cdot b; Llama 2 7B: 512 KiB. Every decode step reads all weights plus the cache, so decode is memory-bound. Share K and V, keep Q: GQA cuts the cache by h / H_{kv}.

Next question → Can we build all of it ourselves and match PyTorch? (Section 14)

14. Build It From Scratch, Then Use It Like a Pro

The question: can we write multi-head causal attention with a KV cache in about 50 lines of numpy, and match PyTorch's F.scaled_dot_product_attention?

The 30-second version: in numpy: project, reshape to (h, T, d_k), score, mask, stable softmax, multiply by V, merge heads, apply W_O. A decode loop that appends to a cache must reproduce the full causal pass exactly. In PyTorch, F.scaled_dot_product_attention dispatches to a fused kernel (FlashAttention, memory-efficient, or plain math). Three gotchas bite, hook Polarity, Placement, Phase: what True means in a boolean mask, how is_causal aligns non-square masks, and the fact that dropout_p ignores model.eval(). First read: run the code once; read the gotchas table.

From scratch in numpy

Every line follows Score, Scale, Mask, Softmax, Sum, then Split, Attend, Concat, Project:

import numpy as np

def softmax(S):                                    # stable: subtract the row max (Section 12)
    S = S - S.max(axis=-1, keepdims=True)
    E = np.exp(S)
    return E / E.sum(axis=-1, keepdims=True)

def split(Y, h):                                   # (T, d) -> (h, T, d_k)
    T, d = Y.shape
    return Y.reshape(T, h, d // h).transpose(1, 0, 2)

def mha_forward(X, Wq, Wk, Wv, Wo, h, mask=None):
    """X: (T, d_model). Wq, Wk, Wv, Wo: (d_model, d_model). mask: (T, T) bool, True = blocked."""
    Q, K, V = split(X @ Wq, h), split(X @ Wk, h), split(X @ Wv, h)       # Ask, Advertise, Deliver
    S = Q @ K.transpose(0, 2, 1) / np.sqrt(Q.shape[-1])                  # Score, Scale
    if mask is not None:
        S = np.where(mask, -np.inf, S)                                   # Mask (the scores, not the weights)
    P = softmax(S)                                                       # Softmax
    O = P @ V                                                            # Sum: (h, T, d_k)
    concat = O.transpose(1, 0, 2).reshape(X.shape[0], -1)                # Concat heads: (T, d_model)
    return concat @ Wo, P                                                # Project

def decode_with_cache(X, Wq, Wk, Wv, Wo, h):
    """Feed tokens one at a time; keep K and V, never Q."""
    K_cache, V_cache, outs = [], [], []
    for t in range(X.shape[0]):
        x = X[t:t + 1]
        K_cache.append(x @ Wk); V_cache.append(x @ Wv)
        q = split(x @ Wq, h)                                             # (h, 1, d_k): only the new query
        K, V = split(np.vstack(K_cache), h), split(np.vstack(V_cache), h)
        P = softmax(q @ K.transpose(0, 2, 1) / np.sqrt(q.shape[-1]))     # no mask needed: the cache holds only the past
        outs.append((P @ V).transpose(1, 0, 2).reshape(1, -1) @ Wo)
    return np.vstack(outs)

# The toy: "the river bank was steep", two heads (sense, subject), W_O = I
X = np.array([[0,0,1,0],[1,0,0,0],[1,1,0,0],[0,0,1,1],[1,0,0,1]], float)
Wq = np.array([[1,0,0,0],[1,-1,0,0],[0,0,0,1],[1,0,1,0]], float)   # [head-1 W_Q | head-2 W_Q]
Wk = np.array([[1,0,1,0],[0,1,1,0],[0,0,0,1],[0,0,-1,0]], float)
Wv = np.array([[1,0,0,1],[-1,1,1,0],[0,0,0,0],[0,0,0,0]], float)
Wo = np.eye(4)
causal = np.triu(np.ones((5, 5), bool), k=1)                       # True above the diagonal = future

out, P = mha_forward(X, Wq, Wk, Wv, Wo, h=2, mask=causal)
print("bank, head 1 weights:", P[0, 2, :3].round(3))                # [0.14  0.576 0.284]
print("bank output:", out[2].round(3))                              # [0.576 0.284 0.333 0.667]
print("toy cache max error:", np.abs(decode_with_cache(X, Wq, Wk, Wv, Wo, 2) - out).max())

rng = np.random.default_rng(0)                                     # a random check: T=16, d=32, h=4
Xr = rng.normal(size=(16, 32)); W = [rng.normal(size=(32, 32)) / np.sqrt(32) for _ in range(4)]
full, _ = mha_forward(Xr, *W, h=4, mask=np.triu(np.ones((16, 16), bool), k=1))
err = np.abs(decode_with_cache(Xr, *W, h=4) - full).max()
assert err < 1e-12, err
print("random cache max error:", err)

split reshapes to (T, h, d_k) and then transposes (reshaping straight to (h, T, d_k) is Section 10's scrambled-heads bug). The decode loop never builds a mask and never stores a query, and it matches the full causal pass within the assert's 10^{-12}. Write this test for any cache you build.

The same thing in PyTorch

Checked against PyTorch 2.14.1 (CPU). The block is guarded so it runs, and prints a note, without PyTorch.

import numpy as np
try:
    import torch
    import torch.nn.functional as F
except ImportError:
    torch = None
    print("PyTorch not installed; skipping this block.")

if torch is not None:
    torch.manual_seed(0)
    B, h, T, d_k = 2, 4, 6, 8
    q, k, v = (torch.randn(B, h, T, d_k) for _ in range(3))

    # 1. SDPA vs the numpy recipe (float32 vs float64, CPU)
    out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
    Qn, Kn, Vn = (t.numpy().astype(np.float64) for t in (q, k, v))
    S = Qn @ Kn.swapaxes(-1, -2) / np.sqrt(d_k)
    S = np.where(np.triu(np.ones((T, T), bool), 1), -np.inf, S)
    P = np.exp(S - S.max(-1, keepdims=True)); P /= P.sum(-1, keepdims=True)
    err = np.abs(out.numpy() - P @ Vn).max()
    assert err < 1e-6, err
    print("SDPA vs numpy:", err)                                     # float32 rounding, ~1e-7

    # 2. Polarity: in SDPA a boolean True means "may attend"
    allowed = torch.ones(T, T, dtype=torch.bool).tril()               # causal, written as "allowed"
    same = F.scaled_dot_product_attention(q, k, v, attn_mask=allowed)
    flipped = F.scaled_dot_product_attention(q, k, v, attn_mask=~allowed)  # nn.MultiheadAttention-style "blocked"
    print("allowed-mask matches is_causal:", torch.allclose(same, out, atol=1e-6))
    print("flipped polarity changes output by:", (flipped - out).abs().max().item())

    # 3. Placement: is_causal is top-left aligned when q_len != k_len
    q_new = torch.randn(B, h, 2, d_k)                                  # 2 new queries over 6 cached keys
    top_left = F.scaled_dot_product_attention(q_new, k, v, is_causal=True)
    bottom_right = torch.ones(2, T, dtype=torch.bool).tril(diagonal=T - 2)
    right = F.scaled_dot_product_attention(q_new, k, v, attn_mask=bottom_right)
    print("top-left == bottom-right?", torch.allclose(top_left, right))  # False: wrong for decoding

    # 4. Phase: dropout_p applies even in eval mode; pass 0.0 yourself at inference
    a = F.scaled_dot_product_attention(q, k, v, dropout_p=0.5)
    b = F.scaled_dot_product_attention(q, k, v, dropout_p=0.5)
    print("dropout made two calls differ:", not torch.allclose(a, b))

    # 5. GQA: 4 query heads share 2 KV heads
    k2, v2 = torch.randn(B, 2, T, d_k), torch.randn(B, 2, T, d_k)
    g = F.scaled_dot_product_attention(q, k2, v2, is_causal=True, enable_gqa=True)
    r = F.scaled_dot_product_attention(q, k2.repeat_interleave(2, 1), v2.repeat_interleave(2, 1), is_causal=True)
    print("enable_gqa == repeat_interleave:", torch.allclose(g, r, atol=1e-6))
    print("torch", torch.__version__)

The signature is F.scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None, enable_gqa=False), with inputs shaped (B, h, T, d_k). It picks a backend for you: a FlashAttention kernel (GPU, fp16/bf16), a memory-efficient kernel, or the plain math path (what runs on CPU). All compute the same thing.

The gotchas: Polarity, Placement, Phase

Remember: Polarity, Placement, Phase. Check what True means, align causal masks bottom-right when decoding with a cache, and turn dropout off yourself. Always test the cached decode against the full causal pass.

Next question → Where does attention sit in a real LLM, and what do modern models actually use? (Section 15)

15. In the Real World: Attention Inside the Block, and What Modern LLMs Use

The question: in Llama 2 7B, attention is 31.9% of the parameters and only 7.7% of a layer's FLOPs at 4K tokens (causal). Yet it gets most of the engineering effort. Where does it sit, and what has each generation of models changed?

The 30-second version: attention mixes information across positions; the feed-forward network (FFN) transforms each position on its own; residual connections carry the running vector (and mitigate rank collapse); normalization keeps the scale in range. A minimal pre-norm block is \mathbf x \leftarrow \mathbf x + \text{Attn}(\text{Norm}(\mathbf x)), then \mathbf x \leftarrow \mathbf x + \text{FFN}(\text{Norm}(\mathbf x)). Across model generations the softmax core is unchanged. What changed is how K and V are shared (MHA → MQA/GQA → latent), where position enters (learned → RoPE), which pairs are allowed (sliding windows), and how logits are bounded (QK-norm, soft-capping). At long context the KV cache, not the FLOPs, dominates serving cost. First read: the block and the model table.

The block

\mathbf x

one token's running vector (the residual stream), updated in place by each sub-block

\leftarrow

replace with: compute the right side, then overwrite \mathbf x

\text{Attn}

multi-head attention over all visible tokens (this chapter)

\text{Norm}

a normalization layer (LayerNorm or RMSNorm) that keeps the vector's scale in range

\text{FFN}

the feed-forward network applied to each token separately

Each piece has a job no other piece does. Attention moves information between tokens but is linear in the values and averaging by nature (Section 7). The FFN adds nonlinearity per token. The residual "+" keeps each token's own information, so the stack doesn't collapse to the mean (Section 11). This is the pre-norm arrangement used since GPT-2. LLM Ch 4 opens the block up: RMSNorm, pre- vs post-norm, SwiGLU.

Our toy did one step of this: bank went from [1, 1, 0, 0] to [1.576, 1.284, 0.333, 0.667] after attention plus the residual.

Where the parameters and FLOPs go

At 4,096 tokens, the causal quadratic part adds 1.07 GFLOP per token on top of the 12.95 GFLOP that the non-embedding weights cost (2N), about 7.7% of Llama 2 7B's forward FLOPs. So why the engineering effort? Because attention's costs are the ones that grow with context: the T^2 memory of training (Section 12) and the per-token KV cache of serving (Section 13). The FFN's cost per token is fixed.

What modern models use

Read down the columns: heads moved from MHA to sharing; positions moved to RoPE; allowed pairs sometimes became windows; logit control appeared. The middle of every row, softmax of QK^\top / \sqrt{d_k} times V, did not change.

Beyond the decoder. Cross-attention lets text queries read keys and values from another source (encoder–decoder models, many vision-language models); retrieval-augmented generation is a hard, external version of the same lookup (LLM Ch 13). Interpretability: the QK circuit decides where a head looks and the OV circuit what it writes (Elhage et al., 2021); induction heads are a well-documented circuit tied to in-context learning (Olsson et al., 2022). Weights alone are still not explanations (Section 4).

Where this connects: the full block (RMSNorm, pre- vs post-norm, SwiGLU, RoPE, sliding windows) is LLM Ch 4; FlashAttention as a GPU algorithm and RoPE math are Ch 5; FLOP counts in scaling laws are Ch 7; paged KV caches and batching are Ch 12; linear attention and state-space models are Ch 15.

Remember: attention mixes positions, the FFN works per position, and residuals keep the stack from collapsing. Modern LLMs keep softmax attention and change how K and V are shared, where position enters, which pairs are allowed, and how logits are bounded.

Next → We have one attention layer, correct, stable and costed. A real LLM stacks it with normalization, a feed-forward network and positions, and every model since GPT-2 has changed those parts. That's LLM Ch 4 (From GPT to LLaMA). First, the chapter on one page (Section 16), then practice.

16. The Chapter on One Page

One map of the whole chapter and every memory hook in one table. Use it before an interview, or to check a second read: if any row feels unfamiliar, go back to its section.

The memory map

Every hook in one place

Remember: if you can redraw the memory map on a blank page and say the hook for every row of the table, you have the chapter. Now practice.

17. Practice

The chapter ends here. Practice lives in two tabs, so this page stays a story you can read straight through.

• LLM Ch 3 · Derivations — 18 derivation cards (from scaled dot-product attention to the online-softmax recurrence), each with its key move from the S-O-F-T toolkit and the common slip; 25 formulas to memorize; and a mastery tracker.
• LLM Ch 3 · Interview Bank — 41 questions in Levels 1–3, each answered in four beats (intuition, answer, why, what they'll ask next); Level 4 with two worked research-level problems (saturation at d_k = 128; exact attention without the T \times T matrix), 20 drills, and a 40-row rapid-fire table for the day before.

Suggested order: finish the chapter's second read → Derivations Cards 1–7 → Interview Bank Levels 1–2 → Cards 8–18 → Levels 3–4.