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.
Detailed answer & concept explanation~6 min readEverything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example.
Everything important, quickly.
Derive the variance step (sum of independent unit variance products), the saturation behavior of softmax at large magnitudes, the Jacobian formula and its degeneration, why √d_k is uniquely the right normalization, and note the empirical robustness in practice.
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.
- Gemma 2 (2024) and Gemma 4 (2026) ship QK-Norm as an alternative stabilization, normalizing Q and K before the dot product instead of (or in addition to) √d_k.
- The Xavier and Kaiming initialization schemes apply the same variance-preservation principle to weights across layers.
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?
Don't say thisRed flags and common mistakes that signal junior thinking. Click to expand.
Red flags and common mistakes 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.
The night-before-the-interview bullets. Scan these on the way to the call.
Primary sources. Skim if you want the original framing.
Same topic, related formats. Practice these next.