Maths for ML / Read an equation

Softmax

Turns a list of scores, one per class, into probabilities that are all positive and add up to 1.

softmax⁡(z)i=ezi∑j=1Kezj\htmlClass{mx-p mx-p-out}{\operatorname{softmax}(\mathbf{z})_{\htmlClass{mx-p mx-p-i}{i}}} = \frac{\htmlClass{mx-p mx-p-top}{e^{z_i}}}{\htmlClass{mx-p mx-p-bot}{\htmlClass{mx-p mx-p-sumj}{\sum_{j=1}^{K}} e^{z_j}}}

The parts

Point at any part of the equation, or pick it from the list.

  • probability of class i · output · shape one number between 0 and 1, for each of the K classes

    The whole output is a vector of K probabilities that add up to 1.

  • which class · index · shape 1 to K

    The class whose probability this line computes. The same formula is used for every i.

  • exp of this class's score · operation · shape one positive number

    Makes the score positive and stretches the gaps: a score 1 higher becomes about 2.7 times bigger.

  • the normalizer · operation · shape one number, the same for every i

    Adds up the exp of every class's score, so dividing by it makes the K outputs sum to 1.

  • sum over all classes · index · shape j counts 1, 2, …, K

    A different letter from i on purpose: i is the one class on top, j runs over every class underneath, including i itself.

Read it with the seven questions

  1. What comes out? A vector of K probabilities, one per class, adding up to 1.
  2. List the parts. Input: the scores z (often called logits), one per class. Indices: i for the class on top, j for the sum. Nothing learned inside softmax itself.
  3. Check the shapes. z has length K, the output has length K. The denominator is one number shared by every class.
  4. Say what each function does. exp makes every score positive and magnifies differences; the sum and the division normalize.
  5. Ask what makes it bigger or smaller. Raising one score raises its own probability and lowers every other one, because the shared denominator grows.
  6. Push the inputs to extremes. All scores equal: every class gets 1/K. One score far above the rest: it takes almost all the probability. Adding the same number to every score changes nothing, which is why code subtracts the largest score first to keep exp from overflowing.
  7. Write it as code. e = np.exp(z - z.max()); p = e / e.sum().

Change it and watch

Move the four scores and watch the probabilities share out. Then try the temperature: dividing every score by T before softmax. Small T makes the winner take all; large T spreads probability evenly. Language models use exactly this knob when they sample text.

A bar chart of four probabilities, for the classes cat, dog, fox and owl, computed by softmax from four score sliders. A temperature slider divides every score by T first: low T makes the largest bar approach 1, high T makes the bars equal.

Now write it in NumPy
import numpy as np

z = np.array([2.0, 1.0, 0.1])
e = np.exp(z - z.max())          # subtracting the max avoids overflow and changes nothing
p = e / e.sum()
print(p.round(3), p.sum())       # [0.659 0.242 0.099] 1.0

You’ll see this in

Symbols on this pagesoftmax softmaxexp exponentialΣ summationxᵢ subscripts & superscript indices