Why do transformers scale attention scores by 1/√d_k, and what breaks if you skip it?

TL;DR: The dot product of two d_k-dimensional vectors carries a standard deviation that rises with √d_k. Without scaling, those logits push the softmax toward a near one-hot output where the gradient disappears, so Q and K stop updating. Dividing by √d_k pulls the logit standard deviation back to roughly 1 and holds the softmax in the range where it stays responsive.

How to approach it. Work out fast whether they are asking for the statistics or the training-dynamics fallout, then cover both: first the variance derivation, then what genuinely breaks in a run that omits it. Say the assumption aloud (query and key components roughly independent, zero mean, unit variance) since the entire argument depends on it.

A strong answer. Attention computes softmax(QKᵀ/√d_k)·V. Take one query-key score q·k = Σ_{i=1}^{d_k} q_i k_i. If the components are independent with mean 0 and variance 1, each term q_i k_i has mean 0 and variance 1, so the sum has mean 0 and variance d_k. The standard deviation therefore scales as √d_k. For d_k = 128, raw logits sit around ±11. Push numbers that large through a softmax and one entry dominates: the output is effectively one-hot. In that saturated region the softmax Jacobian diag(p) − ppᵀ collapses toward zero, so gradients to Q and K nearly vanish exactly early in training when you need them most. Dividing by √d_k rescales the logit standard deviation back to ~1, keeping logits near ±3 where softmax is smooth and informative. This is the original "Attention Is All You Need" rationale.

The variance bookkeeping is the whole argument, so it is worth seeing the chain in one place:

import torch, torch.nn.functional as F
def attention(q, k, v, mask=None):           # q,k,v: (B, heads, T, d_k)
    d_k = q.size(-1)
    scores = q @ k.transpose(-2, -1) / d_k ** 0.5
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))
    return F.softmax(scores, dim=-1) @ v

The insider point: a strong candidate names the failure mode (vanishing gradients through a saturated softmax), not just "it normalizes things."

Key takeaways

• The score variance grows like d_k, so the standard deviation grows like √d_k. You scale by the standard deviation, which is why the constant is √d_k and not d_k.
• The real failure is a saturated softmax with a near-zero Jacobian: gradients to Q and K vanish early in training, not a vague "instability."
• The whole derivation depends on the zero-mean, unit-variance, independent-component assumption. State it, or the variance does not equal d_k.
• Large raw logits also overflow in fp16/bf16; the scale plus softmax max-subtraction (done streaming in FlashAttention) keeps it numerically safe.

What interviewers probe next.

• "Why √d_k and not d_k?" You normalize the standard deviation, which grows like √d_k, not the variance. Dividing by d_k would over-shrink the logits and flatten attention.
• "Does it matter at inference?" The scale is a fixed constant folded into the math, so it is a training-stability concern primarily, but train and serve must use the identical scale or you ship a different function than you trained.
• "How does this interact with fp16/bf16?" Large pre-softmax logits also overflow in low precision. The √d_k scale plus the standard max-subtraction inside softmax keeps it numerically safe; FlashAttention does this max-subtraction in a streaming, IO-aware way.

Common mistakes.

• Saying the scale "normalizes the attention weights." The softmax does that. The scale controls the magnitude of the logits feeding the softmax.
• Quoting d_k instead of √d_k, or being unable to justify the square root.
• Dropping the independence and unit-variance assumption, which is the only reason the variance equals d_k.