Why divide by √d

Dot products of long vectors are big numbers: their spread grows with the square root of the length. Big scores push softmax into giving one token everything, and a saturated softmax stops learning. Dividing the scores by √d keeps them at a sensible size whatever the head size.

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 d\sqrt{d}, where dd is the length of a query (the head size):

A=softmax⁡ ⁣(QK⊤d)A = \operatorname{softmax}\!\Big(\frac{Q K^\top}{\sqrt{d}}\Big)

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 dd products, one per position: q⋅k=q1k1+q2k2+⋯+qdkdq \cdot k = q_1 k_1 + q_2 k_2 + \dots + q_d k_d. 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 dd of them gives a total with variance dd. Its typical size, the standard deviation, is d\sqrt{d}.

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:

length d
variance of raw66.3
theory: d64
std of raw (√d)8.1
variance after ÷ √d1.04
softmax of one query’s scores against 8 keys
from raw scores (top weight 99%)from scaled scores (top weight 43%)
With d = 64, raw scores typically land around ±8 (variance 66, theory 64). Gaps that big make softmax put 99% on one key. Divided by √64 = 8, the variance is back to about 1 (1.04) and the top weight is 43%. Scaled scores look the same whatever d is.
Each sample is a query and a key with d random numbers each (mean 0, variance 1), and their dot product. Pick the vector length d; the top histogram is the raw score q·k, the bottom one the score divided by √d. Below: what softmax does with eight such scores.

The red histogram widens as dd grows. The blue one, the same scores divided by d\sqrt{d}, stays the same width for every dd: dividing a random number by its standard deviation gives it a standard deviation of 1.

example · adding up judges

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 e1≈2.7e^1 \approx 2.7 between their weights; a gap of 8 means e8≈3000e^8 \approx 3000. 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 d\sqrt{d} in the formula is a fixed temperature of d\sqrt{d}, chosen so that the scores have a spread of about 1.

3. A saturated softmax stops learning

The model learns WQW_Q and WKW_K 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 pi(1−pi)p_i(1 - p_i) on the diagonal and −pipj-p_i p_j off it. If one weight is close to 1 and the rest close to 0, every one of those entries is close to 0.

softmax weights
scores: 0.8, 1.7, 0.5, 0.1, 0.2
gradient: how each weight (row) moves when a score (column) is nudged
ThecatsatontheThe0.15−0.09−0.03−0.02−0.02cat−0.090.25−0.07−0.04−0.05sat−0.03−0.070.12−0.01−0.01on−0.02−0.04−0.010.08−0.01the−0.02−0.05−0.01−0.010.09
−0.30.3
largest gradient entry (up) as the scores grow (across)
top weight49%
largest gradient0.250
The weights are spread and the gradient entries are sizeable (largest 0.25). Nudging a score changes the weights, so training can learn which token to look at.
The last “the” of the toy sentence, attending to all five tokens (no mask, to keep the picture simple). Drag the slider to multiply its scores, as if the vectors were longer. Watch the weights and the gradient matrix.

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 dd cause.

example · the √8 in GPT-2

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 dd would shrink the scores too much: their spread would be 1/d1/\sqrt{d}, which for d=64d = 64 is 0.125, and every softmax would be nearly uniform, again with little to learn from. Dividing by d\sqrt{d} is the one choice that keeps the spread at 1 for every head size.

going deeper

The variance calculation

Let q,k∈Rdq, k \in \mathbb{R}^d have independent entries with mean 0 and variance 1. Each product qikiq_i k_i has mean 0 and variance E[qi2] E[ki2]=1\mathbb{E}[q_i^2]\,\mathbb{E}[k_i^2] = 1. The products are independent, so

Var⁡(q⋅k)=∑i=1dVar⁡(qiki)=d,Var⁡ ⁣(q⋅kd)=1.\operatorname{Var}(q \cdot k) = \sum_{i=1}^{d} \operatorname{Var}(q_i k_i) = d, \qquad \operatorname{Var}\!\Big(\frac{q \cdot k}{\sqrt{d}}\Big) = 1.

The softmax Jacobian

For p=softmax⁡(s)p = \operatorname{softmax}(s),

∂pi∂sj=pi (δij−pj),\frac{\partial p_i}{\partial s_j} = p_i\,(\delta_{ij} - p_j),

whose largest possible entry is 14\tfrac14, reached when pi=12p_i = \tfrac12. If pp 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 1/d1/d 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

Related

  • qkv - Attention is three matrix multiplications to make queries, keys and values, one more to compare every query with every key, a softmax to turn scores into weights, and a last multiplication to mix the values. Here is every number.
  • attention-intuition - Attention lets each token pull in information from other tokens. It works like a dictionary lookup that never fails: instead of returning the one entry whose key matches, it returns a weighted average of every entry, weighted by how well each key matches the question.
  • causal-mask - A language model must predict each token without seeing it. The causal mask is a triangle of −∞ added to the attention scores so that every token can only look at itself and earlier tokens. It also lets training check a prediction at every position of a text in one pass.