Softmax
Calculate conditional probabilities, Bayes updates, expectations, variance and softmax probabilities.
THE IDEA
From scores to probabilities
A classifier or a language model ends with a list of real-valued scores, called logits: one per class or per candidate next token. They can be negative, huge, and they do not sum to anything in particular. We want probabilities: positive, summing to 1. Softmax does it in two moves. Exponentiate each score (this makes everything positive and keeps the order), then divide by the total. p_i = exp(z_i) / sum of exp(z_j).
Predict first. Adding the same number, say 100, to every logit: what happens to the probabilities?
Temperature and next-token scores
Language models commonly produce one logit per vocabulary token, and a softmax turns them into the next-token distribution that a sampler draws from. Temperature T divides the logits first: softmax(z / T). Low T widens the gaps and makes the top token dominate; high T shrinks the gaps toward a uniform choice. T = 1 leaves the scores as they are. The scores below are made up.
Next word after "The cat sat on the ..." (made-up scores).
| token | logit | / T | exp(/T - max) | P |
|---|---|---|---|---|
| mat | 3.2 | 3.2 | 1 | 0.686 |
| sofa | 2.1 | 2.1 | 0.333 | 0.228 |
| roof | 1 | 1 | 0.111 | 0.076 |
| moon | -1 | -1 | 0.015 | 0.01 |
| total | 1.459 | 1 |
At T = 1 the four tokens get 68.6%, 22.8%, 7.6%, 1% for mat, sofa, roof, moon. Even a token with a low score (moon) keeps a small chance. Softmax never outputs an exact zero for finite scores, though floating point may round a tiny value to 0.
Why subtract the maximum
Floating-point numbers overflow: exp(710) is already inf in 64-bit floats, and inf / inf is not a number. At the other extreme, exp(-1000) underflows to 0 and a row of such values gives 0 / 0. The shift invariance you just saw is the cure: subtract the largest score before exponentiating. The biggest term becomes exp(0) = 1, nothing overflows, and the answer is mathematically identical.
| score | exp(score) | naive P | stable P |
|---|---|---|---|
| 2 | 7.389 | 0.665 | 0.665 |
| 1 | 2.718 | 0.245 | 0.245 |
| 0 | 1 | 0.09 | 0.09 |
import math
z = [1000, 1000]
[math.exp(x) for x in z] # OverflowError: math range error
m = max(z)
e = [math.exp(x - m) for x in z] # [1.0, 1.0]
[v / sum(e) for v in e] # [0.5, 0.5]Common mistakes
- Softmax outputs sum to 1 but are not automatically calibrated: a model can put 99% on a wrong answer. A high softmax value is not evidence the probability model is right.
- Applying softmax twice, or to values that are already probabilities, flattens the distribution.
- Temperature must be positive. As T goes toward 0 the distribution approaches always picking the top token, but T = 0 itself divides by zero.
- Taking log(softmax(z)) can hit log(0). Use the log-sum-exp form instead (you will write it in the practice).
Chapter context
Probability quantifies uncertainty relative to a model of possible outcomes. Conditioning changes the relevant population. Before inserting numbers into a formula, say which event is given and which event is being predicted.
TRY IT YOURSELF
Part 1. Do it hands-on
Logits are shown in whole tenths on the sliders. Tasks are checked from the exact softmax of the current logits and temperature.
Next word after "The cat sat on the ..." (made-up scores).
| token | logit | / T | exp(/T - max) | P |
|---|---|---|---|---|
| mat | 3.2 | 3.2 | 1 | 0.686 |
| sofa | 2.1 | 2.1 | 0.333 | 0.228 |
| roof | 1 | 1 | 0.111 | 0.076 |
| moon | -1 | -1 | 0.015 | 0.01 |
| total | 1.459 | 1 |
- To do: Without touching the logits, make the model at least 90% sure of “mat”
- To do: Make the distribution nearly flat: no token above 30% (four equal tokens would be 25% each)
- To do: Raise every logit by the same amount so all four logits are positive, keeping the temperature at 1 and the probabilities unchanged
Part 2. Do it in code
- Write softmax(logits, temperature=1.0) the numerically stable way: divide by the temperature, subtract the largest value, exponentiate, then divide by the total. Raise ValueError for an empty list or a temperature that is not positive.
- Write log_softmax(logits): log of the softmax, computed as z - max - log(sum(exp(z - max))) so that tiny probabilities do not become log(0).
- Write sample_frequencies(logits, temperature, n, seed): draw n indices from softmax(logits, temperature) with random.Random(seed).choices and return the fraction of draws for each index.
Stuck, or something not working? You can come back to it any time.
CHECK YOUR UNDERSTANDING
Softmax scores [0,ln 2,ln 3] give what probabilities?
Chapter reference: math, worked example & code
This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.
Math and notation
Bayes' denominator must be nonzero. The softmax normalization runs over the candidate classes.
Worked example
In a fictional population, illness prevalence is 1%, sensitivity is 90% and the false-positive rate is 10%. True positives account for 0.01×0.9=0.009 of the population; false positives account for 0.99×0.1=0.099. Among positives, the illness fraction is 0.009/0.108=1/12≈8.33%. This is an arithmetic example, not medical guidance.
Reference code
The following example uses the Python standard library and does not call a remote service.
from math import exp, log, isclose
posterior = .01*.9 / (.01*.9 + .99*.1)
assert isclose(posterior, 1/12)
scores = [0, log(2), log(3)]
weights = [exp(z-max(scores)) for z in scores]
probs = [w/sum(weights) for w in weights]
assert all(isclose(a,b) for a,b in zip(probs, [1/6,2/6,3/6]))
Python examples are provided for study and local execution; this page does not execute Python.
Common mistakes
P(A given B) and P(B given A) are different quantities. A high uncalibrated model confidence is not evidence of a correct probability model.