Transformers, lesson 3
Why the scores are scaled
Dividing by the square root of the key size keeps softmax from saturating.
Real queries and keys hold far more than two numbers, and longer vectors make bigger scores. This lesson shows why bigger scores break softmax, and the one division that fixes it. You need attention from the last lesson, and slopes and gradients from Neural Networks.
By the end, you can:
Part 1. Longer vectors, bigger scores
Estimate how big a score gets when queries and keys hold d numbers.
The toy queries and keys so far had two numbers each. Real ones have far more: 64 per head in GPT-2 small, 128 in Llama 2 7B. A score adds up one product per number, so longer vectors give bigger totals. How much bigger depends on how the products add up.
Example 1.1
The spread of a score at d = 64
A query and a key each hold 64 random numbers, with an average of 0 and a spread (standard deviation) of 1, all drawn independently. That's the assumption the 2017 transformer paper makes. How big is a typical score q · k?
- Look at one productq₁ × k₁: average 0, spread 1
The product of two independent numbers of spread 1 has spread 1, and is as often negative as positive.
- Add 64 of them64 × 1² = 64
Spreads don't add, but their squares do, for independent numbers. This square is called the variance.
- Go back to a spread√64 = 8
The square root turns the variance back into a spread.
AnswerA spread of 8: scores of 8 or −8 are ordinary.
Why √d and not d? The products are as often negative as positive, so most of them cancel, like d random steps forward or back that end about √d steps from the start. You met the same rule in Neural Networks: a neuron's total grows with the square root of its fan-in.
Variance
The spread squared. For independent random numbers, variances add when the numbers add. So a score that adds d products of variance 1 has variance d, and a spread of √d. The 2017 paper makes this argument in a footnote.
spread of q · k = √d, so √64 = 8
The rule is about random numbers, which is roughly what a model starts with. Training can later push scores up on purpose. The trouble is that at the start, scores grow with d whether or not they mean anything.
Your turn
Queries and keys go from 64 numbers to 1,024 numbers, 16 times as many. About how much bigger do the scores get?
Part 2. Softmax goes all or nothing
Show how bigger scores push softmax to one winner, and why that stalls learning.
Softmax looks only at the gaps between scores, and when scores grow, so do the gaps. At d = 64 a score's spread is 8 times what it is for a single product. So take three moderate scores and make them 8 times bigger.
Example 2.1
The same scores, 8 times bigger
Softmax turns the scores (2, 1, 0) into weights. Then the scores become (16, 8, 0). How do the weights change, and how fast does the top weight respond to its score?
- Softmax the small scoressoftmax(2, 1, 0) = (0.67, 0.24, 0.09)
e^2 = 7.389, e^1 = 2.718 and e^0 = 1, out of 11.107. Every word gets some say.
- Softmax the big scoressoftmax(16, 8, 0) = (0.9997, 0.0003, 0.0000)
e^16 is about 8.9 million and e^8 about 2,981. The top word takes almost everything.
- Compare the slopes0.67 × 0.33 = 0.22, but 0.9997 × 0.0003 = 0.0003
How fast a weight p moves when its own score moves is p × (1 − p). Near 1, that's almost 0.
AnswerAlmost all the weight on one word, and a slope about 660 times flatter.
Training nudges W_Q and W_K by the slopes passed back through softmax. When the slope is nearly 0, the nudges are nearly 0, and the grids barely learn. A head stuck like this is also stiff: it gives one word all the weight, even when the answer needs two.
Saturated
A softmax is saturated when one weight sits near 1 and the rest near 0, so changing the scores a little barely changes the weights. Almost no gradient gets back to the scores. You met the word for tanh and sigmoid neurons in Neural Networks.
slope of a weight p = p × (1 − p)
The order of the scores isn't what softmax cares about. (2, 1, 0) and (16, 8, 0) rank the words the same way, yet give very different weights. Adding the same number to every score changes nothing. Multiplying them all by a number does.
Your turn
Softmax gives the scores (1, 0, 0) the weights (0.58, 0.21, 0.21), and the top weight's slope is 0.24. Multiply the scores by 8, to (8, 0, 0). What are the top weight and its slope now?
Part 3. Divide by √d
Work out scaled dot-product attention by hand, and read its formula.
The fix is one division. Scores grow like √d, so divide every score by √d before the softmax. With d = 64, divide by 8: scores of (16, 8, 0) shrink back to (2, 1, 0), and the weights spread out again. Here's the whole of attention on a tiny case.
Example 3.1
One query, d = 4
The query q = (1, 2, 0, 1) meets three keys: k₁ = (1, 1, 0, 1), k₂ = (0, 1, 1, 0) and k₃ = (1, −1, 2, 1). Their values are (1, 0), (0, 1) and (1, 1). What does attention send out for this query?
- Score the first key1 × 1 + 2 × 1 + 0 × 0 + 1 × 1 = 4
Multiply number by number and add.
- Score the other two0 + 2 + 0 + 0 = 2, 1 − 2 + 0 + 1 = 0
The same for k₂ and k₃, with the products already worked out.
- Divide by √d√4 = 2, so (4, 2, 0) → (2, 1, 0)
d = 4, the number of numbers in each query and key.
- Softmax the scoressoftmax(2, 1, 0) = (0.67, 0.24, 0.09)
The weights from part 2. Without the division, (4, 2, 0) would give (0.87, 0.12, 0.02).
- Blend the values(0.67 + 0.09, 0.24 + 0.09) = (0.76, 0.33)
Each value times its weight, one number at a time: (1, 0) adds to the first, (0, 1) to the second, (1, 1) to both.
Answer(0.76, 0.33).
Dividing by √d puts the spread of the scores back near 1, whatever d is. Softmax then sees scores of the same size in a tiny model and a huge one, so it starts out spread, with slopes big enough to learn from. Divide by d instead and the scores shrink too far: every weight comes out nearly equal.
Scaled dot-product attention
The attention from the 2017 paper "Attention Is All You Need", and PyTorch's default. Q Kᵀ is every query dotted with every key, where Kᵀ is K turned on its side. Divide by √d, softmax each row, then multiply by V.
Attention(Q, K, V) = softmax(Q Kᵀ / √d) V
Which d? The size of one head's queries and keys, not the model's full width. The 2017 paper splits its 512 numbers per token into 8 heads of 64, so it divides by 8. "Heads in parallel", later in this module, shows the split.
Your turn
A head with d = 16 gets the raw scores (8, 4, 0). What weights come out?
Where you'll meet it
"Why divide by √d_k?" is a classic transformer question. Answer in three steps. A score adds d_k products, so its spread grows like √d_k. Bigger gaps saturate softmax, which then passes back almost no gradient. Dividing by √d_k puts the spread back near 1 at any size. The 2017 paper gives the variance argument in a footnote.
PyTorch's F.scaled_dot_product_attention divides by √E by default, where E is the last size of the query, so pass it one head's queries. nanoGPT writes it out by hand: att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))), where k.size(-1) is the head size.
import numpy as np
rng = np.random.default_rng(0)
for d in (4, 64, 1024):
q = rng.standard_normal((10_000, d)) # 10,000 random queries
k = rng.standard_normal((10_000, d)) # and keys, spread 1
scores = (q * k).sum(axis=1)
print(d, scores.std().round(1), (scores / np.sqrt(d)).std().round(2))
# 4 2.0 1.0
# 64 8.1 1.01
# 1024 32.0 1.0The raw spread doubles each time d grows four times: 2, 8, 32. Divided by √d, it sits near 1 at every size, so softmax sees scores of the same size whatever the model.
Check yourself
Four questions from across the lesson. Try each one before you look back.
1
A head's softmax gives one word 0.9997 of the weight. Why does that stall training?
2
Llama 2 7B has 4,096 numbers per token, split across 32 heads of 128 numbers each. What does each head divide its scores by?
3
A model with d = 1,024 divides its scores by d instead of √d. What happens to the weights?
4
In softmax(Q Kᵀ / √d) V, which part holds what each word passes on?
What to remember
1Longer vectors, bigger scores
A score adds d products, so for random numbers of spread 1 its spread is √d: about 8 when d = 64.
2Softmax goes all or nothing
Bigger gaps saturate softmax: (2, 1, 0) gives the top word 0.67, (16, 8, 0) gives it 0.9997, and its slope p × (1 − p) falls from 0.22 to 0.0003.
3Divide by √d
Attention(Q, K, V) = softmax(Q Kᵀ / √d) V divides every score by √d before the softmax: with d = 4, the scores (4, 2, 0) become (2, 1, 0).
Recap and going further
The lesson on one page
- A score adds one product per number
- Its spread grows like √d: 8 at d = 64
- (2, 1, 0) gives the top word 0.67, (16, 8, 0) gives 0.9997
- The slope p × (1 − p) falls to 0.0003
- Tiny slopes, tiny updates to W_Q and W_K
- The head can only look at one word
- Divide every score by √d before softmax
- √64 = 8, √128 ≈ 11.3
- softmax(Q Kᵀ / √d) V
- d is the size of one head's queries and keys
Try this
Measure the growth yourself
In NumPy, draw 10,000 random query and key vectors at sizes 4, 16, 64, 256 and 1,024, with every number from a standard normal distribution. Print the spread (standard deviation) of their dot products, raw and divided by √d. Then score one query against 8 random keys at each size, put the scores through softmax with and without the division, and print the largest weight.
You need: Python with NumPy.
Look for
- Raw spreads near √d, and scaled ones near 1 at every size.
- The size at which the unscaled softmax hands one key nearly all the weight.
- p × (1 − p) for that top weight, scaled and unscaled.
- What happens if you divide by d instead of √d.
The words, in plain language
- Variance
- The spread squared. For independent random numbers, variances add when the numbers add. So a score that adds d products of variance 1 has variance d, and a spread of √d. The 2017 paper makes this argument in a footnote.
- Saturated
- A softmax is saturated when one weight sits near 1 and the rest near 0, so changing the scores a little barely changes the weights. Almost no gradient gets back to the scores. You met the word for tanh and sigmoid neurons in Neural Networks.
- Scaled dot-product attention
- The attention from the 2017 paper "Attention Is All You Need", and PyTorch's default. Q Kᵀ is every query dotted with every key, where Kᵀ is K turned on its side. Divide by √d, softmax each row, then multiply by V.
Transformers overviewSpotted a mistake? Tell us.