Part II — Attention

Self-attention with trainable weights (Q, K, V)

Lesson 04 built attention with no parameters: a token's context vector was a softmax-weighted blend of the others, where the weights came straight from dot-products of the raw embeddings. It worked, but it had no knobs — nothing to learn. This lesson adds the three projection matrices that turn attention into a trainable layer, and derives the single constant, 1/√d_k, that keeps it from strangling its own gradient.

What the last step left broken

In lesson 04, each token used its raw embedding x in all three roles at once: it was the thing doing the looking, the thing being matched against, and the thing handed over. With no free parameters, the model cannot learn what to look for or what to share — the attention pattern is a fixed function of the embeddings. And every token is symmetric: nothing distinguishes "I am searching" from "I am on offer." A trainable layer needs weights, and it needs them to sit in exactly the right three places.

Linear position

Forced by: parameter-free attention (04) has no knobs — it cannot learn what is relevant, because the query, the key, and the value are all the same raw vector.
New idea: project each token into three different learned spaces — a query (what am I looking for?), a key (what do I offer?), and a value (what will I pass on?) — via matrices W_q, W_k, W_v; score with q·kᵀ, then divide by √d_k so the softmax stays in its useful range.
Forces next: this layer still lets every position attend to every other — including the future token it is trying to predict — and it computes only one similarity "view." Lesson 06 must add the causal mask and multiple heads.

The plan

Five moves. (1) See precisely why raw dot-product attention cannot be trained. (2) Split the one role into three — query, key, value — each its own learned projection. (3) Derive 1/√d_k from the variance of a dot product, and see what breaks without it. (4) Assemble the full operation, softmax(QKᵀ/√d_k)·V, with shapes. (5) Drive the scale toggle and watch attention either learn or freeze.

1 · Why parameter-free attention can't be trained

Recall lesson 04. For a query token x_i, the attention weight on token x_j was softmax_j(x_i · x_j), and the context vector was z_i = Σ_j α_ij x_j. Every quantity is a fixed function of the embeddings. There is no matrix anywhere for gradient descent to adjust, so the only way the attention pattern can change is if the embeddings themselves change — but the embeddings also have to serve every other part of the model. Attention has no degrees of freedom of its own.

Worse, the raw scheme collapses three distinct jobs into one vector. When token x_i attends, it plays the searcher: "which tokens are relevant to me?" When some other token is attended to, it plays two more roles: it is the thing being matched (does it satisfy the searcher?) and the thing being retrieved (what content flows back?). A word like "it" should search for a noun to resolve to, advertise itself very differently, and pass along yet another signal. One vector cannot be optimized for three opposing purposes. We need to let the token wear three different hats.

2 · Three roles, three matrices

Give the model three learned linear projections. From each token embedding x (dimension d) we produce three vectors of dimension d_k:

q = W_q x (query) · k = W_k x (key) · v = W_v x (value)

role	question it answers	used as
query q	"What am I looking for?"	the row token when it attends
key k	"What do I offer to matchers?"	matched against every query via dot-product
value v	"What content do I hand over if matched?"	the vectors that get blended into the output

The score between token i (as query) and token j (as key) is q_i · k_j — high when what i seeks aligns with what j advertises. Softmax over j turns the scores into weights; the context vector blends the values, not the raw embeddings:

α_ij = softmax_j(q_i · k_j) → z_i = Σ_j α_ij v_j

Now there are three matrices to learn. The model can discover, for instance, that pronoun-queries should match noun-keys, while the values carry the information the pronoun needs — three separately optimized channels. This is the entire content of "self-attention with trainable weights." One subtlety remains, and it is the difference between a layer that learns and one that freezes on the first step.

3 · Why divide by √d_k

The scores feed a softmax, and softmax is exquisitely sensitive to the scale of its inputs. Track that scale. Assume the entries of q and k are roughly independent with mean 0 and variance 1 (which is what sane initialization arranges). A single score is a dot product over d_k dimensions:

q · k = Σ_{c=1..d_k} q_c k_c

Each term q_ck_c has variance ≈ 1, and there are d_k independent terms, so the variance of the sum is ≈ d_k and its standard deviation is √d_k. In other words, unscaled scores grow like √d_k. For GPT-2's head width d_k = 64, that is a spread of about ±8.

Feed scores of spread 8 into a softmax and the largest one wins by a landslide: a gap of 8 means e⁸ ≈ 3000× more probability mass. The softmax output is essentially one-hot before training has done anything. And a one-hot softmax is flat — its Jacobian is nearly zero — so almost no gradient flows back through the attention weights. The layer is frozen at initialization: it cannot learn where to look because the derivative that would teach it is ~0.

The fix is to pin the score variance back to ≈ 1 regardless of head width. Divide every score by √d_k:

score_ij = (q_i · k_j) / √d_k

Now Var(score) ≈ d_k / (√d_k)² = 1 for any d_k. Softmax stays in its informative regime, gradients flow, and the layer trains. This is one of the rare constants a Transformer derives rather than tunes — it falls straight out of the variance algebra, and the widget below lets you watch attention saturate the moment you turn it off.

The trap

The √d_k is not a decorative normalization you can drop "to keep things simple." Without it, a wide head produces one-hot attention with a dead gradient, and your loss flatlines near its initial value while the attention maps look frozen and spiky from step one. It is a correctness constant, not a style choice.

4 · The full operation, with shapes

Stack all T tokens into a matrix X of shape (T, d). The three projections become matrix multiplies, and the whole layer is four lines:

Q = X @ W_q            # (T, d) @ (d, d_k) -> (T, d_k)   queries
K = X @ W_k            #                              -> (T, d_k)   keys
V = X @ W_v            #                              -> (T, d_k)   values
S = (Q @ K.T) / sqrt(d_k)     # (T, d_k) @ (d_k, T)   -> (T, T)    scores
A = softmax(S, axis=-1)       #                       -> (T, T)    weights (rows sum to 1)
Z = A @ V                     # (T, T) @ (T, d_k)     -> (T, d_k)  context vectors

Read the shapes as the story. Q@Kᵀ produces a T×T table — every token's query dotted with every token's key — the raw compatibility matrix. Softmax along each row makes each query's weights a distribution. A@V then replaces each token with a weighted blend of all values. The output Z is (T, d_k): one enriched context vector per position, ready to be projected back and handed to the rest of the block.

In real code this is a small module holding the three weight matrices; batching just adds a leading B so tensors become (B, T, d) and the matmuls batch over it. Nothing else changes. Note the T×T score table: it is quadratic in sequence length, which becomes the memory wall at long context — the reason FlashAttention exists (see the cross-reference below).

5 · Watch the scale decide everything

The widget uses the six-token sentence "Your journey starts with one step," each token a small embedding, with fixed (seeded) W_q, W_k, W_v. It computes the full T×T attention matrix and shows it as a heatmap — darker means more attention. The ÷√d_k button toggles the scaling. Slide d_k up with scaling off and watch every row collapse toward a single black cell (one-hot, frozen). Turn scaling on and the rows breathe again — attention that can actually learn. The KPIs report the score spread (which should track √d_k unscaled, and stay ≈ 1 scaled) and the peak weight in the highlighted query row.

Failure modes & checklist

Failure modes

Dropping the √d_k scale. Scores grow like √d_k, softmax saturates to one-hot, gradient vanishes. Signal: loss stuck near init; attention maps spiky and frozen from step 1.
Sharing one matrix for Q, K, and V. Collapses the three roles you just separated; the layer loses most of its expressive power.
Scaling by d_k instead of √d_k. Over-shrinks scores toward 0, softmax becomes uniform, and attention averages everything indiscriminately.
Softmax over the wrong axis. Normalizing down columns instead of across each query's keys makes the "weights" meaningless.

Checklist

Three separate matrices W_q, W_k, W_v, each (d, d_k).
Scores = QKᵀ divided by √d_k (not d_k).
Softmax along the key axis so each query row sums to 1.
Context = weights times V (values), not times the raw embeddings.
Output shape (T, d_k) — one context vector per token.

Checkpoint

Try it

Turn the ÷√d_k scaling OFF and drag d_k from 2 up to 64. Watch the "peak weight in row" climb toward 1.00 and the "softmax state" flip to saturated (one-hot) — that is a dead gradient. Now turn scaling ON and drag d_k again: the score spread stays near 1 and the peak weight stays moderate at every width. You have just reproduced, by hand, the reason every Transformer divides by √d_k. Which of the three properties from lesson 00 — accurate, trainable, steerable — does this constant protect?

Where this points next

You now have a genuine trainable attention layer. But it has two problems that a language model cannot tolerate. First, look at the heatmap: token 2 attends freely to tokens 3, 4, 5, 6 — all of which come after it. A model trained to predict the next token would simply read the answer from the future, learn nothing, and be useless at inference time where the future does not exist yet. Second, the layer computes exactly one similarity geometry — one notion of "relevant" — when language needs many at once (grammatical, referential, positional). Lesson 06, Causal & multi-head attention, fixes both: a mask that forbids attending to the future, and multiple heads that attend in parallel.

Takeaway

Trainable self-attention splits each token's single embedding into three learned projections — query (what I seek, W_q x), key (what I offer, W_k x), and value (what I pass on, W_v x) — scores queries against keys with QKᵀ, and blends the values: Z = softmax(QKᵀ/√d_k)·V. The 1/√d_k is forced, not chosen: a d_k-dimensional dot product has standard deviation √d_k, so without the scale a wide head saturates softmax to one-hot and kills its own gradient. Three matrices give the layer knobs to learn what to look for and what to share; one derived constant keeps it learnable.

Interview prompts

Why can't the parameter-free attention of lesson 04 be trained? (§1 — it has no matrices of its own; the pattern is a fixed function of the embeddings, and one vector must serve query, key, and value at once.)
What are the query, key, and value, and what question does each answer? (§2 — query = "what am I looking for," key = "what do I offer," value = "what do I hand over if matched"; each is a learned projection of the embedding.)
Derive why attention scores are divided by √d_k. (§3 — a dot product of unit-variance vectors over d_k dims has variance d_k, std √d_k; unscaled logits saturate softmax to one-hot with ~0 gradient, so the scale pins variance to ≈1.)
What exactly breaks if you omit the scale? (§3 — softmax collapses to one-hot at init, its Jacobian ≈ 0, gradients through attention vanish, and the loss flatlines.)
Why blend the values rather than the raw embeddings? (§2 — the value is a separately learned channel for "what content to pass on," decoupled from "how to be matched" (the key); blending raw x would re-entangle those roles.)
What is the shape of the score matrix and why does it matter at scale? (§4 — (T, T), quadratic in sequence length; it is the memory bottleneck at long context and the reason for FlashAttention.)
Would scaling by d_k instead of √d_k work? (§3/failure modes — no; it over-shrinks scores so softmax goes uniform and attention averages everything, losing selectivity.)

Companion reads: cs336 · 02 The Transformer counts the params and FLOPs of this same operation; gpt_mini · 01 Architecture shows it as runnable code; cs336 · 07 Kernels explains how the T×T scores are computed without ever materializing them.