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 · 5 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.

15:00recommended time

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

30-second answer map

Visual first · depth when needed

Given a multiclass probability vector and one-hot target, interpret p - y as the logit update: gradient descent raises the true-class logit, lowers every competing logit in proportion to its probability, and preserves a zero-sum shift.

Preparing the visual…

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

The L4 candidate states the pairing. The L6 candidate derives the gradient simplification, explains the GLM connection, and discusses why frameworks compute the joint operation rather than the two separately.

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

What does p - y make gradient descent do to each logit?

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
Read it this way: 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.

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.