Scaled dot-product attention & the √d_k

1

Concept

Self-attention begins with a simple compatibility measure: the dot product between a query and a key. If two vectors point in similar directions, their dot product tends to be positive and large. If they are poorly aligned, it is smaller or negative. For a query matrix $Q$ and key matrix $K$, all pairwise scores appear in $QK^\top$. But the raw score scale changes when the vector dimension changes.

Write one score as

For an intuition used in the original Transformer paper, suppose components are independent, have mean zero, and variance one. Each product $q_i k_i$ then has variance near one. Adding $d_k$ such terms gives $\operatorname{Var}(s) \approx d_k$, so the score’s standard deviation grows like $\sqrt{d_k}$. A 256-dimensional head therefore produces wider raw score variation than a 16-dimensional head even when the component statistics are otherwise comparable.

Attention sends each row of scores through softmax. Softmax exponentiates its inputs, so scale matters dramatically. Scores $[0.2,0.1,0]$ produce a fairly distributed result. Scores $[20,10,0]$ produce an almost winner-take-all result. The ranking is identical, yet the second distribution behaves as if the model were extremely certain. When probabilities approach zero or one, many useful gradients become small, which can make learning brittle.

Scaled dot-product attention corrects this dimension effect:

Dividing by $\sqrt{d_k}$ approximately normalizes the score standard deviation under the assumptions above. It does not force exact variance one in a trained network, and it does not normalize each query or key vector. Learned projections, normalization layers, correlations, and real activations violate the tidy independence story. The scaling is still a robust design choice because it removes the most obvious growth with dimension.

The divisor can also be viewed as a fixed temperature. Softmax temperature is usually written $\operatorname{softmax}(s/T)$. Here $T=\sqrt{d_k}$. A larger temperature flattens a distribution; a smaller one sharpens it. In attention the scale is tied to head dimension rather than exposed as a decoding control. A model can still learn sharper or flatter patterns by changing query and key norms.

Numerical stability adds another layer. Practical softmax implementations subtract the row maximum before exponentiating:

Subtracting a constant does not change the ratio of probabilities, but it prevents a large positive score from overflowing the exponential. This technique solves numerical overflow; it does not solve saturation caused by score gaps. The $\sqrt{d_k}$ scaling addresses the latter source at initialization and during early training.

Consider scores from one query against three keys. If the raw dot products are $[8,4,0]$ and $d_k=64$, scaling produces $[1,0.5,0]$. Softmax can assign meaningful mass to all three, allowing gradients to adjust each relationship. Without scaling, the first key would dominate much more strongly before the model had evidence that it deserved such certainty.

The lesson is broader than memorizing a denominator. Whenever a sum collects many roughly independent contributions, its typical magnitude often grows with the number of terms. Passing that sum into a sensitive nonlinearity can change behavior merely because width changed. The Transformer compensates with a scale derived from variance, keeping the routing mechanism in a trainable regime.