Take the attention head from §4.2, make it the size real models use, say $d_k = 512$, initialize it the standard way so that queries and keys have unit-variance entries, compute the scores as $\mathbf{Q}\mathbf{K}^\top$ exactly as that section did, and look at the attention weights before any training. Most rows are nearly one-hot. Each query has already "decided" on a single key, at random, and the gradient flowing back through the softmax is close to zero. The head cannot learn, because it is stuck in the hard lookup that §4.1 worked so hard to escape.

The fix is one division, and it is the reason the formula in every Transformer paper looks the way it does,

This section derives the $\sqrt{d_k}$ from a variance argument, shows exactly how the gradient dies without it, and then does the thing every lab interview eventually asks for. We derive the backward pass of attention by hand and check it against numerical gradients.

The formula, piece by piece

Fix one head. The queries are $\mathbf{Q} \in \R^{T_q\times d_k}$, the keys $\mathbf{K} \in \R^{T_k\times d_k}$, and the values $\mathbf{V} \in \R^{T_k\times d_v}$, as in §4.2 (a batch dimension, and later a head dimension, ride along in front). Scaled dot-product attention computes three matrices,

where the softmax acts on each row of $\mathbf{S}$ separately, so $a_{ij} = e^{s_{ij}}/\sum_l e^{s_{il}}$ and every row of $\mathbf{A}$ sums to 1.

In plain words, score every query against every key, shrink the scores by the square root of the key width, turn each query's scores into weights, and average the values with them.

In the language of §4.1, dividing by $\sqrt{d_k}$ is a soft lookup at temperature $\tau = \sqrt{d_k}$. The question is why that particular temperature.

Why dot products grow with the width

Suppose that at initialization the entries of a query $\mathbf{q}$ and a key $\mathbf{k}$ are independent random variables with mean 0 and variance 1. This is roughly true in a freshly initialized network whose inputs have been normalized (LayerNorm, §5.2) and whose projections use a variance-preserving initialization (§0.7). The score is a sum of $d_k$ products,

Each term has mean $\mathbb{E}[q_m k_m] = \mathbb{E}[q_m]\,\mathbb{E}[k_m] = 0$ by independence. Its variance is

The $d_k$ terms are independent of each other (they involve different coordinates), so their variances add,

In plain words, the typical size of a raw score is $\sqrt{d_k}$. Double the key width and the scores grow by a factor of $\sqrt 2$, for no reason other than that there are more terms in the sum.

Dividing by $\sqrt{d_k}$ undoes exactly that growth, since $\operatorname{Var}(\mathbf{q}^\top\mathbf{k}/\sqrt{d_k}) = d_k/d_k = 1$. The scaled scores have unit variance at every width. (If the entries have variances $\sigma_q^2$ and $\sigma_k^2$ instead of 1, the same steps give $\operatorname{Var}(\mathbf{q}^\top\mathbf{k}) = d_k\sigma_q^2\sigma_k^2$, and the scaling still removes the dependence on $d_k$.) This is the argument in footnote 4 of Vaswani et al. (2017), which also notes that additive attention outperformed unscaled dot-product attention at large $d_k$. It is also Fact 1 of §1.3 in a new role. There it explained why random vectors are nearly orthogonal. Here it explains why their raw dot products are large.

What large scores do to the softmax

Why should we care that scores have standard deviation $\sqrt{d_k}$ rather than 1? Because softmax only sees differences between scores, and when the differences are large the softmax saturates. If the largest score in a row beats the runner-up by 10, the winner's weight is at least $e^{10}/(e^{10} + T - 1)$, which for $T = 32$ keys is 0.9986. The lookup has become hard, and §4.1 told us that a hard lookup has no gradient. To see how little gradient is left, we need the derivative of the softmax as a matrix.

The Jacobian of softmax, and what saturation does to it

Write one row of scores as $\mathbf{s} \in \R^n$ and its softmax as $\mathbf{a}$. §0.2 derived the softmax Jacobian with the quotient rule (and §4.1 used its entries, $\partial a_i/\partial s_j = a_i(\mathbb{1}[i=j] - a_j)$). As a matrix,

Recall its three properties from that section. It is symmetric, its rows sum to zero (shift invariance again), and it is positive semidefinite, because for any $\mathbf{u} \in \R^n$

the variance of the numbers $u_i$ when entry $i$ is chosen with probability $a_i$. Two consequences matter for attention.

The first is a single number for how much gradient can pass. The trace of the Jacobian, the sum of its eigenvalues, is $\sum_i a_i(1 - a_i) = 1 - \lVert\mathbf{a}\rVert^2$. It is largest, $1 - 1/n$, for uniform weights, and it shrinks to 0 as the weights concentrate.

The second is the punchline. If $\mathbf{a}$ is one-hot, then every variance under $\mathbf{a}$ is zero, so the quadratic form vanishes for every $\mathbf{u}$, and a symmetric positive semidefinite matrix with a zero quadratic form is the zero matrix. A saturated softmax has a zero Jacobian. No gradient reaches the scores, and therefore none reaches $\mathbf{W}_Q$ or $\mathbf{W}_K$. The head's attention pattern is frozen wherever it happened to land at initialization, exactly the stuck query of Exercise 4.1e.

In backpropagation we never build this $n\times n$ matrix. We only need its product with the upstream gradient $\mathbf{g} = \partial\mathcal{L}/\partial\mathbf{a}$, the $O(n)$ vector-Jacobian product you wrote in Exercise 0.2c,

In plain words, each score's gradient is its weight times how much its upstream gradient exceeds the weighted average upstream gradient. With a one-hot $\mathbf{a}$, the winner's term is $1\cdot(g_{j^\star} - g_{j^\star}) = 0$ and every other term has weight 0, which is the punchline again, in one line.

Saturation, measured

Now we can measure what scaling buys. For each width, draw random queries and 32 random keys, compute the attention weights with and without the division by $\sqrt{d_k}$, and record the largest weight in each row and the trace of the softmax Jacobian ($1 - \lVert\mathbf{a}\rVert^2$, which is 0 for one-hot rows).

Without scaling, the average largest weight climbs from 0.35 at $d_k = 4$ to 0.96 at $d_k = 1024$, and the gradient capacity (the trace) falls more than tenfold. With scaling, both numbers stay put at every width, about 0.16 and 0.93. The head starts out soft whatever its size, and its sharpness becomes something it learns rather than an accident of $d_k$.

Saturation has a mirror image. If scores are tiny compared to 1, the weights are nearly uniform, $a_j \approx 1/T_k$, and by the formula $\partial\mathbf{o}/\partial s_j = a_j(\mathbf{v}_j - \mathbf{o})$ of §4.1 each score receives a gradient of order $1/T_k$. Unit-variance scores sit between the two extremes, soft enough for every key to receive gradient and sharp enough that the weights are not all the same.

The assumption behind the argument, and when it fails

The variance argument assumes that queries and keys are independent. That is true at initialization and false after training, because the whole point of learning $\mathbf{W}_Q$ and $\mathbf{W}_K$ is to make the query of "it" align with the key of "animal". Suppose a key is correlated with its query, $\mathbf{k} = \rho\mathbf{q} + \sqrt{1-\rho^2}\,\mathbf{z}$ with $\mathbf{z}$ independent noise. Then $\mathbb{E}[\mathbf{q}^\top\mathbf{k}] = \rho\,\mathbb{E}\lVert\mathbf{q}\rVert^2 = \rho d_k$, so the scaled score has mean $\rho\sqrt{d_k}$, which still grows with width (Exercise 4.3d derives the variance too). Aligned vectors have dot products that grow like $d_k$, by the law of large numbers, not like $\sqrt{d_k}$.

This is why the Maximal Update Parametrization ($\mu$P) of Yang et al. (2022) scales attention scores by $1/d_k$ instead of $1/\sqrt{d_k}$. Their analysis of how quantities evolve during training, as width grows, predicts that queries and keys become correlated, and $1/d_k$ is the scaling that keeps the logits of a trained model stable across widths, which is what makes hyperparameters transfer from small models to large ones (§9.5). Other practical answers to growing logits include normalizing queries and keys before the dot product, QK-norm (Henry et al. 2020). Dehghani et al. (2023) adopted it for their 22-billion-parameter ViT after finding that, already at about 8 billion parameters, attention logits grew to huge values, the weights became nearly one-hot, and training diverged after a few thousand steps (§5.7). And some implementations fold the $1/\sqrt{d_k}$ into the initialization of $\mathbf{W}_Q$ rather than into the forward pass. T5's is one. Its query weights are initialized with a standard deviation smaller by a factor of $\sqrt{d_k}$ than its key weights, and the scores are never divided. At initialization the two choices are identical. During training they differ slightly, because an optimizer such as Adam takes steps of a similar size whatever the scale of a weight, so the effective learning rate of the folded-in version is different.

import numpy as np

def scaled_dot_product_attention(Q, K, V, mask=None):
    raise NotImplementedError

A vectorized implementation, and why stability is not optional

Here is the version we will use for the rest of the chapter. Every line is one equation from above.

The naive line fails in float16 for a score of only 12, because $e^{12} \approx 163{,}000$ exceeds the largest float16 number, 65,504 (the cutoff is $\ln 65504 \approx 11.09$). In float32 the cutoff is about 88.7. Subtracting the row maximum fixes overflow completely, because the largest exponent becomes $e^0 = 1$ and the denominator is at least 1. Small exponents can still underflow to zero, but that is harmless, since those weights really are negligible. This is why production kernels subtract the maximum, and why mixed-precision training usually computes the softmax and its normalizer in float32 even when the matrix products run in bfloat16 or float16 (§8.5). FlashAttention goes one step further and computes the maximum and the normalizer incrementally, block by block, so the score matrix never has to exist in memory (§12.2).

The backward pass of attention

Now the derivation that separates people who have used attention from people who understand it. Given the upstream gradient $\mathrm{d}\mathbf{O} = \partial\mathcal{L}/\partial\mathbf{O} \in \R^{T_q\times d_v}$, we want $\mathrm{d}\mathbf{Q}$, $\mathrm{d}\mathbf{K}$, and $\mathrm{d}\mathbf{V}$. (Here $\mathrm{d}\mathbf{X}$ always means $\partial\mathcal{L}/\partial\mathbf{X}$, a matrix of the same shape as $\mathbf{X}$, the convention of the FlashAttention papers.) Write $c = 1/\sqrt{d_k}$ for the scale. The forward pass is three steps,

and we walk back through them in reverse, using index notation so that no step is a leap. Recall from §0.2 the rule for any product $\mathbf{Y} = \mathbf{X}\mathbf{W}$, namely $\mathrm{d}\mathbf{X} = \mathrm{d}\mathbf{Y}\,\mathbf{W}^\top$ and $\mathrm{d}\mathbf{W} = \mathbf{X}^\top\mathrm{d}\mathbf{Y}$. Each step below is that rule, checked by indices.

Step 1, through $\mathbf{O} = \mathbf{A}\mathbf{V}$. Entry by entry, $O_{ik} = \sum_j A_{ij}V_{jk}$. The entry $V_{jk}$ appears in $O_{ik}$ for every query $i$, with coefficient $A_{ij}$, so

The entry $A_{ij}$ appears in $O_{ik}$ for every $k$, with coefficient $V_{jk}$, so

Step 2, through the row-wise softmax. Each row of $\mathbf{A}$ depends only on the same row of $\mathbf{S}$, so we apply the softmax vector-Jacobian product above to every row,

The row sum simplifies nicely. Substituting $\mathrm{d}A_{il} = \sum_k \mathrm{d}O_{ik}V_{lk}$,

so $D_i$ is just the dot product of row $i$ of $\mathrm{d}\mathbf{O}$ with row $i$ of $\mathbf{O}$, and

Step 3, through $\mathbf{S} = c\,\mathbf{Q}\mathbf{K}^\top$. Entry by entry, $S_{ij} = c\sum_m Q_{im}K_{jm}$. So $Q_{im}$ appears in $S_{ij}$ for every key $j$, with coefficient $cK_{jm}$, and $K_{jm}$ appears for every query $i$, with coefficient $cQ_{im}$,

In plain words, the value gradient is the attention-weighted sum of output gradients, the score gradient is the softmax backward applied to $\mathrm{d}\mathbf{O}\mathbf{V}^\top$, and the query and key gradients come from the score gradient exactly as in any matrix product, times the scale.

Four observations make this derivation useful in practice.

• Masks need no special handling. A masked entry has $A_{ij} = 0$, so $\mathrm{d}S_{ij} = 0$ automatically, and no gradient leaks through a masked connection.
• The backward pass costs about twice the forward pass. The forward pass does two large matrix products ($\mathbf{Q}\mathbf{K}^\top$ and $\mathbf{A}\mathbf{V}$). The backward pass does four ($\mathrm{d}\mathbf{O}\mathbf{V}^\top$, $\mathbf{A}^\top\mathrm{d}\mathbf{O}$, $\mathrm{d}\mathbf{S}\mathbf{K}$, $\mathrm{d}\mathbf{S}^\top\mathbf{Q}$), each the same size. This is the attention version of the familiar rule that a backward pass costs twice a forward pass (§5.6).
• The backward pass needs $\mathbf{A}$, a $T_q\times T_k$ matrix per head. Storing it for every layer is what makes naive attention memory-hungry (§4.7). FlashAttention instead stores only $\mathbf{O}$ and one number per row, the log-sum-exp of the scores, and recomputes $\mathbf{A}$ block by block during the backward pass (§12.3).
• The identity $D_i = \mathrm{d}\mathbf{O}_i\cdot\mathbf{O}_i$ matters. Computing $D_i$ from its definition needs a whole row of $\mathbf{A}$ and $\mathrm{d}\mathbf{A}$ at once. Computing it from $\mathrm{d}\mathbf{O}$ and $\mathbf{O}$ needs neither, which is exactly what lets FlashAttention's backward pass work one tile at a time.

import numpy as np

def attention_backward(Q, K, V, dO, mask=None):
    raise NotImplementedError

import numpy as np

def scaled_score_moments(d_k, rho):
    return 0.0, 1.0      # placeholder: the independent case, rho = 0

In PyTorch

PyTorch's scaled_dot_product_attention computes exactly our formula, and dispatches to a fused FlashAttention or memory-efficient kernel when one is available. It also takes dropout_p, which during training zeroes attention weights at random and rescales the rest, as GPT-2 and BERT do at a rate of 0.1. Most recent large language models train with it set to zero. Below we check the function against our NumPy version, then check our hand-derived backward pass against autograd.

Read more

• Vaswani et al. (2017). Attention Is All You Need. Section 3.2.1 and its footnote give the variance argument for the scaling.
• Phuong & Hutter (2022). Formal Algorithms for Transformers. Algorithm 4 states scaled dot-product attention with masks precisely, in the same notation style as this chapter.
• Dao et al. (2022). FlashAttention. Appendix B derives the same backward pass, including the $D_i = \mathrm{d}\mathbf{O}_i\cdot\mathbf{O}_i$ identity, and shows how to run it tile by tile.
• Yang et al. (2022). Tensor Programs V. The $\mu$P paper, which argues for $1/d_k$ attention scaling.
• Henry et al. (2020). Query-Key Normalization for Transformers and Dehghani et al. (2023), Scaling Vision Transformers to 22 Billion Parameters. Normalizing queries and keys, and the logit growth that made it necessary.
• Milakov & Gimelshein (2018). Online normalizer calculation for softmax. Computing a stable softmax in one pass, the seed of FlashAttention.

Watch next

▶

▶