DEV Community

Devanshu Biswas
Devanshu Biswas

Posted on

Softmax and cross-entropy: the clean p y gradient that trains every classifier

#machinelearning #deeplearning #ai #neuralnetworks

Softmax and cross-entropy are almost always taught together, and there is a beautiful reason for it that a lot of tutorials skip over. Individually they are a little fiddly — softmax has a full Jacobian, cross-entropy has a 1/p in its derivative. But fuse them and the entire gradient collapses to one subtraction. I built an interactive page where you drag four logits and watch the probabilities, the loss, and that gradient update live, and the cancellation is the whole point.

Logits are not probabilities

A classifier's last layer emits logits — unbounded real numbers, one per class, that mean nothing on their own. Softmax turns them into a probability distribution: exponentiate, then normalize so they sum to one.

def softmax(z):
    z = z - np.max(z, axis=-1, keepdims=True)   # shift for stability
    e = np.exp(z)                                # largest term is now exp(0)=1
    return e / e.sum(axis=-1, keepdims=True)     # sums to 1
Enter fullscreen mode Exit fullscreen mode

Exponentiating keeps everything positive and preserves the ordering (bigger logit → bigger probability), so softmax never changes which class wins, only how the confidence is expressed. It is a smooth, differentiable stand-in for argmax.

The stability trick you must not skip

exp(1000) is infinity in floating point, and inf/inf is NaN — your loss dies. Real logits reach the hundreds, so this is not hypothetical. The fix is free, because softmax is shift-invariant: subtract any constant from every logit and the result is unchanged, since the constant cancels top and bottom. Choose the max. Now the largest shifted logit is 0, its exp is exactly 1, and nothing overflows. The demo shows this explicitly — it computes the naive exp(max) next to the shifted version so you can see the overflow it dodges. Every production softmax does this.

Cross-entropy is the surprise of the correct answer

Cross-entropy measures how far the predicted distribution is from the truth: L = −Σ yᵢ log(pᵢ). Because the true label is one-hot, every term vanishes except the true class, so it collapses to −log(p_true) — the negative log-likelihood of the right answer.

def cross_entropy(p, y_true):
    return -np.log(p[y_true] + 1e-12)   # only the true class contributes
Enter fullscreen mode Exit fullscreen mode

Assign 100% to the truth and the loss is 0; assign near-zero and it blows up. It punishes confident mistakes brutally, which is exactly the behaviour you want during training.

The punchline: the gradient is just p − y

Here is why the pair is inseparable. Write the loss as L = −z_t + log Σ exp(z_j) and differentiate with respect to the logits. The first term gives −1 only for the true class; the second gives pᵢ for every class. Everything else — the softmax denominator, the exponentials, the 1/p — cancels:

def softmax_ce_grad(z, y_true):
    p = softmax(z)
    onehot = np.zeros_like(p)
    onehot[y_true] = 1.0
    return p - onehot          # the entire gradient, dL/dz
Enter fullscreen mode Exit fullscreen mode

Predicted distribution minus the one-hot target. In the demo you can read it directly off the bars: the true class gets p_true − 1 (negative, so gradient descent pushes that logit up), every wrong class gets +pᵢ (positive, pushed down), and the gradients sum to zero — a pure redistribution of probability toward the truth. The bigger the mistake, the larger the corrective step, and one global learning rate works because no stray factors survive.

Feed raw logits, never pre-softmax

In practice you almost never call softmax before the loss. Frameworks fuse log_softmax + nll_loss using the log-sum-exp trick, which is numerically stable by construction and hands you the exact p − y gradient for free:

loss = F.cross_entropy(logits, target)     # pass RAW logits, target is an index
loss.backward()
print(logits.grad)   # tensor([[-0.354, 0.238, 0.097, 0.020]])  == p - onehot

# DON'T:  F.cross_entropy(F.softmax(logits, -1), target)   # double-softmax bug
Enter fullscreen mode Exit fullscreen mode

Applying softmax yourself first is the classic bug — it double-softmaxes and quietly hurts training.

Where it lives

This one loss is the workhorse of deep learning. Any multi-class classifier ends in logits → softmax → cross-entropy. And a language model is just a giant classifier over the vocabulary: at each position it emits logits over ~50,000 tokens, softmax turns them into next-token probabilities, and cross-entropy against the actual next token is the entire pretraining objective. The p − y gradient is literally what trains GPT-scale models, one token at a time. Learn this single loss and its gradient and you understand how nearly every modern model is trained.

Drag the four logits, pick the true class, and watch the stable softmax, the loss, and the p − y gradient update together:
https://dev48v.infy.uk/dl/day34-softmax-cross-entropy.html

Top comments (0)