Why Softmax and Not Some Other Normalisation
Last timeThe Scaling Factor
Scores have to become weights, and many functions could do that. Softmax is chosen for properties that are specific and checkable, not for convention.
By this point in the formula the hard work looks done. Queries and keys have
produced a matrix of scores. Those scores need to become weights that sum to
one. Softmax does it, and the choice tends to go unexamined.
It should not, because several other functions would also produce weights that
sum to one, and comparing them shows what softmax is actually buying.
What the function has to do
Write down the requirements before looking at any candidate.
- Non-negative outputs. A negative weight would mean subtracting a token's
value, which is not what "attend to" means and makes the output unbounded.
- Outputs summing to one. Without this, the magnitude of the output would
depend on sequence length rather than on content.
- Differentiable everywhere. Gradients have to flow back through this step
to reach the query and key projections.
- Order preserving. A higher score must produce a higher weight, always.
Now test the obvious alternatives.
| Candidate | Non-negative | Sums to one | Differentiable | Order preserving |
|---|---|---|---|---|
| Divide by the sum | no | yes | yes | yes |
| Hard argmax | yes | yes | no | yes |
| Square then normalise | yes | yes | yes | no |
| Softmax | yes | yes | yes | yes |
Each alternative fails exactly one requirement. Dividing by the sum breaks on
negative scores. Hard argmax has zero gradient everywhere it is defined, so
nothing can be learned through it. Squaring maps and to the same
weight, which destroys the ordering the scores were computed to express.
Dividing by the sum is the first thing anyone tries, and it breaks immediately.
Scores are dot products, so they are freely negative. A set of scores summing to
near zero sends the weights to infinity, and a negative score produces a
negative weight.
Hard argmax satisfies everything except differentiability, and that single
failure is fatal: a function whose gradient is zero almost everywhere passes no
learning signal back to the projections that produced the scores.
Squaring fixes the sign problem and keeps differentiability, but it maps
and to the same weight. A strongly negative score, which means an actively
poor match, would be treated as a strong match.
Softmax handles all four by exponentiating first. Exponentiation maps the whole
real line to the positive reals, is smooth everywhere, and is strictly
increasing.
The lesson stops here
9 more paragraphs to go
You have read the opening. The rest of the argument, the problems that check whether it landed, and the lines worth keeping at the end all come with a plan.
The first lesson of every course in the library reads the whole way through, free, so you can see exactly what the rest of them are.
See the planThe contentsThis is the reading half
Starting the course gives you your own copy of it. Every idea on every page has problems standing under it, marked with a reason rather than a tick, and any sentence you do not believe can be opened and argued with. None of that can happen on a page nobody owns.
The contents