Skip to content
mlmentorship

Why is softmax + cross-entropy the right pairing?

The gradient simplifies to (p - y), and that's not a coincidence. The senior answer derives this and connects to GLMs and numerical stability.

Published · 7 min read ·Core ·Advanced

Math / derivation · active recall

Practice before you read

15 minutes. State assumptions, derive cleanly, check dimensions or limiting cases, then interpret the result.

Back to my workbook How practice works

Math / derivation · closed-book attempt

Why is softmax + cross-entropy the right pairing?

State assumptions, derive cleanly, check dimensions or limiting cases, then interpret the result.

Problem

Context. A classifier produces one finite logit per class and observes one correct class index for each example. Softmax defines normalized class probabilities, and the objective is the negative log-probability of the observed class. Task. Derive the gradient with respect to every logit, then explain a numerically stable logits-based computation. Connect the result to the categorical likelihood and compare separate versus fused implementations. Constraints. Begin with one-hot targets and mutually exclusive classes; smoothing, weighting, and multilabel objectives require stated changes. Distinguish algebraic equivalence from floating-point behavior, and do not feed probabilities into an API that expects logits.

15:00recommended time

Closing or reloading clears the scratchpad. Only score, weak rubric dimensions, attempt count, and retry date can be stored locally.

Asked in: ML breadth, math-screen, and LLM internals interviews.

Derive the loss and logit gradient for mutually exclusive classification and explain why implementations combine softmax with cross-entropy.

Problem

Context. A classifier produces one finite logit per class and observes one correct class index for each example. Softmax defines normalized class probabilities, and the objective is the negative log-probability of the observed class.

Task. Derive the gradient with respect to every logit, then explain a numerically stable logits-based computation. Connect the result to the categorical likelihood and compare separate versus fused implementations.

Constraints. Begin with one-hot targets and mutually exclusive classes; smoothing, weighting, and multilabel objectives require stated changes. Distinguish algebraic equivalence from floating-point behavior, and do not feed probabilities into an API that expects logits.

The setup

C-class classification. Logits z = (z_1, ..., z_C). Softmax produces probabilities:

p_i = exp(z_i) / sum_j exp(z_j)

Cross-entropy with one-hot label y:

L = -sum_i y_i * log p_i = -log p_y

(where y is the index of the true class).

The gradient simplification

Compute dL / dz_k:

dL / dz_k = -d log p_y / dz_k
         = -(1 / p_y) * dp_y / dz_k

Two cases for dp_y / dz_k:

  • If k == y: dp_y / dz_k = p_y * (1 - p_y).
  • If k != y: dp_y / dz_k = -p_y * p_k.

Substituting:

  • dL / dz_y = -(1 / p_y) * p_y * (1 - p_y) = -(1 - p_y) = p_y - 1
  • dL / dz_k = -(1 / p_y) * (-p_y * p_k) = p_k

Combining: dL / dz_k = p_k - y_k where y_k = 1 for k = y, 0 otherwise.

The full gradient is p - y (predicted probabilities minus the one-hot true label). Three lines of algebra; the cleanest gradient in deep learning.

Learning objective

Gradient descent subtracts the gradient, raising the underpredicted true-class logit

Descent subtracts the gradient, so the true logit risesFor three equal logits and target class A, gradient is (-2/3,1/3,1/3). A learning rate of 0.3 produces logits (0.2,-0.1,-0.1).Descent subtracts the gradient fromeach logitClass A: gradient -0.6667New logit 0.2Class B: gradient 0.3333New logit -0.1Class C: gradient 0.3333New logit -0.1
Derivation
A three-class softmax cross-entropy gradient and its opposite-signed logit update For predicted probabilities A 0.70, B 0.20, and C 0.10 with B as the true class, subtracting the one-hot target gives gradient components positive 0.70, negative 0.80, and positive 0.10. A gradient-descent step moves in the opposite direction: A left by 0.70 eta, B right by 0.80 eta, and C left by 0.10 eta. The two competitor logits decrease, the true-class logit increases, and all three changes sum to zero. CLASS p y g = p − y LOGIT STEP −ηg A 0.70 0 +0.70 lower A: −0.70η B (true) 0.20 1 −0.80 raise B: +0.80η C 0.10 0 +0.10 −0.10η CHECK Σp = 1 Σy = 1 Σg = 0 Σ(−ηg) = 0 left = lower logit · right = raise logit · η > 0
subtract the one-hot target row by row, then reverse each sign for gradient descent. The two wrong-class logits move left, in proportion to their current probabilities; the true-class logit moves right by 0.80η. The updates sum to zero, matching softmax's invariance to a shared logit shift. These are logit changes, not direct probability changes. Original worked example checked against Goodfellow, Bengio, and Courville and the PyTorch CrossEntropyLoss documentation.
Check: Why does the true class's negative gradient produce a positive logit step under gradient descent?
Read it this way: For three equal logits and target class A, gradient is (-2/3,1/3,1/3). A learning rate of 0.3 produces logits (0.2,-0.1,-0.1).
Why does a negative gradient raise the true-class logit?

Benefits of the fused operation

“Three reasons the joint operation is preferred:

1. Numerical stability. Computing softmax then cross-entropy separately involves taking log(exp(...)), which can overflow or underflow. The joint operation uses log-sum-exp:

log p_y = z_y - log sum_j exp(z_j) = z_y - max_j z_j - log sum_j exp(z_j - max_j z_j)

Subtracting max_j z_j keeps the largest exponentiated argument at zero, avoiding overflow. Frameworks (PyTorch’s nn.CrossEntropyLoss) accept logits directly and apply this internally.

2. Computational efficiency. The joint operation skips computing the explicit probabilities (since the gradient p - y only needs p, computed on demand). Saves memory and a few flops.

3. The gradient is exact and stable. The p - y form has bounded magnitude (each element is in [-1, 1]), so gradients don’t explode at the loss layer.”

The L6 connection: GLM

“Softmax + cross-entropy is the multiclass generalization of sigmoid + binary cross-entropy, both of which are GLMs under their canonical link functions. The gradient simplification (predicted - true) * input is a property of all canonical-link GLMs, not just classification. Linear regression with MSE has the same gradient form (because MSE on a Gaussian noise model is the GLM with identity link).

Modern deep nets use sigmoid + BCE for binary classification and softmax + CE for multiclass classification because the pairings give simple gradients and stable computation.”

Tells that get you a strong-hire vote

  • You derive the gradient cleanly.
  • You name the log-sum-exp trick for numerical stability.
  • You connect to GLMs and canonical links.
  • You explain why frameworks fuse the operations.

Tells that get you down-leveled

  • “It just works” without derivation.
  • Computing softmax explicitly in code (in real systems, you should pass logits to the loss function).
  • No mention of numerical stability.
  • Confusion about which axis softmax operates on.

Common follow-up

“What’s wrong with using MSE for classification?”

The L6 answer:

“Two related problems. (1) Vanishing gradients on confident-wrong predictions: MSE gradient under sigmoid is proportional to (p - y) * p * (1 - p). When the model is very confident and wrong (p ≈ 1 for the wrong class), the p * (1 - p) term vanishes; the model can’t learn its way out. Cross-entropy’s gradient is p - y directly, which stays large precisely when the model is most wrong. (2) MSE assumes Gaussian noise; classification labels are categorical. MLE under the wrong noise model gives the wrong objective. Cross-entropy is MLE under the right (categorical) noise model.”

Related: entropy and mutual information, cross-entropy and softmax, derive logistic regression from MLE, and choose a loss function.

Found an error or stale claim? Report a correction or send a private note.