Md. Asif Uddin

Proposition 322 of 39 in the corpus

The divisor in scaled dot-product attention is what keeps the gradient alive.

Dot products of d-dimensional vectors have variance proportional to d. Left unscaled the logits grow with head width, the softmax saturates, and the gradient through it goes to zero.

Depends on

The same attention logits, scaled and unscaledTwo bar charts of attention weights over eight keys. Scaled, the weights are spread across several keys. Unscaled, almost all the weight lands on one key and the remaining bars are invisible, which is where the gradient dies.scaled by 1/√dk1k2k3k4k5k6k7k8max weight 27%unscaled — logits grow with dk1k2k3k4k5k6k7k8max weight 80%For unit-variance q and k the dot product has variance d, so the logits grow with the head width.Dividing by √dₖ restores it, whatever the width.
Fig. 3 — The same attention logits, scaled and unscaled. Without the divisor the distribution collapses onto one key and the gradient through the softmax goes flat.

Demonstration

Take q and k with independent components of zero mean and unit variance. Their dot product is a sum of d such products, so it has mean zero and variance d, and typical magnitude √d. For a head width of 64 that is around 8; for 128, around 11.

Feed logits of that size to a softmax and the distribution collapses. The largest logit is several units above the rest, exp amplifies the difference exponentially, and the resulting weights are essentially one-hot. That in itself would only be a modelling limitation. The fatal part is the gradient: the Jacobian of the softmax is

∂pᵢ/∂zⱼ = pᵢ(δᵢⱼ − pⱼ)

which is proportional to p(1 − p) on the diagonal. When p is near one or near zero, that product is near zero — the softmax is saturated, and almost no gradient flows back to WQ and WK. The attention pattern is frozen at whatever the initialisation produced.

Dividing by √dk restores unit variance to the logits regardless of head width, and the distribution stays soft enough to receive gradient. It is a one-character change with the property that omitting it does not raise an error, does not produce NaN, and merely results in a model that trains badly for reasons invisible in the loss curve.

The same reasoning explains why head width and head count interact. Splitting a budget of 512 into eight heads of 64 rather than one head of 512 does not only buy multiple attention patterns — it also lowers dk, and with it the variance of the logits before scaling.

Corollary

Temperature is this same knob, exposed at inference. Dividing logits by a value greater than one flattens the distribution; less than one sharpens it. The softmax has no intrinsic scale, so wherever one appears in an architecture, somebody has chosen its temperature — explicitly or by accident.

Sources