Why is the dot product scaled by √d_k in attention? Derive the variance argument.
Why is the attention dot product scaled by √d_k? Walk through the variance argument: what is the distribution of QKᵀ values without scaling, and what specifically goes wrong if you skip the scaling?
QK^T entries are sums of d_k unit variance products, so variance is d_k and std is √d_k. Big scores saturate softmax and kill off-peak gradients. Divide by √d_k to fix it.
Imagine adding up a long list of random numbers. The longer the list, the wider the typical sum becomes, the spread grows roughly like the square root of how many numbers you added. Attention scores work the same way. Bigger `d_k` means wider, more extreme scores. When those scores feed into a "pick a winner" step, one score being far above the rest means that winner takes everything. The learner has no signal that would shift the win to a different choice, because the other choices have basically zero probability. Dividing by `√d_k` undoes the growth, keeping scores in a friendly range where learning still works.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
The √d_k scaling looks like a tiny detail in the attention formula: one division, easy to skim past. It is actually critical for training stability, and the reason ties together three things worth being able to derive on demand: the variance of a sum of independent products, the saturation behavior of softmax, and the structure of the softmax Jacobian.
The failure mode without scaling is not numerical overflow. The dot products themselves are perfectly representable in standard precision (FP16/BF16). The failure is that gradients die. Attention scores grow large with d_k, softmax saturates onto its argmax, and the Jacobian off the peak collapses to zero. Once that happens, the optimizer has no signal to move the attention pattern, and training stalls silently, with a loss curve that looks fine for a while before refusing to improve.
We will walk the variance derivation, the saturation behavior, the Jacobian collapse, why √d_k specifically is the right correction, and how 2026 alternatives (QK-Norm in Gemma 4) achieve the same stability differently.
Variance derivation
Assume Q and K components are independent with mean 0 and variance 1. This is roughly true after standard initialization plus LayerNorm, the standard setup in real transformers.
Each entry of QK^T is the dot product Σ_k q_ik · k_jk, summed over k from 1 to d_k. Each individual term q_ik · k_jk is a product of two independent zero mean unit variance variables. The product has mean 0 (independence) and variance equal to the product of the input variances, so 1 · 1 = 1. Each term in the sum has variance 1.
The sum of d_k independent unit variance variables has variance equal to the sum of their variances, by Bienaymé's identity:
So std is √d_k. For modern head dimensions (d_head = 64 in many models, d_head = 128 in Llama 4 Maverick and Mistral Large 3) that is an std of roughly 8 or 11. A typical score will sit between -22 and +22 (two stds either side), and the largest scores in a row can be much further out.
The variance argument is not specific to attention; it is the same principle behind Xavier and Kaiming weight initialization. Sums of
nindependent unit variance terms have std√n.
Situations where this technique stops working.
2–4 min · Everything important, quickly.
import torch
# Demonstrate the variance argument empirically
torch.manual_seed(0)
for d_k in [8, 64, 128, 512]:
Q = torch.randn(1000, d_k)
K = torch.randn(1000, d_k)
scores = Q @ K.T
print(f'd_k={d_k:4d}: empirical std = {scores.std():.2f}, predicted √d_k = {d_k**0.5:.2f}')
# d_k= 8: std ≈ 2.83 ≈ √8
# d_k= 64: std ≈ 8.00 ≈ √64
# d_k=128: std ≈ 11.31 ≈ √128
# d_k=512: std ≈ 22.63 ≈ √512| Scaling choice | Var of softmax input | Softmax behavior | Training behavior |
|---|---|---|---|
| No scaling, d_k=128 | 128 | Saturated | Stalls |
| Divide by d_k | 1/128 | Never sharpens | Trains but model can't make sharp decisions |
| Divide by √d_k | 1 | Soft, can sharpen with learning | Trains correctly |
Real products, models, and research that use this idea.
- Vaswani et al. 2017 footnote 4 explicitly derives this variance argument and motivates the √d_k scaling.
- Llama 4 Maverick and Mistral Large 3 use d_head = 128, the variance argument applies directly at modern scale.
What an interviewer would ask next. Try answering before peeking at the approach.
QThe unit variance assumption for Q and K isn't exactly true in practice, Q and K are produced by learned projections of LayerNormed inputs. How robust is the √d_k argument?
Empirically very robust because std-normalization is a one-parameter correction, mild deviations from unit variance just shift the effective scaling slightly. The key is that the variance is bounded and doesn't grow with d_k after the projection. LayerNorm + reasonable init keeps Q, K in a regime where √d_k is approximately correct.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Stopping at 'variance is d_k' without connecting to the softmax-saturation to gradient-death chain. The variance argument is the SETUP; the failure mode is the gradient pathology.
60 second bullets to scan on the way to the call.
The Q, K unit variance, independence assumption
Derivation of Var(QKᵀ entry) = d_k step by step
Primary sources. Browse if you want the original framing.
Same topic, related formats. Practice these next.