6 · Softmax + multi-class loss
Turn raw logits into a probability distribution, pair with categorical cross-entropy, and discover why the gradient simplifies beautifully.
Softmax converts a vector of raw scores (logits) into a probability distribution — and is paired with categorical cross-entropy for multi-class classification.
Without this:
Without softmax + cross-entropy you can't train classifiers with more than 2 classes.
In binary classification, a sigmoid output gives one probability p(class=1). For K classes, we need K numbers that:
- Are all non-negative
- Sum to exactly 1 (a valid probability distribution)
Softmax does this:
softmax(z)_i = exp(z_i) / Σ_j exp(z_j)
The network's final layer produces a vector of raw scores called logits — one per class. Softmax converts them to probabilities. The class with the highest probability is the prediction.
Paired with categorical cross-entropy (CCE) loss:
CCE = -Σ_i y_true_i * log(y_pred_i)
This measures how surprised the distribution y_pred is by the true label y_true. When the correct class gets probability 1.0, the loss is 0. When it gets probability near 0, the loss → ∞.
The combination of softmax + CCE has an elegant property that makes it the default: its gradient is simply y_pred - y_true.
Python (in browser)
Python runs entirely in your browser via Pyodide (~6 MB on first Run, cached after).
Python (in browser)
Python runs entirely in your browser via Pyodide (~6 MB on first Run, cached after).
Python (in browser)
Python runs entirely in your browser via Pyodide (~6 MB on first Run, cached after).
The math behind softmax + CCE elegance — and why they must be paired
PyTorch read-along — CrossEntropyLoss expects logits, not probabilities
What does the gradient of softmax + CCE simplify to?
- Softmax converts a logit vector to a valid probability distribution (non-negative, sums to 1) — mandatory before categorical cross-entropy.
- The numerical stability trick `exp(z - z.max())` prevents overflow when logits are large — always use it in custom implementations.
- Softmax + CCE's gradient simplifies to `y_pred - y_true` — one of the most elegant results in deep learning.
- PyTorch's `nn.CrossEntropyLoss` expects raw logits (not softmax probabilities) — double-softmax is a silent bug.
Every classifier with K > 2 classes uses softmax + cross-entropy. ImageNet (1000 classes), language modeling (50k vocab), intent classification all use this.
If you remove it: You'd have no principled way to train a 10-class classifier, and you'd fall back to one-vs-rest binary classifiers — 10× the training cost for a worse result.