Attention: queries, keys, values and masks

Before this lesson: lesson 03 (backpropagation), lesson 05 (LayerNorm and residual connections) and lesson 08 (tokens, embeddings and what attention replaced). The softmax Jacobian is also derived in linear-algebra lesson 16. You need the rule "the variance of a sum of independent terms is the sum of their variances" (probability-and-statistics track). Time: about 70 minutes, plus the notebook.

---

Why this matters #

A recurrent network (lesson 08) carries information from token 1 to token 1000 through 999 sequential steps. Attention lets every token read from every other token in one step, and all tokens are processed in parallel on a GPU. This one operation is the core of large language models (GPT, LLaMA), of BERT, of the vision transformer (lesson 11), of speech models such as Whisper, of AlphaFold 2, and of most weather and climate foundation models. Interviews test it at three levels: write it from a blank file, derive the $1/\sqrt{d_k}$ scaling, and count its cost. The cost grows with the square of the number of tokens, which is the main limit of long contexts.

1. Predict first #

2. The picture #

One query: a soft dictionary lookup #

A Python dictionary {"cat": v1, "dog": v2} does a hard lookup: the query "cat" matches exactly one key and returns its value. Attention does a soft lookup. There are $T$ items. Item $j$ has a key $k_j \in \mathbb{R}^{d_k}$ and a value $v_j \in \mathbb{R}^{d_v}$. A query $q \in \mathbb{R}^{d_k}$ asks for something. Attention does three steps:

1. Score every item: $s_j = q^\top k_j / \sqrt{d_k}$. A large dot product means "this key matches the query".
2. Weights: $a_j = e^{s_j} / \sum_l e^{s_l}$, the softmax of the scores. So $a_j > 0$ and $\sum_j a_j = 1$.
3. Output: the weighted average $o = \sum_j a_j v_j$.

The output is a blend of all values, with more weight on the values whose keys match the query. Because $a_j > 0$ and $\sum_j a_j = 1$, the output is a convex combination: it always lies inside the region spanned by the values. A hard lookup is the limit where one weight is 1 and the others are 0.

Interactive figure: loading…

Self-attention: every token is a query #

A sequence of $T$ tokens is a matrix $X \in \mathbb{R}^{T \times D}$: one row per token, $D$ is the model width. (We use the row convention of PyTorch, as in the whole track.) In self-attention, the queries, keys and values all come from the same sequence. Three learned matrices $W_Q, W_K \in \mathbb{R}^{D \times d_k}$ and $W_V \in \mathbb{R}^{D \times d_v}$ give

Row $i$ of $Q$ is the query of token $i$; row $j$ of $K$ and of $V$ are the key and the value of token $j$. Now every token does the lookup above at the same time:

Entry $(i, j)$ of $QK^\top$ is $q_i^\top k_j$, the score of query $i$ for key $j$. Row $i$ of $A$ holds the attention weights of token $i$: it is positive and sums to 1. Row $i$ of $AV$ is $\sum_j A_{ij} v_j$, the output of token $i$.

The mixing weights $A$ are computed from the input. A convolution (lesson 07) also mixes neighbouring tokens, but with weights that are fixed after training and that depend only on the offset. Attention decides for every input again which tokens to mix.

Interactive figure: loading…

3. Definitions and main results #

Scaled dot-product attention #

Definition. For queries $Q \in \mathbb{R}^{T_q \times d_k}$, keys $K \in \mathbb{R}^{T_k \times d_k}$, values $V \in \mathbb{R}^{T_k \times d_v}$ and a mask $M \in \{0, -\infty\}^{T_q \times T_k}$:

where the softmax acts on each row separately. With no mask, $M = 0$.

Cost. $QK^\top$ needs $2T_qT_kd_k$ floating-point operations (FLOPs; one multiply and one add per term), $AV$ needs $2T_qT_kd_v$, and $A$ needs $T_qT_k$ numbers of memory. For self-attention $T_q = T_k = T$: the cost grows like $T^2$.

Why divide by $\sqrt{d_k}$ #

The figure first, then the proof. The figure draws one query and 16 keys with random entries and shows the softmax weights, with and without the scaling.

Interactive figure: loading…

Theorem (variance of the dot product). Let $q, k \in \mathbb{R}^{d_k}$ be independent, with independent entries of mean 0 and variance 1. Then $\mathbb{E}[q^\top k] = 0$ and $\operatorname{Var}(q^\top k) = d_k$.

Proof.

1. $q^\top k = \sum_{i=1}^{d_k} q_ik_i$. By independence, $\mathbb{E}[q_ik_i] = \mathbb{E}[q_i]\,\mathbb{E}[k_i] = 0$. So $\mathbb{E}[q^\top k] = 0$.
2. $\operatorname{Var}(q_ik_i) = \mathbb{E}[q_i^2k_i^2] - 0^2 = \mathbb{E}[q_i^2]\,\mathbb{E}[k_i^2] = 1 \cdot 1 = 1$.
3. The terms $q_ik_i$ for different $i$ are independent, so their variances add: $\operatorname{Var}(q^\top k) = d_k$. $\blacksquare$

So $q^\top k$ has standard deviation $\sqrt{d_k}$, and $q^\top k / \sqrt{d_k}$ has variance 1 for every head size.

Theorem (softmax Jacobian). For $p = \operatorname{softmax}(s)$, $\partial p_i/\partial s_j = p_i(\delta_{ij} - p_j)$, that is,

Proof.

1. $p_i = e^{s_i}/Z$ with $Z = \sum_l e^{s_l}$, and $\partial Z/\partial s_j = e^{s_j}$.
2. Quotient rule: $\dfrac{\partial p_i}{\partial s_j} = \dfrac{\delta_{ij}e^{s_i}Z - e^{s_i}e^{s_j}}{Z^2} = \delta_{ij}p_i - p_ip_j$. $\blacksquare$

$J$ is symmetric and positive semidefinite, and $J\mathbf{1} = 0$: adding the same constant to all logits changes nothing. Backpropagation needs only the product with the upstream gradient $g = \partial \mathcal{L}/\partial p$:

which costs $O(T)$ instead of $O(T^2)$.

Saturation. Let the largest weight be $p_m = 1 - \varepsilon$. Then every entry of $J$ is at most $\varepsilon$ in absolute value. Why: every entry $J_{ij}$ has a factor $p_i$ or $p_j$ with an index different from $m$, and that factor is at most $\varepsilon$; the one exception is $J_{mm} = p_m(1 - p_m) \le \varepsilon$. So a nearly one-hot softmax passes almost no gradient back to its logits, and from there to $W_Q$ and $W_K$.

Putting the two together. Without the scaling, logits have standard deviation $\sqrt{d_k}$: 8 for $d_k = 64$. The difference of two logits has standard deviation $\sqrt{2d_k} \approx 11$, and $e^{11} \approx 60{,}000$. The softmax then puts almost all weight on one or a few keys. In a simulation with 128 keys and $d_k = 64$, the median of the largest weight is about 0.8 and about 95% of the keys get a weight below $10^{-3}$; with $d_k = 256$ the largest weight is about 0.98. The effect grows with $d_k$. Dividing by $\sqrt{d_k}$ keeps the logits of order 1 at initialization, for every head size.

Backward pass through attention, and the scaling used by μP

With $S = QK^\top/\sqrt{d_k}$, $A = \operatorname{softmax}(S)$ (rows), $O = AV$ and the upstream gradient $dO = \partial \mathcal{L}/\partial O$:

$$\begin{aligned} dV &= A^\top dO, & dA &= dO\,V^\top, \\ dS &= A \odot \big(dA - \operatorname{rowsum}(A \odot dA)\,\mathbf{1}^\top\big), & & \\ dQ &= dS\,K/\sqrt{d_k}, & dK &= dS^\top Q/\sqrt{d_k}. \end{aligned}$$

The line for $dS$ is the softmax VJP above, applied to each row. FlashAttention uses the same formulas, computed tile by tile. The notebook checks them against autograd.

The variance argument assumes that $q$ and $k$ are independent. If they are strongly correlated, for example $q = k$, then $q^\top k = \|q\|^2$ grows like $d_k$, not like $\sqrt{d_k}$. The maximal-update parameterization (μP, Yang et al., Tensor Programs V) therefore scales the logits by $1/d_k$, so that the right learning rate transfers across model widths.

Masks #

A mask is added to the logits before the softmax. An entry $M_{ij} = -\infty$ gives $e^{-\infty} = 0$: query $i$ gives no weight to key $j$.

• Causal mask (decoders, language models): query $i$ may attend to key $j$ only if $j \le i$. So $M_{ij} = 0$ for $j \le i$ and $-\infty$ above the diagonal. Then token $i$ never reads a future token, and training on all positions at once gives the same outputs as generating one token at a time.
• Padding mask: sequences in a batch have different lengths, so short ones are filled with padding tokens. A padded key $j$ gets $M_{ij} = -\infty$ in every row.
• Masks combine by adding them (for boolean "may attend" masks: a logical AND).
• Add the mask before the softmax. Then the masked weights are exactly 0, and the visible weights of each row still sum to 1.
• A fully masked row has every logit equal to $-\infty$. A direct implementation computes $0/0$ and returns NaN, and the NaN spreads through the whole network. Some library kernels return zeros for such rows instead; do not rely on either. Make sure every query sees at least one key.

Interactive figure: loading…

Conventions differ between libraries. In F.scaled_dot_product_attention, a boolean mask entry True means "may attend". In nn.MultiheadAttention, True in attn_mask or key_padding_mask means "blocked". Read the documentation of the function you call.

Multi-head attention #

One head gives each query one probability distribution over the keys. Multi-head attention runs $H$ heads in parallel, each with its own projections, so that a token can attend to several places at once (for example, to the previous token and to a related word far away).

Definition. Let $H$ divide $D$, and let $d_k = d_v = D/H$ be the head size. Take $W_Q, W_K, W_V, W_O \in \mathbb{R}^{D \times D}$ and bias vectors $b_Q, b_K, b_V, b_O \in \mathbb{R}^D$. Head $h \in \{1, \dots, H\}$ uses the columns $(h - 1)d_k + 1$ to $hd_k$ of the first three matrices: $W_Q^h, W_K^h, W_V^h \in \mathbb{R}^{D \times d_k}$, and the same entries of the biases. Then

Here $[\cdot, \dots, \cdot]$ puts the head outputs side by side ($T \times D$), and $W_O$ mixes the heads.

• Implementation: no loop over heads. One matrix product gives $Q = XW_Q + \mathbf{1}b_Q$ of shape $(T, D)$. View it as $(T, H, d_k)$, then move the head axis to the front: $(H, T, d_k)$. Batched matrix products then compute all heads at once. With a batch axis: $(B, T, D) \to (B, T, H, d_k) \to (B, H, T, d_k)$.
• Parameters: $W_Q, W_K, W_V, W_O$ give $4D^2$, the four biases give $4D$: in total $4D^2 + 4D$, for every $H$. Heads split the width; they do not add parameters.
• Cost for $T$ tokens: the four projections need $4 \cdot 2TD^2 = 8TD^2$ FLOPs; the scores and the weighted sums of all heads need $H \cdot 4T^2d_k = 4T^2D$. The two are equal at $T = 2D$; for longer sequences the $T^2$ term dominates. Lesson 11 adds the MLP of the transformer block to this count.
• Low rank: the score of head $h$ is $x_i^\top W_Q^h (W_K^h)^\top x_j / \sqrt{d_k}$, a bilinear form whose $D \times D$ matrix has rank at most $d_k$ (linear-algebra lesson 15).

Interactive figure: loading…

Cross-attention #

In cross-attention, the queries come from one sequence and the keys and values from another: $Q = X_{\text{dec}}W_Q$ with $T_q$ rows, and $K = ZW_K$, $V = ZW_V$ from a second sequence $Z$ with $T_k$ rows (for example, the output of an encoder). Then $A$ is $T_q \times T_k$, and the output has one row per query. Cross-attention uses the padding mask of $Z$, but never a causal mask: the whole second sequence is known. It is the standard way to condition a network on other information, such as a translation model reading the source sentence, or an image generator reading a text prompt.

Permutation equivariance #

A permutation matrix $P \in \mathbb{R}^{T \times T}$ reorders the rows of a matrix: $(PX)_i = X_{\pi(i)}$ for a permutation $\pi$ of $\{1, \dots, T\}$. It is orthogonal: $P^\top P = I$.

Theorem (self-attention is permutation-equivariant). Without a mask, and with $Q = XW_Q$, $K = XW_K$, $V = XW_V$: $\operatorname{Attn}(PX) = P\,\operatorname{Attn}(X)$ for every permutation matrix $P$.

Proof.

1. The input $PX$ gives $Q' = PXW_Q = PQ$, and in the same way $K' = PK$ and $V' = PV$.
2. The scores are $S' = Q'K'^\top/\sqrt{d_k} = PQK^\top P^\top/\sqrt{d_k} = PSP^\top$. In entries: $S'_{ij} = S_{\pi(i)\pi(j)}$. Rows and columns are reordered in the same way.
3. The row-wise softmax commutes with this reordering: $A'_{ij} = \dfrac{e^{S_{\pi(i)\pi(j)}}}{\sum_l e^{S_{\pi(i)\pi(l)}}} = \dfrac{e^{S_{\pi(i)\pi(j)}}}{\sum_{l'} e^{S_{\pi(i)l'}}} = A_{\pi(i)\pi(j)}$, because $l' = \pi(l)$ runs over all indices. So $A' = PAP^\top$.
4. $A'V' = PAP^\top PV = PAV$, because $P^\top P = I$. $\blacksquare$

The same holds for multi-head attention (each head is equivariant, and the concatenation and $W_O$ act on each token separately), and for a whole transformer encoder, because the MLP, LayerNorm and residual connections also act on each token separately. Without positions, the model sees a set of tokens, not a sequence: "dog bites man" and "man bites dog" give the same three output vectors, only in a different order. Lesson 10 adds positions.

A causal mask breaks the symmetry. The masked scores of the input $PX$ are $PSP^\top + M$, and $PMP^\top \ne M$ for $P \ne I$. Query $i$ sees exactly $i$ keys, so a decoder can infer positions even without a positional encoding.

Answers to section 1. (1) $q^\top k$ has standard deviation $\sqrt{64} = 8$, and the difference of two such logits has standard deviation $8\sqrt{2} \approx 11$. The softmax is then nearly one-hot (almost all weight on one or a few keys), most entries of its Jacobian are close to 0, and the gradients to $W_Q$ and $W_K$ become very small. Dividing by $\sqrt{d_k} = 8$ fixes this. (2) The outputs are shuffled in the same way and are otherwise unchanged (permutation equivariance). A pooled output, such as the mean over the tokens, does not change at all.

4. Worked example #

One head, three tokens. Take $T = 3$, $d_k = d_v = 2$ and $W_Q = W_K = W_V = I$, so $Q = K = V = X$:

No mask. Row 1: $e^{0.707} = 2.028$ and $e^0 = 1$. The sum is $2.028 + 1 + 2.028 = 5.056$, so the weights are $(2.028, 1, 2.028)/5.056 = (0.401, 0.198, 0.401)$. The other rows in the same way:

Check row 1 of $O$: $0.401\,(1, 0) + 0.198\,(0, 1) + 0.401\,(1, 1) = (0.802, 0.599)$.

Causal mask. Row 1 sees only key 1: weights $(1, 0, 0)$, output $(1, 0)$. Row 2 sees keys 1 and 2, with logits $(0, 0.707)$: $e^0 = 1$, $e^{0.707} = 2.028$, weights $(1, 2.028)/3.028 = (0.330, 0.670)$, output $0.330\,(1, 0) + 0.670\,(0, 1) = (0.330, 0.670)$. Row 3 sees all keys and does not change: $(0.752, 0.752)$. The output of token 1 no longer depends on tokens 2 and 3.

Jacobian of row 2 (causal). $p = (0.330, 0.670)$, so

using $p_1 - p_1^2 = p_1(1 - p_1) = p_1p_2$.

Saturation in numbers. Logits $(2, 1, 0)$ give $p = (0.665, 0.245, 0.090)$ and $\max_{ij}|J_{ij}| = 0.223$. The same logits times 8, which is what $d_k = 64$ gives without the scaling, are $(16, 8, 0)$: $p = (0.99966, 0.00034, 0.0000001)$ and $\max_{ij}|J_{ij}| = 3.4 \cdot 10^{-4}$. The largest Jacobian entry is about 660 times smaller.

Counting. Multi-head attention with $D = 512$, $H = 8$ (so $d_k = 64$) and biases has $4(512^2 + 512) = 1{,}050{,}624$ parameters. For $T = 512$ tokens: projections $8TD^2 = 8 \cdot 512^3 \approx 1.07$ GFLOP; scores and weighted sums $4T^2D = 4 \cdot 512^3 \approx 0.54$ GFLOP. At $T = 2D = 1024$ the two parts cost the same.

5. In code #

PyTorch checks the worked example and the claims of section 3. The notebook builds the same functions from NumPy.

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)
f64 = torch.float64

# (a) the worked example, without and with the causal mask
X = torch.tensor([[1., 0.], [0., 1.], [1., 1.]], dtype=f64)
print(F.scaled_dot_product_attention(X, X, X).numpy().round(3))
# [[0.802 0.599] [0.599 0.802] [0.752 0.752]]
print(F.scaled_dot_product_attention(X, X, X, is_causal=True).numpy().round(3))
# [[1.    0.   ] [0.33  0.67 ] [0.752 0.752]]

# (b) Var(q·k) = d_k, and Var(q·k / sqrt(d_k)) = 1
for d_k in (16, 64, 256):
    q, k = torch.randn(100_000, d_k, dtype=f64), torch.randn(100_000, d_k, dtype=f64)
    s = (q * k).sum(-1)
    print(d_k, round(s.var().item(), 1), round((s / d_k**0.5).var().item(), 2))
# 16 16.0 1.0 / 64 63.8 1.0 / 256 255.3 1.0   (Monte Carlo: close to d_k and 1)

# (c) softmax Jacobian = diag(p) - p pᵀ, and it is tiny when p is nearly one-hot
for z in ([2., 1., 0.], [16., 8., 0.]):
    z = torch.tensor(z, dtype=f64)
    J = torch.autograd.functional.jacobian(lambda s: torch.softmax(s, -1), z)
    p = torch.softmax(z, -1)
    print(torch.allclose(J, torch.diag(p) - torch.outer(p, p)), f"{J.abs().max().item():.2e}")
# True 2.23e-01
# True 3.35e-04

# (d) 4D² + 4D parameters for every H, and permutation equivariance
print([sum(p.numel() for p in nn.MultiheadAttention(512, H).parameters()) for H in (1, 8, 64)])
# [1050624, 1050624, 1050624]
mha = nn.MultiheadAttention(64, 4, batch_first=True, dtype=f64)
attn = lambda x, mask=None: mha(x, x, x, attn_mask=mask, need_weights=False)[0]
x, perm = torch.randn(1, 10, 64, dtype=f64), torch.randperm(10)
print(torch.allclose(attn(x)[:, perm], attn(x[:, perm])))                  # True
causal = nn.Transformer.generate_square_subsequent_mask(10, dtype=f64)
print(torch.allclose(attn(x, causal)[:, perm], attn(x[:, perm], causal)))  # False

F.scaled_dot_product_attention divides by $\sqrt{d_k}$ itself, with $d_k$ the last dimension of the queries. generate_square_subsequent_mask returns the additive mask $M$ (0 and $-\infty$).

6. Where it shows up #

• Language models. GPT and LLaMA are stacks of causal self-attention (decoder-only). BERT uses self-attention without a causal mask (encoder-only). At generation time the keys and values of past tokens are stored in a KV cache; multi-query and grouped-query attention share key and value heads between query heads to make this cache smaller (lesson 11).
• Long inputs. The $T^2$ term limits the context length. FlashAttention computes exact attention in tiles with an online softmax: the extra memory becomes $O(T)$, while the FLOPs stay $O(T^2)$.
• Conditioning by cross-attention. Machine translation (the original transformer), speech recognition (Whisper: the text decoder reads the audio encoder), and text-to-image latent diffusion (the image tokens of the denoiser read the text tokens).
• Many inputs, few queries. The Perceiver lets a small set of learned query vectors cross-attend to a very large input, so the cost grows like $T$ instead of $T^2$. ClimaX merges the physical variables of each weather patch by cross-attention with one learned query.
• Science. AlphaFold 2 applies attention along the rows and columns of a sequence alignment and over pairs of residues. Attention restricted to the neighbours of each node is graph attention (lesson 13). Attention without positions is a natural layer for sets, such as point clouds or particles.

7. Common mistakes #

• Masking after the softmax. The visible weights then sum to less than 1, and they depend on the masked keys through the normalizer. Add $-\infty$ to the logits before the softmax.
• A fully masked row. It gives NaN (or silent zeros). Check padding masks for rows with no visible key.
• A flipped mask convention. True means "attend" in one API and "blocked" in another.
• A wrong head split. Reshaping $(B, T, D)$ directly to $(B, H, T, d_k)$ mixes tokens and features. View as $(B, T, H, d_k)$, then transpose.
• Scaling by $\sqrt{D}$ instead of $\sqrt{d_k}$, the head size.
• "More heads, more parameters." The count is $4D^2 + 4D$ for every $H$ that divides $D$.
• A causal mask in an encoder or in cross-attention. Only decoder self-attention is causal.
• Reading row $j$ of $A$ as "who attends to token $j$". Row $i$ holds the weights of query $i$; column $j$ shows how much each query reads from token $j$.

8. Check yourself #

1. Cross-attention with batch $B = 2$, $T_q = 5$ queries, $T_k = 7$ keys, $D = 64$ and $H = 4$ heads. Give the shapes of the per-head queries, the keys, the attention weights and the output.
   Answer

   $d_k = 64/4 = 16$. Queries $(2, 4, 5, 16)$; keys and values $(2, 4, 7, 16)$; weights $(2, 4, 5, 7)$; per-head outputs $(2, 4, 5, 16)$; after merging the heads and $W_O$: $(2, 5, 64)$, one row per query.
2. Logits $s = (3, 0)$. Compute $p$ and the softmax Jacobian.
   Answer

   $p = (e^3, 1)/(e^3 + 1) = (0.953, 0.047)$. $J = p_1p_2\begin{pmatrix}1&-1\\-1&1\end{pmatrix}$ with $p_1p_2 = 0.045$.
3. The entries of $q$ and $k$ are independent with mean 0 and variance $\sigma^2$. What is $\operatorname{Var}(q^\top k/\sqrt{d_k})$? What keeps $\sigma$ close to 1 in a network?
   Answer

   $\sigma^4$: each term $q_ik_i$ has variance $\sigma^2 \cdot \sigma^2$, and there are $d_k$ terms, divided by $d_k$. The scaling works only if $q$ and $k$ have entries of unit scale. A normalized input to the attention layer (LayerNorm, lesson 05) and a variance-preserving initialization of $W_Q, W_K$ (lesson 04) give this at initialization; qk-norm keeps it during training.
4. Why is the mask added before the softmax and not applied after it?
   Answer

   Before: masked weights are exactly 0 and the visible weights are normalized to sum to 1. After: the normalizer still contains the masked (for example, future) keys, so the visible weights do not sum to 1 and they depend on the masked tokens. For a causal model this is a leak of future information.
5. A model has $D = 1024$. For which number of tokens $T$ do the attention scores and weighted sums ($4T^2D$ FLOPs) cost more than the four projections ($8TD^2$)?
   Answer

   $4T^2D > 8TD^2 \iff T > 2D = 2048$.
6. Can a decoder-only model without any positional encoding know token positions?
   Answer

   Yes, it can infer them. Under the causal mask, query $i$ averages over exactly $i$ tokens. For example, a head with equal weights gives the first token the weight $1/i$, which encodes $i$. The mask breaks permutation equivariance.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. Write scaled dot-product attention with a mask.

   Show answer

2. Scores $(0, \ln 3)$ for two keys. What are the attention weights?

   Show answer

3. Why do we divide the logits by $\sqrt{d_k}$?

   Show answer

4. What is the Jacobian of $p = \operatorname{softmax}(s)$?

   Show answer

5. What does a causal mask do, and how is it built?

   Show answer

6. What happens to an attention row where every key is masked?

   Show answer

7. How many parameters does multi-head attention with width $D$ and $H$ heads have?

   Show answer

8. How do you split $Q$ of shape $(B, T, D)$ into $H$ heads?

   Show answer

9. How does the cost of self-attention grow with the number of tokens $T$?

   Show answer

10. What does permutation equivariance of self-attention mean, and what breaks it?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Implement multi-head attention with a causal mask." — Write the shapes first; view then transpose for the heads; add the mask before the softmax; compare with F.scaled_dot_product_attention.
• "Why divide by $\sqrt{d_k}$?" — $q^\top k$ has variance $d_k$ at initialization; large logits saturate the softmax, its Jacobian $\operatorname{diag}(p) - pp^\top$ goes to 0, and learning stops.
• "What is the cost of self-attention?" — $8TD^2 + 4T^2D$ FLOPs per layer and $O(T^2)$ memory for the weights unless the computation is tiled (FlashAttention).
• "Why several heads, and does it add parameters?" — Several attention patterns at once; no, the count stays $4D^2 + 4D$.
• "Self-attention vs cross-attention?" — Same formula; in cross-attention the queries come from one sequence and the keys and values from another.
• "What does attention do if you shuffle the tokens?" — Without positions or a causal mask, the outputs are shuffled in the same way (permutation equivariance).

10. Practice #

The notebook for this lesson runs in Colab. It uses NumPy for the exercises and compares with PyTorch when it is installed (it is on Colab). Five exercises, each with a check cell and a folded solution:

1. Scaled dot-product attention with a boolean mask, checked on the worked example and against F.scaled_dot_product_attention.
2. A Monte Carlo check of $\operatorname{Var}(q^\top k) = d_k$, and a plot of the largest weight with and without the scaling.
3. The softmax Jacobian and its VJP, checked by finite differences, and its size as the logits grow.
4. Multi-head attention from scratch (view and transpose), matching nn.MultiheadAttention with copied weights.
5. A permutation-equivariance test, with and without the causal mask.

11. Watch and read #

Attention in transformers, step-by-step | Deep Learning Chapter 63Blue1Brown · Deep learning, chapter 6 · 26 min

Queries, keys and values as pictures; the parts on masking and multiple heads match section 3 of this lesson.

The matrix math behind transformer neural networks, one step at a time!!!StatQuest with Josh Starmer · 24 min

The matrix form QKᵀ, softmax and AV with small numbers, as in the worked example.

Let's build GPT: from scratch, in code, spelled out.Andrej Karpathy · 1 h 56 min

From 1:02:00, self-attention in code, the causal mask, the scaling by the head size and multi-head attention.

Attention Is All You NeedYannic Kilcher · 27 min

A walk through the original paper.

• Vaswani et al., Attention Is All You Need (2017): §3.2 (scaled dot-product and multi-head attention; footnote 4 is the variance argument) and Table 1 (cost per layer).
• Simon Prince, Understanding Deep Learning (free PDF), chapter 12: §12.2–12.3 (dot-product self-attention, the scaling, multiple heads), §12.7 (the causal mask in a decoder) and §12.8 (encoder–decoder and cross-attention).
• Rush et al., The Annotated Transformer, sections "Attention" and "Applications of Attention in our Model": the paper as runnable code.
• Zhang, Lipton, Li, Smola, Dive into Deep Learning, chapter 11 (attention as queries, keys and values; multi-head and self-attention), with code.
• Stanford CS224N, the lecture "Self-Attention and Transformers" (slides and notes on the course page).