Answers · How models work

Why do we divide the attention scores by the square root of the key dimension?

In shortThe square root of the key dimension is the typical size of an unscaled attention score, so dividing by it stops softmax from becoming a hard pick-one.

The square root of the key dimension is the standard deviation of an unscaled attention score, so dividing by it holds the spread of the scores at one however wide the head happens to be. It is a derived number, not a tuning knob: it falls out of the variance of a dot product between independent vectors. Any other constant would fail at some head width, and the failure is not cosmetic, because softmax converts a change of scale into a change of sharpness.

Where the spread comes from

Fix one head with key dimension $d_k$ , and suppose the components of a query $q$ and a key $k$ are independent random variables with mean zero and variance one. That is roughly what a fresh initialisation gives, and roughly what normalisation maintains. A single attention score is their dot product, a sum of $d_k$ products:

q \cdot k = \sum_{i=1}^{d_k} q_i k_i

Each term $q_i k_i$ is a product of two independent mean-zero variables, so its mean is $0 \cdot 0 = 0$ and the whole score is centred. For the variance of one term, independence lets the expectation of the product of squares factorise, so $\mathbb{E}[q_i^2 k_i^2] = \mathbb{E}[q_i^2]\,\mathbb{E}[k_i^2] = 1 \cdot 1 = 1$ . Every term therefore contributes exactly one unit of variance. The terms are independent of one another, and independent variances add, so $d_k$ of them give a variance of $d_k$ .

FIG 1The variance of one attention score

: the dimension of the query and key vectors inside one head
: variance, the average squared distance from the mean
: standard deviation, the square root of the variance, measured in score units

Each of the terms adds one unit of variance, and the variance of a sum of independent terms is the sum of the variances. The standard deviation, which is in the same units as the score itself, grows like the square root of the width.

Read the numbers off the formula. With $d_k = 16$ a typical score sits about $4$ from zero, at the $d_k = 64$ of the original transformer about $8$ , and at $d_k = 256$ about $16$ . Dividing each score by $\sqrt{d_k}$ brings all three of those back to $1$ . That is the whole content of the denominator: it restores a fixed unit of spread to a quantity whose spread was drifting with the width of the head.

What a score of 8 does to softmax

Softmax sees only differences between scores, and it turns a difference $\Delta$ into a ratio of weights of $e^{\Delta}$ . Multiply every score by a constant $c$ and every difference becomes $c\Delta$ , so every ratio is raised to the power $c$ : the distribution sharpens. This is exactly a temperature. With raw scores of spread $8$ , take five of them, $12$ , $8$ , $4$ , $0$ , $-4$ , and exponentiate: $e^{12} = 162754$ , $e^{8} = 2981$ , $e^{4} = 54.6$ , $e^{0} = 1$ , $e^{-4} = 0.018$ , summing to $165791$ . The top weight is $162754/165791 = 0.98$ . The second token is only four raw points behind, and gets $2981/165791 = 0.018$ , a factor of $e^{4} = 54.6$ smaller.

FIG 2The same five scores, before and after dividing by 8

step	raw score	weight unscaled	scaled score	weight scaled	what happened
1	12	0.98	1.5	0.43	Top token. Unscaled it takes 98 per cent of the attention.
2	8	0.018	1	0.26	Four raw points behind, yet 55 times smaller in weight.
3	4	0.0003	0.5	0.16	Still a plausible match, effectively erased.
4	0	6e-06	0	0.1	Neutral.
5	-4	1e-07	-0.5	0.06	A poor match, downweighted in both columns.

5 steps

The ordering never changes; only the sharpness does. Scaling makes a near-argmax into a graded mixture.

Divide by $8$ instead and the scores become $1.5$ , $1$ , $0.5$ , $0$ , $-0.5$ , with exponentials $4.48$ , $2.72$ , $1.65$ , $1$ , $0.61$ summing to $10.46$ . The weights are $0.43$ , $0.26$ , $0.16$ , $0.10$ , $0.06$ . The top two are now $0.43/0.26 = 1.65$ apart, which is $e^{0.5}$ , exactly as it must be, since the gap shrank from $4$ to $0.5$ . The information in the second and third tokens survives.

Why a smaller number is not enough

The natural guess is that the division is a general trick for keeping numbers small, so any constant that brings the magnitudes down would serve, or that $8$ was found by experiment. Suppose you fixed the divisor at $8$ regardless of width. At $d_k = 256$ the raw spread is $16$ , so a gap of one standard deviation becomes $2$ after dividing, and the weight ratio is $e^{2} = 7.39$ . At $d_k = 16$ the raw spread is $4$ , the same standard-deviation gap becomes $0.5$ , and the ratio is $e^{0.5} = 1.65$ . The head is saturated at one width and washed out at another, and the constant that was correct at $d_k = 64$ is wrong at both. Dividing by $\sqrt{d_k}$ gives a gap of exactly $1$ in both cases and a ratio of $e = 2.72$ . The scale factor has to track the width because the standard deviation does; no single number does that.

The damage from getting it wrong is not confined to ugly weights. The derivative of softmax output $a_i$ with respect to input $s_j$ is a product of probabilities, so it is largest when the weights are moderate and collapses when they are extreme. On the diagonal the derivative is $a_i(1-a_i)$ : at the scaled top weight of $0.43$ that is $0.43 \times 0.57 = 0.25$ , while at the unscaled $0.98$ it is $0.98 \times 0.02 = 0.02$ . Off the diagonal it is $-a_i a_j$ , so the gradient reaching the third token's score is $-0.43 \times 0.16 = -0.068$ when scaled against $-0.98 \times 0.0003 = -0.0003$ when not, roughly two hundred times smaller. Almost no signal reaches the projections, and the head cannot learn to attend differently because the confidence has switched off the mechanism by which it would learn.

FIG 3Why a saturated softmax stops learning

: off the diagonal this is minus the product of two probabilities, tiny unless both are moderate
: on the diagonal, largest at a_i = 1/2 and tending to zero at both 0 and 1

Every partial derivative is a product involving the output probabilities, so it collapses at both ends of the range and peaks in the middle.

What the derivation assumes

The argument needed two things: that query and key components are roughly independent, and that they have roughly unit variance. Neither is exactly true in a trained network. Learned projections introduce correlations between components, and the actual variance depends on where the block sits and how it is normalised. If the components have variance $\sigma^2$ rather than $1$ , the score variance becomes $d_k \sigma^4$ and the correct divisor is $\sigma^2\sqrt{d_k}$ , so a fixed $\sqrt{d_k}$ under-corrects as the variance drifts upward with depth. That is why the placement of layer normalisation changes how stable training is, and why variants that normalise queries and keys directly before the dot product exist at all: both are repairs to the assumption rather than to the algebra. The constant is worth knowing precisely because its assumptions are visible; when they fail, the constant is the part you adjust.

Dividing by $\sqrt{d_k}$ leaves the score distribution with unit spread whatever the head width, which keeps softmax graded, which keeps gradients alive. A wide head without it does not attend badly so much as commit to an answer before it has learned anything.

Worth keeping

The variance of an unscaled attention score is the key dimension, so its standard deviation is the square root of that dimension.
Softmax is exponential in score gaps, so multiplying all scores by a constant raises every weight ratio to that constant and changes sharpness without changing order.
When one softmax weight approaches 1, every partial derivative is a product of probabilities and collapses, so an unscaled head stops receiving gradient.

Where this goes next

This question is one step of The Mathematics of Attention, a course that works through its subject from the beginning. The page that teaches this part of it is The Scaling Factor, and it reads the whole way through whether or not you have an account.

What an account changes is the practice: problems standing under every idea, marked with a reason rather than a tick, in a copy of the course that is yours.

Read that page The whole course