Softmax & Cross-Entropy Loss

A classifier's last layer produces raw scores, one per class. Softmax turns them into probabilities, and cross-entropy scores how much probability went to the right class. After this lesson your program computes both without overflowing, and also the gradient that starts backpropagation.

Softmax: scores to probabilities

Given logits z[0..n-1]:

softmax(z)[i] = e^z[i] / ( e^z[0] + e^z[1] + ... + e^z[n-1] )

Every output is positive, the outputs sum to 1, and the biggest logit gets the biggest probability. Adding the same constant to every logit changes nothing, because e^(z+c) = e^c * e^z and the factor cancels between top and bottom. That freedom is the whole stability trick.

The overflow trap and the max-shift

e^1000 is far beyond the largest 64-bit float, so logits of [1000, 1000] produce infinity divided by infinity: nan. Subtract the largest logit first. Then the biggest exponent is e^0 = 1, nothing overflows, and the denominator is at least 1. For [1000, 1000] the shifted logits are [0, 0], the exponentials [1, 1], and the answer [0.5, 0.5].

Cross-entropy: how surprised was the model?

With the true class c and probabilities p, the loss is

loss = -ln( p[c] )

It is 0 when the model gave the right class probability 1 and grows without bound as that probability approaches 0. The natural log gives units of nats; every framework uses it.

A worked trace

Logits z = [0.5, 1.5, 3.0], true class 1.

shift by max (3.0):   -2.5   -1.5    0.0
exponentials:          0.0821  0.2231  1.0       sum = 1.3052
probabilities:         0.0629  0.1710  0.7662
loss = -ln(0.1710) = 1.7664

The model put 77% on class 2 and only 17% on the true class 1, so the loss is large.

Compute the loss without forming the probability

If a logit is far below the rest, its probability underflows to exactly 0 and -ln(0) is infinite, even though the true loss is a perfectly finite number. Rewrite the loss to avoid ever taking the log of a tiny probability:

loss = ( max + ln( sum of e^(z[i] - max) ) ) - z[c]

The bracket is the log-sum-exp of the logits, computed on the shifted values. For z = [0, -1000] and true class 1 this gives exactly 1000, whereas "softmax, then log" gives infinity.

The gradient is the prediction error

Differentiate softmax and cross-entropy together and the messy Jacobians cancel:

d loss / d z[i]  =  p[i] - (1 if i == c else 0)

For the trace above: [0.0629, -0.8290, 0.7662]. The true class gets pushed up (negative gradient), every other class pushed down in proportion to its probability. This one subtraction is the starting point of the backward pass for every classifier, which is why frameworks fuse the two functions into one operation. If an example can carry several labels at once ("dog" and "outdoor"), use one independent sigmoid per label instead; softmax forces the labels to compete.

Common mistakes

• Skipping the max-shift and getting nan for large logits.
• Taking -ln(p[c]) from rounded or underflowed probabilities.
• Using base-10 or base-2 logs.
• Forgetting that the gradient at the true class is p - 1, which is negative.
• Computing the max over the wrong vector after a new LOGITS line.

Your exercise

Commands set LOGITS and LABEL, then ask for PROBS, LOSS or GRAD. The starter parses and prints. Implement softmax, loss (use the log-sum-exp form) and grad.