Softmax

Only the gaps matter

Shift every score by the same amount and softmax gives the same answer. The scores [1, 3] and [11, 13] both give 0.119 and 0.881.

Think of a race. If every runner's time had 10 seconds added, the gaps between them would stay the same. So would the order they finished in. Softmax only looks at those gaps.

import math
def softmax(xs):
    e = [math.exp(x) for x in xs]
    return [round(v / sum(e), 3) for v in e]
print(softmax([1, 3]) == softmax([11, 13]))

The overflow trap and the max trick

A computer stores each number in a fixed amount of space. $e^{1000}$ is far too big to fit. It overflows: the computer gives up and calls it infinity. Real logits can be large, so this does happen. Infinity divided by infinity isn't a number, so NumPy prints nan ("not a number").

The fix uses the rule you just saw. Only the gaps matter, so subtract the largest score from every score first. The biggest score becomes 0, and $e^0 = 1$. Every other score becomes negative. $e$ to a negative power is a small number between 0 and 1. Nothing overflows, and the answer is unchanged.

import numpy as np

z = np.array([1000.0, 1001.0, 1002.0])
naive = np.exp(z) / np.exp(z).sum()          # overflow: inf / inf
stable = np.exp(z - z.max()) / np.exp(z - z.max()).sum()
print("naive: ", naive)
print("stable:", stable.round(4))

The naive line gives nan for every option. NumPy may also print an overflow warning. The stable line gives 0.09, 0.2447 and 0.6652. Those are the probabilities for the scores [0, 1, 2]. They're also the worked example's numbers in reverse order. Every real softmax implementation uses this trick.

Temperature

Temperature is a dial for how bold the choices are. You divide every score by a number $T$, the temperature, before softmax.

Dividing by a number below 1 spreads the scores further apart. Dividing by a number above 1 squeezes them together. And bigger gaps give more lopsided probabilities.

Here $p_i$ is the probability of option $i$.

In plain words: divide every score by the temperature, then do softmax as usual.

Take the worked example's scores [2, 1, 0] again:

• $T = 1$ leaves the scores alone. The probabilities are 0.665, 0.245 and 0.090.
• $T = 0.5$ doubles the scores to [4, 2, 0]. The gaps grow, and "cat" jumps to 0.867.
• $T = 2$ halves the scores to [1, 0.5, 0]. The gaps shrink, giving 0.506, 0.307 and 0.186.

So the general rule is:

• $T$ below 1 stretches the gaps. The top option gets even more probability. The output is sharper and more predictable.
• $T$ above 1 shrinks the gaps. The probabilities even out. The output is flatter and more varied.
• As $T$ gets close to 0, the top option takes nearly all the probability. That comes close to always picking the top option, which is called greedy choice.

Temperature never changes the order. The top score stays on top at every temperature.

The code below defines a reusable softmax. It applies the temperature and the max trick.

import numpy as np

def softmax(z, T=1.0):
    s = np.asarray(z) / T
    e = np.exp(s - s.max())
    return e / e.sum()

z = [2.0, 1.0, 0.1, -1.0]
for T in [0.3, 1.0, 3.0]:
    print(f"T={T}: {softmax(z, T).round(3)}")

At T=0.3 the top option gets 0.964. At T=3 the four probabilities sit much closer together: 0.382, 0.274, 0.203 and 0.141.

Where softmax appears

Anywhere a model must share out 100% between options, you'll find softmax.

A classifier sorts things into fixed groups, such as "spam" or "not spam". A language model scores every token it knows to pick the next one. Attention comes in Phase 5. For now, know that each word gives every other word a score, and softmax turns those scores into shares of attention.

With only two options, softmax is the S-curve (the sigmoid) from the functions and graphs lesson. Give one option the score $z$ and the other the score 0. The first option's probability is then the sigmoid of $z$. For $z = 2$, both come to 0.881.

import math

z = 2.0
print(round(math.exp(z) / (math.exp(z) + math.exp(0)), 3))
print(round(1 / (1 + math.exp(-z)), 3))