- Key idea
- The variance of q·k grows with d; dividing by √d brings it back to 1, so softmax stays soft and gradients keep flowing
- Real world
- A vote where each judge gives a score. Add up 4 judges and totals differ by a few points; add up 256 judges and totals differ by dozens. If the winner is picked by “highest total takes all”, more judges means a landslide every time, unless you divide by the right amount first.
- In an LLM
- In GPT-2 small each query and key has d = 64 numbers. Without scaling, every score would be 8 times larger, softmax would be far sharper, and the query and key matrices would get much weaker training signals.
The step in question
In the query, key, value recipe, step 3 divides every score by , where is the length of a query (the head size):
It looks like a small detail. It isn’t: without it, attention in a model with a large head size barely trains. This page shows why, in three steps.
1. Dot products grow with length
A dot product adds up products, one per position: . If the numbers are random with mean 0 and variance 1 (roughly what a freshly initialised network gives), each product is a random number with variance 1, and adding up of them gives a total with variance . Its typical size, the standard deviation, is .
So a head with 64 numbers produces scores around ±8, and one with 256 numbers scores around ±16, before the model has learned anything. Try it:
The red histogram widens as grows. The blue one, the same scores divided by , stays the same width for every : dividing a random number by its standard deviation gives it a standard deviation of 1.
Each judge scores a contestant somewhere between −1 and +1, at random. With 4 judges, two contestants’ totals typically differ by one or two points. With 256 judges, by around 10. If you turn totals into a share of the prize with softmax, the 256-judge contest nearly always gives one contestant everything, for no better reason than there were more judges. Divide each total by √(number of judges) and the contests become comparable again.
2. Big scores saturate softmax
Softmax only cares about the differences between scores. A gap of 1 between two scores means a ratio of between their weights; a gap of 8 means . When scores are spread by ±8, the largest one wins almost all the weight and attention stops being a soft average. It becomes a hard pick, and early in training that pick is decided by random noise in the weights.
Multiplying all the scores by a number is the same as changing the softmax : multiply by 10 and it behaves like temperature 0.1. The in the formula is a fixed temperature of , chosen so that the scores have a spread of about 1.
3. A saturated softmax stops learning
The model learns and by gradient descent, which needs to know how the weights would change if a score changed. That sensitivity is a matrix of derivatives (the Jacobian) with entries on the diagonal and off it. If one weight is close to 1 and the rest close to 0, every one of those entries is close to 0.
At the toy model’s real scale (×1) the largest gradient entry is about 0.25, the biggest it can ever be. At ×4 “cat” takes 97% and the largest entry falls to 0.027; at ×10 it is about 0.00005. Whatever error the model makes, almost none of it gets back through the softmax to the query and key matrices. This is a vanishing gradient, and it is exactly what unscaled scores with a large cause.
GPT-2 small has a head size of 64, so its attention divides every score by 8. Llama 2 7B has a head size of 128 and divides by about 11.3. The model can still learn sharp attention, and real heads do (the previous-token head puts over 99% on one token), but it gets there by learning larger query and key vectors on purpose, not by accident at initialisation.
Why √d and not d
Dividing by would shrink the scores too much: their spread would be , which for is 0.125, and every softmax would be nearly uniform, again with little to learn from. Dividing by is the one choice that keeps the spread at 1 for every head size.
The variance calculation
Let have independent entries with mean 0 and variance 1. Each product has mean 0 and variance . The products are independent, so
The softmax Jacobian
For ,
whose largest possible entry is , reached when . If approaches a one-hot vector, every entry approaches 0.
Variants
Some models pick a different scale. T5 folds the factor into its weight initialisation and uses no explicit division. μP (maximal update parametrisation) uses with a matching initialisation so that hyperparameters transfer between model widths. Many recent models also normalise queries and keys (QK-norm) to stop scores growing during training, a fix for training instability in large models.
In papers
- The scaling and its explanation (footnote 4): Vaswani et al., 2017, section 3.2.1.
- μP and attention scaling: Yang et al., 2022.
- QK-norm against attention logit growth: Dehghani et al., 2023.