Why Scaled Dot-Product Attention Divides by the Square Root of the Key Dimension

The attention formula carries a constant that looks arbitrary until it is derived:

$$\mathrm{Attn}(Q,K,V) = \mathrm{softmax}!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$$

The division by $\sqrt{d_k}$ is not a learning-rate-like hyperparameter that happened to work. It is a variance correction, and the argument for it is short enough to state in three lines. It also explains a class of failures that appear when the factor is removed, retuned, or applied inconsistently across layers.

A dot product accumulates variance linearly in dimension

Take a single query row $q$ and a single key row $k$, both of length $d_k$. Suppose their entries are independent, zero mean, and unit variance. This is not an exotic assumption: after a linear projection of normalized activations, entries are approximately centered with a spread that the projection’s initializer was chosen to control. Under it,

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

Each term $q_i k_i$ has mean zero and variance one, because the product of two independent unit-variance zero-mean variables has variance $1 \cdot 1 = 1$. The sum of $d_k$ independent such terms therefore has variance $d_k$ and standard deviation $\sqrt{d_k}$.

With $d_k = 64$, a common per-head width, the logits entering the softmax have a standard deviation around 8. Across a row of $n$ keys, the largest logit sits roughly $\sigma\sqrt{2\ln n}$ above the mean. For $n = 512$ that is about $3.5\sigma$, so the top logit exceeds the mean by nearly 28. The softmax of such a row is a near one-hot vector: the largest probability is within a rounding error of 1 and the rest are $e^{-28}$ or smaller.

Dividing the logits by $\sqrt{d_k}$ restores unit standard deviation, which keeps the softmax operating in the regime where several keys receive non-negligible mass. The factor is exactly the standard deviation of the unscaled product, so the scaled version has the same distribution regardless of head width. That invariance is the point: a model can change $d_k$ without silently changing the effective temperature of every attention layer.

import numpy as np

rng = np.random.default_rng(0)
for d_k in (16, 64, 256):
    q = rng.normal(size=(1, d_k))
    k = rng.normal(size=(512, d_k))
    raw = (q @ k.T)[0]                    # 512 logits, variance ~ d_k
    scaled = raw / np.sqrt(d_k)           # variance ~ 1 regardless of d_k
    p_raw = np.exp(raw - raw.max()); p_raw /= p_raw.sum()
    p_sc = np.exp(scaled - scaled.max()); p_sc /= p_sc.sum()
    print(d_k, raw.std(), scaled.std(), p_raw.max(), p_sc.max())
# raw std grows as sqrt(d_k); p_raw.max approaches 1; p_sc.max stays moderate

Saturation shows up as a vanishing gradient, not as a wrong answer

A saturated attention row still produces a defensible output. If one key dominates, the layer returns that key’s value vector, which is a plausible thing to do. The damage is in the derivative.

The Jacobian of the softmax with respect to its logits is

$$\frac{\partial p_i}{\partial z_j} = p_i(\delta_{ij} - p_j)$$

Every entry of this matrix carries a factor of $p_i$. When the row is one-hot, the winning entry has $p_i \approx 1$ and all others have $p_i \approx 0$. The diagonal entry for the winner is $p_i(1-p_i) \approx 0$; every off-diagonal entry is a product of two numbers, at least one of which is near zero. The entire Jacobian collapses toward the zero matrix.

Since the logits are $q \cdot k / \sqrt{d_k}$, a vanishing Jacobian with respect to the logits means a vanishing gradient with respect to both $q$ and $k$. The value projection and the output projection still receive gradient, because $V$ enters the formula linearly after the softmax. So the layer continues to learn how to mix values while the mechanism that decides which values to mix stops moving. In a stack, that leaves the queries and keys frozen at their initialization scale while everything downstream adapts around them.

The same argument runs in reverse. If the logits are too small, through a large $d_k$ with unit-variance inputs and no correction or through a temperature that is too high, the softmax approaches the uniform distribution $1/n$ over the row. The output becomes the mean of all value vectors, identical for every query position. The layer is then a constant function of the query, and its gradient with respect to $q$ is small for a different reason: the distribution is flat, so moving a logit slightly changes nothing. Both ends of the temperature range are degenerate, and the scaling factor is what keeps the operating point away from both.

There is a useful way to see the two regimes together. The entropy of the row distribution is maximal at $\ln n$ when the logits are equal and falls toward zero as the spread grows. The scaled dot product does not fix the entropy at any particular value, because the actual spread of $q$ and $k$ depends on the layer’s weights. It fixes the scale of the logits relative to what the weights would otherwise produce, so that a single initialization scheme behaves similarly at every width.

The correction assumes more than the formula admits

The derivation rests on three assumptions that real layers violate, and knowing which ones matters more than knowing the derivation.

The first is independence between $q$ and $k$. In self-attention both come from the same token stream through different projection matrices, so their entries are correlated through the shared input. A positive correlation between corresponding coordinates inflates the variance of the product beyond $d_k$; a negative one deflates it. Cross-attention between unrelated sequences is closer to the independent case. The correction is therefore a first-order fix, accurate when the projections decorrelate the two sides and approximate when they do not.

The second is that entries have unit variance. Nothing enforces this. Layer normalization placed before the projections makes the input to each projection roughly unit-variance per coordinate, and a projection initialized with variance $1/d_{\text{model}}$ preserves that, so the assumption holds at initialization and drifts during training. The drift is not necessarily harmful: the learned temperature can be absorbed into the norms of the query and key weights, since scaling $W_q$ by a constant $c$ multiplies every logit by $c$. The scaling factor sets a starting point; gradient descent is free to move away from it. This is why removing the factor does not always destroy a model: the weights can compensate, given enough training signal and a warmup schedule gentle enough to survive the initial saturation.

The third assumption is that $d_k$ is the right dimension to normalize by. The formula uses the per-head width, not the model width. A model with $d_{\text{model}} = 512$ and 8 heads has $d_k = 64$, so the divisor is 8 rather than roughly 22.6. Using the model width instead would over-flatten the logits by a factor of $\sqrt{8}$, which is a large change in effective temperature. Implementations that parameterize attention with a single projection of shape $d_{\text{model}} \times d_{\text{model}}$ and then split it into heads must still divide by the per-head width after the split, or the temperature changes with the head count.

Reading the factor as a temperature

The cleanest mental model is that the whole expression is a softmax with a learned, dimension-dependent temperature:

$$\mathrm{softmax}\left(\frac{q \cdot k}{T}\right), \qquad T = \sqrt{d_k}$$

Writing it this way makes the interaction with other design choices visible. Doubling the head count at fixed model width halves $d_k$ and therefore reduces $T$ by $\sqrt{2}$, sharpening every head. Doubling the model width at fixed head count doubles $d_k$ and softens every head. Neither change is neutral for optimization, which is why head count and width are usually scaled together and why per-head width stays within a narrow band rather than growing with the model.

It also makes the relationship to the rest of the block explicit. The residual stream carries activations whose scale is controlled by normalization, and the attention output is added back into that stream. A saturated attention layer returns a value vector whose norm is that of a single token’s projection rather than an average of many, so the scale of the block’s contribution changes with the temperature even though the residual path does not. Keeping the logits near unit variance keeps that contribution’s scale stable across layers, which is the same reason normalization is applied before the projections rather than after.

Failure modes at the boundaries

Three cases are worth naming because they appear as bugs rather than as bad hyperparameters.

A row with no valid keys produces a uniform distribution over an empty set, which is undefined. Masking a full row to negative infinity makes every logit $-\infty$; the softmax then computes $0/0$ and returns NaN, which propagates through the whole batch on the next backward pass. Implementations avoid this either by substituting a large finite negative value or by detecting fully masked rows and zeroing their output. The trade-off is that a large finite value is not a true zero: with a mask value of $-10^{9}$ in 16-bit arithmetic, the value overflows the representable range and behaves as $-\infty$ anyway, while in 32-bit it produces a probability of exactly zero only by underflow.

Mixed-precision arithmetic changes where the saturation boundary sits. The softmax is usually computed in 32-bit even when the projections run in 16-bit, because the exponential of a logit near 30 is not representable in half precision and the running maximum subtraction is what keeps the exponent argument small. An implementation that computes the softmax entirely in 16-bit will overflow for logits above roughly 11, which corresponds to a modest $d_k$ and unremarkable weight norms. This is a precision constraint on the mechanism, not a numerical nicety.

Finally, the factor is sometimes folded into the query projection at initialization by scaling $W_q$ by $d_k^{-1/4}$ so that the runtime division disappears. That is algebraically identical and saves one multiply per score, but it changes what the optimizer sees: the gradient with respect to $W_q$ is scaled by the same constant, so a learning rate tuned for the unscaled parameterization is no longer equivalent. Folding the constant is safe when the learning rate is set afterwards; folding it into a pretrained checkpoint and continuing training at the original rate is not.

The mechanics of what a single head can express once its logits are well conditioned are covered in multi-head attention as subspace projection, and the masking that decides which logits are allowed to compete at all is the subject of causal masking and mask placement.