The Attention Equation, Term by Term

Lesson 3 of 4 in Self-Attention: Every Token Looks at Every Token.

Here is the whole mechanism in one line, exactly as the transformer paper wrote it (Vaswani et al., 2017):

Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) · V

Read it term by term, inside the parentheses first:

  • QK^T — every query dotted against every key. For n tokens this produces an n × n table of raw Attention scores: row i holds token i’s query scored against every token’s key. A big entry means ‘your index card matches my question’.
  • / sqrt(d_k) — divide every score by the square root of the key dimension (Head dimension). A gentle rescaling with an outsized effect on training — the reason is below.
  • softmax — applied to each row separately: exponentiate every score, divide by the row’s total. Each row becomes positive weights that sum to exactly 1 — a budget of attention that the row must spend across the visible tokens. Softmax is the same normalize-into-weights move the model later uses on Logits.
  • · V — multiply the weights back onto the value vectors and sum. Each token’s output is a weighted average of everyone’s values, dominated by the tokens it scored highest.

A 3 by 3 heatmap of attention weights. Rows are the querying tokens the, cat, sat; columns are the attended tokens. Row values: the attends 0.40 to the, 0.20 to cat, 0.40 to sat; cat attends 0.20, 0.40, 0.40; sat attends 0.25, 0.25, 0.50. Every row sums to 1; the strongest cell is sat attending to itself at 0.50.

The attention matrix from this lesson’s worked example — a toy, three-token sequence with 2-dimensional heads, computed by hand in the deep dive below. Each row is one token’s attention budget and sums to 1; weights rounded to two decimals. Toy example: real models use hundreds of dimensions, but the arithmetic is exactly this. (calculated — source: Worked example in this lesson; equation from Vaswani et al. (2017))

Why divide by sqrt(d_k)? Dot products grow with vector length: sum d_k products of roughly unit-scale numbers and the result’s typical size grows like sqrt(d_k). With d_k in the hundreds, raw scores get large — and softmax is an exponential amplifier. Feed it scores like 20 versus 5 and it returns weights of nearly 1 and nearly 0: a saturated, almost one-hot distribution. Saturated softmax has near-zero gradients, so during training almost no learning signal flows back through the attention pattern. Dividing by sqrt(d_k) cancels the growth, keeps early-training scores in softmax’s responsive range, and is cited in the original paper as the fix for exactly this effect — it is the ‘scaled’ in scaled dot-product attention.

Everything above is the plain story; the deep dive below runs the full computation on three tokens with 2-dimensional vectors, small enough to check with a pencil.

One worked example, by hand

Three tokens — ‘the’, ‘cat’, ‘sat’ — and d_k = 2, so everything fits on paper. Assume the projections have already produced these vectors (hand-picked so the arithmetic stays clean):

  • Queries: q_the = [1, 0], q_cat = [0, 1], q_sat = [1, 1]
  • Keys: k_the = [1, 0], k_cat = [0, 1], k_sat = [1, 1]
  • Values: v_the = [1, 0], v_cat = [0, 2], v_sat = [2, 2]

Step 1 — scores (QK^T). Dot each query with each key. For ‘sat’: q_sat · k_the = 1·1 + 1·0 = 1, q_sat · k_cat = 0 + 1 = 1, q_sat · k_sat = 1 + 1 = 2. Doing all nine:

  • row ‘the’: [1, 0, 1]
  • row ‘cat’: [0, 1, 1]
  • row ‘sat’: [1, 1, 2]

Step 2 — scale by 1/sqrt(d_k) = 1/sqrt(2) ≈ 0.71. Row ‘sat’ becomes [0.71, 0.71, 1.41]. Tiny effect here; at d_k = 128 it is the difference between a usable softmax and a saturated one.

Step 3 — softmax each row. Exponentiate, then divide by the row sum. For row ‘sat’: e^0.71 ≈ 2.03, e^0.71 ≈ 2.03, e^1.41 ≈ 4.11; the sum is ≈ 8.17, so the weights are [0.25, 0.25, 0.50] (rounded to two decimals). Rows ‘the’ and ‘cat’ work out to [0.40, 0.20, 0.40] and [0.20, 0.40, 0.40] the same way. Stack the three rows and you have exactly the heatmap above — each row a budget summing to 1.

Step 4 — blend the values. The output for ‘sat’ is its weight row times the value vectors:

0.25·[1, 0] + 0.25·[0, 2] + 0.50·[2, 2] = [0.25, 0] + [0, 0.5] + [1, 1] = [1.25, 1.5]

That [1.25, 1.5] (using the two-decimal weights; the unrounded blend is ≈ [1.26, 1.50]) is the attention output for ‘sat’ — half of it the token’s own value, the rest drawn from its neighbours. In a real model the output joins the Residual stream and the same computation runs for every token, in every head, in every layer, with d_k in the tens-to-hundreds instead of 2 — but not one step you did here changes.

One thing to notice for the next lesson: row ‘the’ spent weight on ‘cat’ and ‘sat’ — tokens that come after it. Full self-attention allows that; a text generator cannot. That fix is causal masking.

Interactive checkpoint quiz (2 questions) — open this page in a browser to take it.