Softmax and cross-entropy loss

Step 1 Understand it

The problem it solves. A classifier's last layer outputs arbitrary numbers, but you want probabilities, and a loss that strongly punishes confident wrong answers. Softmax plus cross-entropy is the standard pairing for every multi-class classifier and every language model, which predicts the next token from a softmax over tens of thousands of options.

Watch it happen

The idea

Softmax turns raw scores (logits) into probabilities that sum to 1: exp(zᵢ) / Σ exp(zⱼ). Cross-entropy loss is −log(probability of the correct class). Together, their gradient is simply: probabilities − one-hot label.

How it works, step by step

1. Start with logits, one per class: z = [2, 1, 0.1].
2. Exponentiate: e^z = [7.39, 2.72, 1.11]. All positive, and the leader pulls further ahead.
3. Divide by the total (11.21): p = [0.66, 0.24, 0.10]. They add up to 1.
4. For stability, subtract the largest logit first. The answer is identical, but exp never overflows.
5. If the true class is the first one, the loss is −log 0.66 ≈ 0.42.
6. The gradient with respect to the logits is p − y = [0.66 − 1, 0.24, 0.10] = [−0.34, 0.24, 0.10]: raise the right score, lower the others in proportion.

Picture it

flowchart LR
  Z["logits<br/>[2, 1, 0.1]"] -->|"exp"| E["[7.39, 2.72, 1.11]"] -->|"÷ sum 11.2"| P["probabilities<br/>[0.66, 0.24, 0.10]"]
  P -->|"right answer is class 0"| L["loss = −log 0.66 ≈ 0.42"]

Analogy

Turning vote counts into vote shares.

Where the analogy breaks

Real vote shares are proportional. Softmax is exponential, so a small lead in score becomes a big lead in probability.

Worked example

Logits [2, 1, 0.1] become about [0.66, 0.24, 0.10]. If the correct class is 0, loss = −log(0.66) ≈ 0.42, and the gradient is [0.66 − 1, 0.24, 0.10] = [−0.34, 0.24, 0.10].

z = z - z.max(axis=1, keepdims=True)                 # for numerical stability
p = np.exp(z) / np.exp(z).sum(axis=1, keepdims=True)
loss = -np.log(p[np.arange(len(y)), y]).mean()

Must-know points

• softmax(z)ᵢ = exp(zᵢ) / Σⱼ exp(zⱼ).
• Adding the same number to every logit doesn't change the softmax, so subtracting the max is free and safe.
• Loss = −log p(correct class). Its gradient on the logits is p − one_hot(y).
• In PyTorch, F.cross_entropy takes raw logits and applies log-softmax itself. Don't softmax first.
• Scaling logits by 1/T (temperature) makes the distribution sharper (T < 1) or flatter (T > 1). Used in sampling (Phase 4).
• At random initialisation with K classes, the loss starts near ln K.

Common mistakes

• Applying softmax and then F.cross_entropy, which applies it again.
• Computing exp of large logits without subtracting the max: overflow, inf, NaN.
• Taking log(softmax) in two steps instead of using log-softmax, which loses precision.

Step 2 Go further (optional)

Optional. The lesson above is written to be enough on its own. If you want a second explanation, start with the one marked Worth it.

• Watch

  Neural Networks Part 5: ArgMax and SoftMax

  StatQuest (Josh Starmer)

• Watch

  Neural Networks Part 6: Cross Entropy

  StatQuest (Josh Starmer)

• Read

  CS231n: Linear classification (the softmax classifier)

  Stanford

Step 3 Practise

Try each exercise yourself first: on paper, out loud or in code. Then open the answer and compare it with yours.

1. ExplainSay it in plain words, with no notes.

   Explain why softmax exaggerates the leader, with the vote-share picture.

   Show answer

   Softmax turns raw scores into shares that add up to 100%, like turning vote counts into vote shares. But it first raises e to the power of each score, and exponentials grow fast. So a small lead in scores becomes a big lead in shares: scores [2, 1, 0.1] give shares of about [66%, 24%, 10%], and scores [4, 2, 0.2] give about [87%, 12%, 1%].

   So softmax exaggerates the leader: a candidate ahead by a couple of points of score walks away with most of the vote. That makes the model's top choice clear and gives useful gradients, but it also means models can sound very confident even when the underlying scores are only slightly apart.

2. DeriveWrite the maths, or the steps, from memory.

   Derive dL/dz for softmax followed by cross-entropy, and show that it equals p − y. No notes.

   Show answer

   Softmax: pᵢ = e^(zᵢ) / Σₖ e^(zₖ). Cross-entropy with a one-hot y: L = −Σᵢ yᵢ log pᵢ.

   Step 1, the softmax derivative: ∂pᵢ/∂zⱼ = pᵢ(δᵢⱼ − pⱼ), where δᵢⱼ is 1 if i = j and 0 otherwise. (Raising zⱼ raises pⱼ and lowers every other share.)

   Step 2, the chain rule:
   ∂L/∂zⱼ = Σᵢ (∂L/∂pᵢ)(∂pᵢ/∂zⱼ) = Σᵢ (−yᵢ/pᵢ) · pᵢ(δᵢⱼ − pⱼ) = −Σᵢ yᵢ(δᵢⱼ − pⱼ) = −yⱼ + pⱼ·Σᵢ yᵢ.

   Step 3: the labels add up to 1, so Σᵢ yᵢ = 1, and ∂L/∂zⱼ = pⱼ − yⱼ.

   So the gradient is just "predicted probabilities minus the truth". For logits [2, 1, 0.1] and correct class 0: p ≈ [0.66, 0.24, 0.10], gradient ≈ [−0.34, 0.24, 0.10].

3. BuildCode it from scratch in NumPy, with no ML library.

   Implement a numerically stable softmax with cross-entropy, and train softmax regression on MNIST.

   Show answer

   Stable softmax subtracts each row's maximum before exp. For the loss, compute log-probabilities directly with log-sum-exp: log pᵢ = zᵢ − log Σ e^(zₖ), which never takes log of 0.

   Softmax regression is just one linear layer, 784 → 10, trained with gradient (p − Y) / n. On MNIST, five passes over the training set with mini-batches of 128 reach about 92% test accuracy, which is the classic number for a linear model. An MLP gets about 98%.

   pythonCopy

   import numpy as np
   from torchvision import datasets              # only to download MNIST
   
   train, test = datasets.MNIST("data", train=True, download=True), datasets.MNIST("data", train=False, download=True)
   Xtr, ytr = train.data.numpy().reshape(-1, 784) / 255.0, train.targets.numpy()
   Xte, yte = test.data.numpy().reshape(-1, 784) / 255.0, test.targets.numpy()
   
   def log_softmax(z):
       z = z - z.max(axis=1, keepdims=True)
       return z - np.log(np.exp(z).sum(axis=1, keepdims=True))
   
   rng = np.random.default_rng(0)
   W, b = np.zeros((784, 10)), np.zeros(10)
   for epoch in range(5):
       for idx in np.array_split(rng.permutation(len(Xtr)), len(Xtr) // 128):
           logp = log_softmax(Xtr[idx] @ W + b)
           loss = -logp[np.arange(len(idx)), ytr[idx]].mean()
           g = np.exp(logp); g[np.arange(len(idx)), ytr[idx]] -= 1; g /= len(idx)   # p − y
           W -= 0.5 * Xtr[idx].T @ g
           b -= 0.5 * g.sum(axis=0)
       acc = (np.argmax(Xte @ W + b, axis=1) == yte).mean()
       print(f"epoch {epoch + 1}: last batch loss {loss:.3f}, test accuracy {acc:.1%}")

4. BreakShow when it fails, and why.

   Compute softmax of [1000, 1000] without subtracting the max. Then compute log(softmax) directly and with log-sum-exp. Where does each version break?

   Show answer

   Softmax of [1000, 1000] without subtracting the max: exp(1000) overflows to inf, and inf / (inf + inf) = nan. The right answer is obviously [0.5, 0.5]. Subtracting the max turns it into softmax of [0, 0], which is exactly [0.5, 0.5].

   log(softmax) computed directly breaks the other way. For logits [0, −1000], softmax gives [1, 0] (e⁻¹⁰⁰⁰ underflows to 0), and log(0) = −inf. If the correct class is the second one, the loss becomes inf and the gradients nan. Log-sum-exp computes log p = z − logsumexp(z) = [0, −1000] exactly, with no inf anywhere.

   That is why frameworks fuse them: PyTorch's CrossEntropyLoss takes raw logits and does the stable log-softmax inside. Never apply softmax yourself before it.

   pythonCopy

   import numpy as np
   
   z = np.array([1000.0, 1000.0])
   with np.errstate(over="ignore", invalid="ignore"):
       print(np.exp(z) / np.exp(z).sum())                    # [nan nan]
   print(np.exp(z - z.max()) / np.exp(z - z.max()).sum())     # [0.5 0.5]
   
   z = np.array([0.0, -1000.0])
   p = np.exp(z - z.max()) / np.exp(z - z.max()).sum()
   with np.errstate(divide="ignore"):
       print(np.log(p))                                       # [  0. -inf]
   print(z - (z.max() + np.log(np.exp(z - z.max()).sum())))   # [    0. -1000.]

Step 4 Prove it

Check you understood

Tick each check when you can do it without notes