The softmax function, worked by hand

The problem

At the end of each step, a language model gives a score to every token it knows. These scores are called logits raw, unscaled scores: any number, including negative ones. Logits are not probabilities. They do not add up to anything in particular. We need a fixed procedure that turns them into probabilities without changing their order.[1]

Softmax does not change the order of the scores. It changes what they mean.

The formula

For a list of n scores, written z, the probability for score number i is:

EQ 01 // SoftmaxDisplay maths

\operatorname{softmax}(z)_i = \frac{e^{z_i}}{\sum_{j=1}^{n} e^{z_j}}

Plainly

Raise e to the power of one score. Divide by the same thing done to every score, added up.

The number e is about 2.718.*1 Raising it to a power always gives a positive result, so every probability is above zero. Dividing by the total makes them add up to exactly 1.

Worked by hand

Take three scores. A calculator is enough.

Write down the scores: 2.0, 1.0, 0.1.
Raise e to each one: 7.389, 2.718, 1.105.
Add them up: 7.389 + 2.718 + 1.105 = 11.212.
Divide each by the total: 0.659, 0.242, 0.099.
Check the result. The three numbers add up to 1.000, and the order has not changed.

Fig. 01 // Softmax, step by step01 / 03

Step 1 // Start with the scores

#z1

2.0

#z2

1.0

#z3

0.1

Raw scores. They add up to 3.1.

Fig. 01

Three scores in, three probabilities out. The largest score still wins.

In code

The same procedure, in Python.

Listing 01 // softmax.py

import math

def softmax(scores):
  exps = [math.exp(z) for z in scores]
  total = sum(exps)
  return [x / total for x in exps]

print(softmax([2.0, 1.0, 0.1]))
# ≈ [0.659, 0.242, 0.099]

In practice, libraries subtract the largest score from every score first. This keeps the numbers small enough to compute safely and gives the same answer.[2]