Why do we scale the dot-product attention scores by the square root of d_k?

How to think about it

We divide the dot product q · k by sqrt(d_k) before softmax, the function that turns raw scores into weights summing to one. Here q is a query vector, k is a key vector, and d_k is their number of coordinates; without the division, the score spread grows with sqrt(d_k), softmax saturates toward one-hot weights, and useful learning signals become weak.

Why the dot product grows with dimension

An attention score measures how well one query matches one key:

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

Imagine, as a simple initialization model, that every coordinate of q and k has mean zero and variance one, and that coordinates are independent. Variance is a measure of squared spread around the mean. Each product q_i k_i then has variance roughly one. Adding d_k independent products adds their variances:

Var(q · k) = d_k

The standard deviation, which is the typical distance from the mean, is therefore sqrt(d_k).

The expected dot product is still zero. That is the detail people often skip. The score does not necessarily become more positive as the vector gets wider; its typical absolute size becomes larger.

For d_k = 64, the unscaled dot product has a standard deviation of about 8. Dividing by sqrt(64), which is 8, brings that standard deviation back to about 1. For d_k = 256, the corresponding figures are 16 before scaling and 1 after scaling.

The exact variance in a trained model is not guaranteed to equal d_k. Learned projections, correlations, normalization, and changing activation scales all affect it. The point is to remove the predictable dependence on vector width.

What goes wrong inside softmax

Softmax is sensitive to differences between scores. Large differences produce a nearly one-hot distribution: one key receives almost all the attention and the others receive almost none.

Consider one query comparing three keys. Suppose the raw scores are:

Scores sent to softmax	Attention weights
[8, 0, -8]	[0.9997, 0.0003, 0.0000001]
[1, 0, -1]	[0.6652, 0.2447, 0.0900]

The first row is what can happen when a typical score scale is around eight. The second row is the same score pattern after division by sqrt(64) = 8.

Softmax can represent a sharp choice when the model has earned one. The problem is becoming sharp before the model has learned which key is correct. At that point, small changes to most scores barely change the weights.

For a softmax probability p_i, the derivative of that probability with respect to its own logit is p_i(1 - p_i). At p_i = 0.9997, that derivative is only about 0.0003. The full softmax Jacobian has similarly small terms around a saturated distribution. Learning signals that need to redistribute attention can therefore become weak.

This does not mean every possible loss has exactly zero gradient. For example, cross-entropy can still produce a direct signal for a badly ranked target. The practical issue is that the softmax mapping itself has become poorly conditioned: changing query and key representations produces very little change in most attention weights.

Where the factor appears

Scaled dot-product attention is usually written as:

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

Q is a matrix of query vectors, K is a matrix of key vectors, and V is a matrix of value vectors. The matrix product QK^T produces one score for every query-key pair. Softmax is applied across the keys for each query. The resulting weights select and mix rows of V.

A standard PyTorch-style implementation looks like this:

import math
import torch

# Q: [batch, heads, query_length, d_k]
# K: [batch, heads, key_length, d_k]
# V: [batch, heads, key_length, value_dim]

scores = Q @ K.transpose(-2, -1)
scores = scores / math.sqrt(d_k)
weights = torch.softmax(scores, dim=-1)
output = weights @ V

In a real Transformer, a causal or padding mask is also applied to scores before softmax. The important order is that the dot product is scaled before softmax sees it.

In multi-head attention, d_k is normally the key dimension of one head, not the total model width. If the model width d_model is 512 and there are 8 equally sized heads, each head commonly has d_k = 64, so the denominator is 8. The number of tokens is unrelated to this denominator.

Why the square root, not d_k

The square root is the standard-deviation correction:

Var(q · k) = d_k

and therefore:

Var((q · k) / sqrt(d_k)) = 1

under the simplifying assumptions above.

Dividing by d_k would overcorrect. The resulting variance would be 1 / d_k, and the standard deviation would be 1 / sqrt(d_k). With d_k = 64, that means a standard deviation of only 0.125. Scores would cluster close to zero, making attention too close to uniform. The model would struggle to focus on the useful tokens.

So the choice is not arbitrary. Dividing by the square root preserves an order-one score scale. Dividing by the full dimension makes the scores unnecessarily small.

The senior-level nuance

This scaling is a variance-control heuristic, not a promise that every attention logit will have variance exactly one. In trained networks, query and key coordinates may not be independent or unit variance. Even so, the correction is a strong baseline because it prevents the score scale from growing merely because a head has more coordinates.

It also acts like a temperature. Softmax with logits divided by a temperature T becomes smoother as T increases. Here, sqrt(d_k) is an implicit temperature that grows with head width. Larger heads therefore do not automatically produce much sharper distributions just because they contain more terms in their dot products.

Do not confuse this with normalizing every query and key to unit length. Cosine attention removes the raw growth of the dot product by constraining its range, but it changes the geometry: vector magnitude can no longer express part of the match strength. Cosine-based systems often need a learned temperature to avoid distributions becoming too flat.

Likewise, additive attention uses a learned function to produce a score rather than directly summing coordinate-wise products. It has a different scale-control mechanism. Applying the Transformer formula blindly to every attention variant is cargo cult with better notation.

Failure modes you can diagnose

If the scaling is missing, the first visible symptom is often attention entropy near zero from the beginning: each head points almost entirely at one token, often an arbitrary one. Training may plateau, become sensitive to initialization, or degrade when the head dimension increases. A naïve exponential implementation may also encounter overflow with extreme logits, although numerically stable softmax implementations usually subtract the largest score before exponentiating.

If the score is divided by d_k instead of sqrt(d_k), the symptom reverses. Attention weights remain diffuse and similar across queries. Outputs become bland weighted averages, and the model may underfit because it cannot select relevant keys sharply enough.

Another common bug is scaling twice, once inside an attention helper and again in a wrapper. The resulting weights are usually too uniform. Scaling after softmax is not equivalent: saturation has already happened, and changing the weights afterward does not restore the lost softmax sensitivity.

In production, inspect the distribution of pre-softmax logits and the spread of attention weights by head. When changing the number or size of heads, check whether the score scale and attention entropy change dramatically. Those checks catch this bug faster than staring at the final loss curve.

What they will ask next

Is d_k the same as d_model?

Not necessarily. d_model is the total hidden width, while d_k is the query and key width used for one attention head. In a single-head design they may be equal. In multi-head attention, d_k is often d_model / number_of_heads, but the architecture can choose other dimensions.

Does scaling guarantee that the logits have variance exactly one?

No. Variance one follows only from the simplified assumptions about independent, mean-zero, unit-variance coordinates. In a trained model, the actual variance can differ. The scaling removes the predictable growth with d_k; it does not replace monitoring or normalization.

What if the model uses cosine attention or a learned temperature?

Then the model may already have a different mechanism for controlling score scale. Adding the square-root factor can make the distribution unnecessarily flat. I would inspect the attention definition and the learned temperature before applying the standard Transformer scaling.

Say this in the interview: We divide query-key dot products by sqrt(d_k) because summing d_k coordinate products makes their standard deviation grow as sqrt(d_k); the division keeps softmax out of saturation, so attention remains selective without becoming one-hot before training has learned what to attend to.