Softmax and cross-entropy loss, worked through

From scores to probabilities

Suppose an image classifier with three classes outputs the logits $z = (2, 1, 0.1)$ for cat, dog and bird. Logits can be any real numbers, and we want a probability distribution: non-negative numbers that sum to 1, with larger logits getting more probability. Softmax does this by exponentiating and normalising:

Exponentiating makes every value positive and turns differences in logits into ratios of probabilities: if $z_i$ is 1 more than $z_j$, class $i$ is $e \approx 2.718$ times as likely.

Useful properties

• Adding a constant to every logit changes nothing. $e^{z_i + c} / \sum_j e^{z_j + c} = e^{z_i} / \sum_j e^{z_j}$, since $e^c$ cancels. Only differences between logits matter.
• Numerical stability. That property is how softmax is computed safely. $e^{1000}$ overflows a floating-point number, so implementations subtract the largest logit first. The widget below does this; the “z / T − max” column shows the shifted values.
• It is a soft version of argmax. The largest logit always gets the largest probability, but the others keep some, and the result changes smoothly with the logits, which is what gradient descent needs.

Temperature

Dividing the logits by a temperature $T > 0$ before softmax controls how peaked the distribution is: $p_i \propto e^{z_i / T}$. With $T = 0.5$ the example becomes (0.864, 0.117, 0.019), sharper; with $T = 2$ it becomes (0.502, 0.304, 0.194), flatter. As $T \to 0$ softmax approaches a one-hot argmax; as $T \to \infty$ it approaches uniform. Language models expose temperature as a sampling setting, and a single temperature fitted after training is a simple, effective way to make a classifier’s confidence better calibrated.

Cross-entropy loss

During training we know the true class $y$. The cross-entropy loss is the negative log of the probability the model assigned to it:

If the model is certain and right, $p_y = 1$ and the loss is 0. As $p_y$ shrinks the loss grows without bound, so confident mistakes are punished hard. In the example, if the image really is a cat the loss is $-\ln$ 0.6590 = 0.4170; if it is a bird the loss is $-\ln$ 0.0986 = 2.3170.

Why the logarithm? Minimising $-\log p_y$ summed over a dataset is the same as maximising the probability the model assigns to all the correct labels together, so it is maximum-likelihood training. In information-theory terms it is the number of nats (or bits, with $\log_2$) needed to encode the true label using the model’s probabilities. A model that guesses uniformly among $k$ classes scores $\ln k$, a useful baseline: about 1.099 for three classes.

The gradient is p − y

Write $y$ as a one-hot vector (1 for the true class, 0 elsewhere). Differentiating $\mathcal{L} = -z_y + \log \sum_j e^{z_j}$ with respect to $z_i$:

For the cat example that is (−0.3410, 0.2424, 0.0986). Gradient descent subtracts this, so it raises the logit of the true class and lowers the others, each in proportion to how much probability it wrongly holds. The gradient never explodes (every entry is between $-1$ and 1) and it vanishes only when the prediction is exactly right. This clean form is a big reason softmax and cross-entropy are always used together, and why libraries fuse them into one numerically stable function.

Things to try: press “All equal” to see the uniform baseline loss of $\ln 3$; press “Confidently wrong” and watch the loss climb; then drag the true class’s logit up and watch the gradient entries shrink towards zero.

References

1. J. S. Bridle, “Probabilistic interpretation of feedforward classification network outputs, with relationships to statistical pattern recognition”, in Neurocomputing, NATO ASI Series F68, Springer, 1990. The paper that named softmax.
2. C. Guo, G. Pleiss, Y. Sun and K. Q. Weinberger, “On calibration of modern neural networks”, ICML 2017. arXiv:1706.04599
3. I. Goodfellow, Y. Bengio and A. Courville, Deep Learning, MIT Press, 2016. Section 6.2.2.