AIM-02-004-P · Primer — Softmax & Cross-Entropy

The idea in one sentence

Softmax turns a list of raw, unbounded scores into a proper probability distribution that sums to one, and cross-entropy measures how much probability that distribution put on the correct answer.

Why it exists

A classifier's last linear layer produces logits — arbitrary real numbers like \([2.0,\ 1.0,\ 0.1]\) — with no guarantee they're positive or sum to anything meaningful. To train with likelihood-based methods and to report a calibrated "how sure am I" number, you need to convert those logits into a valid probability distribution (non-negative, sums to 1) and a loss that gets steeper the more confidently wrong the model is. Think of it like converting horse-race scores into betting odds: every horse needs a share of the total pool that adds up to 100%, and the favorite (highest score) should get an outsized share, not just a proportionally larger one — exponentiating the scores before normalizing is exactly what creates that "winner takes disproportionately more" effect. Cross-entropy then plays bookmaker: it charges you -log(probability you assigned to what actually happened), so being 99% confident and right costs almost nothing, but being 99% confident and wrong costs a fortune.

How it works

The equation:

\[\text{softmax}(z)_i = \frac{\exp(z_i - \max(z))}{\sum_j \exp(z_j - \max(z))}\]

Subtracting the max before exponentiating doesn't change the result (it cancels in the ratio) but keeps every exponent \(\le 0\), so \(\exp(\cdot)\) never overflows even for logits in the thousands. Cross-entropy for a single example with true class index \(y\) is then just \(L = -\log(p_y)\) — the negative log of the probability the model assigned to the correct class.

Worked example: logits \(= [2.0,\ 1.0,\ 0.1]\), true class = index \(0\). Subtract the max (\(2.0\)): shifted \(= [0.0,\ -1.0,\ -1.9]\). Exponentiate: \([1.0,\ 0.3679,\ 0.1496]\). Sum \(\approx 1.5174\). Divide each by the sum: probs \(\approx [0.6590,\ 0.2424,\ 0.0986]\) — they sum to \(1.0\). Cross-entropy \(= -\log(0.6590) \approx 0.417\) nats. If instead the true class were index \(2\) (the model's least-favorite guess), the loss would be \(-\log(0.0986) \approx 2.317\) — over five times larger, because the model was confidently wrong.

One more mechanism worth knowing: when softmax feeds directly into cross-entropy, the gradient with respect to the logits collapses to the remarkably clean form \((p - y)\), where \(y\) is the one-hot true label. That's why frameworks fuse the two operations into a single numerically-stable op rather than composing softmax and log separately — it's both faster and avoids taking \(\log(0)\) if a probability underflows to exactly zero.

See it — Logits become probabilities, then get scored against the true label

flowchart TD
    A["Logits<br/>(raw scores)"] --> B["Subtract max<br/>(for stability)"]
    B --> C["Exponentiate<br/>each value"]
    C --> D["Normalize<br/>(divide by sum)"]
    D --> E["Probabilities<br/>(sum to 1)"]
    E --> F{"True class<br/>index y"}
    F --> G["Take -log(p_y)"]
    G --> H["Cross-entropy<br/>loss"]

The one gotcha

The single most common bug: applying softmax yourself and then feeding those probabilities into a framework's built-in loss. PyTorch's nn.CrossEntropyLoss (and TensorFlow's from_logits=True losses) expect raw logits — they compute softmax internally, fused with the log for numerical stability. Pre-softmaxing before calling them applies softmax twice, silently flattening gradients and crippling training without ever throwing an error.

The one-line takeaway

Softmax normalizes; cross-entropy punishes confident wrongness logarithmically.

Where you'll meet it

PyTorch nn.CrossEntropyLoss / nn.LogSoftmax + nn.NLLLoss; TensorFlow/Keras tf.nn.softmax_cross_entropy_with_logits and SparseCategoricalCrossentropy(from_logits=True).

Named concepts: softmax cross entropy