FLAGSHIP FOUNDATION

Why softmax cross-entropy produces p − y

Download the CPU example

The problem

A classifier emits logits: unrestricted scores, one per class. We need one loss that rewards the target score, compares it with every competing score, and gives a gradient for all scores at once.

Consider logits z = [2, 1, 0] and target class 1. The largest score belongs to class 0, so this example is wrong. A useful gradient must lower class 0 relative to class 1. It should also say what to do with class 2.

Applying softmax alone gives probabilities but no training objective. Applying a logarithm directly to raw logits is invalid because logits can be negative and do not sum to one. Softmax and cross-entropy solve different parts of the problem and simplify when combined.

Work through a small example

Subtract the largest logit before exponentiating. This does not change softmax, because adding the same constant to every logit cancels between numerator and denominator.

1. Shift: [2, 1, 0] − 2 = [0, −1, −2].
2. Exponentiate: [1, 0.3679, 0.1353].
3. Normalize by 1.5032: p = [0.6652, 0.2447, 0.0900].
4. Encode class 1 as y = [0, 1, 0].
5. Take the negative log probability: L = −log(0.2447) = 1.4076.

The target component is negative: 0.2447 − 1 = −0.7553. Gradient descent subtracts this negative value, so it raises the target logit. The two non-target components are positive, so gradient descent lowers them.

The three components sum to zero. This is expected: shifting every logit by the same amount cannot change the probabilities or loss, so the gradient has no component in the all-ones direction.

The general rule

Let zi be logit i. Let yi be the target distribution, with non-negative entries that sum to one. Softmax defines:

pi = exp(zi) / Σk exp(zk)

Cross-entropy is L = −Σi yi log pi. Substitute the softmax logarithm:

L = log(Σk exp(zk)) − Σi yizi

The simplification uses Σ yi = 1. Differentiate with respect to one logit zj.

1. The derivative of log-sum-exp is exp(zj) / Σ exp(zk) = pj.
2. The derivative of −Σ yizi is −yj.
3. Therefore ∂L / ∂zj = pj − yj.

This is the derivative with respect to logits, not model weights. Backpropagation still multiplies it by the Jacobian of the logits with respect to each earlier activation or parameter.

Implement and inspect

The complete program computes the loss without taking log(softmax(...)) as two fragile operations. It then compares the derived gradient with autograd.

import torch
from softmax_p_minus_y import softmax_cross_entropy

logits = torch.tensor([[2.0, 1.0, 0.0]], requires_grad=True)
targets = torch.tensor([1])
loss, probabilities, manual = softmax_cross_entropy(logits, targets)
(automatic,) = torch.autograd.grad(loss, logits)

print(probabilities)
print(loss)
print(manual)
torch.testing.assert_close(manual, automatic)

Run from the repository root with python examples/flagship/softmax_p_minus_y.py. The tested output is approximately:

probabilities: [0.6652, 0.2447, 0.09]
loss: 1.4076
p - y: [0.6652, -0.7553, 0.09]
matches autograd: True
gradient sum: 0.0

For a batch shaped (B, C), probabilities and gradients also have shape (B, C). If the loss uses a batch mean, the returned gradient is (p − y) / B. A sum reduction does not include that factor.

Engineering checks

• Use logits as input. PyTorch cross_entropy expects raw logits. Applying softmax first changes the function and usually weakens numerical stability.
• Check shift invariance. Adding 10,000 to every class score should not change the probabilities, loss, or gradient.
• Check the reduction. Doubling batch size can halve each row's gradient under mean even when the examples are copied.
• Check target shape and meaning. Integer targets name one class. Distribution targets must be normalized and use the same class order as the logits.
• Mask before reducing. For padded sequence targets, remove ignored positions from both the loss sum and the denominator.

A finite-difference check is useful for a custom loss. Perturb one logit by a small ε and compare (L(z+ε) − L(z−ε)) / (2ε) with the analytic component. Use float64 for this test and try more than one ε.

Go deeper

Soft labels keep the same form

Label smoothing replaces a one-hot vector with a distribution such as [0.05, 0.90, 0.05]. The gradient remains p − y because the target still sums to one. The update no longer asks for target probability exactly 1.

Class weighting changes the scale. Temperature-scaled softmax with p = softmax(z/τ) introduces a factor 1/τ when differentiating with respect to z. Do not repeat the plain shortcut without accounting for the modified definition.

The curvature has structure

The Hessian with respect to logits is diag(p) − ppT. It is positive semidefinite, so cross-entropy is convex as a function of the logits. It has a zero direction corresponding to adding the same constant to all logits.

Each logit-gradient component lies between −1 and 1 for a single example, but this does not guarantee small parameter gradients. Earlier Jacobians, batch reduction, sequence length, mixed precision, and loss scaling all affect the final parameter update.

Know where the result does not apply

Independent multi-label classification uses one sigmoid per label and binary cross-entropy. Its component gradient is also prediction minus target, but there is no shared softmax normalization across classes. Structured losses, focal loss, sampled softmax, and contrastive objectives have different derivatives.

Practice and recap

Question: For two logits [0, 0] and target class 0, calculate p, y, and p − y. Which direction does gradient descent move each logit?

References and limits

• PyTorch cross-entropy documentation for current API shapes and reduction behavior.
• PyTorch log-sum-exp documentation for the stable primitive used in the example.

This article derives one loss at the logit level. It does not claim that cross-entropy alone produces calibrated probabilities, solves class imbalance, or guarantees useful features.