4. Attention from first principles

Attention lets each token build a new vector as a weighted average of information from other tokens, with weights computed from content. It is a soft, differentiable dictionary lookup, and it is the only place in a transformer where tokens exchange information.

Start from a dictionary

A Python dict maps keys to values: give it a query, it finds the matching key and returns its value. That is hard lookup - exact match or nothing - and it has no gradient, so it can’t be learned.

Make it soft:

1. Compare the query with every key, giving a similarity score.
2. Turn the scores into weights that are positive and sum to 1 (softmax).
3. Return the weighted average of all the values.

If one key matches far better than the rest, the result approaches hard lookup; if several match, you get a blend. Every step is differentiable, so gradient descent can learn what to look for and what to offer.

Queries, keys and values

In self-attention every token plays all three roles, each through its own learned linear projection of the token’s vector x:

• Query q = x W_Q: “what am I looking for?”
• Key k = x W_K: “what do I contain that others might look for?”
• Value v = x W_V: “what do I hand over if someone attends to me?”

Separate roles matter. A pronoun’s query can look for “noun naming a person” while that noun’s key advertises exactly this, and its value carries the details the pronoun needs. One vector couldn’t play all three roles well.

The formula

Stack all T tokens into X of shape (T, d). Then Q = XW_Q, K = XW_K, V = XW_V, with Q and K of shape (T, d_k) and V of shape (T, d_v).

Step	Operation	Shape	Meaning
1	S = QKᵀ	(T, T)	S[i, j] = how well token i’s query matches token j’s key
2	S / √d_k	(T, T)	rescale so the softmax doesn’t saturate
3	+ M	(T, T)	mask: add −∞ where attention is forbidden
4	A = row-wise softmax	(T, T)	each row is a probability distribution over positions
5	A V	(T, d_v)	each output is a weighted average of value vectors

Row i of A answers “where does token i read from?” This is the attention pattern people visualize.

A worked example

Three tokens, “the cat sat”, with d_k = 2. We compute the output for “sat”.

• Query of “sat”: q = [1, 0]
• Keys: the = [0, 1], cat = [2, 0], sat = [1, 1]
• Values, one-hot for clarity: the = [1, 0, 0], cat = [0, 1, 0], sat = [0, 0, 1]

1. Scores q·k: the → 0, cat → 2, sat → 1.
2. Divide by √2: 0, 1.41, 0.71.
3. Softmax: e⁰ = 1, e^1.41 ≈ 4.11, e^0.71 ≈ 2.03; the sum is about 7.14, so the weights are about 0.14, 0.58 and 0.28.
4. Output = 0.14·[1,0,0] + 0.58·[0,1,0] + 0.28·[0,0,1] = [0.14, 0.58, 0.28].

“sat” now carries mostly information from “cat”, its subject. Training adjusts W_Q, W_K and W_V so that useful matches like this one score high.

Why divide by √d_k

Suppose the components of q and k are independent with mean 0 and variance 1. Their dot product sums d_k such products, so its variance is d_k. With d_k = 128 the scores have a standard deviation of about 11, and a softmax over numbers spread that far is nearly one-hot, with almost zero gradient. Dividing by √d_k restores variance 1 and keeps softmax responsive.

Softmax details worth knowing

• Adding a constant to every score in a row doesn’t change the result, so implementations subtract the row maximum first to avoid overflow.
• Softmax must put its weight somewhere. Heads with nothing useful to do often park weight on the first token, an “attention sink”, which matters for streaming and long-context tricks.
• Dividing scores by a temperature controls sharpness: the same knob as in sampling (Part 9).

Masking

• Causal mask: in language modeling, token i may attend only to positions j ≤ i. M is −∞ above the diagonal, so those weights become exactly 0. Training can then compute predictions for all T positions in one parallel pass while each prediction sees only the past (“teacher forcing”). The same structure enables the KV cache during generation (Part 9).
• Padding mask: when sequences of different lengths share a batch, padded positions are masked out.

Self-attention vs cross-attention

In self-attention, Q, K and V all come from the same sequence. In cross-attention, queries come from one sequence and keys and values from another: a translation decoder reading the source sentence, or text tokens reading image features. Same formula, different inputs.

Attention has no sense of order

Shuffle the input tokens and, without a mask, the outputs shuffle identically: attention treats its input as a set. “Dog bites man” and “man bites dog” would look the same. Position must be injected separately (Part 6).

What it costs

Per layer, the score matrix QKᵀ and the product AV each take about 2·T²·d FLOPs, and a naive implementation stores the (T × T) matrix. The projections and the MLP cost about 24·T·d² FLOPs per layer. So the quadratic term dominates only when T exceeds roughly 6·d - beyond about 25k tokens for d = 4096. Below that, MLPs and projections are most of the compute. Memory for the T × T matrix was the first wall, removed by FlashAttention (Part 9).

Questions you should now be able to answer

• What are the shapes of Q, K, V, QKᵀ and the output for T = 10, d = 512 and d_k = d_v = 64?
• Why use separate W_Q and W_K instead of computing XXᵀ directly? (Hint: symmetry, and each token’s similarity to itself.)
• Derive why the dot product of two random vectors has variance d_k, and what that does to softmax.
• How does the causal mask let us train on all positions in parallel?
• At what sequence length does attention’s quadratic cost overtake the rest of the layer when d = 8192?