Chapter 3 — Derivations

Attention is one line of linear algebra, and interviewers probe that line one symbol at a time. Why is the scale $\sqrt{d_k}$ and not $d_k$ ? What is the softmax Jacobian, and why does it vanish when a row saturates? How does a gradient reach $W_Q$ ? How many bytes does Llama 2 7B's KV cache take per token, and how does FlashAttention stay exact without the $T \times T$ matrix? Each answer is a short derivation, and each one began as a fix for a problem someone hit: overflow in fp16, rows that went one-hot, a decoder that peeked at the future, a GPU that ran out of memory at long context. This tab has the four algebra moves those derivations use, the chapter's 25 key formulas, and 18 cards. Each card opens with the problem that produced the result, states the result, derives it with every step shown, checks it on the toy sentence "the river bank was steep" or on a published model configuration, and names the slip candidates make and the way the question is asked. Every identity below was checked numerically (finite differences, autograd or a direct computation) in sims/check_derivations.py. Practice by covering everything below a card's Statement, deriving it on a blank page against a timer, then checking. The protocol at the end says how.

The toolkit

Every card checks itself on the running example from Sections 3–4: "the river bank was steep", one head with $d_k = 2$ . Under the causal mask, bank may see the, river and itself. Its values live in a two-dimensional [geography, money] space.

\mathbf q_{\text{bank}} = (2,\,-1); \quad \mathbf k_{\text{the}},\ \mathbf k_{\text{river}},\ \mathbf k_{\text{bank}} = (0, 0),\ (1, 0),\ (1, 1); \quad \mathbf v_{\text{the}},\ \mathbf v_{\text{river}},\ \mathbf v_{\text{bank}} = (0, 0),\ (1, 0),\ (0, 1); \quad \mathbf p = (0.140,\ 0.576,\ 0.284), \quad \mathbf o_{\text{bank}} = (0.576,\ 0.284)

$\mathbf q_{\text{bank}}$: bank's query: what bank is looking for (bold = vector)
$\mathbf k_{\text{the}}$: the key of the: zero, so it advertises nothing
$\mathbf k_{\text{river}}$: river's key: how river advertises itself
$\mathbf k_{\text{bank}}$: bank's own key
$\mathbf v_{\text{the}}$: the value of the: zero, it delivers nothing
$\mathbf v_{\text{river}}$: river's value: pure geography
$\mathbf v_{\text{bank}}$: bank's own value before attention: pure money
$\mathbf p$: bank's attention weights on the, river and bank; they sum to 1
$\mathbf o_{\text{bank}}$: bank's new vector: 0.576 geography, 0.284 money

The raw scores are $\mathbf q_{\text{bank}}^\top\mathbf k = 0, 2, 1$ ; divided by $\sqrt 2$ they become 0, 1.414, 0.707, and the softmax turns them into $\mathbf p$ . Bigger checks use Llama 2 7B (32 layers, width 4,096, 32 heads of 128) and the 4,096-key simulation from Section 5. Every derivation in this tab uses some combination of four moves. Hook: S-O-F-T.

S · Shapes tally. Write the shape of every factor first. Multiplying an $m \times k$ matrix by a $k \times n$ matrix costs $2mnk$ FLOPs (one multiply and one add per term), and storing a tensor costs its element count times the bytes per number.

(m \times k)(k \times n) \to (m \times n): \quad \text{FLOPs} = 2mnk, \qquad \text{bytes} = (\text{number of elements}) \times b

$m$: the number of rows of the left matrix
$k$: the shared inner dimension, summed over in every dot product
$n$: the number of columns of the right matrix
$b$: bytes per number: 2 for fp16 or bf16, 4 for fp32

O · One-hot minus p. The derivative of a log-softmax is the one-hot vector minus the probabilities, so every softmax gradient reads "what you got minus what you expected".

p_i = \frac{e^{s_i}}{\sum_{r=1}^{T} e^{s_r}} \;\Rightarrow\; \frac{\partial\log p_i}{\partial s_j} = \delta_{ij} - p_j, \qquad \frac{\partial p_i}{\partial s_j} = p_i(\delta_{ij} - p_j)

$p_i$: the softmax weight on key $i$
$s_i$: the score (logit) of key $i$
$e$: Euler's number, 2.718…; $e^{x}$ is always positive
$\sum_{r=1}^{T}$: add up over all $T$ keys $r$ in the row (the normalizer)
$T$: the number of keys in the row
$s_r$: the score of key $r$
$\frac{\partial\log p_i}{\partial s_j}$: how fast $\log p_i$ changes as you nudge score $s_j$ , holding the rest fixed: a slope, not a fraction to divide
$\delta_{ij}$: 1 if $i = j$ , else 0: the one-hot part
$p_j$: the softmax weight on key $j$
$\frac{\partial p_i}{\partial s_j}$: how fast weight $p_i$ changes as you nudge score $s_j$
$i$: the key whose weight we watch
$j$: the key whose score we nudge

F · Factor out the max. Softmax ignores a constant added to every score. So pull out the largest score $m$ : no exponential can overflow, and partial sums computed with different maxima can be moved to a common base by one multiplication.

\operatorname{softmax}(\mathbf s + c\mathbf 1) = \operatorname{softmax}(\mathbf s), \qquad \log\sum_j e^{s_j} = m + \log\sum_j e^{s_j - m}, \qquad m = \max_j s_j

$\mathbf s$: a row of scores, one per key
$c$: any constant added to every score
$\mathbf 1$: the all-ones vector
$\sum_j$: add up over the keys $j$ in the row
$s_j$: the score of key $j$
$m$: the largest score in the row

T · Total variance of a sum. Independent terms add their variances, and the product of two independent zero-mean numbers has variance equal to the product of their variances.

\text{Var}\Big(\sum_{i=1}^{n} X_i\Big) = \sum_{i=1}^{n}\text{Var}(X_i) \;\;(\text{independent}), \qquad \text{Var}(XY) = \sigma_X^2\sigma_Y^2 \;\;(X \perp Y,\ \mathbb E X = \mathbb E Y = 0)

$\text{Var}$: variance: the average squared distance from the mean
$\sum_{i=1}^{n}$: add up over the $n$ terms
$X_i$: the $i$ -th random term, e.g. one product $q_ik_i$
$n$: the number of terms in the sum
$i$: here: which term
$X$: a random number with mean 0
$Y$: another, independent of $X$ ( $\perp$ = independent)
$\sigma_X^2$: the variance of $X$
$\sigma_Y^2$: the variance of $Y$
$\mathbb E$: expectation: the average over repeated random draws

Each card names the move that does the work.

Formula sheet

The operator

Scaled dot-product attention (Card 1)

\text{Attn}(Q, K, V) = \operatorname{softmax}\Big(\frac{QK^\top}{\sqrt{d_k}} + M\Big)V, \qquad Q = XW_Q,\;\; K = XW_K,\;\; V = XW_V

$Q$: the queries, one row per token: what each token looks for ( $T \times d_k$ )
$K$: the keys, one row per token: how each token advertises itself ( $T \times d_k$ )
$V$: the values, one row per token: what each token delivers ( $T \times d_v$ )
$QK^\top$: the $T \times T$ score matrix: every query dotted with every key (⊤ = transpose)
$d_k$: the width of each query and key: 2 in the toy, 128 in Llama 2 7B
$M$: the mask: 0 where attention is allowed, $-\infty$ where it is blocked
$X$: the input token vectors, one row per token ( $T \times d$ )
$W_Q$: the learned query map ( $d \times d_k$ )
$W_K$: the learned key map ( $d \times d_k$ )
$W_V$: the learned value map ( $d \times d_v$ )

One row (Card 1)

\mathbf o_i = \sum_j p_{ij}\,\mathbf v_j, \qquad p_{ij} = \frac{e^{s_{ij}}}{\sum_r e^{s_{ir}}}, \qquad s_{ij} = \frac{\mathbf q_i^\top\mathbf k_j}{\sqrt{d_k}} + M_{ij}

$\mathbf o_i$: token $i$ 's output: its new vector
$\sum_j$: add up over the keys $j$
$p_{ij}$: the share of token $i$ 's attention that goes to token $j$ (≥ 0; each row sums to 1)
$\mathbf v_j$: token $j$ 's value
$e$: Euler's number, 2.718…
$s_{ij}$: the score of key $j$ for query $i$
$\sum_r$: add up over every key $r$ in row $i$ (the normalizer)
$s_{ir}$: the score of key $r$ for query $i$
$\mathbf q_i$: token $i$ 's query (row $i$ of $Q$ )
$\mathbf k_j$: token $j$ 's key (row $j$ of $K$ )
$M_{ij}$: the mask entry: 0 if $i$ may see $j$ , $-\infty$ if not
$i$: the query token (the one being updated)
$j$: a key token (one being read)
$r$: any key in the row, used inside the normalizer

Only $W_QW_K^\top$ reaches the scores (Section 3)

S = \frac{XW_QW_K^\top X^\top}{\sqrt{d_k}}, \qquad (W_QA)(W_KA^{-\top})^\top = W_QW_K^\top \;\text{ for any invertible } A

$S$: the score matrix before masking
$W_QW_K^\top$: the "QK circuit": a $d \times d$ matrix of rank at most $d_k$ , all the scores ever see
$X^\top$: $X$ transposed
$A$: any invertible $d_k \times d_k$ matrix
$A^{-\top}$: the inverse of $A$ , transposed

Softmax

Shift invariance, logsumexp, temperature (Card 2)

\operatorname{softmax}(\mathbf s + c\mathbf 1) = \operatorname{softmax}(\mathbf s), \qquad \log\sum_j e^{s_j} = m + \log\sum_j e^{s_j - m}, \qquad \operatorname{softmax}(\mathbf s/\tau) \to \begin{cases}\text{one-hot on the top score} & \tau \to 0\\ \text{uniform } 1/T & \tau \to \infty\end{cases}

$\mathbf s$: one row of scores
$c$: any constant
$\mathbf 1$: the all-ones vector
$s_j$: the score of key $j$ in this row
$m$: the largest score in the row
$\tau$: the temperature: divides every score; small = sharp, large = flat
$T$: the number of keys in the row

Entropy and effective keys (Card 3)

H(\mathbf p) = -\sum_j p_j\log p_j \;\le\; \ln T, \qquad \text{effective keys} = e^{H}

$H(\mathbf p)$: the entropy of a weight row, in nats: 0 = one-hot, $\ln T$ = uniform
$\mathbf p$: one row of attention weights
$p_j$: the weight on key $j$
$H$: the entropy, written without its argument

Scaling

Variance of a dot product (Card 3)

\text{Var}(\mathbf q^\top\mathbf k) = d_k\,\sigma_q^2\sigma_k^2 \;\;(\text{i.i.d. components, mean } 0), \qquad \text{Var}\Big(\frac{\mathbf q^\top\mathbf k}{\sqrt{d_k}}\Big) = 1 \;\;(\sigma_q = \sigma_k = 1)

$\text{Var}$: variance: the average squared spread around the mean
$\mathbf q^\top\mathbf k$: the raw dot product of a query and a key
$\mathbf q$: a query vector with random components
$\mathbf k$: a key vector with random components
$\sigma_q^2$: the variance of each query component
$\sigma_k^2$: the variance of each key component
$\sigma_q$: the standard deviation of each query component
$\sigma_k$: the standard deviation of each key component

Entropy of Gaussian logits (Card 3)

H \approx \ln T - \frac{\sigma_s^2}{2} \;\;(\sigma_s \ll \sqrt{2\ln T}), \qquad \sigma_s = 1,\ T = 4{,}096:\;\; H \approx 7.818

$\sigma_s$: the spread (standard deviation) of the logits across keys
$\sigma_s^2$: its square, the logit variance

Two keys, exactly (Card 3)

\Pr\big(\max_j p_j > 0.99\big) = \Pr\big(\lvert s_1 - s_2\rvert > \ln 99\big), \qquad s_1 - s_2 \mid \mathbf q \sim \mathcal N\big(0,\ 2\alpha^2\lVert\mathbf q\rVert^2\big), \qquad \lVert\mathbf q\rVert^2 \sim \chi^2_{d_k}

$\max_j p_j$: the top weight in the row
$s_1$: the first key's score
$s_2$: the second key's score
$\mathcal N$: a normal distribution, written $\mathcal N(\text{mean}, \text{variance})$ ; $\sim$ = is distributed as; $\mid$ = given
$\alpha$: the scale applied to the dot products: 1 (none) or $1/\sqrt{d_k}$
$\lVert\mathbf q\rVert^2$: the query's squared length
$\chi^2_{d_k}$: chi-square with $d_k$ degrees of freedom: a sum of $d_k$ squared standard normals

Aligned queries and keys (Card 16)

\mathbf k = \rho\,\mathbf q + \sqrt{1 - \rho^2}\,\mathbf z:\quad \mathbb E\Big[\frac{\mathbf q^\top\mathbf k}{\sqrt{d_k}}\Big] = \rho\sqrt{d_k},\;\; \text{sd} = \sqrt{1 + \rho^2}; \qquad \mathbb E\Big[\frac{\mathbf q^\top\mathbf k}{d_k}\Big] = \rho

$\rho$: how aligned key and query are: 0 = independent, 1 = identical
$\mathbf z$: a standard-normal vector independent of $\mathbf q$
$\mathbb E$: expectation: the average over random draws

Gradients

Softmax Jacobian (Card 4)

J = \frac{\partial\mathbf p}{\partial\mathbf s} = \operatorname{diag}(\mathbf p) - \mathbf p\mathbf p^\top, \qquad J\mathbf 1 = \mathbf 0, \qquad 0 \le \lambda(J) \le 2\max_i p_i(1 - p_i) \le \tfrac12

$J$: the softmax Jacobian: how each weight moves when each score is nudged
$\frac{\partial\mathbf p}{\partial\mathbf s}$: the matrix of slopes $\partial p_i/\partial s_j$ : slopes, not fractions to divide
$\operatorname{diag}(\mathbf p)$: the diagonal matrix with the weights on its diagonal
$\mathbf p\mathbf p^\top$: the outer product of the weights with themselves
$\mathbf 0$: the zero vector
$\lambda(J)$: any eigenvalue of $J$
$p_i$: the weight on key $i$
$i$: any key in the row

Backward pass (Card 5)

dV = P^\top dO, \quad dP = dO\,V^\top, \quad dS = P \odot \big(dP - \operatorname{rowsum}(P \odot dP)\big), \quad dQ = \frac{dS\,K}{\sqrt{d_k}}, \quad dK = \frac{dS^\top Q}{\sqrt{d_k}}

$dV$: the gradient of the loss with respect to the values (same shape as $V$ )
$P^\top$: the attention weights, transposed
$P$: the $T \times T$ matrix of attention weights
$dO$: the upstream gradient arriving at the outputs
$dP$: the gradient with respect to the weights
$V^\top$: the values, transposed
$dS$: the gradient with respect to the scores
$\odot$: entry-by-entry product
$dQ$: the gradient with respect to the queries
$dK$: the gradient with respect to the keys
$dS^\top$: $dS$ transposed

One row, one key (Card 5)

\frac{\partial\mathcal L}{\partial s_j} = p_j\big(\mathbf g^\top\mathbf v_j - \mathbf g^\top\mathbf o\big), \qquad \frac{\partial\mathcal L}{\partial\mathbf q} = \sum_j \frac{\partial\mathcal L}{\partial s_j}\,\frac{\mathbf k_j}{\sqrt{d_k}}, \qquad \operatorname{rowsum}(P \odot dP) = \operatorname{rowsum}(dO \odot O)

$\frac{\partial\mathcal L}{\partial s_j}$: how fast the loss changes as you nudge score $s_j$ : a slope, not a fraction
$\mathcal L$: the training loss
$\mathbf g$: the upstream gradient at this row's output, $\partial\mathcal L/\partial\mathbf o$
$\mathbf o$: this row's output
$\frac{\partial\mathcal L}{\partial\mathbf q}$: the gradient with respect to this row's query
$O$: the matrix of outputs, one row per token

Masks and positions

The causal mask (Card 6)

M_{ij} = \begin{cases}0 & j \le i\\ -\infty & j > i\end{cases} \;\Rightarrow\; p_{ij} = 0 \text{ for } j > i, \qquad \frac{\#\text{allowed pairs}}{T^2} = \frac{T + 1}{2T}

$T^2$: the number of (query, key) pairs

Permutation equivariance and the position leak (Card 7)

\text{Attn}(\Pi X) = \Pi\,\text{Attn}(X) \;\;(\text{no mask, no positions}), \qquad \text{causal, equal scores: share of BOS at position } t = \frac1t

$\Pi$: a permutation matrix: it reorders the tokens
$t$: a position in the sequence (1 = first)

Nadaraya–Watson is attention (Card 8)

\hat f(\mathbf x) = \frac{\sum_j \kappa_\lambda(\mathbf x, \mathbf x_j)\,y_j}{\sum_j \kappa_\lambda(\mathbf x, \mathbf x_j)}, \qquad -\frac{\lVert\mathbf q - \mathbf k_j\rVert^2}{2\lambda^2} = \frac{\mathbf q^\top\mathbf k_j}{\lambda^2} - \frac{\lVert\mathbf q\rVert^2}{2\lambda^2} - \frac{\lVert\mathbf k_j\rVert^2}{2\lambda^2}

$\hat f$: the kernel-smoothed estimate (hat = estimated from data)
$\mathbf x$: the point where we want an estimate
$\kappa_\lambda$: the kernel: a similarity weight that falls with distance
$\mathbf x_j$: data point $j$
$y_j$: the value observed at data point $j$
$\lVert\mathbf q - \mathbf k_j\rVert^2$: the squared distance between query and key $j$
$\lambda$: the kernel bandwidth (usually written $h$ ; here $h$ counts heads)
$\lVert\mathbf k_j\rVert^2$: key $j$ 's squared length

The rank bound (Card 10)

\operatorname{rank}(QK^\top) \le d_k, \qquad d_k = 1:\;\; s_{ij} = q_ik_j \;\Rightarrow\; \arg\max_j s_{ij} \in \big\{\arg\max_j k_j,\ \arg\min_j k_j\big\}

$q_i$: token $i$ 's query, a single number when $d_k = 1$
$k_j$: token $j$ 's key, a single number when $d_k = 1$

Multi-head and cost

Multi-head attention and its parameters (Card 9)

\big[\text{head}_1, \ldots, \text{head}_h\big]W_O = \sum_{n=1}^{h}\text{head}_n\,W_O^{(n)}, \qquad \#\text{params} = 3h\,d\,d_k + h\,d_v\,d = 4d^2 \;\;(h\,d_k = h\,d_v = d), \qquad \text{GQA: } 2d^2 + 2d\,H_{\text{kv}}d_h

$h$: the number of heads
$W_O$: the output projection, mixing the concatenated heads back to width $d$
$\sum_{n=1}^{h}$: add up over the heads $n$
$W_O^{(n)}$: the block of rows of $W_O$ that reads head $n$
$n$: which head
$d$: the model width $d_{\text{model}}$ : 4 in the toy, 4,096 in Llama 2 7B
$d_v$: the width of each head's values
$H_{\text{kv}}$: the number of key/value heads (GQA shares each across several query heads)
$d_h$: the head width, $d_k = d_v$ (128 in every Llama in this chapter)

FLOPs per layer and the crossover (Card 11)

\underbrace{2Td(2d + 2H_{\text{kv}}d_h)}_{\text{projections; } 8Td^2 \text{ for MHA}} + \underbrace{4T^2d}_{QK^\top \text{ and } PV} + \underbrace{6Td\,d_{\text{ff}}}_{\text{SwiGLU FFN}}, \qquad T^\star = 2d + 1.5\,d_{\text{ff}}

$T$: the context length (tokens in the sequence)
$d_{\text{ff}}$: the feed-forward hidden width (11,008 in Llama 2 7B)
$T^\star$: the context length where the quadratic part equals everything else (star = the crossover)

Memory: naive scores vs online (Card 12)

\text{naive} = B\,h\,T^2\,b \;\text{bytes per layer}, \qquad \text{online} = B\,h\,T \times 4 \;\text{bytes per layer (one fp32 logsumexp per row)}

$B$: the batch size (sequences processed together)
$b$: bytes per number: 2 for fp16 or bf16

Inference

KV cache (Card 13)

\text{KV bytes} = 2 \cdot L \cdot H_{\text{kv}} \cdot d_h \cdot b \cdot T \cdot B

$L$: the number of layers

Sharing keys and values (Card 14)

\frac{\text{cache}_{\text{MHA}}}{\text{cache}_{\text{GQA}}} = \frac{h}{H_{\text{kv}}}, \qquad \text{decode attention intensity} = \frac{h \cdot 4Td_h}{H_{\text{kv}} \cdot 2Td_h\,b} = \frac{h}{H_{\text{kv}}} \;\text{FLOP/byte} \;\;(b = 2)

Online softmax and the merge (Card 15)

m_n = \max\big(m_{n-1},\, \max_{j \in \mathcal B_n} s_j\big), \quad \ell_n = e^{m_{n-1} - m_n}\ell_{n-1} + \sum_{j \in \mathcal B_n} e^{s_j - m_n}, \quad \mathbf u_n = e^{m_{n-1} - m_n}\mathbf u_{n-1} + \sum_{j \in \mathcal B_n} e^{s_j - m_n}\mathbf v_j, \quad \mathbf o = \frac{\mathbf u_N}{\ell_N}

$m_n$: the running max after block $n$
$n$: here: the block number
$m_{n-1}$: the running max before block $n$
$\mathcal B_n$: the set of keys in block $n$
$\ell_n$: the running normalizer: the sum of $e^{s_j - m_n}$ over keys seen so far
$\ell_{n-1}$: the normalizer before block $n$ , on the old base $m_{n-1}$
$\mathbf u_n$: the running unnormalized output: weighted sum of values seen so far
$\mathbf u_{n-1}$: the unnormalized output before block $n$
$\mathbf u_N$: the unnormalized output after the last block
$\ell_N$: the normalizer after the last block
$N$: the number of blocks

m = \max(m_a, m_b), \qquad \ell = e^{m_a - m}\ell_a + e^{m_b - m}\ell_b, \qquad \mathbf u = e^{m_a - m}\mathbf u_a + e^{m_b - m}\mathbf u_b, \qquad \mathbf o = \mathbf u/\ell

$m$: here: the merged max
$m_a$: partial result $a$ 's max (e.g. GPU 1's half of the keys)
$m_b$: partial result $b$ 's max
$\ell$: the merged normalizer
$\ell_a$: partial $a$ 's normalizer
$\ell_b$: partial $b$ 's normalizer
$\mathbf u$: the merged unnormalized output
$\mathbf u_a$: partial $a$ 's unnormalized output
$\mathbf u_b$: partial $b$ 's unnormalized output

Uniform heads and rank collapse (Card 17)

\text{Var}\Big(\frac1T\sum_j v_j\Big) = \frac{\sigma_v^2}{T}, \qquad r_{l+1} \le C\,r_l^3 \;\Rightarrow\; r_L \le C^{-1/2}\big(\sqrt C\,r_0\big)^{3^L}

$v_j$: one coordinate of value $j$ , random with variance $\sigma_v^2$
$\sigma_v^2$: the variance of a single value coordinate
$r_{l+1}$: how far the tokens are from all being identical after layer $l + 1$
$r_l$: the same distance after layer $l$
$l$: the layer number
$C$: a constant set by the size of the weight matrices
$r_L$: the distance after the last layer
$r_0$: the distance at the input

Sinks and QK-norm (Card 18)

w = \frac{e^{a}}{e^{a} + T - 1} \iff a = \ln\frac{w(T - 1)}{1 - w}, \qquad s_{ij} = g\,\hat{\mathbf q}_i^\top\hat{\mathbf k}_j \;\Rightarrow\; \max_j p_{ij} \le \frac{e^{g}}{e^{g} + (T - 1)e^{-g}}

$w$: the share of attention the sink token takes
$a$: the sink's logit advantage over each of the other $T - 1$ keys
$g$: the QK-norm gain: the logit of a perfectly aligned pair
$\hat{\mathbf q}_i$: query $i$ scaled to length 1 (hat = normalized)
$\hat{\mathbf k}_j$: key $j$ scaled to length 1

Cards

Stars: ★ = asked often, ★★ = asked at senior level, ★★★ = research level.

Card 1 · Scaled dot-product attention in matrix form, with shapes ★

The story. Bahdanau, Cho and Bengio (2014) scored each source word with a small feed-forward network inside an RNN translator. Vaswani et al. (2017) replaced that network with a dot product and wrote every token's lookup as one matrix product, noting that dot-product attention is much faster in practice because it runs on highly optimized matrix-multiplication code.

Statement. For $T$ tokens, attention is two matrix products around a row-wise softmax, and every output row is a convex combination (a weighted average with non-negative weights summing to 1) of the value rows.

O = \operatorname{softmax}\Big(\frac{QK^\top}{\sqrt{d_k}} + M\Big)V: \qquad \underbrace{(T \times d_k)(d_k \times T)}_{S:\; T \times T} \;\to\; \underbrace{(T \times T)(T \times d_v)}_{O:\; T \times d_v}

$O$: the outputs, one new vector per token ( $T \times d_v$ )
$Q$: the queries, one row per token (Ask)
$K$: the keys, one row per token (Advertise)
$V$: the values, one row per token (Deliver)
$QK^\top$: every query dotted with every key: the $T \times T$ raw scores (⊤ = transpose)
$d_k$: the query and key width: 2 in the toy, 128 in Llama 2 7B
$M$: the mask: 0 where allowed, $-\infty$ where blocked
$T$: the number of tokens: 5 in the toy
$d_v$: the value width: 2 in the toy
$S$: the scaled, masked score matrix

Derivation.

Start with one token. Token $i$ wants a blend of the values, weighted by how well its query matches each key. The weights must be non-negative and sum to 1, and a softmax of the scores gives exactly that.

\mathbf o_i = \sum_{j=1}^{T} p_{ij}\,\mathbf v_j, \qquad p_{ij} = \frac{e^{s_{ij}}}{\sum_{r=1}^{T} e^{s_{ir}}}, \qquad s_{ij} = \frac{\mathbf q_i^\top\mathbf k_j}{\sqrt{d_k}} + M_{ij}

$\mathbf o_i$: token $i$ 's new vector
$\sum_{j=1}^{T}$: add up over all $T$ keys $j$
$p_{ij}$: the share of token $i$ 's attention given to token $j$
$\mathbf v_j$: token $j$ 's value
$e$: Euler's number, 2.718…
$s_{ij}$: the score of key $j$ for query $i$
$\sum_{r=1}^{T}$: add up over every key $r$ (the normalizer)
$s_{ir}$: the score of key $r$ for query $i$
$\mathbf q_i$: token $i$ 's query (row $i$ of $Q$ )
$\mathbf k_j$: token $j$ 's key (row $j$ of $K$ )
$M_{ij}$: the mask entry for the pair $(i, j)$
$i$: the query token being updated
$j$: a key token being read
$r$: any key, inside the normalizer

Stack the scores. Entry $(i, j)$ of $QK^\top$ is row $i$ of $Q$ dotted with row $j$ of $K$ , so one matrix product computes all $T^{2}$ scores at once.

(QK^\top)_{ij} = \sum_{c=1}^{d_k} Q_{ic}K_{jc} = \mathbf q_i^\top\mathbf k_j \;\Rightarrow\; S = \frac{QK^\top}{\sqrt{d_k}} + M \in \mathbb R^{T \times T}

$\sum_{c=1}^{d_k}$: add up over the $d_k$ components $c$
$Q_{ic}$: component $c$ of token $i$ 's query
$K_{jc}$: component $c$ of token $j$ 's key
$c$: a component (coordinate) index
$\mathbb R^{T \times T}$: the set of $T \times T$ matrices of real numbers

Softmax acts on each row separately and gives the weight matrix $P$ . Row $i$ of $PV$ is $\sum_j p_{ij}\mathbf v_j^\top$ , which is exactly $\mathbf o_i^\top$ .

P = \operatorname{softmax}_{\text{row}}(S), \qquad (PV)_{i\cdot} = \sum_j p_{ij}\,\mathbf v_j^\top = \mathbf o_i^\top \;\Rightarrow\; O = PV

$P$: the $T \times T$ attention weights; each row sums to 1
$(PV)_{i\cdot}$: row $i$ of the product $PV$
$\sum_j$: add up over the keys $j$

Each output is a convex combination. The weights are non-negative and sum to 1, so every coordinate of $\mathbf o_i$ lies between the smallest and largest value of that coordinate. Attention can blend values; it cannot extrapolate beyond them.

\min_j v_{jc} \;\le\; o_{ic} = \sum_j p_{ij}\,v_{jc} \;\le\; \max_j v_{jc}

$v_{jc}$: coordinate $c$ of value $j$
$o_{ic}$: coordinate $c$ of output $i$

Batches and heads add leading axes. Only the last two axes are multiplied, so $B$ and $h$ are loop dimensions that run in parallel.

(B, h, T, d_k) \times (B, h, d_k, T) \to (B, h, T, T), \qquad (B, h, T, T) \times (B, h, T, d_v) \to (B, h, T, d_v)

$B$: the batch size
$h$: the number of heads

Check on the toy. Bank's causal row puts 0.140, 0.576 and 0.284 on the, river and bank. Both output coordinates lie in [0, 1], the range of the values. For Llama 2 7B the same algebra makes a 4,096 × 4,096 score matrix per head from $d_k = 128$ .

\mathbf o_{\text{bank}} = 0.140\,(0, 0) + 0.576\,(1, 0) + 0.284\,(0, 1) = (0.576,\ 0.284)

$\mathbf o_{\text{bank}}$: bank's new vector in [geography, money] space

Key move. S (shapes tally): write the shape of each factor and the formula writes itself.

Common slip. Taking the softmax over the wrong axis. With dim=0 each key's weights sum to 1 over the queries, so the outputs stop being averages. Two more: writing $Q^\top K$ (a $d_k \times d_k$ matrix), and multiplying by a 0/1 mask instead of adding $-\infty$ (Card 6).

How it's asked. "Write attention in matrix form and give the shape of every intermediate." (Q13, Q16.) Strong answer: the shapes with batch and heads, the $T \times T$ intermediate named as the memory problem (Card 12), and the one-line summary "linear in $V$ , nonlinear in $Q$ and $K$ ".

Card 2 · Softmax: shift invariance, temperature, stable logsumexp ★

The story. In fp16 the largest number is 65,504, so $e^{x}$ overflows once $x$ passes 11.09. Every production softmax therefore subtracts the row max before exponentiating, and the same trick, applied block by block, is what lets FlashAttention stay exact (Card 15).

Statement. Adding a constant to every score leaves the softmax unchanged, the log-normalizer can be computed with the max factored out, and dividing the scores by a temperature moves the softmax between one-hot and uniform.

\operatorname{softmax}(\mathbf s + c\mathbf 1) = \operatorname{softmax}(\mathbf s), \qquad \log\sum_j e^{s_j} = m + \log\sum_j e^{s_j - m}, \qquad \lim_{\tau \to 0}\operatorname{softmax}(\mathbf s/\tau) = \mathbf e_{j^\star}, \qquad \lim_{\tau \to \infty}\operatorname{softmax}(\mathbf s/\tau) = \tfrac1T\mathbf 1

$\mathbf s$: one row of scores
$c$: here: any constant added to every score
$\mathbf 1$: the all-ones vector
$s_j$: the score of key $j$
$m$: the largest score in the row
$\tau$: the temperature: every score is divided by it
$\mathbf e_{j^\star}$: the one-hot vector on the top key
$j^\star$: the key with the largest score (assumed unique)

Derivation.

Shift: multiply numerator and denominator by the same positive number $e^{c}$ .

\frac{e^{s_i + c}}{\sum_j e^{s_j + c}} = \frac{e^{c}\,e^{s_i}}{e^{c}\sum_j e^{s_j}} = \frac{e^{s_i}}{\sum_j e^{s_j}}

$s_i$: the score of key $i$

Choose $c = -m$ . Every exponent is now at most 0, so each term lies in (0, 1] and the top key contributes exactly 1. The sum lies between 1 and $T$ : nothing overflows, and the denominator can never be 0.

1 \;\le\; \sum_j e^{s_j - m} \;\le\; T

The log of the normalizer (the logsumexp): factor $e^{m}$ out of the sum.

\log\sum_j e^{s_j} = \log\Big(e^{m}\sum_j e^{s_j - m}\Big) = m + \log\sum_j e^{s_j - m}

Check: scores (1000, 999). The naive $e^{1000}$ overflows even in fp64; the stable form is exact.

\log\big(e^{1000} + e^{999}\big) = 1000 + \log\big(1 + e^{-1}\big) = 1000.313

Temperature. Write the weight at temperature $\tau$ with the max factored out. For every key but the top one, $s_j - m < 0$ , so its exponent heads to $-\infty$ as $\tau \to 0$ and its weight to 0. As $\tau \to \infty$ every exponent heads to 0 and every weight to $1/T$ .

p_j(\tau) = \frac{e^{(s_j - m)/\tau}}{\sum_r e^{(s_r - m)/\tau}}: \qquad \tau \to 0:\; p_j \to \begin{cases}1 & j = j^\star\\ 0 & \text{otherwise}\end{cases}, \qquad \tau \to \infty:\; p_j \to \frac{e^{0}}{T\,e^{0}} = \frac1T

$p_j(\tau)$: the weight on key $j$ at temperature $\tau$
$p_j$: the weight on key $j$
$\sum_r$: add up over every key $r$ in the row
$s_r$: the score of key $r$

Check on the toy. Adding 10 to bank's scores leaves 0.140 / 0.576 / 0.284 unchanged. Dividing by $\sqrt{d_k}$ is a temperature of $\sqrt 2$ applied to the raw scores 0, 2, 1, and it softens the row.

\operatorname{softmax}(0,\ 2,\ 1) = (0.090,\ 0.665,\ 0.245), \qquad \operatorname{softmax}\big((0,\ 2,\ 1)/\sqrt 2\big) = (0.140,\ 0.576,\ 0.284)

The O move, used in Cards 4 and 5: $\log p_i = s_i - \log\sum_r e^{s_r}$ , and the derivative of the logsumexp is the softmax itself.

\frac{\partial}{\partial s_j}\log\sum_r e^{s_r} = \frac{e^{s_j}}{\sum_r e^{s_r}} = p_j \;\Rightarrow\; \frac{\partial\log p_i}{\partial s_j} = \delta_{ij} - p_j

$\frac{\partial}{\partial s_j}$: how fast the expression after it changes as you nudge score $s_j$ : a slope, not a fraction to divide
$\frac{\partial\log p_i}{\partial s_j}$: how fast $\log p_i$ changes as you nudge $s_j$
$p_i$: the weight on key $i$
$\delta_{ij}$: 1 if $i = j$ , else 0
$i$: the key whose weight we watch

Key move. F (factor out the max).

Common slip. Calling softmax "scale-invariant". It is shift-invariant; scaling is a temperature and changes the answer (step 6). Also: subtracting the max in the numerator only, or subtracting the minimum, which can still overflow.

How it's asked. "Your custom kernel's softmax returns inf and NaN in fp16 on long prompts but not short ones. Why, and what's the fix?" (Q23; also Q3.) Strong answer: subtract the row max, accumulate the sum in fp32, and keep the logsumexp, because that single number per row is what FlashAttention stores for its backward pass (Cards 12 and 15).

Card 3 · Var(q·k) = d_k, hence 1/√d_k ★

The story. Vaswani et al. (2017) noted that, without scaling, dot-product attention falls behind additive attention at large $d_k$ . They wrote that they "suspect" large dot products push the softmax into regions with extremely small gradients, and a footnote gave the variance argument below.

Statement. If the components of $\mathbf q$ and $\mathbf k$ are independent with mean 0, the dot product's variance grows linearly in $d_k$ , so dividing by $\sqrt{d_k}$ restores unit variance. Over $T$ keys, logits with spread $\sigma_s$ give an entropy of about $\ln T - \sigma_s^{2}/2$ .

\text{Var}(\mathbf q^\top\mathbf k) = d_k\,\sigma_q^2\sigma_k^2, \qquad \text{Var}\Big(\frac{\mathbf q^\top\mathbf k}{\sqrt{d_k}}\Big) = 1 \;\text{ when } \sigma_q = \sigma_k = 1, \qquad H \approx \ln T - \frac{\sigma_s^2}{2}

$\text{Var}$: variance: the average squared spread around the mean
$\mathbf q^\top\mathbf k$: the raw dot product of a random query and key
$\mathbf q$: a query with independent random components
$\mathbf k$: a key with independent random components
$\sigma_q^2$: the variance of each query component
$\sigma_k^2$: the variance of each key component
$\sigma_q$: the standard deviation of each query component
$\sigma_k$: the standard deviation of each key component
$H$: the entropy of a weight row, in nats: 0 = one-hot, $\ln T$ = uniform
$\sigma_s$: the spread (standard deviation) of the logits across keys
$\sigma_s^2$: the logit variance

Derivation.

The dot product is a sum of $d_k$ products.

\mathbf q^\top\mathbf k = \sum_{c=1}^{d_k} q_c k_c

$q_c$: component $c$ of the query
$k_c$: component $c$ of the key
$c$: here: a component index, 1 to $d_k$

Each product has mean 0 and variance $\sigma_q^{2}\sigma_k^{2}$ , because the two factors are independent with mean 0.

\mathbb E[q_ck_c] = \mathbb E[q_c]\,\mathbb E[k_c] = 0, \qquad \text{Var}(q_ck_c) = \mathbb E[q_c^2k_c^2] - 0^2 = \mathbb E[q_c^2]\,\mathbb E[k_c^2] = \sigma_q^2\sigma_k^2

$\mathbb E$: expectation: the average over random draws

Products for different $c$ are independent, so their variances add.

\text{Var}(\mathbf q^\top\mathbf k) = \sum_{c=1}^{d_k}\text{Var}(q_ck_c) = d_k\,\sigma_q^2\sigma_k^2

Dividing by a constant divides the variance by its square, so $1/\sqrt{d_k}$ brings unit-variance components back to unit-variance logits. At $d_k = 128$ the standard deviation drops from 11.3 to 1 (simulated variance 128.4 against 128). With unequal component variances, $\sigma_q = 1.5$ and $\sigma_k = 0.5$ at $d_k = 64$ , theory gives 36 and simulation 35.8: the "right" scale is $1/(\sqrt{d_k}\,\sigma_q\sigma_k)$ , and normalization layers are what keep $\sigma$ near 1.

\text{Var}\Big(\frac{\mathbf q^\top\mathbf k}{\sqrt{d_k}}\Big) = \frac{d_k\,\sigma_q^2\sigma_k^2}{d_k} = \sigma_q^2\sigma_k^2, \qquad d_k = 128:\;\; \text{sd} = \sqrt{128} = 11.3 \;\to\; 1

Why the spread matters: the entropy of a row. Start from $H = -\sum_j p_j\log p_j$ with $\log p_j = s_j - \log\sum_r e^{s_r}$ . For many keys with Gaussian logits of spread $\sigma_s$ , the sums become averages: $\sum_j e^{s_j} \approx T\,\mathbb E[e^{s}] = T e^{\sigma_s^{2}/2}$ and $\mathbb E[s\,e^{s}] = \sigma_s^{2}e^{\sigma_s^{2}/2}$ .

H = \log\sum_j e^{s_j} - \sum_j p_js_j \;\approx\; \ln T + \frac{\sigma_s^2}{2} - \frac{\mathbb E[s\,e^{s}]}{\mathbb E[e^{s}]} = \ln T + \frac{\sigma_s^2}{2} - \sigma_s^2 = \ln T - \frac{\sigma_s^2}{2}

$\sum_j$: add up over the keys $j$
$p_j$: the weight on key $j$
$s_j$: the logit of key $j$
$e$: Euler's number, 2.718…
$s$: one random logit, normal with mean 0 and spread $\sigma_s$

Plug in $T = 4{,}096$ . Scaled logits ( $\sigma_s = 1$ ) give 7.818; the simulated median is 7.83. The formula holds only while $\sigma_s$ is well below $\sqrt{2\ln T} = 4.08$ : at $\sigma_s = 2$ it gives 6.32 against a simulated 6.41, and at $\sigma_s = 4$ it predicts 0.32 against 3.17. Unscaled logits at $d_k = 128$ ( $\sigma_s = 11.3$ ) are far past that edge. The row is nearly one-hot: median top weight 0.868, entropy 0.49 nats, 1.63 effective keys, against 2,511 effective keys when scaled.

H \approx \ln 4{,}096 - \tfrac12 = 8.318 - 0.5 = 7.818 \;\;(\sigma_s = 1), \qquad \sigma_s = 11.3 \;\gg\; \sqrt{2\ln 4{,}096} = 4.08

The exact two-key version. With two keys, the top weight exceeds 0.99 exactly when the score gap exceeds $\ln 99 = 4.595$ . Given $\mathbf q$ , the gap $\alpha\,\mathbf q^\top(\mathbf k_1 - \mathbf k_2)$ is normal with variance $2\alpha^{2}\lVert\mathbf q\rVert^{2}$ , and $\lVert\mathbf q\rVert^{2}$ is chi-square with $d_k$ degrees of freedom. Average the two-sided tail over it.

\Pr\big(\max_j p_j > 0.99\big) = \mathbb E_{\lVert\mathbf q\rVert^2 \sim \chi^2_{d_k}}\Big[2\,\Phi\Big(-\frac{\ln 99}{\alpha\sqrt{2\lVert\mathbf q\rVert^2}}\Big)\Big] = \begin{cases}0.773 & \alpha = 1\\ 0.00135 & \alpha = 1/\sqrt{128}\end{cases}

$\max_j p_j$: the larger of the two weights
$\mathbb E_{\lVert\mathbf q\rVert^2 \sim \chi^2_{d_k}}$: the average over the random squared length of the query
$\lVert\mathbf q\rVert^2$: the query's squared length
$\chi^2_{d_k}$: chi-square with $d_k$ degrees of freedom: a sum of $d_k$ squared standard normals
$\Phi$: the standard normal cumulative distribution function
$\alpha$: the scale on the dot products: 1 (none) or $1/\sqrt{d_k}$

Replacing $\lVert\mathbf q\rVert^{2}$ by its mean $d_k$ gives the Gaussian approximation, 0.774 and 0.00116 (with scaling, $\ln 99$ is 3.25 standard deviations of the gap).

The gradient consequence (Card 4). A saturated row sends its gradient almost only to the winner, and that gradient shrinks as the row saturates further: at $d_k = 128$ , 22% of unscaled rows put more than 0.99 on one key (46% at $d_k = 512$ ). The size of the Jacobian is not monotone in $d_k$ . Unscaled, its Frobenius norm peaks at $d_k = 32$ , and at $d_k = 128$ it is larger than the scaled one, but it is concentrated on about 1.6 keys.

\lVert J\rVert_F = 0.181 \;\text{(unscaled, on 1.63 effective keys)} \quad\text{vs}\quad 0.025 \;\text{(scaled, on 2,511)} \qquad (d_k = 128)

$\lVert J\rVert_F$: the size of the softmax Jacobian: square root of the sum of its squared entries (median over rows)

Key move. T (total variance of a sum).

Common slip. Dividing by $\sqrt{d_{\text{model}}}$ , or by $d_k$ ("to normalize the dot product"). With $1/d_k$ at initialization the logits have spread $1/\sqrt{d_k}$ and every row is almost uniform: entropy 8.314 against the maximum 8.318 at $T = 4{,}096$ . Second slip: saying unscaled gradients simply "vanish". They concentrate on the winner (step 8). Third: forgetting the assumption. The argument is about independent $\mathbf q$ and $\mathbf k$ at initialization (Card 16).

How it's asked. "Show why the scale is $\sqrt{d_k}$ , not $d_k$ or $\sqrt{d_{\text{model}}}$ , and state your assumptions." (Q4, Q14, Worked Problem 1, Drill 2.) Strong answer: the three-line variance argument, "too big makes rows one-hot, too small makes them uniform", the exact two-key numbers, and the caveat that training aligns $\mathbf q$ and $\mathbf k$ , which is why modern stacks also bound logits directly (Cards 16 and 18).

Card 4 · The softmax Jacobian ★★

The story. Backpropagating through attention needs to know how every weight in a row moves when any one score moves. The answer also explains why a saturated row stops learning.

Statement. The softmax Jacobian is $\operatorname{diag}(\mathbf p) - \mathbf p\mathbf p^\top$ . It is symmetric and positive semi-definite, its rows sum to 0, it has rank $T - 1$ when every weight is positive, its largest eigenvalue is at most 1/2, and it goes to 0 as the row becomes one-hot.

\frac{\partial p_i}{\partial s_j} = p_i(\delta_{ij} - p_j), \qquad J = \operatorname{diag}(\mathbf p) - \mathbf p\mathbf p^\top:\;\; J = J^\top,\;\; J \succeq 0,\;\; J\mathbf 1 = \mathbf 0,\;\; \operatorname{rank}J = T - 1,\;\; \lVert J\rVert_2 \le \tfrac12

$\frac{\partial p_i}{\partial s_j}$: how fast weight $p_i$ changes as you nudge score $s_j$ : a slope, not a fraction to divide
$J$: the softmax Jacobian: entry $(i, j)$ is $\partial p_i/\partial s_j$
$\operatorname{diag}(\mathbf p)$: the diagonal matrix with the weights on its diagonal
$\mathbf p$: the row of weights
$\mathbf p\mathbf p^\top$: the outer product of the weights: entry $(i, j)$ is $p_ip_j$
$J^\top$: $J$ transposed
$\succeq$: "is positive semi-definite": no direction has negative curvature
$\mathbf 0$: the zero vector
$\lVert J\rVert_2$: the largest stretch $J$ applies to any vector (its top eigenvalue here)

Derivation.

Use the log-derivative from Card 2 (the O move) and the identity $\partial p = p\,\partial\log p$ .

\frac{\partial p_i}{\partial s_j} = p_i\,\frac{\partial\log p_i}{\partial s_j} = p_i(\delta_{ij} - p_j)

$\frac{\partial\log p_i}{\partial s_j}$: how fast $\log p_i$ changes as you nudge $s_j$

In matrix form the diagonal is $p_i(1 - p_i)$ and the off-diagonal entries are $-p_ip_j$ , so $J$ is symmetric. Its rows sum to 0 because the weights sum to 1: raising every score by the same amount moves nothing (Card 2 again).

J\mathbf 1 = \mathbf p - \mathbf p\,(\mathbf p^\top\mathbf 1) = \mathbf p - \mathbf p = \mathbf 0

$\mathbf p^\top\mathbf 1$: the sum of the weights, which is 1

Positive semi-definite: the quadratic form is a variance, the spread of $\mathbf u$ 's entries when key $i$ is drawn with probability $p_i$ .

\mathbf u^\top J\mathbf u = \sum_i p_iu_i^2 - \Big(\sum_i p_iu_i\Big)^2 = \text{Var}_{\mathbf p}(\mathbf u) \;\ge\; 0

$\mathbf u$: any vector with one entry per key (a direction to nudge the scores)
$\mathbf u^\top J\mathbf u$: the quadratic form of $J$ along $\mathbf u$
$\sum_i$: add up over the keys $i$
$u_i$: entry $i$ of $\mathbf u$
$\text{Var}_{\mathbf p}(\mathbf u)$: the variance of $\mathbf u$ 's entries under the weights $\mathbf p$

Rank $T - 1$ . When every $p_i > 0$ , the variance is zero only if all entries of $\mathbf u$ are equal. For a PSD matrix, $J\mathbf u = \mathbf 0$ exactly when $\mathbf u^\top J\mathbf u = 0$ , so the null space is the single direction $\mathbf 1$ .

J\mathbf u = \mathbf 0 \iff \mathbf u^\top J\mathbf u = 0 \iff u_1 = \cdots = u_T \iff \mathbf u \propto \mathbf 1

$u_1$: the first entry of $\mathbf u$
$u_T$: the last entry of $\mathbf u$

The 1/2 bound, by Gershgorin's disc theorem: every eigenvalue lies within some row's diagonal entry plus the absolute sum of its off-diagonal entries. Row $i$ has diagonal $p_i(1 - p_i)$ and off-diagonal sum $p_i\sum_{j \ne i}p_j = p_i(1 - p_i)$ .

\lambda \;\le\; \max_i\big[p_i(1 - p_i) + p_i(1 - p_i)\big] = 2\max_i p_i(1 - p_i) \;\le\; 2 \cdot \tfrac14 = \tfrac12

$\lambda$: any eigenvalue of $J$

Equality holds at $\mathbf p = (\tfrac12, \tfrac12)$ , where the eigenvalue is exactly 1/2; the largest value over 20,000 random rows was 0.500.

Saturation. As $\mathbf p$ approaches one-hot, every $p_i(1 - p_i)$ goes to 0, so by step 5 the whole Jacobian goes to 0: no score change can move the weights. For two keys the nonzero eigenvalue is $2p(1 - p)$ .

\mathbf p = (0.9,\ 0.1):\;\; J = \begin{pmatrix}0.09 & -0.09\\ -0.09 & 0.09\end{pmatrix},\;\; \lambda = 0,\ 0.18; \qquad \mathbf p = (0.99,\ 0.01):\;\; \lambda_{\max} = 0.0198

$\lambda_{\max}$: the largest eigenvalue of $J$

You never form $J$ . Backprop only needs $J$ times a vector, and that costs $O(T)$ , not $O(T^{2})$ .

J\mathbf y = \mathbf p \odot \mathbf y - \mathbf p\,(\mathbf p^\top\mathbf y) = \mathbf p \odot \big(\mathbf y - (\mathbf p^\top\mathbf y)\,\mathbf 1\big)

$\mathbf y$: the incoming gradient with respect to the weights, one entry per key
$\odot$: entry-by-entry product
$\mathbf p^\top\mathbf y$: the weighted average of $\mathbf y$ : one number

Check on bank's row, $\mathbf p = (0.140, 0.576, 0.284)$ . The eigenvalues are 0.393, 0.175 and 0, and the Gershgorin bound is $2 \times 0.576 \times 0.424 = 0.488$ . A finite-difference Jacobian on a random 7-key row matches to $3.8 \times 10^{-11}$ .

J_{\text{bank}} = \begin{pmatrix}0.120 & -0.081 & -0.040\\ -0.081 & 0.244 & -0.164\\ -0.040 & -0.164 & 0.203\end{pmatrix}, \qquad \lambda = 0.393,\ 0.175,\ 0

$J_{\text{bank}}$: the Jacobian of bank's causal row (order: the, river, bank)

Key move. O (one-hot minus p).

Common slip. Treating softmax as elementwise, like a sigmoid, and keeping only the diagonal $p_i(1 - p_i)$ . The off-diagonal terms are what make the rows sum to 0. Second slip: quoting the bound as 1/4, which is the sigmoid's largest slope; softmax's bound is 1/2.

How it's asked. "Derive $\partial\mathbf p/\partial\mathbf s$ for one attention row and say when it vanishes." (Q15, Drill 3.) Strong answer: $\operatorname{diag}(\mathbf p) - \mathbf p\mathbf p^\top$ , the variance reading of the quadratic form, the 1/2 bound, saturation sending it to 0, and the $O(T)$ vector product of step 7 that Card 5 uses.

Card 5 · Backpropagation through attention ★★

The story. Dao et al. (2022) needed a backward pass that never stores the $T \times T$ weights. It works because the softmax part of the gradient needs only one number per row, $D_i = dO_i \cdot \mathbf o_i$ , which is derived here.

Statement. With $S = QK^\top/\sqrt{d_k} + M$ , $P = \operatorname{softmax}(S)$ and $O = PV$ , the gradients are as below. For one row with upstream gradient $\mathbf g$ , a key's score gradient is its weight times how much its value beats the row's average output.

dV = P^\top dO, \quad dP = dO\,V^\top, \quad dS = P \odot \big(dP - \operatorname{rowsum}(P \odot dP)\big), \quad dQ = \frac{dS\,K}{\sqrt{d_k}}, \quad dK = \frac{dS^\top Q}{\sqrt{d_k}}; \qquad \frac{\partial\mathcal L}{\partial s_j} = p_j\big(\mathbf g^\top\mathbf v_j - \mathbf g^\top\mathbf o\big)

$dV$: the gradient of the loss with respect to $V$ (same shape as $V$ )
$P^\top$: the weights, transposed
$P$: the $T \times T$ attention weights
$dO$: the upstream gradient arriving at the outputs, $\partial\mathcal L/\partial O$
$dP$: the gradient with respect to the weights
$V^\top$: the values, transposed
$dS$: the gradient with respect to the scores
$\odot$: entry-by-entry product
$dQ$: the gradient with respect to the queries
$dK$: the gradient with respect to the keys
$dS^\top$: $dS$ transposed
$\frac{\partial\mathcal L}{\partial s_j}$: how fast the loss changes as you nudge score $s_j$ in one row
$\mathcal L$: the training loss
$p_j$: this row's weight on key $j$
$\mathbf g$: the upstream gradient at this row's output
$\mathbf o$: this row's output

Derivation.

Write the forward pass with every intermediate named.

S = \frac{QK^\top}{\sqrt{d_k}} + M, \qquad P = \operatorname{softmax}_{\text{row}}(S), \qquad O = PV, \qquad dO = \frac{\partial\mathcal L}{\partial O}

$S$: the scaled, masked scores
$O$: the outputs
$\frac{\partial\mathcal L}{\partial O}$: the matrix of slopes of the loss with respect to each output entry

$O = PV$ is a matrix product. For $Y = AB$ the rule is $dA = dY\,B^\top$ and $dB = A^\top dY$ (check the shapes: the S move).

dV = P^\top dO \;\;(T \times T)(T \times d_v), \qquad dP = dO\,V^\top \;\;(T \times d_v)(d_v \times T)

The softmax acts row by row, so row $i$ of $dS$ is $J_i$ times row $i$ of $dP$ , which by Card 4 step 7 is $\mathbf p_i \odot (dP_i - \mathbf p_i^\top dP_i)$ . Stacked, that is one line.

dS = P \odot \big(dP - \operatorname{rowsum}(P \odot dP)\big)

The row sum simplifies. Since $dP_{ij} = dO_i \cdot \mathbf v_j$ , the weighted sum over $j$ is $dO_i$ dotted with the row's output. So the backward pass needs one number per row, computed from $O$ and $dO$ ; FlashAttention calls it $D_i$ .

\sum_j p_{ij}\,dP_{ij} = \sum_j p_{ij}\,dO_i^\top\mathbf v_j = dO_i^\top\mathbf o_i \;\Rightarrow\; \operatorname{rowsum}(P \odot dP) = \operatorname{rowsum}(dO \odot O)

$p_{ij}$: the weight query $i$ puts on key $j$
$dP_{ij}$: the gradient with respect to that weight
$dO_i$: row $i$ of $dO$ : the upstream gradient at output $i$
$\mathbf v_j$: value $j$
$\mathbf o_i$: output $i$
$i$: the query row
$j$: a key

The scores are bilinear in $Q$ and $K$ , and the mask is a constant. Then chain once more through $Q = XW_Q$ (and likewise $K$ , $V$ ) to reach the weights.

dQ = \frac{dS\,K}{\sqrt{d_k}}, \qquad dK = \frac{dS^\top Q}{\sqrt{d_k}}, \qquad dW_Q = X^\top dQ, \quad dW_K = X^\top dK, \quad dW_V = X^\top dV

$dW_Q$: the gradient reaching the query map $W_Q$
$dW_K$: the gradient reaching the key map
$dW_V$: the gradient reaching the value map
$X^\top$: the layer's input, transposed

The per-row form, from the Jacobian's symmetry: the loss sees weight $p_i$ through $\mathbf g^\top\mathbf v_i$ .

\frac{\partial\mathcal L}{\partial s_j} = \sum_i \mathbf g^\top\mathbf v_i\;p_i(\delta_{ij} - p_j) = p_j\Big(\mathbf g^\top\mathbf v_j - \sum_i p_i\,\mathbf g^\top\mathbf v_i\Big) = p_j\big(\mathbf g^\top\mathbf v_j - \mathbf g^\top\mathbf o\big)

$\mathbf v_i$: value $i$ in this row
$p_i$: the row's weight on key $i$
$\delta_{ij}$: 1 if $i = j$ , else 0
$i$: here: a key in the row (summed over)

The score gradients sum to 0 across the row, and a key's entry is large only if it has weight and its value differs from the current average. That is the hook: push values by weight, pull scores by surprise.

Check on the toy, with the chapter's convention (Section 8): the loss falls as bank reads more geography, so $\mathbf g = \partial\mathcal L/\partial\mathbf o_{\text{bank}} = (-1, 0)$ . Then $\mathbf g^\top\mathbf v = (0, -1, 0)$ for the, river, bank, and $\mathbf g^\top\mathbf o = -0.576$ .

\frac{\partial\mathcal L}{\partial\mathbf s} = \big(0.140\,(0 + 0.576),\;\; 0.576\,(-1 + 0.576),\;\; 0.284\,(0 + 0.576)\big) = (0.081,\ -0.244,\ 0.164)

$\frac{\partial\mathcal L}{\partial\mathbf s}$: the gradient with respect to bank's three scores

Into the query and the values. The query gradient sums the score gradients times the keys over $\sqrt 2$ ; each value's gradient is its weight times $\mathbf g$ .

\frac{\partial\mathcal L}{\partial\mathbf q_{\text{bank}}} = \frac{0.081\,(0, 0) - 0.244\,(1, 0) + 0.164\,(1, 1)}{\sqrt 2} = (-0.057,\ 0.116), \qquad \frac{\partial\mathcal L}{\partial\mathbf v_j} = p_j\,\mathbf g

$\frac{\partial\mathcal L}{\partial\mathbf q_{\text{bank}}}$: the gradient with respect to bank's query
$\frac{\partial\mathcal L}{\partial\mathbf v_j}$: the gradient with respect to value $j$

A finite-difference check of the query gradient agrees to $2.4 \times 10^{-12}$ . Gradient descent steps against these. River's score gradient is negative, so the step raises river's score and lowers the other two; the query moves along $(0.057, -0.116)$ , which raises its score on river's key $(1, 0)$ and lowers it on its own key $(1, 1)$ , so bank reads more geography, as this loss wants. (With $\mathbf g = (1, 0)$ , a loss that grows with geography, every sign flips and descent turns the query away from river.) On a random causal $T = 6$ problem, the matrix formulas match PyTorch autograd in float64 to $1.1 \times 10^{-16}$ and finite differences to $4.4 \times 10^{-10}$ .

Key move. O (one-hot minus p), with S to get every transpose right.

Common slip. Dropping the $-\operatorname{rowsum}$ term (treating softmax as elementwise). Also: forgetting the $1/\sqrt{d_k}$ in $dQ$ and $dK$ , writing $dV = dO\,P^\top$ (wrong shape unless $d_v = T$ ), and reading $\partial\mathcal L/\partial s_j > 0$ as "this key should gain weight". Descent moves the score the other way.

How it's asked. "Derive $\partial\mathcal L/\partial S$ for one attention row and backpropagate to $Q$ , $K$ and $V$ ." (Q15, Drill 4.) Strong answer: the five matrix lines, the $D_i = dO_i \cdot \mathbf o_i$ simplification, and the consequence: the backward pass needs only $Q, K, V, O, dO$ and one logsumexp per row, so $P$ can be recomputed instead of stored (Card 12).

Card 6 · Causal mask correctness ★

The story. Vaswani et al. (2017) trained the decoder on whole sentences in one pass and had to "prevent leftward information flow". They did it inside the softmax, by setting the illegal scores to $-\infty$ .

Statement. Adding $-\infty$ to every score with $j > i$ makes those weights exactly 0, so output $i$ depends on tokens 1 to $i$ only. The same fact makes the KV cache valid. A fully masked row is 0/0, and zeroing weights after the softmax both breaks the sum-to-1 and leaks the future through the normalizer.

M_{ij} = -\infty \;\;(j > i) \;\Rightarrow\; p_{ij} = 0 \;\;(j > i) \;\Rightarrow\; \mathbf o_i = f(\mathbf x_1, \ldots, \mathbf x_i)

$M_{ij}$: the mask entry for query $i$ and key $j$
$f$: some function: the point is which inputs it may depend on
$\mathbf x_1$: the first token's input vector
$\mathbf x_i$: token $i$ 's input vector

Derivation.

Since $e^{-\infty} = 0$ , blocked keys vanish from the numerator and from the normalizer.

p_{ij} = \frac{e^{s_{ij}}\,\mathbb 1[j \le i]}{\sum_{r \le i} e^{s_{ir}}}, \qquad \mathbf o_i = \sum_{j \le i} p_{ij}\,\mathbf v_j

$\mathbb 1[j \le i]$: 1 if key $j$ is at or before query $i$ , else 0
$\sum_{r \le i}$: add up over the visible keys $r = 1, \ldots, i$
$\sum_{j \le i}$: add up over the visible keys $j$

Everything left in $\mathbf o_i$ is built from $\mathbf x_1, \ldots, \mathbf x_i$ . In a stack, induct over layers: if every layer- $\ell$ vector at positions up to $i$ depends only on tokens up to $i$ , so does layer $\ell + 1$ , because the residual and the feed-forward network act on each position alone.

\mathbf o_i = \sum_{j \le i}\operatorname{softmax}_{j \le i}\Big(\frac{(W_Q^\top\mathbf x_i)^\top(W_K^\top\mathbf x_j)}{\sqrt{d_k}}\Big)\,W_V^\top\mathbf x_j

$W_Q^\top\mathbf x_i$: token $i$ 's query, written from its input
$W_K^\top\mathbf x_j$: token $j$ 's key
$W_V^\top\mathbf x_j$: token $j$ 's value
$\mathbf x_j$: token $j$ 's input vector

The test that catches every causal-mask bug: change the future and the past must not move. In the toy, replacing steep with (0, 3, 0, 1) leaves the causal outputs of the, river, bank and was unchanged to the last bit. Without the mask, bank's output moves by 0.217.

\max_{i \le 4}\lVert\mathbf o_i' - \mathbf o_i\rVert_\infty = 0 \;\;(\text{causal}), \qquad \lVert\mathbf o_{\text{bank}}' - \mathbf o_{\text{bank}}\rVert_\infty = 0.217 \;\;(\text{no mask})

$\lVert\mathbf o_i' - \mathbf o_i\rVert_\infty$: the largest change in any coordinate of output $i$ (‖·‖∞ = largest absolute entry)
$\mathbf o_i'$: output $i$ after the last token is changed
$\mathbf o_{\text{bank}}'$: bank's output after the change

The same fact proves the KV cache valid. Keys and values at positions up to $t$ never change when later tokens arrive, so decoding one token at a time only needs the new token's query, key and value. In the toy, the cached decode matches the full causal pass exactly.

\mathbf o_t = \sum_{j \le t} p_{tj}\,\mathbf v_j \;\text{ uses only }\; \{\mathbf k_j, \mathbf v_j\}_{j \le t} \;\Rightarrow\; \text{cache them; compute only } \mathbf q_t, \mathbf k_t, \mathbf v_t

$t$: the decode step: the position of the newest token
$\mathbf o_t$: the newest token's output
$p_{tj}$: the newest token's weight on key $j$
$\mathbf k_j$: cached key $j$
$\mathbf q_t$: the newest token's query
$\mathbf k_t$: the newest token's key
$\mathbf v_t$: the newest token's value

A fully masked row has no allowed key. With $-\infty$ a hand-written softmax is 0/0, which is NaN (and so is nn.MultiheadAttention with its default need_weights=True; PyTorch 2.14.1's F.scaled_dot_product_attention on CPU returned zeros instead, and this has varied by version and backend). An off-by-one triu(ones, 0) masks the diagonal, so the first row is fully masked and turns NaN; filling with $-10^{9}$ hides the bug by returning a uniform row instead.

p_{1j} = \frac{e^{-\infty}}{\sum_r e^{-\infty}} = \frac00 = \text{NaN}, \qquad \text{fill } -10^9:\;\; p_{1j} = \frac{e^{-10^9}}{5\,e^{-10^9}} = 0.2

$p_{1j}$: the first token's weight on key $j$

Masking after the softmax. Zero the future weights of the unmasked softmax: the kept weights no longer sum to 1 (bank's sum to 0.583), and the denominator still contains the future scores. Changing steep moves bank's output by 0.162: the future leaks.

\tilde p_{ij} = \frac{e^{s_{ij}}\,\mathbb 1[j \le i]}{\sum_{r=1}^{T} e^{s_{ir}}}: \qquad \sum_j\tilde p_{\text{bank},j} = 0.583, \qquad \text{denominator still contains } e^{s_{\text{bank,steep}}}

$\tilde p_{ij}$: the weight after zeroing post-softmax (tilde = the buggy version)
$\tilde p_{\text{bank},j}$: bank's buggy weight on key $j$
$s_{\text{bank,steep}}$: bank's score for the future token steep

Dividing by the new row sum afterwards does give the right answer (the full normalizer cancels), but it computes and throws away every future exponential.

Count what the mask allows: $T(T + 1)/2$ of the $T^{2}$ pairs. Causal attention needs about half the work, but only if the kernel actually skips masked blocks.

\frac{T(T + 1)/2}{T^2} = \frac{T + 1}{2T}:\qquad \frac{15}{25} = 0.6 \;\;(T = 5), \qquad 0.500122 \;\;(T = 4{,}096)

Key move. F: $e^{-\infty} = 0$ inside the normalizer.

Common slip. Masking the weights instead of the scores (step 6); triu with k=0 instead of k=1; filling fully masked rows with -inf (NaN) when padding can blank a whole row. The safe fill is the dtype's most negative finite number, plus zeroing those rows' outputs.

How it's asked. "A new language model reaches a training loss of 0.05 in an hour and generates garbage. What's your first test?" (Q6; also Q17, Q18, Drill 9.) Strong answer: the change-the-future unit test, mask the scores not the weights, the fully masked row, and "the same property is why the KV cache is exact".

Card 7 · Permutation equivariance without positions ★★

The story. Vaswani et al. (2017) noted that their model "contains no recurrence and no convolution", so to use word order they had to "inject some information about the relative or absolute position of the tokens". This card proves why.

Statement. Without a mask or positional information, reordering the input tokens reorders the outputs the same way, and pooling over tokens erases the order completely. A causal mask breaks the symmetry, and even with identical tokens the first token's share at position $t$ is $1/t$ .

\text{Attn}(\Pi X) = \Pi\,\text{Attn}(X), \qquad \frac1T\mathbf 1^\top\text{Attn}(\Pi X) = \frac1T\mathbf 1^\top\text{Attn}(X), \qquad \text{causal, equal scores: } p_{t,\text{BOS}} = \frac1t

$\Pi$: a permutation matrix: each row has a single 1, so $\Pi X$ reorders the rows of $X$
$X$: the input token vectors, one row per token
$\frac1T\mathbf 1^\top$: average over the $T$ tokens (mean pooling)
$\mathbf 1^\top$: the all-ones row vector: sums over tokens
$p_{t,\text{BOS}}$: the weight position $t$ puts on the first (beginning-of-sequence) token

Derivation.

Permuting the rows of $X$ permutes the rows of $Q$ , $K$ and $V$ .

(\Pi X)W_Q = \Pi Q, \qquad (\Pi X)W_K = \Pi K, \qquad (\Pi X)W_V = \Pi V

$W_Q$: the query map
$W_K$: the key map
$W_V$: the value map

The scores are permuted on both sides: rows (queries) and columns (keys).

(\Pi Q)(\Pi K)^\top = \Pi\,QK^\top\,\Pi^\top

$\Pi^\top$: the transpose of $\Pi$ , which is also its inverse

A row-wise softmax doesn't care about order. Reordering rows reorders the output rows, and reordering the entries inside a row reorders that row's weights.

\operatorname{softmax}_{\text{row}}(\Pi S\,\Pi^\top) = \Pi\,\operatorname{softmax}_{\text{row}}(S)\,\Pi^\top

Multiply by the permuted values and use $\Pi^\top\Pi = I$ .

\Pi P\,\Pi^\top\,\Pi V = \Pi P\,(\Pi^\top\Pi)\,V = \Pi P V = \Pi O, \qquad \Pi^\top\Pi = I

$I$: the identity matrix: undoing a reordering leaves the tokens where they were

Pooling: summing over tokens doesn't care about their order, because $\mathbf 1^\top\Pi = \mathbf 1^\top$ . So a pooled representation is invariant, while the per-token outputs are equivariant.

\mathbf 1^\top\,\Pi O = (\Pi^\top\mathbf 1)^\top O = \mathbf 1^\top O

$\mathbf 1$: the all-ones column vector

The causal mask is not permutation-invariant, so step 3 fails once it is added. Reversing the toy sentence permutes the unmasked outputs to $1.1 \times 10^{-16}$ but changes the causal ones by up to 0.426.

\Pi(S + M)\Pi^\top = \Pi S\,\Pi^\top + \Pi M\,\Pi^\top \;\ne\; \Pi S\,\Pi^\top + M

Position leaks through the mask. Give the first token value 1, every other token value 0, and make all scores equal. Position $t$ sees $t$ keys and spreads its weight evenly, so its output is $1/t$ : a position signal from nothing. Without the mask every position gets $1/T$ , and there is no signal.

\mathbf o_t = \frac1t\sum_{j=1}^{t}\mathbf v_j = \frac1t \cdot 1 + \frac{t - 1}{t} \cdot 0 = \frac1t

$t$: a position in the sequence (1 = first)
$\mathbf o_t$: the output at position $t$
$\sum_{j=1}^{t}$: add up over the $t$ visible tokens

Key move. S: track where each permutation lands.

Common slip. Saying "permutation-invariant". The outputs move with the tokens (equivariant); only a pooled summary is invariant. Second slip: "a decoder without positional encodings has no idea of order". The causal mask leaks it (Haviv et al., 2022), but weakly, which is why real models add positions.

How it's asked. "An engineer shuffles the input tokens of an encoder and finds the outputs are just shuffled. Bug?" (Q9, Drill 12.) Strong answer: equivariant by construction, so "dog bites man" and "man bites dog" get the same vectors rearranged; pooling makes them identical; the four ways real models add order (Section 9).

Card 8 · Attention as Nadaraya–Watson kernel smoothing ★★★

The story. In 1964 Nadaraya and Watson independently proposed estimating a curve by averaging nearby observations with kernel weights. Tsai et al. (2019) showed that Transformer attention can be read as the same estimator with a learned kernel.

Statement. A Gaussian-kernel Nadaraya–Watson smoother is exactly dot-product attention at temperature $\lambda^{2}$ plus a per-key bias $-\lVert\mathbf k_j\rVert^{2}/2\lambda^{2}$ . When all keys have the same length, the bias drops out and the two are identical.

\hat f(\mathbf q) = \frac{\sum_j \kappa_\lambda(\mathbf q, \mathbf k_j)\,\mathbf v_j}{\sum_j \kappa_\lambda(\mathbf q, \mathbf k_j)}, \quad \kappa_\lambda(\mathbf q, \mathbf k_j) = e^{-\lVert\mathbf q - \mathbf k_j\rVert^2/2\lambda^2} \;\Rightarrow\; \hat f(\mathbf q) = \sum_j\operatorname{softmax}_j\Big(\frac{\mathbf q^\top\mathbf k_j}{\lambda^2} - \frac{\lVert\mathbf k_j\rVert^2}{2\lambda^2}\Big)\mathbf v_j

$\hat f(\mathbf q)$: the smoothed estimate at query point $\mathbf q$ (hat = estimated from data)
$\hat f$: the smoothed estimate
$\kappa_\lambda(\mathbf q, \mathbf k_j)$: the Gaussian kernel: similarity between $\mathbf q$ and data point $\mathbf k_j$
$\mathbf q$: the point where we want an estimate (attention: the query)
$\lVert\mathbf q - \mathbf k_j\rVert^2$: the squared distance from the query to key $j$
$\lambda$: the kernel bandwidth: how far "nearby" reaches (usually written $h$ ; here $h$ counts heads)
$\lVert\mathbf k_j\rVert^2$: key $j$ 's squared length

Derivation.

Nadaraya–Watson weights are positive and sum to 1. Any such weights are a softmax of their logarithms.

\frac{\kappa_\lambda(\mathbf q, \mathbf k_j)}{\sum_r\kappa_\lambda(\mathbf q, \mathbf k_r)} = \operatorname{softmax}_j\big(\log\kappa_\lambda(\mathbf q, \mathbf k_j)\big)

$\kappa_\lambda(\mathbf q, \mathbf k_r)$: the kernel weight of key $r$
$\mathbf k_r$: key $r$

Expand the squared distance in the Gaussian log-kernel.

-\frac{\lVert\mathbf q - \mathbf k_j\rVert^2}{2\lambda^2} = \frac{\mathbf q^\top\mathbf k_j}{\lambda^2} - \frac{\lVert\mathbf q\rVert^2}{2\lambda^2} - \frac{\lVert\mathbf k_j\rVert^2}{2\lambda^2}

$\lVert\mathbf q\rVert^2$: the query's squared length

The query term is the same for every key, so it cancels in the softmax (F). What remains is dot-product attention at temperature $\lambda^{2}$ , plus a bias that penalizes long keys.

\operatorname{softmax}_j\Big(\frac{\mathbf q^\top\mathbf k_j}{\lambda^2} - \frac{\lVert\mathbf q\rVert^2}{2\lambda^2} - \frac{\lVert\mathbf k_j\rVert^2}{2\lambda^2}\Big) = \operatorname{softmax}_j\Big(\frac{\mathbf q^\top\mathbf k_j}{\lambda^2} - \frac{\lVert\mathbf k_j\rVert^2}{2\lambda^2}\Big)

Equal key lengths make the bias a constant too. Then Gaussian kernel smoothing is plain dot-product attention, and the $1/\sqrt{d_k}$ scale is a bandwidth.

\lVert\mathbf k_j\rVert = \text{const} \;\Rightarrow\; \hat f(\mathbf q) = \sum_j\operatorname{softmax}_j\Big(\frac{\mathbf q^\top\mathbf k_j}{\lambda^2}\Big)\mathbf v_j, \qquad \frac{1}{\lambda^2} = \frac{1}{\sqrt{d_k}} \iff \lambda = d_k^{1/4}

Check. On 40 noisy points of a sine curve with bandwidth 0.5, Nadaraya–Watson and biased attention agree to $5.6 \times 10^{-16}$ ; dropping the bias is off by up to 1.09. With equal-length keys, plain attention matches to $6.9 \times 10^{-17}$ , and a random three-dimensional case matches to $3.3 \times 10^{-16}$ .

\max\lvert\hat f_{\text{NW}} - \hat f_{\text{attn + bias}}\rvert = 5.6 \times 10^{-16}, \qquad \max\lvert\hat f_{\text{NW}} - \hat f_{\text{attn, no bias}}\rvert = 1.09

$\hat f_{\text{NW}}$: the Nadaraya–Watson estimate
$\hat f_{\text{attn + bias}}$: dot-product attention with the key-length bias
$\hat f_{\text{attn, no bias}}$: plain dot-product attention

Consequences of reading attention this way. The output stays inside the convex hull of the values (Card 1). A weighted average cannot extrapolate. With plain dot-product scores, a zero query gives equal scores and therefore the plain mean: the has $\mathbf q = \mathbf 0$ and its unmasked output is the mean value (0.4, 0.2).

\mathbf q = \mathbf 0:\;\; \operatorname{softmax}(\mathbf 0) = \tfrac1T\mathbf 1 \;\Rightarrow\; \mathbf o = \frac1T\sum_j\mathbf v_j = (0.4,\ 0.2)

$\mathbf 0$: the zero vector
$T$: the number of keys: 5
$\mathbf o$: the output of the

What attention adds to the 1964 estimator: the kernel is learned ( $W_Q$ and $W_K$ ), anisotropic ( $W_QW_K^\top$ is not the identity), and asymmetric, and the values are learned too.

Key move. F: the query's own length cancels in the softmax.

Common slip. "Attention is a Gaussian kernel": not without the key-length bias, and not with $\mathbf q \ne \mathbf k$ maps. Second slip: keeping the $\lVert\mathbf q\rVert^{2}$ term. It cancels.

How it's asked. "A statistician on your team says attention is 'just Nadaraya–Watson'. Defend or refute, precisely." (Q24.) Strong answer: the expansion in three lines, the key-length bias, equal-length keys (which QK-norm enforces) making it exact with bandwidth $1/\sqrt g$ (Card 18), and what is learned that the 1964 estimator fixed by hand.

Card 9 · Multi-head parameter count ★

The story. Vaswani et al. (2017) split their width of 512 into 8 heads of 64 and noted that, because each head is smaller, the total cost is similar to single-head attention at full width. Finance partners still ask why 32 heads don't cost 32 times as much.

Statement. With $h$ heads of width $d_k = d_v = d/h$ , multi-head attention has $4d^{2}$ weights, whatever $h$ is. With grouped-query attention it has $2d^{2} + 2d\,H_{\text{kv}}d_h$ .

\#\text{params} = \underbrace{3h\,d\,d_k}_{W_Q,\,W_K,\,W_V} + \underbrace{h\,d_v\,d}_{W_O} = 4d^2 \;\;(h\,d_k = h\,d_v = d), \qquad \text{GQA: } 2d^2 + 2d\,H_{\text{kv}}d_h

$h$: the number of heads
$d$: the model width $d_{\text{model}}$ : 4 in the toy, 4,096 in Llama 2 7B
$d_k$: each head's query and key width
$W_Q$: all heads' query maps, side by side ( $d \times h\,d_k$ )
$W_K$: all heads' key maps
$W_V$: all heads' value maps
$d_v$: each head's value width
$W_O$: the output projection ( $h\,d_v \times d$ )
$H_{\text{kv}}$: the number of key/value heads
$d_h$: the head width, $d_k = d_v$

Derivation.

Each head $n$ has a query, key and value map from width $d$ .

W_Q^{(n)}, W_K^{(n)} \in \mathbb R^{d \times d_k}, \quad W_V^{(n)} \in \mathbb R^{d \times d_v} \;\Rightarrow\; h\,(2d\,d_k + d\,d_v) \;\text{parameters}

$W_Q^{(n)}$: head $n$ 's query map
$W_K^{(n)}$: head $n$ 's key map
$W_V^{(n)}$: head $n$ 's value map
$n$: which head, 1 to $h$
$\mathbb R$: the real numbers; $\mathbb R^{d \times d_k}$ = real matrices of that shape

The output projection maps the concatenated heads, width $h\,d_v$ , back to $d$ .

W_O \in \mathbb R^{h d_v \times d} \;\Rightarrow\; h\,d_v\,d \;\text{parameters}

With $h\,d_k = h\,d_v = d$ , the head count disappears. The toy has $4 \cdot 4^{2} = 64$ , and so does one big head with $d_k = 4$ . Llama 2 7B has 67,108,864 per layer, whether it used 1 head or 64.

3d\,(h\,d_k) + (h\,d_v)\,d = 3d^2 + d^2 = 4d^2: \qquad 4 \cdot 4^2 = 64, \qquad 4 \cdot 4{,}096^2 = 67{,}108{,}864

Biases add $3h\,d_k + d = 4d$ . GPT-2 small, $d = 768$ , has biases. The toy with biases has $64 + 16 = 80$ .

4d^2 + 4d = 4 \cdot 768^2 + 4 \cdot 768 = 2{,}359{,}296 + 3{,}072 = 2{,}362{,}368

Grouped-query attention keeps $h$ query heads but only $H_{\text{kv}}$ key and value heads, so the $K$ and $V$ maps shrink to $d \times H_{\text{kv}}d_h$ . Llama 2 70B ( $d = 8{,}192$ , 64 query heads, 8 KV heads of 128):

2d^2 + 2d\,H_{\text{kv}}d_h = 2 \cdot 8{,}192^2 + 2 \cdot 8{,}192 \cdot 8 \cdot 128 = 134{,}217{,}728 + 16{,}777{,}216 = 150{,}994{,}944

Full MHA would be $4d^{2}$ = 268,435,456 (GQA saves 43.75%), and MQA ( $H_{\text{kv}} = 1$ ) 136,314,880.

Why $W_O$ is not decoration: multiplying the concatenation by $W_O$ is the same as each head multiplying its own block of rows, then adding. Heads write additively into the residual stream.

\big[\text{head}_1, \ldots, \text{head}_h\big]\,W_O = \sum_{n=1}^{h}\text{head}_n\,W_O^{(n)}

$\sum_{n=1}^{h}$: add up over the heads
$W_O^{(n)}$: the $d_v$ rows of $W_O$ that read head $n$

Key move. S (shapes tally).

Common slip. $4h\,d^{2}$ (giving every head full width), forgetting $W_O$ , or saying more heads cost more parameters.

How it's asked. "A finance partner budgeting a model asks why going from 8 to 32 heads doesn't add parameters. And how many are in GPT-2 small's attention per layer?" (Q8, Q19.) Strong answer: $4d^{2}$ in two lines; 2,362,368 with biases; then what heads do cost: each score matrix has rank at most $d_k$ (Card 10), and the KV cache grows with the number of KV heads (Card 13).

Card 10 · The rank bound and the d_k = 1 head ★★★

The story. Bhojanapalli et al. (2020) argued that shrinking each head to $d/h$ creates a "low-rank bottleneck" in what a head can express, and proposed setting the head size independently of the number of heads.

Statement. The score matrix has rank at most $d_k$ . At $d_k = 1$ , every query's top key is either the key with the largest value or the one with the smallest, so a one-dimensional head can sharply select at most two keys.

\operatorname{rank}(QK^\top) \le d_k, \qquad d_k = 1:\;\; \arg\max_j s_{ij} \in \Big\{\arg\max_j k_j,\ \arg\min_j k_j\Big\}

$s_{ij}$: the score of key $j$ for query $i$
$k_j$: key $j$ , a single number when $d_k = 1$

Derivation.

$QK^\top$ is a sum of $d_k$ outer products, one per component, and each has rank 1.

QK^\top = \sum_{c=1}^{d_k} Q_{\cdot c}\,K_{\cdot c}^\top \;\Rightarrow\; \operatorname{rank}(QK^\top) \le d_k

$c$: here: a component index, 1 to $d_k$
$Q_{\cdot c}$: column $c$ of $Q$ : component $c$ of every query
$K_{\cdot c}^\top$: column $c$ of $K$ , laid on its side

The bound is on the logits, not the weights. The exponential is nonlinear, so $P$ can be full rank; but $\log P$ is the low-rank $S$ minus one logsumexp per row, so its rank is at most $d_k + 1$ . In the simulation with $T = 64$ and $d_k = 8$ , the score matrix has rank 8 and the weight matrix rank 64.

\log P = S - \boldsymbol\ell\,\mathbf 1^\top \;\Rightarrow\; \operatorname{rank}(\log P) \le d_k + 1, \qquad \text{while } \operatorname{rank}(P) \text{ can be } T

$\log P$: the log of each weight, entry by entry
$\boldsymbol\ell$: the vector of row logsumexps, one per query
$\mathbf 1^\top$: the all-ones row vector

At $d_k = 1$ each score is a product of two numbers. A positive query is maximized by the largest key, a negative query by the smallest, and a zero query gives a uniform row.

q_i > 0:\;\; \arg\max_j q_ik_j = \arg\max_j k_j, \qquad q_i < 0:\;\; \arg\max_j q_ik_j = \arg\min_j k_j

$q_i$: query $i$ , a single number when $d_k = 1$

So however many queries there are, a $d_k = 1$ head has at most two possible top keys. The simulation confirms it in 100% of 2,000 random instances. With $d_k = 2$ , 12 keys spread on a circle can each be some query's top choice: a query pointing at a key selects it.

\big\lvert\{\arg\max_j s_{ij} : i = 1, \ldots, T\}\big\rvert \le 2 \;\;(d_k = 1), \qquad 12 \text{ of } 12 \text{ keys on a circle } (d_k = 2)

Key move. S: a rank count from the shapes.

Common slip. "The attention weights have rank at most $d_k$ ." The weights can be full rank; the logits are what is low-rank. The opposite slip is to claim a head can express any weight pattern. It cannot: the pattern of logits is constrained.

How it's asked. "A head has $d_k = 1$ . Which keys can ever be its top choice?" (Drills 7–8; Q40 follow-up.) Strong answer: the outer-product argument, logits vs weights, and the link to linear attention, where the weight matrix itself is a product of $T \times r$ feature matrices and so has rank at most $r$ , which is part of why it struggles with sharp retrieval.

Card 11 · Attention FLOPs and the crossover ★★

The story. Vaswani et al. (2017) compared layer types by their cost per layer, $O(n^{2} d)$ for self-attention against $O(n d^{2})$ for recurrence, and argued attention is cheaper whenever the sequence is shorter than the width. That held for 2017 sentences; long contexts reversed it.

Statement. Per layer and forward pass, the projections cost $2Td(2d + 2H_{\text{kv}}d_h)$ , which is $8Td^{2}$ for MHA; the two quadratic products cost $4T^{2}d$ ; and a SwiGLU feed-forward network costs $6Td\,d_{\text{ff}}$ . The quadratic part overtakes everything else at $T^\star = 2d + 1.5\,d_{\text{ff}}$ .

\text{FLOPs per layer} = \underbrace{2Td(2d + 2H_{\text{kv}}d_h)}_{\text{Q, K, V, O projections}} + \underbrace{4T^2d}_{QK^\top\text{ and } PV} + \underbrace{6Td\,d_{\text{ff}}}_{\text{SwiGLU FFN}}, \qquad T^\star = 2d + 1.5\,d_{\text{ff}}

$T$: the context length in tokens
$d_{\text{ff}}$: the feed-forward hidden width (11,008 in Llama 2 7B)
$T^\star$: the crossover context length (star = the point where the two costs meet)

Derivation.

The $2mnk$ rule: each of the $mn$ outputs is a dot product of length $k$ , one multiply and one add per term.

(m \times k)(k \times n):\;\; mn \;\text{outputs} \times k \;\text{multiply-adds} \times 2 = 2mnk

$m$: here: rows of the left matrix
$k$: here: the shared inner dimension
$n$: here: columns of the right matrix

Projections. $Q$ and the output projection are $(T \times d)(d \times d)$ . $K$ and $V$ are $(T \times d)(d \times H_{\text{kv}}d_h)$ .

2Td^2 + 2Td\,(H_{\text{kv}}d_h) + 2Td\,(H_{\text{kv}}d_h) + 2Td^2 = 2Td\,(2d + 2H_{\text{kv}}d_h) = 8Td^2 \;\text{ when } H_{\text{kv}}d_h = d

The quadratic part. Per head, $QK^\top$ is $(T \times d_h)(d_h \times T)$ and $PV$ is $(T \times T)(T \times d_h)$ ; there are $h$ heads with $h\,d_h = d$ . Under a causal mask, only the fraction $(T + 1)/2T$ is needed, and only if the kernel skips the masked blocks. The softmax's own exponentials are also $O(T^{2})$ but are not matrix products, so they are left out of this count.

h\,(2T^2d_h + 2T^2d_h) = 4T^2\,(h\,d_h) = 4T^2d

$h$: the number of query heads

The SwiGLU network has three $d \times d_{\text{ff}}$ matrices (gate, up, down).

3 \times 2Td\,d_{\text{ff}} = 6Td\,d_{\text{ff}}

Llama 2 7B at $T = 4{,}096$ ( $d = 4{,}096$ , $d_{\text{ff}} = 11{,}008$ ). The quadratic part is 14.2% of the layer (7.66% causal). At $T = 512$ it is 2.03%; at 32,768 it is 57.0%.

8Td^2 = 549.8, \qquad 4T^2d = 274.9, \qquad 6Td\,d_{\text{ff}} = 1{,}108.1 \;\;\text{GFLOP}; \qquad \frac{274.9}{549.8 + 274.9 + 1{,}108.1} = 14.2\%

The crossover: set the quadratic part equal to the rest and divide by $4Td$ .

4T^2d = 8Td^2 + 6Td\,d_{\text{ff}} \;\Rightarrow\; T^\star = 2d + 1.5\,d_{\text{ff}} = 8{,}192 + 16{,}512 = 24{,}704

Causal: the quadratic part becomes $2Td(T + 1)$ . Solving gives 49,407, which the large- $T$ shortcut "double $T^\star$ " rounds to 49,408.

2Td\,(T + 1) = 8Td^2 + 6Td\,d_{\text{ff}} \;\Rightarrow\; T = 4d + 3d_{\text{ff}} - 1 = 49{,}407

The same comparison per token (divide by $T$ ): the weights cost a constant per token, the pairs cost $4Td$ , which grows with the context.

\frac{\text{FLOPs}}{T} = 8d^2 + 6d\,d_{\text{ff}} + 4Td = 134\text{M} + 271\text{M} + 67\text{M} \;\;(T = 4{,}096)

Key move. S (shapes tally, $2mnk$ ).

Common slip. Counting a multiply-add as one FLOP; saying "doubling the context quadruples the cost" (only the quadratic part does, and it is 14% at 4K); claiming causal attention halves cost on a kernel that computes the masked blocks anyway.

How it's asked. "An executive slide says doubling the context quadruples the cost. True for Llama 2 7B?" (Q12, Drill 15.) Strong answer: the three terms, the 14.2% share at 4K, the crossover at 24,704, and the hook " $d^{2}$ for the weights, $T^{2}$ for the pairs".

Card 12 · Attention memory: naive vs online ★★

The story. Rabe and Staats (2021) titled their paper "Self-attention Does Not Need $O(n^{2})$ Memory". The $T \times T$ matrix is an intermediate, not an output, and nothing forces anyone to keep it.

Statement. Storing the weights costs $B\,h\,T^{2}\,b$ bytes per layer. Keeping only one fp32 logsumexp per row costs $B\,h\,T \times 4$ bytes, and the backward pass recomputes $P$ from it.

\text{naive} = B\,h\,T^2\,b \;\text{bytes per layer}; \qquad \text{online} = B\,h\,T \times 4 \;\text{bytes per layer, plus recomputing } P \text{ in the backward pass}

$B$: the batch size
$b$: bytes per number: 2 for fp16 or bf16

Derivation.

Elements times bytes, for one head of Llama 2 7B at $T = 4{,}096$ in fp16.

T^2\,b = 4{,}096^2 \times 2 = 33{,}554{,}432 \;\text{bytes} = 32 \;\text{MiB per head}

Times 32 heads per layer, and times 32 layers if every layer's weights are kept for the backward pass.

32 \;\text{MiB} \times 32 \;\text{heads} = 1 \;\text{GiB per layer}, \qquad \times\, 32 \;\text{layers} = 32 \;\text{GiB}

Quadratic growth: Llama 3 8B at 131,072 tokens.

131{,}072^2 \times 2 \;\text{bytes} = 32 \;\text{GiB per head}, \qquad \times\, 32 \;\text{heads} = 1 \;\text{TiB per layer}

The backward pass uses $P$ only inside $dS = P \odot (dP - D)$ (Card 5). $P$ can be rebuilt from $S$ and one logsumexp per row, and $D$ from $O$ and $dO$ .

p_{ij} = e^{s_{ij} - \text{lse}_i}, \qquad \text{lse}_i = \log\sum_r e^{s_{ir}}, \qquad D_i = dO_i^\top\mathbf o_i

$p_{ij}$: weight $(i, j)$ , recomputed in the backward pass
$e$: Euler's number, 2.718…
$s_{ij}$: score $(i, j)$ , recomputed from $Q$ and $K$
$\sum_r$: add up over the keys $r$ in row $i$
$s_{ir}$: the score of key $r$ in row $i$
$D_i$: the per-row number from Card 5
$dO_i$: the upstream gradient at output $i$
$\mathbf o_i$: output $i$

So the stored state is $T$ numbers per head instead of $T^{2}$ .

32 \times 4{,}096 \times 4 \;\text{bytes} = 0.5 \;\text{MiB per layer}, \qquad \times\, 32 \;\text{layers} = 16 \;\text{MiB} \;\;\text{vs}\;\; 32 \;\text{GiB}

For Llama 3 8B at 131,072 tokens it is 16 MiB per layer, against 1 TiB.

The price is compute: recomputing $QK^\top$ in the backward pass adds $2T^{2}d$ FLOPs per layer, half the forward quadratic part. On a GPU that trade wins, because attention is limited by memory traffic, not arithmetic (LLM Ch 5).

\text{extra backward FLOPs} = 2T^2d = \tfrac12 \times 4T^2d

$d$: the model width
$T$: the context length

Key move. S (elements × bytes).

Common slip. Confusing this training-time $O(T^{2})$ score memory with the inference-time $O(T)$ KV cache (Card 13). Second slip: saying FlashAttention reduces FLOPs. It does the same FLOPs plus a recomputation, and wins on memory traffic.

How it's asked. "A team extends Llama 3 8B to 128K tokens and runs out of memory in training even at batch 1. Where did it go?" (Q33, Drill 16.) Strong answer: 32 GiB per head for the scores; the online softmax keeps 16 MiB per layer; then the KV cache (16 GiB at 128K for inference) and activation checkpointing for the rest.

Card 13 · KV-cache size ★

The story. Generation adds one token at a time. Recomputing every past key and value at each step would redo the whole prefix, so decoders cache them, and at long context the cache, not the weights, becomes the memory bill.

Statement. The KV cache stores one key and one value per layer, per KV head, per token.

\text{KV bytes} = \underbrace{2}_{K,\,V} \cdot \underbrace{L}_{\text{layers}} \cdot H_{\text{kv}} \cdot d_h \cdot b \;\times\; T \times B

$L$: the number of layers
$H_{\text{kv}}$: the number of key/value heads per layer
$d_h$: the head width

Derivation.

Per token and layer, each KV head stores a key and a value of width $d_h$ . Queries are never cached: each is used once, at its own step (Card 6).

\text{per token, per layer} = 2\,H_{\text{kv}}\,d_h \;\text{numbers}

Times $L$ layers and $b$ bytes. Llama 2 7B: 32 layers, 32 KV heads, $d_h = 128$ , fp16.

2 \cdot 32 \cdot 32 \cdot 128 \cdot 2 = 524{,}288 \;\text{bytes} = 512 \;\text{KiB per token}

Times the context and the batch.

524{,}288 \times 4{,}096 = 2 \;\text{GiB per sequence}, \qquad \times\, 8 \;\text{sequences} = 16 \;\text{GiB}

Other models. Llama 3 8B's shape reaches 16 GiB at 131,072 tokens (Llama 3.1 8B's context); an int8 cache halves Llama 2 7B's to 256 KiB per token.

\text{GPT-2 small: } 2 \cdot 12 \cdot 12 \cdot 64 \cdot 2 = 36 \;\text{KiB}, \qquad \text{Llama 2 70B: } 2 \cdot 80 \cdot 8 \cdot 128 \cdot 2 = 320 \;\text{KiB}, \qquad \text{Llama 3 8B: } 2 \cdot 32 \cdot 8 \cdot 128 \cdot 2 = 128 \;\text{KiB}

The serving budget. Llama 2 7B's fp16 weights take 13.48 GB of an 80 GB GPU; the rest holds at most 30 sequences of 4,096 tokens. This is an upper bound: it ignores activations, fragmentation and framework overhead.

\Big\lfloor\frac{80 \times 10^9 - 13{,}476{,}831{,}232}{2 \times 2^{30}}\Big\rfloor = 30

The decode speed limit. Each new token reads every weight once, plus the whole cache, so memory bandwidth sets the pace. On an A100 (2,039 GB/s) at batch 1: 6.61 ms for the weights (at most 151 tokens/s), plus 1.05 ms for a 4K cache.

t_{\text{token}} \;\ge\; \frac{\text{weight bytes} + \text{KV bytes}}{\text{bandwidth}} = \frac{13.48 \;\text{GB} + 2 \;\text{GiB}}{2{,}039 \;\text{GB/s}} = 6.61 + 1.05 \;\text{ms} \;\Rightarrow\; \le 130 \;\text{tokens/s}

$t_{\text{token}}$: the time to generate one token

Key move. S (elements × bytes).

Common slip. Dropping the 2 (keys and values); using the number of query heads $h$ instead of $H_{\text{kv}}$ ; caching queries; mixing GiB ( $2^{30}$ , used for caches) with GB ( $10^{9}$ , used by GPU vendors).

How it's asked. "How much memory does Llama 2 7B's KV cache take for a 4,096-token conversation in fp16, and how many fit on one 80 GB GPU?" (Q22, Q30, Drills 13 and 20.) Strong answer: 512 KiB per token, 2 GiB per sequence, at most 30 sequences, the bandwidth-bound decode, and the levers (GQA, an int8 cache, paging).

The story. Shazeer (2019) found that incremental decoding was slowed by reloading the large keys and values at every step, and proposed sharing one K/V head across all query heads (multi-query attention). Ainslie et al. (2023) took the middle path, groups of query heads sharing a K/V head, and showed an existing MHA checkpoint can be uptrained to GQA with about 5% of its original pre-training compute.

Statement. In grouped-query attention, each group of $G = h/H_{\text{kv}}$ query heads shares one key head and one value head. The cache and the K/V parameters shrink by $h/H_{\text{kv}}$ , and decode attention does $h/H_{\text{kv}}$ FLOPs per byte read. $H_{\text{kv}} = h$ is MHA and $H_{\text{kv}} = 1$ is MQA.

\text{query head } n \text{ reads KV head } \lfloor n/G\rfloor; \qquad \frac{\text{cache}_{\text{MHA}}}{\text{cache}_{\text{GQA}}} = \frac{h}{H_{\text{kv}}}, \qquad \text{decode intensity} = \frac{h}{H_{\text{kv}}} \;\text{FLOP/byte}

$n$: here: which query head, counted from 0
$G$: the group size: query heads per KV head, $h/H_{\text{kv}}$
$h$: the number of query heads

Derivation.

GQA is MHA in which each key and value head is repeated across its group. The simulation's GQA output is identical to MHA with repeated K/V heads (difference exactly 0).

K^{(n)} = K_{\text{shared}}^{(\lfloor n/G\rfloor)}, \qquad V^{(n)} = V_{\text{shared}}^{(\lfloor n/G\rfloor)}, \qquad G = \frac{h}{H_{\text{kv}}}

$K^{(n)}$: the keys query head $n$ attends over
$K_{\text{shared}}^{(\lfloor n/G\rfloor)}$: the stored key head for that group
$V^{(n)}$: the values query head $n$ reads
$V_{\text{shared}}^{(\lfloor n/G\rfloor)}$: the stored value head for that group

The cache is proportional to $H_{\text{kv}}$ (Card 13), so the saving is $h/H_{\text{kv}}$ . Llama 2 70B with 64 query heads and 8 KV heads:

\frac{2 \cdot 80 \cdot 64 \cdot 128 \cdot 2}{2 \cdot 80 \cdot 8 \cdot 128 \cdot 2} = \frac{2{,}560 \;\text{KiB}}{320 \;\text{KiB}} = 8, \qquad \text{MQA: } 40 \;\text{KiB per token}

What that buys in serving. Llama 2 70B's fp16 weights take 137.95 GB of 2 × 80 GB, leaving 22.05 GB. A 4K sequence needs 1.25 GiB with GQA and would need 10 GiB with MHA.

\Big\lfloor\frac{22{,}046{,}703{,}616}{1.25 \times 2^{30}}\Big\rfloor = 16 \;\;(\text{GQA}), \qquad \Big\lfloor\frac{22{,}046{,}703{,}616}{10 \times 2^{30}}\Big\rfloor = 2 \;\;(\text{MHA})

Arithmetic intensity of one decode step over $T$ cached tokens. Each query head does $\mathbf q K^\top$ and $\mathbf p V$ : $4Td_h$ FLOPs. Each KV head's keys and values must be read: $2Td_h\,b$ bytes.

\frac{h \cdot 4Td_h}{H_{\text{kv}} \cdot 2Td_h\,b} = \frac{2h}{b\,H_{\text{kv}}} = \frac{h}{H_{\text{kv}}} \;\;(b = 2)

$T$: the number of cached tokens
$b$: bytes per number (2 for fp16)

That is 1 for Llama 2 7B, 4 for Llama 3 8B and 8 for Llama 2 70B.

Compare with the hardware. An A100 can do 153 FLOPs in the time it reads one byte. Decode attention sits far below that ridge, so it is limited by memory traffic, and GQA speeds it up by reading less, not by computing less.

\text{ridge} = \frac{312 \times 10^{12} \;\text{FLOP/s}}{2{,}039 \times 10^9 \;\text{B/s}} = 153 \;\text{FLOP/byte} \;\gg\; 1,\ 4,\ 8

Parameters (Card 9): Llama 2 70B's attention per layer is 151.0M with GQA against 268.4M for MHA.

1 - \frac{150{,}994{,}944}{268{,}435{,}456} = 0.4375

Key move. S (shapes and bytes).

Common slip. "GQA reduces the number of query heads" (it keeps all $h$ ); "GQA cuts attention FLOPs" (they are unchanged; the bytes fall); "MQA is free" (Ainslie et al. report quality between MQA and MHA for GQA, with MQA lowest).

How it's asked. "Your 7B MHA model is too expensive to serve at long context. Leadership asks whether to retrain it with GQA." (Q31; also Q32, Drill 14.) Strong answer: cache ÷ $h/H_{\text{kv}}$ , intensity $h/H_{\text{kv}}$ vs the ridge, uptraining by mean-pooling the K/V heads at about 5% of pre-training compute instead of retraining, and the choice of group count.

Card 15 · The online-softmax (FlashAttention) recurrence ★★★

The story. Milakov and Gimelshein (2018) showed that a softmax's maximum and normalizer can be computed together in one pass. Dao et al. (2022) fused that recurrence with the weighted sum of values, so exact attention runs block by block without ever forming the $T \times T$ matrix.

Statement. Process the keys in blocks $\mathcal B_1, \ldots, \mathcal B_N$ , keeping a running max $m$ , a running normalizer $\ell$ and an unnormalized output $\mathbf u$ . Start from $m_0 = -\infty$ , $\ell_0 = 0$ , $\mathbf u_0 = \mathbf 0$ . The result is exactly softmax attention.

m_n = \max\big(m_{n-1},\, \max_{j \in \mathcal B_n} s_j\big), \quad \ell_n = e^{m_{n-1} - m_n}\ell_{n-1} + \sum_{j \in \mathcal B_n} e^{s_j - m_n}, \quad \mathbf u_n = e^{m_{n-1} - m_n}\mathbf u_{n-1} + \sum_{j \in \mathcal B_n} e^{s_j - m_n}\mathbf v_j, \quad \mathbf o = \frac{\mathbf u_N}{\ell_N}

$m_n$: the running max after block $n$
$n$: here: the block number
$m_{n-1}$: the running max before block $n$
$\mathcal B_n$: the set of keys in block $n$
$s_j$: the score of key $j$ for this query row
$\ell_n$: the running normalizer: the sum of $e^{s_j - m_n}$ over the keys seen so far
$\ell_{n-1}$: the normalizer before block $n$ , on the old base $m_{n-1}$
$\mathbf u_n$: the running unnormalized output: weighted sum of the values seen so far
$\mathbf u_{n-1}$: the unnormalized output before block $n$
$\mathbf u_N$: the unnormalized output after the last block
$\ell_N$: the normalizer after the last block
$N$: the number of blocks
$\mathbf o$: the row's output

Derivation.

The claim to prove (the invariant): after block $n$ , $\ell_n$ and $\mathbf u_n$ are the direct sums over every key seen so far, measured relative to the current max.

\ell_n = \sum_{j \in \mathcal B_{1:n}} e^{s_j - m_n}, \qquad \mathbf u_n = \sum_{j \in \mathcal B_{1:n}} e^{s_j - m_n}\,\mathbf v_j

$\mathcal B_{1:n}$: all keys in blocks 1 to $n$

Base case. With $m_0 = -\infty$ the rescale factor $e^{m_0 - m_1}$ is 0, so block 1's sums start from scratch.

\ell_1 = 0 \cdot \ell_0 + \sum_{j \in \mathcal B_1} e^{s_j - m_1}, \qquad \mathbf u_1 = \sum_{j \in \mathcal B_1} e^{s_j - m_1}\,\mathbf v_j

$\ell_1$: the normalizer after block 1
$\ell_0$: the starting normalizer, 0
$\mathcal B_1$: the keys of block 1
$m_1$: the max of block 1
$\mathbf u_1$: the unnormalized output after block 1

Inductive step. The factor $e^{m_{n-1} - m_n}$ moves the old sums from base $m_{n-1}$ to base $m_n$ (the F move); adding block $n$ 's terms, already on base $m_n$ , restores the invariant. The same line works for $\mathbf u$ .

e^{m_{n-1} - m_n}\,\ell_{n-1} = \sum_{j \in \mathcal B_{1:n-1}} e^{s_j - m_{n-1}}\,e^{m_{n-1} - m_n} = \sum_{j \in \mathcal B_{1:n-1}} e^{s_j - m_n}

$\mathcal B_{1:n-1}$: all keys in blocks 1 to $n - 1$

At the end the common factor $e^{-m_N}$ cancels, leaving exact softmax attention. The row's logsumexp, $m_N + \log\ell_N$ , is the one number saved for the backward pass (Card 12).

\frac{\mathbf u_N}{\ell_N} = \frac{e^{-m_N}\sum_j e^{s_j}\,\mathbf v_j}{e^{-m_N}\sum_j e^{s_j}} = \sum_j p_j\,\mathbf v_j

$m_N$: the max over all keys
$\sum_j$: add up over all keys
$p_j$: the exact softmax weight on key $j$

Safety and memory. Every exponent is at most 0 and the normalizer is at least 1, so nothing overflows. The state per query row is $d_v + 2$ numbers, independent of $T$ .

s_j - m_n \le 0 \;\Rightarrow\; e^{s_j - m_n} \le 1, \qquad \ell_n \ge 1, \qquad (m_n, \ell_n, \mathbf u_n):\; d_v + 2 \;\text{numbers}

$d_v$: the value width

Bank's row, one key at a time (scaled scores 0, 1.414, 0.707). River raises the max, so the past is rescaled by $e^{0 - 1.414} = 0.243$ ; bank does not.

(m, \ell, \mathbf u):\;\; \big(0,\ 1,\ (0, 0)\big) \;\to\; \big(1.414,\ 0.243 \cdot 1 + 1,\ (1, 0)\big) = \big(1.414,\ 1.243,\ (1, 0)\big) \;\to\; \big(1.414,\ 1.243 + 0.493,\ (1,\ 0.493)\big)

$m$: here: the running max
$\ell$: the running normalizer
$\mathbf u$: the running unnormalized output

The output is $(1, 0.493)/1.736 = (0.576, 0.284)$ , exactly the direct answer.

Merging two partial results (two GPUs, or split-K decoding) is the same rescale with $m = \max(m_a, m_b)$ . Drill 11's partials (3, 1.1353, 2.1353) and (5, 1.0498, 4.1494) merge to 3.6881, the exact answer, with logsumexp 5.1852.

\mathbf o = \frac{e^{m_a - m}\mathbf u_a + e^{m_b - m}\mathbf u_b}{e^{m_a - m}\ell_a + e^{m_b - m}\ell_b}: \qquad \frac{e^{-2}(2.1353) + 4.1494}{e^{-2}(1.1353) + 1.0498} = 3.6881

$m_a$: partial $a$ 's max
$m_b$: partial $b$ 's max
$\mathbf u_a$: partial $a$ 's unnormalized output
$\mathbf u_b$: partial $b$ 's unnormalized output
$\ell_a$: partial $a$ 's normalizer
$\ell_b$: partial $b$ 's normalizer

At scale: one row of 4,096 keys in 32 blocks of 128 matches the direct computation to $7.8 \times 10^{-16}$ in fp64 and $1.7 \times 10^{-7}$ in fp32. Block sizes from 1 to 4,096 all agree to about $10^{-15}$ , and merging two halves agrees to $1.0 \times 10^{-15}$ .

\max\big\lvert\mathbf o_{\text{online}} - \mathbf o_{\text{direct}}\big\rvert = 7.8 \times 10^{-16} \;\;(\text{fp64}), \qquad 1.7 \times 10^{-7} \;\;(\text{fp32})

$\mathbf o_{\text{online}}$: the output from the recurrence
$\mathbf o_{\text{direct}}$: the output from the full softmax

Key move. F (factor out the max), applied repeatedly.

Common slip. Rescaling $\ell$ but not $\mathbf u$ ; using $e^{m_n - m_{n-1}}$ (wrong sign); calling FlashAttention "an approximation". It is exact; only floating-point rounding differs.

How it's asked. "You may store only $O(T)$ numbers per head. Derive a one-pass algorithm that is exactly softmax attention, prove it, and show how to combine results from two GPUs that each saw half the keys." (Worked Problem 2, Drills 10–11.) Strong answer: the invariant and its induction, $d_v + 2$ numbers of state, the merge, and "the FLOPs are unchanged, so the speed comes from memory traffic" (LLM Ch 5).

Card 16 · When 1/√d_k is wrong: alignment and μP's 1/d_k ★★★

The story. Yang et al. (2022, μP) argued that after training, queries and keys are correlated, so $\mathbf q^\top\mathbf k$ grows like $d_k$ (a law-of-large-numbers sum) rather than $\sqrt{d_k}$ (a central-limit sum), and scaled attention by $1/d_k$ so that hyperparameters transfer across widths.

Statement. If the key is partly aligned with the query, the scaled logit's mean grows like $\sqrt{d_k}$ while its noise stays fixed. Dividing by $d_k$ keeps the mean fixed and shrinks the noise.

\mathbf k = \rho\,\mathbf q + \sqrt{1 - \rho^2}\,\mathbf z:\;\; \mathbb E\Big[\frac{\mathbf q^\top\mathbf k}{\sqrt{d_k}}\Big] = \rho\sqrt{d_k},\;\; \text{sd} = \sqrt{1 + \rho^2}; \qquad \mathbb E\Big[\frac{\mathbf q^\top\mathbf k}{d_k}\Big] = \rho,\;\; \text{sd} = \sqrt{\frac{1 + \rho^2}{d_k}}

$\mathbf k$: here: a key partly aligned with the query
$\rho$: the alignment: 0 = independent (initialization), 1 = identical
$\mathbf z$: a standard-normal vector independent of $\mathbf q$
$\mathbb E$: expectation: the average over random draws

Derivation.

The mean. $\mathbb E[\mathbf q^\top\mathbf z] = 0$ by independence, and $\mathbb E\lVert\mathbf q\rVert^{2} = d_k$ for standard-normal components.

\mathbb E[\mathbf q^\top\mathbf k] = \rho\,\mathbb E\lVert\mathbf q\rVert^2 + \sqrt{1 - \rho^2}\;\mathbb E[\mathbf q^\top\mathbf z] = \rho\,d_k

$\mathbb E\lVert\mathbf q\rVert^2$: the average squared length of the query: $d_k$
$\mathbf q^\top\mathbf z$: the query dotted with the independent noise

The variance (the T move). $\lVert\mathbf q\rVert^{2}$ is chi-square with variance $2d_k$ ; $\mathbf q^\top\mathbf z$ has variance $d_k$ (Card 3); the two are uncorrelated, because flipping the sign of $\mathbf z$ flips one and not the other.

\text{Var}(\mathbf q^\top\mathbf k) = \rho^2 \cdot 2d_k + (1 - \rho^2)\,d_k = d_k\,(1 + \rho^2)

Divide by each candidate scale.

\frac{\rho\,d_k}{\sqrt{d_k}} = \rho\sqrt{d_k}, \quad \frac{\sqrt{d_k(1 + \rho^2)}}{\sqrt{d_k}} = \sqrt{1 + \rho^2}; \qquad \frac{\rho\,d_k}{d_k} = \rho, \quad \frac{\sqrt{d_k(1 + \rho^2)}}{d_k} = \sqrt{\frac{1 + \rho^2}{d_k}}

Check at $d_k = 128$ . Fully aligned ( $\rho = 1$ ): the mean of $\mathbf q^\top\mathbf k/\sqrt{d_k}$ is 11.3 (simulated 11.3) against 1.00 with $1/d_k$ . Half aligned ( $\rho = 0.5$ ): mean 5.66 and spread 1.118, simulated 5.66 and 1.114.

\rho = 1:\;\; \sqrt{128} = 11.3 \;\text{ vs }\; 1, \qquad \rho = 0.5:\;\; 0.5\sqrt{128} = 5.66,\;\; \sqrt{1.25} = 1.118

Reconcile the two. At initialization $\rho \approx 0$ : $1/\sqrt{d_k}$ gives unit-variance logits (Card 3), while $1/d_k$ gives spread $1/\sqrt{d_k}$ , and rows that are almost uniform (entropy 8.314 against the maximum 8.318 over 4,096 keys). Neither scale is right at every stage: $\sqrt{d_k}$ is an argument about initialization, $d_k$ an argument about trained alignment. Production LLMs mostly keep $1/\sqrt{d_k}$ and bound the logits directly (Card 18).

\rho = 0:\;\; \text{sd}\Big(\frac{\mathbf q^\top\mathbf k}{d_k}\Big) = \frac{1}{\sqrt{d_k}} = \frac{1}{\sqrt{128}} \;\Rightarrow\; H = 8.314 \approx \ln 4{,}096 = 8.318

$H$: the row entropy in nats (maximum $\ln T$ )

Key move. T (total variance), with a mean term this time.

Common slip. Declaring one scale universally right; treating μP's $1/d_k$ as a temperature to tune by hand rather than part of a width-scaling rule.

How it's asked. "A μP paper divides by $d_k$ instead of $\sqrt{d_k}$ . Which is right?" (Q26, Drill 6.) Strong answer: both, under different assumptions; the aligned calculation; and what models do in practice (keep $1/\sqrt{d_k}$ , add QK-norm or soft-capping).

Card 17 · Uniform attention and rank collapse ★★★

The story. Dong, Cordonnier and Loukas (2021) titled their paper "Attention is not all you need": stacked self-attention without skip connections or feed-forward layers drives every token toward the same vector, doubly exponentially fast in depth.

Statement. A uniform head returns the mean, which shrinks the variance by $1/T$ . For a pure attention layer, the distance from "all tokens identical" falls at least as fast as its cube (Dong et al., 2021, prove a bound of this form for self-attention; the constant depends on the weight norms), which unrolls to a doubly exponential collapse.

\text{uniform: } \text{Var}\Big(\frac1T\sum_{j=1}^{T} v_j\Big) = \frac{\sigma_v^2}{T}; \qquad r_{l+1} \le C\,r_l^3 \;\Rightarrow\; r_L \le C^{-1/2}\big(\sqrt C\,r_0\big)^{3^L}

$\sum_{j=1}^{T}$: add up over all $T$ values
$v_j$: one coordinate of value $j$ , random with variance $\sigma_v^2$
$\sigma_v^2$: the variance of a single value coordinate
$r_{l+1}$: how far the tokens are from all being identical, after layer $l + 1$
$r_l$: the same distance after layer $l$
$l$: the layer number
$C$: a constant set by the size of the weight matrices
$r_L$: the distance after the last layer
$L$: here: the depth of the stack
$r_0$: the distance at the input

Derivation.

Uniform weights give the plain mean, and independent values add their variances (T). At $T = 4{,}096$ that is $2.44 \times 10^{-4}$ (simulated $2.45 \times 10^{-4}$ ).

\text{Var}\Big(\frac1T\sum_{j=1}^{T}v_j\Big) = \frac{1}{T^2}\sum_{j=1}^{T}\text{Var}(v_j) = \frac{T\sigma_v^2}{T^2} = \frac{\sigma_v^2}{T}

Split the tokens into their mean and a residual. Collapse means the residual goes to 0.

X = \mathbf 1\bar{\mathbf x}^\top + R, \qquad \mathbf 1^\top R = \mathbf 0^\top, \qquad r = \lVert R\rVert

$\bar{\mathbf x}$: the average token vector (bar = mean)
$R$: the residual: each token minus the average; its rows sum to 0
$\mathbf 0^\top$: a row of zeros
$r$: here: the size of the residual
$\lVert R\rVert$: the size (Frobenius norm) of $R$

Split each logit three ways, with $W = W_QW_K^\top/\sqrt{d_k}$ and $\mathbf r_i$ the rows of $R$ . The query-only term is constant along a row and cancels (F). The key-only term is the same in every row. Only the cross term makes rows differ, and it is quadratic in the residual.

s_{ij} = \underbrace{(\bar{\mathbf x} + \mathbf r_i)^\top W\bar{\mathbf x}}_{\text{query only: cancels}} + \underbrace{\bar{\mathbf x}^\top W\mathbf r_j}_{\text{key only: same in every row}} + \underbrace{\mathbf r_i^\top W\mathbf r_j}_{O(r^2)}

$W$: the combined score matrix $W_QW_K^\top/\sqrt{d_k}$
$\mathbf r_i$: row $i$ of $R$ : token $i$ 's deviation from the average
$\mathbf r_j$: token $j$ 's deviation
$O(r^2)$: a quantity that shrinks like $r^2$ (big-O: ignoring constant factors)

If every row of $P$ were the same distribution $\boldsymbol\pi$ , the output would be one vector repeated: no residual at all. So write $P = \mathbf 1\boldsymbol\pi^\top + \Delta$ with $\Delta = O(r^{2})$ . Since $\Delta\mathbf 1 = \mathbf 0$ , the mean part of $X$ drops out, and the output's residual is quadratic times linear.

P = \mathbf 1\boldsymbol\pi^\top + \Delta,\;\; \Delta\mathbf 1 = \mathbf 0 \;\Rightarrow\; \text{res}(PXW_V) = \text{res}(\Delta X W_V) = \text{res}(\Delta R\,W_V) = O(r^2) \cdot O(r) = O(r^3)

$\boldsymbol\pi$: the weight row all rows would share without the cross term
$\Delta$: how much the rows of $P$ differ from that common row
$W_V$: the value map
$P$: the layer's attention weights

A direct check: shrink the input residual by 10× and the output residual shrinks by $10^{3}$ (measured slopes 3.00 and 3.00 per decade).

Unroll the cubic recursion with $u = \sqrt C\,r$ .

u_{l+1} = \sqrt C\,r_{l+1} \le \sqrt C \cdot C\,r_l^3 = u_l^3 \;\Rightarrow\; u_L \le u_0^{3^L}

$u_{l+1}$: the rescaled residual $\sqrt C\,r$ after layer $l + 1$
$u_l$: the rescaled residual after layer $l$
$u_L$: the rescaled residual after the last layer
$u_0$: the rescaled residual at the input

Once $u_0 < 1$ , the exponent triples every layer: doubly exponential.

The simulation (32 tokens, width 64, 20 random stacks). Pure attention: 0.783 → 0.273 → 0.0126 → $2.2 \times 10^{-6}$ over layers 1–4, and the ratios $r_{l+1}/r_l^{3}$ stay of order 1, as a cubic law predicts. With residual connections the distance is still 0.746 at layer 5 and 0.103 at layer 12: residuals mitigate collapse at random initialization; they do not abolish it.

\frac{r_{l+1}}{r_l^3} = 0.82,\ 0.57,\ 0.62,\ 1.10 \;\;(l = 0, \ldots, 3), \qquad r_{12} = 0.103 \;\;\text{with residuals}

$r_{12}$: the distance after 12 layers

Key move. T for the uniform head; F to cancel the query-only term.

Common slip. Saying residuals prevent collapse; confusing rank collapse (all tokens become the same) with entropy collapse (each row becomes one-hot, Section 11); claiming the mean is the "best" summary when a head is uniform. It is the most diluted one.

How it's asked. "Defend or attack: a stack of attention layers alone is a good language model." (Q41, Drill 5.) Strong answer: the logit split, the cubic, doubly exponential collapse, and why the block has residuals and a feed-forward network (Section 15).

Card 18 · Attention sinks: the logit needed to park mass ★★

The story. Xiao et al. (2023) found that window attention, which evicts the oldest tokens, breaks down once the first tokens leave the cache, because many heads park their attention there. Their fix, StreamingLLM, keeps the first four tokens as "attention sinks".

Statement. A sink whose logit beats each of $T - 1$ equal keys by $a$ takes share $w = e^{a}/(e^{a} + T - 1)$ , so taking a share $w$ needs only $a = \ln(w(T - 1)/(1 - w))$ . QK-norm with gain $g$ caps every weight.

w = \frac{e^{a}}{e^{a} + T - 1} \iff a = \ln\frac{w(T - 1)}{1 - w}; \qquad s_{ij} = g\,\hat{\mathbf q}_i^\top\hat{\mathbf k}_j \;\Rightarrow\; \max_j p_{ij} \le \frac{e^{g}}{e^{g} + (T - 1)e^{-g}}

$w$: the share of the row's attention the sink takes
$a$: the sink's logit advantage over each other key
$g$: the QK-norm gain: the logit of a perfectly aligned query–key pair
$\hat{\mathbf q}_i$: query $i$ scaled to length 1 (hat = normalized)
$\hat{\mathbf k}_j$: key $j$ scaled to length 1
$\max_j p_{ij}$: the largest weight in row $i$

Derivation.

By shift invariance, set the other $T - 1$ logits to 0 and the sink's to $a$ .

w = \frac{e^{a}}{e^{a} + (T - 1)\,e^{0}}

Solve for $a$ .

w\,e^{a} + w(T - 1) = e^{a} \;\Rightarrow\; e^{a} = \frac{w(T - 1)}{1 - w} \;\Rightarrow\; a = \ln\frac{w(T - 1)}{1 - w}

The numbers. The required advantage grows only like $\ln T$ , so a sink is cheap: an advantage of 5 already takes 12.7% of a 1,024-token row.

a_{90\%} = \ln(9 \times 1{,}023) = 9.13 \;\;(T = 1{,}024), \qquad 7.04 \;\;(T = 128), \qquad 10.51 \;\;(T = 4{,}096)

$a_{90\%}$: the advantage needed for the sink to take 90%

Why sinks exist. A softmax row must sum to 1, so a head with nothing useful to read must still put its weight somewhere. The cheapest place is a token every query can see, with a near-zero value. In the toy, the ( $\mathbf k = \mathbf v = \mathbf 0$ ) takes 14% of bank's attention and contributes nothing.

\sum_j p_{ij} = 1 \;\text{always}; \qquad \mathbf v_{\text{sink}} \approx \mathbf 0 \;\Rightarrow\; \mathbf o_i \approx \sum_{j \ne \text{sink}} p_{ij}\,\mathbf v_j \;\;\text{with total weight } 1 - w

$\mathbf v_{\text{sink}}$: the sink token's value
$\sum_{j \ne \text{sink}}$: add up over every key except the sink

QK-norm caps any single weight. With unit-length queries and keys every logit lies in $[-g, g]$ . A weight is largest when its own logit is $g$ and every other is $-g$ .

p_{ij} = \frac{e^{s_{ij}}}{\sum_r e^{s_{ir}}} \;\le\; \frac{e^{g}}{e^{g} + (T - 1)\,e^{-g}}

With $T = 4{,}096$ : $g = 5$ caps the top weight at 0.843, and $g = 1/\sqrt{128}$ (unit vectors with the usual scale) caps it at $2.9 \times 10^{-4}$ , so no row can select anything. Setting the cap to 0.9 and solving gives the smallest gain that allows 90%:

g = \tfrac12\ln\frac{0.9 \times 4{,}095}{0.1} = 5.26

The kernel view (Card 8). For unit vectors the squared distance is $2 - 2\hat{\mathbf q}^\top\hat{\mathbf k}$ , so the gain is a Gaussian bandwidth $\lambda = 1/\sqrt g$ : 3.36 at $g = 1/\sqrt{128}$ .

g\,\hat{\mathbf q}^\top\hat{\mathbf k} = g - \frac{g}{2}\lVert\hat{\mathbf q} - \hat{\mathbf k}\rVert^2 = \text{const} - \frac{\lVert\hat{\mathbf q} - \hat{\mathbf k}\rVert^2}{2\lambda^2}, \qquad \lambda = \frac{1}{\sqrt g}

$\hat{\mathbf q}$: a unit-length query
$\hat{\mathbf k}$: a unit-length key
$\lambda$: here: the equivalent Gaussian bandwidth

Key move. F (shift invariance sets the other logits to 0).

Common slip. Treating a sink as a bug to suppress, or reading its large weight as "the first token is important". Its value is near zero; it is a no-op. The operational slip: an eviction policy that drops the oldest tokens first, which removes the sinks.

How it's asked. "Interpretability finds heads putting 80% of their mass on the first token, and a cache-eviction policy that drops the oldest tokens wrecks quality. Explain both." (Q34, Drills 18–19.) Strong answer: sum-to-1 forces a parking spot, $a = \ln(w(T - 1)/(1 - w))$ is small, StreamingLLM keeps the first four tokens, and QK-norm's gain sets how sharp any head can be.

Practice protocol

Read only the Statement. Cover the rest of the card.
Derive on a blank page against a timer: 5 minutes for ★, 10 for ★★, 15 for ★★★. Talk out loud as if an interviewer were listening; the narration is half the grade.
Land the check number. Every card has one (bank's 0.140 / 0.576 / 0.284, the gradient $(-0.057, 0.116)$ , 512 KiB per token, 24,704, the online trace 0.243 → 1.736). If your derivation doesn't reproduce it, the error is in your algebra.
Compare step by step. Mark the first step you missed and name the toolkit move it needed (S, O, F or T). Read the Common slip and check whether you made it.
Answer "How it's asked" cold, including the strong-answer extension. Senior and research rounds live there.
Re-derive after 48 hours, then after a week, without looking. A card is mastered when you can do it twice in a row within time and hit the check.
Run the chains. Several cards are one argument continued; practice them in one sitting:

The softmax chain: Card 2 → 4 → 5 → 15 (shift → Jacobian → backward → online recurrence).
The scale chain: Card 3 → 16 → 18 (variance at init → alignment in training → bounding logits directly).
The cost chain: Card 11 → 12 → 13 → 14 (FLOPs → score memory → KV cache → sharing K and V).
The expressiveness chain: Card 1 → 8 → 10 → 17 (convex combination → kernel smoother → rank → collapse).
The order chain: Card 6 → 7 (the mask's guarantee → the symmetry the mask breaks).

Suggested order: Cards 1–7 after the chapter's second read, then Interview Bank Levels 1–2, then Cards 8–18, then Levels 3–4.

Mastery checklist

[ ] Card 1 · Attention from one row to matrix form; shapes with batch and heads; convex combination; bank's output (0.576, 0.284) by hand.
[ ] Card 2 · Shift invariance by multiplying through by $e^{c}$ ; logsumexp with the max factored out (1000.313); temperature limits; $\partial\log p_i/\partial s_j = \delta_{ij} - p_j$ .
[ ] Card 3 · $\text{Var}(\mathbf q^\top\mathbf k) = d_k\sigma_q^{2}\sigma_k^{2}$ in three lines; entropy $\ln T - \sigma_s^{2}/2$ (7.818) and where it breaks (4.08); two-key exact 0.773 vs 0.00135; the gradient concentrates rather than simply vanishing.
[ ] Card 4 · $J = \operatorname{diag}(\mathbf p) - \mathbf p\mathbf p^\top$ ; PSD as a variance; rank $T - 1$ ; Gershgorin bound 1/2; bank's eigenvalues 0.393, 0.175, 0; $J\mathbf y$ in $O(T)$ .
[ ] Card 5 · The five backward lines; $\operatorname{rowsum}(P \odot dP) = \operatorname{rowsum}(dO \odot O)$ ; toy $\partial\mathcal L/\partial\mathbf s = (0.081, -0.244, 0.164)$ and $\partial\mathcal L/\partial\mathbf q = (-0.057, 0.116)$ with $\mathbf g = (-1, 0)$ ; descent turns the query toward river.
[ ] Card 6 · $e^{-\infty} = 0$ in the normalizer; change-the-future test (0 vs 0.217); KV cache validity; fully masked row = NaN in a hand softmax (SDPA on CPU, torch 2.14.1: zeros); post-softmax zeroing sums to 0.583 and leaks 0.162; $(T + 1)/2T$ .
[ ] Card 7 · $\text{Attn}(\Pi X) = \Pi\,\text{Attn}(X)$ in four lines; pooling invariant; causal mask breaks it (0.426); BOS share $1/t$ .
[ ] Card 8 · Expand $-\lVert\mathbf q - \mathbf k\rVert^{2}/2\lambda^{2}$ ; the query term cancels; the key-length bias; equal norms make it exact ( $5.6 \times 10^{-16}$ vs 1.09 without the bias).
[ ] Card 9 · $4d^{2}$ whatever $h$ is (64; 67,108,864); GPT-2 small 2,362,368 with biases; GQA 150,994,944 vs 268,435,456; concat · $W_O$ = sum over heads.
[ ] Card 10 · $\operatorname{rank}(QK^\top) \le d_k$ by outer products; $\log P$ rank $\le d_k + 1$ but $P$ full rank (8 vs 64); $d_k = 1$ picks only argmax or argmin $k$ .
[ ] Card 11 · $2mnk$ ; $8Td^{2} + 4T^{2}d + 6Td\,d_{\text{ff}}$ ; 549.8 / 274.9 / 1,108.1 GFLOP and 14.2% at 4K; $T^\star = 24{,}704$ ; causal 49,407.
[ ] Card 12 · 32 MiB per head, 1 GiB per layer, 32 GiB per model at 4K; 32 GiB per head at 128K; logsumexp 16 MiB; recompute $P$ from $\text{lse}$ and $D$ .
[ ] Card 13 · $2 \cdot L \cdot H_{\text{kv}} \cdot d_h \cdot b$ ; 512 KiB per token, 2 GiB per 4K sequence, 30 sequences on 80 GB; 36 / 320 / 128 KiB for the other models; 6.61 + 1.05 ms per token.
[ ] Card 14 · GQA = MHA with repeated K/V heads; cache ratio 8 (2,560 vs 320 KiB); 16 vs 2 sequences; intensity $h/H_{\text{kv}}$ against a ridge of 153.
[ ] Card 15 · The $(m, \ell, \mathbf u)$ recurrence; invariant proved by induction; bank's trace (rescale 0.243, $\ell$ = 1.243 → 1.736); merge to 3.6881; $7.8 \times 10^{-16}$ in fp64.
[ ] Card 16 · Aligned mean $\rho d_k$ and variance $d_k(1 + \rho^{2})$ ; $1/\sqrt{d_k}$ gives $\rho\sqrt{d_k}$ (11.3 at $\rho = 1$ ), $1/d_k$ gives $\rho$ ; why both scales exist.
[ ] Card 17 · Uniform head: $\sigma_v^{2}/T$ ; the three-way logit split; output residual $O(r^{3})$ (slope 3.00); unroll to $u_0^{3^L}$ ; 0.783 → $2.2 \times 10^{-6}$ pure vs 0.103 at layer 12 with residuals.
[ ] Card 18 · $a = \ln(w(T - 1)/(1 - w))$ : 9.13 at $T = 1{,}024$ ; why sinks exist; QK-norm cap $e^{g}/(e^{g} + (T - 1)e^{-g})$ (0.843 at $g = 5$ ; $g = 5.26$ for 0.9); bandwidth $1/\sqrt g$ .

Chapter 3 — Derivations

The toolkit

Formula sheet

The operator

Softmax

Scaling

Gradients

Masks and positions

Multi-head and cost

Inference

Cards

Card 1 · Scaled dot-product attention in matrix form, with shapes ★

Card 2 · Softmax: shift invariance, temperature, stable logsumexp ★

Card 3 · Var(q·k) = d_k, hence 1/√d_k ★

Card 4 · The softmax Jacobian ★★

Card 5 · Backpropagation through attention ★★

Card 6 · Causal mask correctness ★

Card 7 · Permutation equivariance without positions ★★

Card 8 · Attention as Nadaraya–Watson kernel smoothing ★★★

Card 9 · Multi-head parameter count ★

Card 10 · The rank bound and the d_k = 1 head ★★★

Card 11 · Attention FLOPs and the crossover ★★

Card 12 · Attention memory: naive vs online ★★

Card 13 · KV-cache size ★

Card 14 · MQA/GQA: what sharing saves ★★

Card 15 · The online-softmax (FlashAttention) recurrence ★★★

Card 16 · When 1/√d_k is wrong: alignment and μP's 1/d_k ★★★

Card 17 · Uniform attention and rank collapse ★★★

Card 18 · Attention sinks: the logit needed to park mass ★★

Practice protocol

Mastery checklist