# Module 6: Attention

> Scaled dot-product attention worked by hand: queries, keys and values, the proof behind dividing by √d_k, causal masks, and multi-head attention.

- Part II: The Transformer · core track · Intermediate · 5 hours
- By the end you can compute attention by hand, prove why scores are divided by √d_k, and explain what multiple heads add.
- Builds on: [Module 5: Sequences Before Transformers](https://transformers-and-math.drmhse.com/md/05-sequences-before-transformers.md)
- Web page (interactive figures, labs and equation explainers): https://transformers-and-math.drmhse.com/modules/06-attention
- From *Transformers & Math* by Mike Chumba: https://transformers-and-math.drmhse.com

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 Hand

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:

$$
\operatorname{Attention}(Q, K, V) = \operatorname{Softmax}\left(\frac{QK^\top}{\sqrt{d_k}} + M\right) V
$$

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.

```mermaid
graph TD
    Q["Queries Q, shape N x d_k"] --> MatMul1["Multiply: Q · K^T"]
    K["Keys K, shape M x d_k"] --> MatMul1
    MatMul1 --> Scores["Raw scores, shape N x M"]
    Scores --> Scale["Divide by √d_k"]
    Scale --> Mask["Add mask M (e.g. causal: -∞ on future keys)"]
    Mask --> Smax["Softmax each row (safe form)"]
    Smax --> AttnW["Attention weights A, each in [0, 1], shape N x M"]
    AttnW --> MatMul2["Multiply: A · V"]
    V["Values V, shape M x d_v"] --> MatMul2
    MatMul2 --> Out["Output, shape N x d_v"]
```

> **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:

| | query | key | value |
|---|---|---|---|
| token 1 | $[3, -3]$ | $[-3, 3]$ | $[0, -1]$ |
| token 2 | $[1, 1]$ | $[1, 1]$ | $[1, 2]$ |

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.

> **Interactive figure** `attention-stepper`, on the web page: https://transformers-and-math.drmhse.com/modules/06-attention#figure-attention-stepper
>
> **Try:** Step to the softmax with the causal mask switched off. Where does token 1 put almost all of its weight, and why does the very negative score $-12.73$ make the answer so lopsided?
>
> **What happens:** On token 2. Softmax reacts to the *gap* between scores: a gap of about 12.7 makes one weight $e^{12.7} \approx 330{,}000$ times the other. A single large gap is enough to turn a weighted average into a copy of one value, which is exactly what section 6.3 is about.

> **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 Width

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

$$
z = \mathbf{q}^\top \mathbf{k} = \sum_{i=1}^{d_k} q_i k_i
$$

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$:

$$
\mathbb{E}[z] = \sum_{i=1}^{d_k} \mathbb{E}[q_i k_i] = \sum_{i=1}^{d_k} \mathbb{E}[q_i] \mathbb{E}[k_i] = 0
$$

2. **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:

$$
\operatorname{Var}(q_i k_i) = \mathbb{E}[(q_i k_i)^2] - (\mathbb{E}[q_i k_i])^2 = \mathbb{E}[q_i^2] \mathbb{E}[k_i^2] - 0 = (1)(1) = 1
$$

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

$$
\operatorname{Var}(z) = \sum_{i=1}^{d_k} \operatorname{Var}(q_i k_i) = \sum_{i=1}^{d_k} 1 = d_k
$$

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

$$
\sigma_z = \sqrt{\operatorname{Var}(z)} = \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$:

$$
\operatorname{Var}\left(\frac{z}{\sqrt{d_k}}\right) = \frac{1}{d_k} \operatorname{Var}(z) = \frac{d_k}{d_k} = 1
$$

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

> **Interactive figure** `sqrt-dk-histogram`, on the web page: https://transformers-and-math.drmhse.com/modules/06-attention#figure-sqrt-dk-histogram
>
> **Try:** Turn scaling off and raise $d_k$ from 16 to 256. Predict what happens to the histogram of scores and to the share of trials in which one key takes more than 99% of the weight. Then turn scaling on.
>
> **What happens:** Unscaled, the histogram widens like $\sqrt{d_k}$ and the saturated share climbs steeply, because wider scores mean bigger gaps. With scaling on, the histogram is the same width at every $d_k$ and saturation all but disappears. The proof above, made visible.

> **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 Masks

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$:

$$
M_{\text{causal}} = \begin{bmatrix} \cdot & -\infty & -\infty & -\infty \\ \cdot & \cdot & -\infty & -\infty \\ \cdot & \cdot & \cdot & -\infty \\ \cdot & \cdot & \cdot & \cdot \end{bmatrix}
$$

- **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 Order

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.

> **Interactive figure** `permutation-demo`, on the web page: https://transformers-and-math.drmhse.com/modules/06-attention#figure-permutation-demo
>
> **Try:** Shuffle the tokens. Before looking at the output panel, predict whether any output vector changes its *values*, or only its *position* in the list.
>
> **What happens:** Only its position. Each token's output moves with the token and is otherwise identical, to the last digit. Switch on position vectors and shuffle again: now the outputs change, because a token's vector depends on where it sits.

> **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 Attention

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:

$$
\operatorname{MultiHead}(Q, K, V) = \left[\operatorname{head}_1; \, \operatorname{head}_2; \, \dots; \, \operatorname{head}_h\right] W^O
$$

$$
\operatorname{head}_i = \operatorname{Attention}\left(Q W_i^Q, \, K W_i^K, \, V W_i^V\right)
$$

$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.

```mermaid
graph TD
    Input["Token vectors X, shape N x d_model"] --> H1["Head 1, e.g. subject–verb links"]
    Input --> H2["Head 2, e.g. pronoun → the noun it refers to"]
    Input --> H3["Head 3, e.g. punctuation and boundaries"]
    Input --> H8["Head 8, e.g. who did what to whom"]
    H1 --> Concat["Concatenate [head_1; head_2; ...; head_8]"]
    H2 --> Concat
    H3 --> Concat
    H8 --> Concat
    Concat --> WO["Output projection W^O, 512 x 512"]
    WO --> Output["Multi-head output, shape N x d_model"]
```

> **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.

> **Interactive figure** `multi-head`, on the web page: https://transformers-and-math.drmhse.com/modules/06-attention#figure-multi-head
>
> **Try:** Pick the query word *it* and look at the four heads. Before you do, guess which head's weights would let the model work out what *it* refers to.
>
> **What happens:** The head that sends the pronoun to its antecedent puts most of the weight of *it* on *animal*. The others look at the previous word, the verb's subject, or the clause boundary. No single head does everything, which is why the outputs of all four are joined and mixed.

> **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):

$$
\text{FLOPs}_{\text{MHA}} = 4 N d_{\text{model}}^2 + 2 N^2 d_{\text{model}}
$$

(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

**Exercise (check).** In one sentence each, what are the query, the key and the value of a token for?

*Solution:* The query describes what the token is looking for; the key describes what the token can be found by; the value is the information the token hands over to whichever tokens attend to it.

**Exercise (apply).** Scores for one query are $[2, 0, 0]$ (already scaled). Compute the attention output when the values are the scalars $[1, 4, 7]$.

*Answer:* 1.959 (within 0.01)

*Solution:* Softmax of $[2, 0, 0]$: $e^2 = 7.389$ and the total is $7.389 + 1 + 1 = 9.389$, so the weights are $0.787$, $0.1065$ and $0.1065$. The output is $0.787 \cdot 1 + 0.1065 \cdot 4 + 0.1065 \cdot 7 = 0.787 + 1.172 \approx 1.959$.

**Exercise (apply).** Using the rules of section 3.2, prove that $\operatorname{Var}(z / \sqrt{d_k}) = 1$ when $z$ is a sum of $d_k$ independent products, each with mean 0 and variance 1.

*Solution:* By rule 4, $\operatorname{Var}(z) = d_k \cdot 1 = d_k$. By rule 3, $\operatorname{Var}(z / \sqrt{d_k}) = \frac{1}{d_k}\operatorname{Var}(z) = 1$.

**Exercise (apply).** Write the causal mask for three tokens, and say what the softmax weights of the first row must be whatever the scores are.

*Solution:* $\begin{bmatrix} 0 & -\infty & -\infty \\ 0 & 0 & -\infty \\ 0 & 0 & 0 \end{bmatrix}$. The first row allows only key 1, so its weights are $[1, 0, 0]$: the first token can only attend to itself.

**Exercise (apply).** For $d_{\text{model}} = 4096$, at what sequence length does the $2N^2 d_{\text{model}}$ term of multi-head attention's cost overtake the $4N d_{\text{model}}^2$ term?

*Answer:* 8192 (within 0)

*Solution:* The quadratic term is larger when $N > 2 d_{\text{model}} = 8192$ tokens.

**Exercise (extend).** What does a head compute if its softmax weights are exactly uniform over all $N$ tokens? Why is that not useless?

*Solution:* Every token's output is the plain average of all value vectors, the same for every token. It is not useless: it gives every position a summary of the whole context, and a head can learn to fall back on near-uniform attention when nothing specific is relevant. (Section 16.6 shows a related learned habit, attention sinks.)

**Exercise (extend).** Show that self-attention without a mask or positions is permutation-equivariant: if the rows of $X$ are reordered by a permutation matrix $P$, the output rows are reordered the same way.

*Solution:* With $X' = PX$, the queries, keys and values are $PQ$, $PK$, $PV$. The scores become $PQ(PK)^\top = P(QK^\top)P^\top$: the same table with rows and columns reordered. Softmax works row by row, so it commutes with reordering: $\operatorname{softmax}(PSP^\top) = P\operatorname{softmax}(S)P^\top$. Multiplying by $PV$ gives $P\operatorname{softmax}(S)P^\top P V = P\operatorname{softmax}(S)V$, since $P^\top P = I$. The output is $P$ times the original output.

### Core Academic References
- **Vaswani et al. (2017)**: *Attention Is All You Need*. NeurIPS 2017. [arXiv:1706.03762](https://arxiv.org/abs/1706.03762)
- **Bahdanau, Cho & Bengio (2014)**: *Neural Machine Translation by Jointly Learning to Align and Translate*. ICLR 2015. [arXiv:1409.0473](https://arxiv.org/abs/1409.0473)
- **Lee et al. (2019)**: *Set Transformer: A Framework for Attention-based Permutation-Invariant Neural Networks*. ICML 2019. [arXiv:1810.00825](https://arxiv.org/abs/1810.00825)