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

TL;DR: Dot products of two d_k-dimensional vectors have standard deviation that grows like √d_k. Left unscaled, the logits saturate the softmax into a near one-hot distribution where its gradient vanishes, so Q and K stop learning. Dividing by √d_k renormalizes the logit standard deviation back to ~1 and keeps the softmax in its responsive regime.

How to approach it. Decide quickly whether they want the statistics or the training-dynamics consequence, then give both: the variance derivation, then what actually fails in a run without it. State the assumption out loud (query and key components roughly independent, zero mean, unit variance) because the whole argument rests 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.