The √d_k scaling exists for one specific training-stability reason: without it, the model can't learn. The failure mode is not numerical: it's gradient flow through softmax.
Here's the chain of reasoning step by step.
First, the dot product Q_i · K_j is a sum of d_k products. If Q and K components are roughly unit variance, the variance of the sum is d_k. So the typical magnitude of an entry in QKᵀ grows as √d_k. For d_head = 128 (typical for modern LLMs), a typical entry is around 11 in absolute value.
Second, feeding large values into softmax produces a peaked output. softmax({0, 0, 11, 0}) is roughly {0.00002, 0.00002, 0.99994, 0.00002}. The largest entry takes almost all the probability mass; the rest collapse toward zero.
Third, look at what happens in the backward pass. The softmax Jacobian at a peaked output is near-zero off the peak: nudging a small entry barely changes its tiny probability. So gradients can't flow back to non-peak positions.
Fourth, the model is stuck. It can't learn to attend somewhere new because the gradient signal to do so is essentially zero. Training stalls.
Dividing by √d_k brings the variance back to 1, keeping softmax inputs in a range where the output stays soft and gradients flow to all positions:
$$\text{Attn}(Q,K,V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$$
That's why every modern attention implementation includes the scaling.
The trap here is that the failure is a gradient flow problem, not a numerical one, people reach for FP16 overflow and get the wrong mechanism.
The variance of a QKᵀ entry scales linearly with d_k, so the standard deviation scales as √d_k. For typical head dimensions like 64 or 128, that's large enough that unscaled scores saturate softmax: one position grabs essentially all the probability mass.
Once softmax is saturated, the softmax Jacobian goes near-zero off the peak. Gradients can't redistribute attention to other positions, because nudging a near-zero probability barely changes it. Training stalls almost immediately on any non-trivial head dim.
The scaling normalizes the variance back to 1:
$$\text{Attn}(Q,K,V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$$
The other options miss the mechanism. FP16 overflow is secondary: gradient failure happens long before NaNs. Softmax always normalizes by construction, so the output is still a probability distribution, just a degenerate one. And multi-head parallelism is completely orthogonal to scaling. The correct answer is the gradient flow story.
Quick reference
TL;DR
Without √d_k, dot product magnitudes grow with d_k, softmax saturates onto one position, the softmax Jacobian collapses off the peak, and gradients can't redistribute attention. Training stalls.
Imagine a class voting on which movie to watch. If everyone scores movies between 1 and 5, the vote is balanced and every movie has a real shot. But suppose the scores get amplified to 100 to 500. Now the top movie wins with essentially 100% of the weight and the rest get zero. Worse, if you wanted the class to change its mind, nudging a near-zero option a little still leaves it near zero: there's no mechanism for the runners-up to gain ground. The √d_k division is the rule that keeps scores in the 1-to-5 range no matter how big the input vectors get, so the vote stays competitive and the model can keep learning.
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.
Custom attention kernels where QK^T is normalized externally (QK-Norm, pre-attention layer norm) If you've already constrained Q and K variance via normalization, applying √d_k on top double-shrinks the scores into the under-saturated regime.
Very small head dimensions like d_k ≤ 4 in toy models The variance argument only matters when d_k is large enough to saturate softmax; at very small d_k the scaling is a no-op and can be omitted for clarity, though it doesn't hurt.
Linearized attention variants that replace softmax with a kernel approximation Performer, Linear Attention, and related methods don't use softmax at the QK^T step, so the saturation failure mode doesn't exist; whether to apply √d_k depends on the kernel's own variance behavior.
Hand-tuned attention biases where you intentionally want a near one hot pattern If saturation is the goal (e.g. hard-routing in some MoE variants), the scaling works against you; you'd want to remove it or use a learned temperature instead.
Quick scan · 2–4 min
Everything important, quickly.
Derive Var(Q · K) = d_k from independence, show why standard deviation scales as √d_k, walk the softmax Jacobian to show off-peak gradients vanish under saturation, and rule out alternative normalizations (d_k, log d_k) as over- or under-correcting.
Code example
python
# Demonstrating softmax saturation when scaling is omitted
import torch
import torch.nn.functional as F
d_k = 128
q = torch.randn(1, d_k)
k = torch.randn(64, d_k) # 64 candidate positions
scores_unscaled = q @ k.T # std ~= sqrt(d_k) ~= 11
scores_scaled = scores_unscaled / (d_k ** 0.5)
attn_unscaled = F.softmax(scores_unscaled, dim=-1).squeeze()
attn_scaled = F.softmax(scores_scaled, dim=-1).squeeze()
print(f'unscaled max prob: {attn_unscaled.max().item():.4f}')
print(f'scaled max prob: {attn_scaled.max().item():.4f}')
# unscaled max ~= 0.99 (saturated); scaled max ~= 0.05 (soft)
Real-world examples
Real products, models, and research that use this idea.
1Vaswani et al. 2017 footnote 4 explicitly motivates the scaling: 'We suspect that for large values of d_k, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients.'
2Modern decoder LLMs like Llama 4 Maverick and Gemma 4 use d_head = 64 to 128 per head, exactly the regime where unscaled attention would saturate immediately.
3PyTorch's torch.nn.functional.scaled_dot_product_attention applies the √d_k scaling internally; users almost never implement it by hand anymore.
4Some 2024-2025 architectures use QK-Norm (layer-norm applied to Q and K before the dot product) as a complementary variance-control mechanism for very large head dimensions.
Interviewer follow-ups
What an interviewer would ask next. Try answering before peeking at the approach.
QDoes removing √d_k still hurt at inference, where there's no gradient to compute?
Sign in or go Pro to see the approach for this follow-up.
QWhy √d_k specifically and not d_k or log(d_k)?
Sign in or go Pro to see the approach for this follow-up.
QWhat is QK-Norm and how does it relate to √d_k scaling?
Sign in or go Pro to see the approach for this follow-up.
Don't say this
Red flags and common mistakes that signal junior thinking. Click to expand.
Most common mistake
Choosing the FP overflow answer. Numerical overflow can happen for extreme values, but the primary failure is the gradient signal, and training will stall well before any FP16 NaN appears.