Why does attention scale the logits by √d_k?
What interviewers are testing
Interviewers use this question to separate candidates who memorized the Transformer diagram from those who understand why the architecture resists saturation. The tell is whether you can connect d_k, the variance of a dot product, and the softmax Jacobian into one causal chain instead of reciting "to keep gradients stable". Follow-ups about QK normalization, temperature, and numerically stable softmax reveal whether the mental model generalizes beyond the original paper.
Mental model
Every attention logit is a dot product of two roughly independent random vectors, so its variance grows linearly with d_k and its standard deviation grows with √d_k. Large logits push softmax into saturation, where the largest weight approaches 1, the rest approach 0, and the Jacobian collapses so almost no gradient reaches the queries and keys. Dividing by √d_k restores unit variance and keeps softmax in its responsive, well-conditioned regime.
Step-by-step solution
Step 1 of 5
Variance grows with d_k
Start with a single attention logit: the dot product of one query row against one key column, written as a sum of d_k products. If query and key entries are roughly independent with zero mean and modest variance, the variance of the sum grows linearly with d_k, so its standard deviation grows with the square root of d_k. At d_k = 64 the logits already have a standard deviation near eight; at d_k = 512 it approaches twenty-three. Softmax is exponential, so it amplifies exactly the gaps the scaling leaves behind. Drag the slider in the animation to change d_k and watch the unscaled logit histogram widen while the largest softmax weight climbs toward one. Dividing the logits by √d_k rescales their variance back to roughly one, which keeps softmax in the range where it can represent differences instead of saturating on them.
Animation — Variance grows with d_k
Why divide attention scores by √d_k?
Each attention score is a sum of d_k products, so its standard deviation grows with √d_k. Dividing by √d_k keeps the logits at σ = 1.
softmax over logits with σ ≈ √d_k saturates: nearly all weight collapses onto one token and gradients vanish.
Edge cases & traps
- Scaling by the model dimension instead of the per-head dimension: the denominator must be d_k = d_model / num_heads, or the logits are shrunk too aggressively and attention flattens toward uniform.
- Applying the scale after softmax: dividing the weights by √d_k breaks their sum-to-one constraint — the factor belongs on the logits before the softmax.
- Trusting the numerically stable softmax (row-max subtraction) to fix saturation: it prevents overflow but leaves the relative logit gaps untouched, so one-hot weights and vanishing gradients remain.
- Hard-coding the scale from sequence length: the factor depends only on d_k and is constant per head — recomputing it from context length is a conceptual bug.
- Dismissing the issue because small heads still train: the effect is mild for tiny d_k, which is why the failure only shows up as head size or depth grows.