Why do we scale by sqrt(d_k) in scaled dot-product attention?

How to think about it

Short answer: We divide by sqrt(d_k) because a query-key dot product gets a standard deviation of roughly sqrt(d_k), where d_k is the number of components in each key and query vector within one attention head. The division keeps the pre-softmax scores, or logits, at roughly the same scale as the head dimension changes, so softmax does not become accidentally too sharp.

Why the dot product grows

An attention score is a dot product between a query and a key. A query represents what the current token is looking for; a key represents what a candidate token offers for matching. For one pair, write the score as:

q · k = q_1 k_1 + q_2 k_2 + ... + q_{d_k} k_{d_k}

Assume, as a useful initialization model, that each q_i and k_i has mean zero and variance one, and that the components are independent.

Each product q_i k_i then has mean zero and variance one. The products add, so their variances add:

Var(q · k) = d_k

Variance is not the same thing as typical size. Standard deviation is the square root of variance, so:

Std(q · k) = sqrt(d_k)

That square root is the whole reason for the denominator. The terms do not all point in the same direction; positive and negative terms partly cancel. If they did not cancel, the mean could grow more like d_k. Under the usual zero-mean assumption, the spread grows like sqrt(d_k).

Head dimension d_k	Raw score standard deviation	After dividing by sqrt(d_k)
16	4	1
64	8	1
256	16	1
512	about 22.6	1

Without scaling, changing the head dimension silently changes the effective temperature of attention. A model with 512-dimensional heads does not merely have larger vectors. It has much wider logits entering the same softmax.

Why wide logits are a problem

Softmax turns scores into nonnegative weights that sum to one:

p_j = e^{s_j} / sum_l e^{s_l}

The exponential makes softmax very sensitive to score differences. A gap of 1 gives a ratio of about 2.7 between two scores. A gap of 10 gives a ratio of about 22,026.

For example:

softmax([0, 10, 0, 0])

is approximately:

[0.000045, 0.999864, 0.000045, 0.000045]

The largest score receives virtually all the attention. That may look decisive, but at initialization it usually means the model is choosing a nearly arbitrary winner based on random projections.

The gradient problem is more precise than “softmax stops learning.” The softmax Jacobian contains terms of the form:

∂p_i / ∂s_j = p_i (delta_ij - p_j)

where delta_ij is one when the two indices match and zero otherwise. If one probability is nearly one and the rest are nearly zero, these derivatives are tiny. Changing a losing score barely changes its already negligible probability. Consequently, the path from the attention output back to the query and key projections carries a weak signal.

One important interview nuance: not every gradient in the entire network becomes exactly zero. If softmax is used directly with cross-entropy, the derivative with respect to a logit can still be p - y, which may be large for a wrong prediction. The narrower claim is that a saturated attention softmax has a nearly flat response to changes in its scores. That makes learning the query-key matching function poorly conditioned.

Dividing by sqrt(d_k) restores the score standard deviation to roughly one. It does not force attention to be uniform. It simply prevents the vector dimension from deciding how sharp the initial distribution will be.

A concrete attention example

Imagine a query from the phrase “forgot my password” comparing four candidate keys: password, account, reset, and today. Suppose the unscaled dot products in a head with d_k = 64 are:

[8, 4, 0, -4]

Without scaling, softmax gives approximately:

[0.9817, 0.0180, 0.00033, 0.000006]

The model has effectively committed to password. The other three positions receive almost no attention and provide almost no useful competition during training.

Since sqrt(64) = 8, scaling gives:

[1, 0.5, 0, -0.5]

Softmax now gives approximately:

[0.4551, 0.2760, 0.1674, 0.1015]

The correct candidate is still preferred, but the model can adjust all four scores. That is a much healthier starting point for learning.

import math
import torch
import torch.nn.functional as F

raw = torch.tensor([8.0, 4.0, 0.0, -4.0])
d_k = 64

unscaled = F.softmax(raw, dim=0)
scaled = F.softmax(raw / math.sqrt(d_k), dim=0)

print(unscaled)
print(scaled)

The scale is applied to the query-key scores, not to the values. In the usual pattern, the computation is conceptually:

Attention(Q, K, V) = softmax((Q K^T) / sqrt(d_k)) V

Q contains queries, K contains keys, and V contains the value vectors whose weighted combination becomes the output.

A causal or padding mask is normally applied after scaling and before softmax. Disallowed positions receive a very negative score, effectively negative infinity, so they receive zero attention.

The production nuance

The denominator is based on the per-head key dimension, not automatically on the model’s total hidden size.

Suppose d_model = 512 and there are 8 heads. Each head typically has:

d_k = 512 / 8 = 64

The correct scale is therefore sqrt(64) = 8.

Using sqrt(512), which is about 22.6, would over-scale the scores. If the raw per-head score standard deviation is about 8, dividing by 22.6 leaves a standard deviation of only about 0.35. Attention becomes too flat. With four candidates, the weights may be much closer to uniform, and the head has difficulty selecting relevant positions.

The reverse mistake is dividing by d_k rather than sqrt(d_k). For d_k = 64, that would reduce an expected standard deviation of 8 to only 0.125. The logits would be so small that softmax would be nearly uniform. You would have traded saturated attention for indecisive attention.

The derivation is also an approximation. If query components have variance sigma_q^2 and key components have variance sigma_k^2, the score variance is roughly:

d_k sigma_q^2 sigma_k^2

Scaling by sqrt(d_k) removes the dimension-dependent part, but it does not fix exploding query or key variance. Layer normalization, initialization, residual connections, and training dynamics still matter.

This is also why the factor is not a universal law for every attention variant. In cosine attention, queries and keys are usually normalized by their vector norms, so their dot product is bounded between negative one and one. Applying the standard sqrt(d_k) divisor on top of that can make logits too small. Such systems often use a separate fixed or learned temperature. The right question is always: what is the scale of the scores entering softmax?

A failure mode you can recognize

If the scaling is accidentally removed, inspect the attention weights from the first few batches. A common symptom is that many heads have a maximum attention weight around 0.99 or 1.0 immediately, often pointing to different and seemingly arbitrary tokens. Attention entropy, a measure of how spread out the weights are, is unusually low.

For comparison, a perfectly uniform distribution over 128 tokens has entropy log(128), about 4.85 nats. A one-hot distribution has entropy zero. Real attention need not sit at either extreme, but an abrupt near-zero entropy pattern at initialization is a strong clue that the logits are too large.

If the opposite bug is present and the code divides by the wrong, much larger dimension, the first symptom is often the reverse: attention weights remain almost uniform, heads fail to focus, and training improves slowly. In both cases, check the actual per-head logit standard deviation rather than guessing from the model name.

What they’ll ask next

Why not divide by d_k?

Because the dot product’s variance grows as d_k, while its standard deviation grows as sqrt(d_k). Dividing by d_k makes the post-scaling standard deviation 1 / sqrt(d_k), which becomes too small as the head gets wider and pushes attention toward uniform weights.

Why use the head dimension instead of the total model dimension?

Each attention head computes its own dot products using vectors of width d_k. The variance calculation applies to those vectors. With a 512-dimensional model split into eight heads, each head has dimension 64, so the denominator is 8 rather than about 22.6.

Does scaling guarantee that gradients will not vanish?

No. It reduces one predictable source of softmax saturation. A trained model can still produce very large score gaps, and gradients can also vanish because of masking, poor initialization, numerical issues, or other parts of the network. Scaling gives every head a dimension-independent starting temperature; it is not a guarantee of healthy optimization.

Say this in the interview

“Because a zero-mean query-key dot product has variance d_k and standard deviation sqrt(d_k), dividing by sqrt(d_k) keeps softmax logits at a stable scale and prevents attention from becoming dimension-dependent and saturated.”