Softmax, Cross-Entropy & Optimization

In the previous lesson, a model predicted one number: next-token latency in milliseconds. Its loss could say "increase the prediction." A code-review classifier now faces a different job. Given a short review note, should it label the issue bug, docs, or security?

Suppose the evidence says docs is correct, but the model currently favors bug. We need a score for how bad that choice is, and a gradient that says which competing scores to raise or lower. Softmax turns raw scores into probability shares. Cross-entropy measures how little probability reached the correct choice. Together they form the standard categorical output-and-loss pair taught for neural classifiers and language models.^{[1]Reference 1Deep Learning.https://www.deeplearningbook.org/}

A decision begins as three raw scores

The final layer of a classifier emits one number per possible label. These numbers are logits: unconstrained scores, not probabilities.

A larger logit means the model prefers a label relative to its competitors. It doesn't mean 3.0 is 300%, and a logit can be negative without creating a negative probability.

This first snippet makes the mistake visible: the scores select a favorite, but don't satisfy probability rules.

Softmax assigns probability shares

Softmax takes every logit $z_i$, exponentiates it, and divides by the sum of all exponentiated logits:

$$p_i = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}}$$

Here $p_i$ is the probability assigned to label $i$, $z_i$ is its raw logit, and $K$ is the number of choices. Exponentiation makes each share positive. Dividing by their total makes the shares sum to one.

For our three class scores, the arithmetic is small enough to do by hand:

The model places only 0.114 probability on the correct docs label. That's the quantity the loss must expose.

Use NumPy to reproduce the table. This version is already stable because it subtracts the largest logit before taking exponentials.

Stable arithmetic matters before scale does

Softmax depends on differences between logits, not on their absolute offset. Adding or subtracting the same constant from every score leaves every probability unchanged:

$$\frac{e^{z_i-c}}{\sum_j e^{z_j-c}} = \frac{e^{z_i}e^{-c}}{e^{-c}\sum_j e^{z_j}} = \frac{e^{z_i}}{\sum_j e^{z_j}}$$

That cancellation gives us a safety rule: subtract the largest logit before exponentiating. The largest exponent becomes $e^0 = 1$, while the relative differences stay intact.

This failure case uses the same score gaps shifted upward by 997. Direct exponentiation overflows; max-shifting produces the original distribution.

Cross-entropy is most stable when computed straight from logits. For a correct class index $y$, the loss can be written without first materializing a tiny probability:

$$L = \log\left(\sum_j e^{z_j}\right) - z_y$$

The first term is called log-sum-exp. Its stable implementation uses the same maximum $m = \max_j z_j$:

$$\log\left(\sum_j e^{z_j}\right) = m + \log\left(\sum_j e^{z_j-m}\right)$$

This code computes the loss for the correct docs label directly from logits. Notice that it works for ordinary and extremely large offsets.

Cross-entropy grades the correct label

The supervised label for this example is docs. Written as a one-hot target vector in the same label order, it's:

$$y = [0,\ 1,\ 0]$$

One-hot means that exactly one class receives target weight one. Cross-entropy compares this target with the predicted probability vector:

$$L = -\sum_i y_i \log(p_i)$$

Every term except docs is multiplied by zero, so the loss reduces to:

$$L = -\log(p_{\text{docs}}) = -\log(0.114) \approx 2.170$$

Had the label been the model's preferred bug, the loss would have been only $-\log(0.844) \approx 0.170$. The loss is high because the supervised answer received little probability, not because the top choice merely happened to be wrong.

For a shorter arithmetic exercise, consider a second example with logits [1.0, 0.0, -2.0] in the same label order. bug still leads, but by a smaller margin. Compare the loss under two possible labels.

Now compute it:

Why squared probability error is a weak fit here

The earlier latency lesson used half squared error for a continuous target in milliseconds. We could attach softmax to a classifier and square its probability error, but the softmax derivative then shrinks the signal when a wrong class is already saturated near probability one. Cross-entropy paired with softmax avoids that extra shrinkage: its logit gradient will be $p-y$.^{[1]Reference 1Deep Learning.https://www.deeplearningbook.org/}

This example makes the difference visible. The logits are extremely sure about the wrong bug label while the target is docs.

Cross-entropy isn't the only possible classification objective, but it gives a direct categorical likelihood objective and a useful correction even when a wrong option dominates.

The gradient says exactly what to correct

For one labeled example, combine stable softmax with cross-entropy:

$$L = \log\left(\sum_j e^{z_j}\right) - z_y$$

Differentiate with respect to any logit $z_k$:

$$\frac{\partial L}{\partial z_k} = \frac{e^{z_k}}{\sum_j e^{z_j}} - \mathbf{1}[k=y] = p_k - y_k$$

This derivative is commonly summarized component-wise as $p_i - y_i$: predicted probability share minus target share.

The indicator $\mathbf{1}[k=y]$ equals one only for the correct class. For our docs example:

Logits themselves are outputs, not normally the parameters an optimizer stores. Updating them below is a diagnostic shortcut: it shows the direction that upstream weight updates are trying to produce on the next forward pass.

As in the previous lesson, a finite-difference test catches sign mistakes in a new loss implementation. Nudging each logit by a tiny amount should agree with $p-y$.

PyTorch receives logits directly

Training code shouldn't implement this loss from scratch unless you're testing your understanding or building a custom variant. In PyTorch, nn.CrossEntropyLoss receives raw logits for class-index targets and computes the equivalent of log-softmax followed by negative log-likelihood internally. Passing already-softmaxed probabilities changes the function being optimized.^{[2]Reference 2Optimizing Model Parameters.https://docs.pytorch.org/tutorials/beginner/basics/optimization_tutorial.html}

First, confirm that PyTorch returns the same loss and $p-y$ gradient as our hand calculation:

Now reproduce a common bug. The incorrect call passes probabilities where the loss expects logits, effectively normalizing an already normalized output again.

Shapes won't warn you about this bug. Keep the model's final classification layer linear during training and feed those raw logits into cross-entropy.

Temperature changes a distribution, not a label

Softmax can be made sharper or flatter by dividing logits by a positive temperature $T$:

$$p_i(T) = \frac{e^{z_i/T}}{\sum_j e^{z_j/T}}$$

When $T$ is below one, logit differences grow and the leading option takes more probability. When $T$ is above one, differences shrink and more probability remains on alternatives. Hinton, Vinyals, and Dean use this operation to reveal softer class relationships during knowledge distillation.^{[3]Reference 3Distilling the Knowledge in a Neural Network.https://arxiv.org/abs/1503.02531} In text generation, the same mathematical scaling can shape the distribution a decoding policy samples from. Softmax still doesn't choose a token by itself.

Use the same stable function to check each distribution:

A sequence supplies many supervised positions

Our classifier example produced one loss. A language model trains on the same categorical idea many times at once: after each context position, it produces logits for the next token, and the observed next token supplies the class label. A sequence model still needs a mechanism for history, which is the question the next lesson will address.

To make the language-model version literal, use a tiny vocabulary: bug, docs, and today. Suppose a two-token reply should be docs today. At each supervised position, the model produces a new logit vector over that same vocabulary:

This final NumPy example applies the same stable loss across two supervised reply positions. Each row holds logits over the same tiny vocabulary; averaging yields the one scalar sent backward.

Debugging checklist

Mastery check

Key concepts

• Logits are unconstrained scores whose differences determine probability shares.
• Stable softmax subtracts the maximum logit before exponentiation.
• Cross-entropy with a class label is the negative log probability of that label.
• Softmax plus cross-entropy produces the logit gradient $p-y$.
• PyTorch categorical cross-entropy receives raw logits, not probabilities.
• Temperature rescales logits; decoding decides how a token is selected.
• Sequence training averages this same categorical loss across supervised positions.

Evaluation rubric

• Foundational: Compute the stable softmax probabilities and cross-entropy loss for the docs example by hand.
• Intermediate: Explain why max shifting leaves probabilities unchanged and diagnose an overflowing implementation.
• Intermediate: Derive $p-y$ and explain the sign of each label gradient when bug is wrongly preferred.
• Advanced: Verify the gradient numerically, use the PyTorch interface correctly, and connect per-position loss to next-token sequence training.

Common pitfalls

• Symptom: A logit is reported as a confidence percentage. Cause: Raw output scores were confused with normalized probabilities. Fix: Apply softmax for interpretation and treat calibration as a separate evaluation question.
• Symptom: Custom loss returns NaN on large scores. Cause: Direct exponentiation overflowed. Fix: Compute max-shifted softmax or stable log-sum-exp.
• Symptom: nn.CrossEntropyLoss runs but learns poorly. Cause: Model probabilities were passed as if they were logits. Fix: Pass the final linear layer's raw scores.
• Symptom: Lower temperature reinforces an incorrect option. Cause: Temperature sharpened existing score order; it doesn't learn a correction. Fix: Use training gradients to repair logits and tune decoding separately.
• Symptom: Mean sequence loss looks acceptable despite a serious wrong token. Cause: Easy token positions dilute one high-loss event. Fix: Inspect unreduced position losses during debugging.

Follow-up questions