---
title: "Softmax & Cross-Entropy"
description: "Turn logits into probabilities, then score them against the correct class."
source: https://vizlearn.app/ai/neural-networks/softmax-cross-entropy
access: premium (interactive player requires Lifetime; this doc is free)
---

# Softmax & Cross-Entropy

Softmax turns a vector of raw scores into a probability distribution by exponentiating each score and dividing by their sum, so the ten outputs land in `[0, 1]` and sum to 1, up to floating-point rounding. Cross-entropy then scores that distribution against the correct label as `−log p_y`, and its gradient with respect to the scores collapses to `p − y`.

That last sentence is the whole reason the two are commonly paired, and it is easy to read past. Softmax is a soft `argmax`: it takes a lead in score space and turns it into a lead in probability space, exponentially. Cross-entropy is a fine that grows without bound as the true class gets less likely. Bolt them together and something remarkable happens on the way back down — the derivative of the loss with respect to the scores is one subtraction, with no trace of softmax left in it. Swap cross-entropy for the mean-squared error you learned in a regression class and that cancellation disappears, taking most of your gradient with it. On the digit this lesson runs — a handwritten 3 the network calls a 2 at 99.86% confidence — cross-entropy pushes back 361 times harder than MSE, from identical probabilities.

## How Softmax and Cross-Entropy Work

The last layer of a classifier emits one raw score per class, called a logit. Logits are unbounded and signed: on the digit this lesson runs, eight of the ten are negative and the whole vector sums to `−29.52`. They are not probabilities and no amount of squinting makes them probabilities.

Softmax fixes that in three moves. Exponentiate every score, which makes them all positive and blows small differences into large ratios. Add them up. Divide. The result is non-negative and sums to 1 by construction, because you literally divided by the sum.

There is a fourth move hiding in front of all of them: subtract the largest score from every score first. This changes nothing — multiplying numerator and denominator by the same `e^(−c)` cancels — so `softmax(s) = softmax(s − c)` for any constant `c`. What it buys you is that every exponent is now at most 0, so `np.exp` can never overflow. Without the shift, a logit of 800 (rare, but reachable in a badly initialized network) gives you `inf / inf = NaN` and a training run that dies for reasons that look nothing like the actual cause.

Cross-entropy reads the finished distribution and asks one question: how much probability did you assign to the class that was actually right? The general form `−Σ yᵢ log pᵢ` is a sum over all ten classes, but nine of those terms are multiplied by a zero in the one-hot target, so it collapses to `−log p_y`. One number in, one number out. As `p_y` slides toward zero, the loss climbs without limit; there is no worst possible score.

Then comes the part worth memorizing. Differentiate the loss with respect to the *scores softmax was handed* — through the log, through the division, through the exponentials — and every softmax term cancels:

`∂L/∂s = p − y`

The predicted distribution minus the one-hot target. That is the entire backward pass through this layer.

`s` is whatever you feed softmax. With no temperature knob that is just the logits, which is why you will usually meet this written `∂L/∂z = p − y`. Divide by a temperature `T` first, as this lesson lets you, and the gradient with respect to the raw logits picks up a factor of `1/T` — the shape of the answer does not change, only its scale.

## Softmax and Cross-Entropy Step by Step

The default run is a real MNIST test image of a handwritten **3** that the shipped 784→25→25→10 network classifies as a **2**. Every number below comes from that forward pass at temperature `T = 1.0`.

**1. The logits.** Ten raw scores: `z₂ = 10.00` towers over everything, `z₃ = 3.42` is the runner-up, and the rest sit between `−1.79` and `−18.15`.

**2. The shift.** Subtract `10.00` from all ten. Class 2 becomes 0, class 3 becomes `−6.578`, everything else goes further negative. The probabilities are unchanged; the overflow risk is gone.

**3. Exponentiate.** `e⁰ = 1` for class 2. For class 3, `e^(−6.578) = 0.0014` — the 6.578-point lead has become a factor of **719**. Eight of the ten exponentials are now below `1e-5`. This is the beat where "soft argmax" stops being a slogan.

**4. Sum.** The ten exponentials add to `Σ = 1.0014`. One class already carries essentially all the mass, so the normalizer barely does anything. Hold on to that: a normalizer this close to 1 is the visible signature of a saturated softmax, and saturation is what breaks the last act.

**5. Divide.** `p₂ = 0.9986`, `p₃ = 0.0014`, and the other eight classes are worth `1.57e-5` between them. Non-negative, sums to 1, done.

**6. Predict, and compare.** `argmax p = 2` at 99.86% confidence. The label is 3. This is not a near miss; the network is confidently, expensively wrong.

**7. Both losses.** Cross-entropy reads one number: `−log(0.0014) = 6.5794` nats. MSE reads all ten: `½Σ(pᵢ − yᵢ)² = 0.9972`. Against a one-hot target MSE cannot exceed 1.0, so it is reporting 99.7% of the worst score it is capable of producing. Both losses identify a poor prediction, but cross-entropy has no maximum and can score much worse ones.

**8. One operator, two roads.** Both gradients reach the scores through the same matrix: `∂L/∂s = J · ∂L/∂p`, where `J = diag(p) − p pᵀ` is softmax's Jacobian. What differs is what enters it. Cross-entropy differentiates to `−1/p₃ = −720.11` at the true class and to exactly zero at the other nine, so it touches one column of `J` — and every column of `J` carries a factor of the probability it belongs to. Here that factor is `p₃`, and `−1/p₃` is its exact reciprocal: `(−1/p₃) · p₃(1 − p₃) = p₃ − 1 = −0.9986`. The saturation cancels. MSE arrives at the same matrix carrying `p₃ − 1 = −0.9986`, an ordinary number with no `1/p` anywhere in it, so all ten columns contribute and the answer comes out as `p₃ · ((p₃ − 1) − pᵀ(p − y))` — with the `p₃` still standing.

**9. Both gradients, on one axis.** Cross-entropy hands the optimizer `−0.9986` at the true class and `+0.9986` at the impostor: push class 3 up as hard as possible, push class 2 down just as hard. MSE hands back `−0.0028`. Same probabilities, same label, same everything — and a push **361× weaker** on the one example that most needed correcting.

That is the previous step's bracket, evaluated: `0.0014 × (−1.9944) = −0.0028`. Both losses met the same Jacobian. Cross-entropy walked out with `p − y` whole because it arrived holding the reciprocal of the very factor the matrix applies; MSE arrived holding a number of ordinary size and paid the full `p₃`.

## Why Cross-Entropy and Not MSE

The usual explanation is that cross-entropy is "the right loss for probabilities" and MSE is "for regression", which is true and teaches nothing. The operational answer is that MSE's gradient vanishes precisely where you need it most.

Look at what MSE reports on this digit versus a digit the network gets right at `p = 1.0`. The **losses** separate cleanly: `0.9972` against `0.0000`. The **gradients** do not: `0.0028` against `9.6e-22`. On any axis that shows cross-entropy's push, both are a flat line. MSE can tell you that you are wrong; it cannot tell your optimizer what to do about it. A confidently wrong prediction can produce a small but nonzero update, despite needing substantial correction.

You can prove to yourself that this is about saturation rather than about MSE being cursed. Raise the temperature to `T = 3.0`, which divides the logits before softmax sees them and loosens the distribution without retraining anything. `p₃` climbs from `0.0014` to `0.0935`, and the gap between the two gradients falls from 361× to **6.4×**. Push the other way, to a digit the network is maximally confident about, and raising the temperature *opens* a gap where there was none. The rule is that the gap tracks saturation, not that temperature shrinks it.

There is a cost argument too, and it is worth getting right, because it is usually overstated. As written in the code above, `grad_mse = J @ (p - t)` builds a dense `C × C` matrix — the outer product `p pᵀ` fills in every entry, so `J` is not diagonal — and that line is `O(C²)` in both time and memory. At ten classes nobody notices. At a 50,000-token vocabulary the matrix alone is 2.5 billion entries — about 20 GB in float64 for every example.

But the matrix is avoidable. `J(p − y)` expands to `p ⊙ ((p − y) − pᵀ(p − y))`, which is `O(C)` and never forms `J` at all — that identity is how this lesson computes the MSE gradient. So the quadratic cost belongs to the naive implementation, not to MSE. The vanishing gradient does belong to MSE, and that one you cannot rewrite your way out of.

## Edge Cases

- **`p_y` underflows to zero.** Then `−log(0)` is `+inf` and the run is over. Logit-based losses avoid computing softmax and then taking its log; they fuse the two into a log-sum-exp form (PyTorch's `cross_entropy`, TensorFlow's `softmax_cross_entropy_with_logits`) that stays finite. If you are hand-rolling this pair, prefer the fused version.
- **Deleting the max-shift.** The probabilities do not change — the shift is exact, not an approximation. On this example it would not even overflow: the largest score is `10.00` and `e^10 ≈ 22026`, perfectly ordinary. The guard matters near `e^710`, where float64 runs out of exponent.
- **A correct, maximally confident prediction.** At `T = 1` the confident 6 has `p_y ≈ 1 − 2.5e-11`, which prints as `1.0000`; cross-entropy's loss and gradient are both about `2.5e-11`. That is the honest answer: there is almost nothing to learn from an example you already have right. When both gradients fall below `1e-6` the lesson suppresses the CE/MSE ratio, because a quotient of two negligible gradients says nothing useful about the losses.
- **Non-one-hot targets.** Label smoothing and distillation replace `y` with a soft target. `∂L/∂s = p − y` still holds for a normalized target distribution; only `y` changes. The `−log p_y` shortcut does not, since more than one term survives.

## Common Mistakes

- **Applying softmax twice.** PyTorch's cross-entropy and interfaces explicitly configured for logits take raw scores and handle softmax internally; other APIs may expect probabilities. Feeding it probabilities runs softmax on a softmax, which flattens the distribution and silently weakens every gradient. The model still trains, badly, which is the worst kind of bug.
- **Reading cross-entropy in bits.** `np.log` is the natural log, so the loss comes out in nats. A "loss of 2.3" is `ln 10`, meaning the model is doing about as well as guessing among ten classes. Divide by `ln 2` if you want bits.
- **Treating the logit gap as the probability ratio.** A 6.578-point lead is not a 6.6× lead. Exponentiation turns differences into ratios, and 6.578 points is 719×. This is why small changes in the last layer's weights can flip a confident prediction.
- **Auto-scaling the probability axis when you plot it.** A saturated distribution and a flat one look identical if the y-axis fits itself to the data. Pin the axis to `[0, 1]` — the shape of the distribution is the thing you are trying to see.
- **Concluding that MSE is broken.** It is not; it is mismatched. MSE is a common loss for linear-output regression. Paired with softmax it inherits a Jacobian that annihilates the gradient in a saturated regime; how often a classifier visits it depends on training.

The animation simplifies in three ways worth naming: it scores one example rather than a mini-batch, it stops the gradient at the logits instead of continuing into the weights (that is [backpropagation](https://vizlearn.app/ai/neural-networks/backpropagation)'s job), and its temperature knob rescales fixed logits rather than retraining the network. The arithmetic, the probabilities, and both gradient vectors are real, computed from a real trained model on a real test image.