Attention

In June 2017, Vaswani and colleagues at Google published "Attention Is All You Need". Before it, sequence models read text with recurrence (section 5.2) or convolutions (small filters slid along the text, each seeing a few neighbours; section 25.1). The Transformer dropped both: every token looks directly at every other token through self-attention. Almost every large language model today is built on that design.

This module builds attention by hand on small numbers, proves why scores are divided by $\sqrt{d_k}$, explains masks, shows what attention cannot do alone (tell word order), and ends with multi-head attention and its cost.

6.1 The Question Attention Answers

For each token, attention answers: "which other tokens should I take information from, and how much?" It uses the three projections of section 1.6:

• Queries $Q \in \mathbb{R}^{N \times d_k}$: one row per asking token, describing what it looks for.
• Keys $K \in \mathbb{R}^{M \times d_k}$: one row per token that can be looked at, describing what it offers.
• Values $V \in \mathbb{R}^{M \times d_v}$: one row per key, the information handed over if chosen.

6.2 Scaled Dot-Product Attention, by HandRun this section

In self-attention all three come from the same sentence, so $N = M$; they differ in cross-attention (section 8.1). The whole computation is one line:

Read it from the inside out:

1. $QK^\top$ is an $N \times M$ table of scores: entry $(i, j)$ is query $i$ dotted with key $j$ (section 1.4).
2. Dividing by $\sqrt{d_k}$ keeps the scores from growing with the vector width (section 6.3).
3. The mask $M$ holds $0$ where attention is allowed and $-\infty$ where it is not, so blocked pairs get weight $e^{-\infty} = 0$ (section 6.4).
4. Softmax turns each row into weights that sum to 1 (section 4.3).
5. Multiplying by $V$ gives each query a weighted average of the value rows.

The letter $M$ does double duty: as a size ($N \times M$) it is the key count; as a matrix ($+M$) it is the $N \times M$ mask. $d_k$, the query and key width, is usually $64$ or $128$; $d_v$ is the value width.

In words. Each query is scored against every key; the scores are scaled, masked and softmaxed into weights, which average the values. Each query gets one output vector, built from the tokens it was allowed to see.

The Whole Computation on Two Tokens

Finish section 1.6's example. Its $W^Q$ and $W^K$, plus a value matrix $W^V$ with rows $[1, 0]$, $[0, 1]$, $[0, 0]$, $[1, 1]$, give:

1. Scores. Every query against every key: $\mathbf{q}_1 \cdot \mathbf{k}_1 = -9 - 9 = -18$, $\mathbf{q}_1 \cdot \mathbf{k}_2 = 0$, $\mathbf{q}_2 \cdot \mathbf{k}_1 = 0$, $\mathbf{q}_2 \cdot \mathbf{k}_2 = 2$.
2. Scale by $\sqrt{d_k} = \sqrt{2} \approx 1.414$: the rows become $[-12.73, \, 0]$ and $[0, \, 1.414]$.
3. Mask. With a causal mask, token 1 may not look at token 2, so its second score becomes $-\infty$: row 1 is $[-12.73, \, -\infty]$.
4. Softmax each row. Row 1 has only one allowed entry, so its weights are $[1, 0]$. Row 2 is $[e^0, e^{1.414}] = [1, 4.113]$, divided by their sum $5.113$: weights $[0.196, \, 0.804]$.
5. Weighted sum of values. Token 1's output is just $\mathbf{v}_1 = [0, -1]$. Token 2's output is $0.196 \cdot [0, -1] + 0.804 \cdot [1, 2] = [0.804, \, 1.413]$.

Each output is a mix of value vectors chosen by the scores. And the mask changed the answer: without it, nearly all of token 1's weight would go to token 2, which it is not supposed to see.

In words. A stepper walks the two-token example from queries and keys through scores, scaling, mask and softmax to the output, highlighting the entries involved. Mask and scaling can be toggled, and later steps update.

6.3 Why Divide by the Square Root of the Key WidthRun this section

The division is not a tuning trick; it follows from a short calculation with the rules of section 3.2. At the start of training, treat each query and key entry as independent with mean 0 and variance 1: $q_i, k_i \sim \mathcal{N}(0, 1)$. One raw score is

a sum of $d_k$ products. Find its mean and variance.

1. The mean is zero, by rules 1 and 2, since each $\mathbb{E}[q_i] = 0$:

1. Each product has variance 1. Its mean is 0, and by independence $\mathbb{E}[q_i^2 k_i^2] = \mathbb{E}[q_i^2]\,\mathbb{E}[k_i^2]$, each factor being a variance of 1:

1. The variances add (rule 4), over $d_k$ terms:

1. So a typical score has size $\sqrt{d_k}$:

At $d_k = 64$ a typical score is about $\pm 8$; at $d_k = 128$, about $\pm 11.3$.

What Goes Wrong Without the Scaling

A gap of $g$ between two scores makes one weight $e^{g}$ times the other. At $d_k = 128$, gaps of 20 are common, and $e^{20}$ is about 500 million, so softmax puts almost all weight on one key. This saturation stalls learning: section 4.3 showed that the softmax Jacobian is then nearly zero, so the query and key weights barely change.

By rule 3, dividing by $\sqrt{d_k}$ divides the variance by $d_k$:

Scaled scores have variance 1 at any head width, so softmax starts training where its gradients are healthy.

In words. A histogram of random query-key scores at a chosen head width, with or without scaling. A second panel runs softmax trials over 64 keys and reports how often the largest weight exceeds 0.99, and how small p times one minus p becomes when it does.

Measured (lab 2b, experiments/variance_scaling.py). The variance result holds at every head width. Over 200,000 random draws per width, $\operatorname{Var}(\mathbf{q}^\top\mathbf{k}) / d_k$ comes out at $0.9987$, $1.0040$, $1.0029$ and $1.0038$ for $d_k = 32, 64, 128, 256$. The scaled variance is $1.0$ throughout.

Measured (lab 2b). Over 500 random trials at $d_k = 128$ with $64$ keys, $35\%$ of the unscaled trials push $p_{\max}$ above $0.99$. With scaling, $0\%$ do. In the trials that do saturate, the gradient term $p_{\max}(1 - p_{\max})$ is $0.0020$, close to zero. The result is a fraction of trials because saturation depends on the random draw: one draw clears $0.99$ while the next can fall short.

6.4 MasksRun this section

A mask added to the scores puts $-\infty$ on forbidden pairs, so softmax spreads each row's weight over the allowed keys only. Three kinds cover every model in the course:

• No mask. Every token sees every token. Encoders use this (section 8.2), so a token's meaning can depend on both sides.
• Causal mask. Token $i$ sees tokens $1$ to $i$ only. Every text-writing model uses this (section 8.4), since token $i+1$ must not be visible while predicting it. For four tokens, with $\cdot$ for $0$:

• Padding mask. Short sequences in a batch are filled out with placeholder tokens; the mask puts $-\infty$ in their columns so no real token attends to them.

Masks combine by adding. Section 16.3 meets a fourth kind, the sliding window. Every row must keep at least one allowed key, or the softmax computes $0/0$ = NaN.

6.5 Attention Cannot See Word OrderRun this section

A score $\mathbf{q}_i \mathbf{k}_j^\top$ is built from two tokens' vectors and nothing else; no position appears anywhere. Shuffle the tokens and, with no mask, the outputs come back shuffled the same way and otherwise unchanged. This is permutation-equivariance: for any permutation matrix $P$ (one that reorders rows), $\operatorname{Attention}(P X, P X, P X) = P \operatorname{Attention}(X, X, X)$. In "The dog bit the man" and "The man bit the dog", "dog" gets exactly the same output vector. The model cannot tell who bit whom.

In words. Self-attention on a short sentence, showing every token's output vector. Shuffling the tokens reorders the outputs without changing their values. A switch adds position vectors, and then shuffling changes the outputs themselves.

The fix is to give each token its position before attention: section 7.1 shows the classic way, and section 15.1 the modern way, which builds position into the query–key score itself.

6.6 Multi-Head AttentionRun this section

One attention computation gives one set of weights per query, so linking a verb to its subject and a pronoun to its noun must blend into one average. Instead, run several smaller attention computations, heads, side by side, each with its own projections, and join the results:

$h$ is the number of heads, and $[\,;\,]$ means concatenate along the feature axis. $W_i^Q, W_i^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$ and $W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$ make head $i$'s queries, keys and values; the output projection $W^O \in \mathbb{R}^{(h \cdot d_v) \times d_{\text{model}}}$ maps the joined outputs back to model width. Usually $d_k = d_v = d_{\text{model}}/h$: the paper uses $h = 8$ and $d_{\text{model}} = 512$, so $d_k = 64$, and eight heads cost about the same as one head of width 512.

In words. Each head projects the tokens to its own 64 features and runs full attention there. The labels are patterns heads in trained models have been seen to learn; nobody assigns them. The outputs are joined and mixed by one output projection.

In words. Four heads read "The animal didn't cross the street because it was tired", with hand-set weights mimicking trained models: one looks at the previous word, one sends "it" to "animal", one sends each verb to its subject, and one attends to the sentence start and "because". The four tables stack into one tensor, a layer per head.

Why the Output Projection $W^O$ Is Needed

1. It mixes the heads. After joining, each head occupies its own 64 features; every output feature of $W^O$ combines all 512, so it can use what every head found.
2. It writes back into the shared vector. Each layer adds its output to the token's vector (the residual stream, section 7.2), and $W^O$ lets each head write in any direction of it, not just its own slice. It also fixes the shape when $h \cdot d_v \neq d_{\text{model}}$.

The Cost of Multi-Head Attention

Costs are counted in FLOPs (floating-point operations). For $N$ tokens, counting multiply-adds:

1. Projecting queries, keys and values for all heads: three $N \times d_{\text{model}}$ by $d_{\text{model}} \times d_{\text{model}}$ products, $3 \cdot N \cdot d_{\text{model}}^2$.
2. The scores $Q K^\top$: each head multiplies $N \times d_k$ by $d_k \times N$, costing $N \cdot d_k \cdot N$. Over $h$ heads that is $h \cdot N^2 \cdot d_k = N^2 \cdot d_{\text{model}}$, because $h \cdot d_k = d_{\text{model}}$.
3. The weighted sum $A V$: likewise $h \cdot (N \cdot N \cdot d_v) = N^2 \cdot d_{\text{model}}$.
4. The output projection $W^O$: $N \cdot d_{\text{model}}^2$.

Adding up (the softmax is small and left out):

(Counting multiply and add separately doubles both terms; the comparison is unchanged.) The first term grows with $N$, the second with $N^2$, and the second wins when $2N^2 d_{\text{model}} > 4N d_{\text{model}}^2$, that is, $N > 2d_{\text{model}}$. Past that length the $N \times N$ score table dominates compute and memory. This is why long contexts are expensive.

Exercises

Core Academic References

• Vaswani et al. (2017): Attention Is All You Need. NeurIPS 2017. arXiv:1706.03762
• Bahdanau, Cho & Bengio (2014): Neural Machine Translation by Jointly Learning to Align and Translate. ICLR 2015. arXiv:1409.0473
• Lee et al. (2019): Set Transformer: A Framework for Attention-based Permutation-Invariant Neural Networks. ICML 2019. arXiv:1810.00825