Softmax, Cross-Entropy & Optimization

The previous lesson predicted one continuous scalar: extra decode latency in milliseconds. Half squared error could say "raise the prediction by four units." A review triage bot faces a different challenge: pick exactly one category among bug, docs, and security.

Suppose the incident note reads, "This docs page is stale," while the model emits raw scores bug=3.0, docs=1.0, security=0.0. It picks bug, missing the correct label. What training signal should backpropagation send back through the network?

A binary right-or-wrong check can't distinguish a narrow near-miss from an arrogant blunder. Assigning docs an 11% probability should cost significantly more than assigning it 49%, even if bug wins the top spot in both runs. We need a mathematical bridge that turns unbounded classification scores into honest probability shares, converts missed confidence into an exact penalty, and sends back a gradient that pulls the correct score upward while pushing wrong favorites down.

A decision begins as three raw scores

The final linear layer of a neural network outputs one real number per possible class. These numbers are logits: unbounded scores produced by taking dot products between internal representations and class weight vectors. They don't represent probabilities yet, and they aren't constrained between zero and one.

A larger logit means the network prefers that label over competitors. A score of 3.0 doesn't mean 300% confidence. Read the table as an uncalibrated ranking: bug leads docs by two score points, but we haven't assigned any percentage share to any option.

Logits can be negative, zero, or wildly positive without breaking mathematics. Probabilities, by contrast, must satisfy two strict Kolmogorov axioms: every probability $p_i$ must lie in $[0, 1]$, and the entire set of mutually exclusive choices must sum to exactly $1$.

Predict the three checks in inspect-raw-logits.py before running it. bug takes the highest score, the raw sum is $4.0$, and the list fails the probability test. We've got a ranking, but not a normalized distribution that an information-theoretic loss can evaluate.

The raw scores identify a winner, but they don't yet say how much probability the correct label received.

Softmax turns score gaps into probability mass

Softmax maps an unconstrained vector of $K$ real numbers into a valid categorical probability distribution. It operates in two deliberate stages: exponentiation and normalization.

First, it raises Euler's constant $e \approx 2.71828$ to each logit. Because $e^z > 0$ for every real number $z \in (-\infty, +\infty)$, exponentiation eliminates negative values and guarantees positivity. At the same time, exponentiation magnifies differences: a constant score lead of $+2.0$ expands into an exponential multiplier of $e^{2.0} \approx 7.389$.

Second, it sums those positive weights into a single shared denominator, called the partition function:

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

Dividing each exponential weight by this sum forces the output shares to sum to $1.0$.

For our three triage scores, walk through the arithmetic: $e^{3.0} \approx 20.086$, $e^{1.0} \approx 2.718$, and $e^{0.0} = 1.000$. The shared sum is $23.804$.

The general equation applies this two-step transformation to any vector $z \in \mathbb{R}^K$:

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

Here $p_i$ is the probability assigned to class $i$, $z_i$ is its logit, and $K$ is the number of competing classes. Notice the distribution: the model assigns 0.844 to the incorrect bug label, leaving only 0.114 for the correct docs label. That value is the handoff to the loss function.

Before running stable-softmax.py, confirm the invariants: the class ranking stays preserved, every entry is strictly positive, and the values sum to 1.0.

The correct label is docs, but its probability is only 0.114. Why can't the model fix this by simply tweaking the post-softmax numbers directly?

Answer

Probabilities are consequences of logits and compete for a fixed budget of 1.0. Learning updates parameters, which changes the logits. Pushing the docs logit upward relative to bug is what increases its probability on the next forward pass.

Stable arithmetic prevents floating-point overflow

Our gentle logits [3.0, 1.0, 0.0] fit easily inside standard floating-point variables. What happens when a deep network produces raw scores like [1000.0, 998.0, 997.0]?

In IEEE 754 single-precision floating-point (float32), numbers overflow once they exceed approximately $3.4 \times 10^{38}$. Because $e^{88.72} \approx 3.4 \times 10^{38}$, any logit greater than $88.72$ overflows to +inf. In half-precision formats (float16 and bfloat16), the dynamic range is even tighter: float16 overflows at $65504$, which corresponds to $e^{11.09}$. Logits as modest as $12.0$ trigger +inf.

Once an exponential overflows to +inf, the denominator evaluates to +inf. Dividing inf / inf yields NaN (Not a Number), poisoning every downstream weight tensor in your model.

Softmax possesses a key algebraic property: shift invariance.^{[1]Reference 1Deep Learning.https://www.deeplearningbook.org/} Adding or subtracting any arbitrary scalar constant $c$ from every logit leaves the resulting probabilities completely unchanged:

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

The common factor $e^{-c}$ factors out of the sum and cancels between numerator and denominator. This identity unlocks the classic numerical stabilization trick: choose $c = m = \max_j(z_j)$, the maximum logit across the vector.

Subtracting $m$ shifts the entire logit vector so that the largest value becomes exactly zero ($z_{\text{max}} - m = 0$). Every other shifted logit is strictly negative ($z_j - m \le 0$). Exponentiating non-positive numbers guarantees:

1. The largest exponential is $e^0 = 1.0$. Overflow is mathematically impossible.
2. Every other exponential lies in $(0, 1]$.
3. The denominator sum is at least $1.0$, preventing division by zero.

Very negative shifted logits (such as $-1000$) will underflow to floating-point 0.0. That underflow is benign during the forward pass: a weight of 0.0 simply means that candidate receives zero probability share.

NumPy's naive exponential returns inf, followed by NaN when dividing by the infinite sum. Shifting by the maximum value first eliminates overflow while returning the exact same probability distribution.

Cross-entropy prices missed predictions through surprise

Softmax gave us a probability distribution where docs received $p_{\text{docs}} \approx 0.114$. How do we score that prediction?

Cross-entropy measures the quality of a predicted distribution against ground truth. It's grounded in Maximum Likelihood Estimation (MLE) and information theory.

In information theory, observing an event with probability $p$ yields an information content, or surprise, of:

$$I(p) = -\log(p)$$

We use natural logarithms throughout, measuring surprise in nats (using base-two logarithms would measure surprise in bits). If an event is certain ($p = 1.0$), observing it brings zero surprise: $-\ln(1.0) = 0$. If an event is extremely rare ($p \to 0$), observing it brings near-infinite surprise: $-\ln(0.001) \approx 6.91$ nats.

When training a supervised model, the ground truth for an input is a target distribution $y$. For single-label classification, $y$ is a one-hot vector where the correct class index $c$ gets weight $1.0$ and all competing classes get $0.0$. For our triage note:

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

Cross-entropy calculates the expected surprise across all classes under the true distribution:

$$L_{\text{CE}} = -\sum_{k=1}^K y_k \log(p_k)$$

Because $y_k = 0$ for all non-target classes and $y_c = 1$ for the target, every term except the correct class multiplies by zero. Cross-entropy collapses to the Negative Log-Likelihood (NLL) of the true target class:

$$L = -\log(p_c)$$

For our incident note where $c = \text{docs}$:

$$L = -\log(0.114195\ldots) \approx 2.170\text{ nats}$$

If the network had assigned 0.844 to docs, the loss would be only $-\log(0.844) \approx 0.170$ nats. If it had placed 0.999 on docs, the loss would plummet to $0.001$ nats.

The log-sum-exp formulation prevents underflow crashes

Computing p = softmax(z) and then calling math.log(p[c]) introduces a subtle failure mode. If a logit is very negative (for example, logits = [1000.0, 0.0, -1.0] with target docs), $z_{\text{docs}} - m = 0.0 - 1000.0 = -1000.0$. In floating-point arithmetic, $e^{-1000}$ underflows to 0.0. Calling log(0.0) produces -inf, crashing training.

The loss can be expressed directly in terms of raw logits, bypassing intermediate probabilities entirely.^{[1]Reference 1Deep Learning.https://www.deeplearningbook.org/} Expanding $p_c$:

$$L = -\log\left(\frac{e^{z_c}}{\sum_j e^{z_j}}\right) = \log\left(\sum_{j=1}^K e^{z_j}\right) - z_c$$

The first term is log-sum-exp. We stabilize it using the same maximum shift $m = \max_j(z_j)$:

$$\log\left(\sum_{j=1}^K e^{z_j}\right) = \log\left(e^m \sum_{j=1}^K e^{z_j - m}\right) = m + \log\left(\sum_{j=1}^K e^{z_j - m}\right)$$

Substituting this back into the loss expression:

$$L = (m - z_c) + \log\left(\sum_{j=1}^K e^{z_j - m}\right)$$

Grouping $(m - z_c)$ first prevents catastrophic cancellation. When $z_c = 0.0$ and $m = 1000.0$, the first term evaluates cleanly to $1000.0$, while the log-sum evaluates to $\log(1.0 + \dots) \approx 0.0$. The loss returns $1000.0$ without ever taking the logarithm of zero.

Another sample gets 0.80 probability on its correct class. Which sample contributes more loss: that sample or our docs sample at 0.114?

Answer

Our docs sample contributes far more loss. With its unrounded probability, $-\log(0.114)$ is roughly $2.170$ nats, while $-\log(0.80)$ is only $0.223$ nats. Lower target probabilities create exponentially harsher penalties.

Entropy is reference uncertainty and KL divergence is prediction error

Deepen the distinction between entropy, cross-entropy, and Kullback-Leibler (KL) divergence.

Entropy $H(P)$ measures the inherent uncertainty already present inside a reference probability distribution $P$:

$$H(P) = -\sum_k P(k) \log P(k)$$

Cross-entropy $H(P, Q)$ measures the total cost of representing distribution $P$ using model distribution $Q$. Kullback-Leibler divergence $D_{\text{KL}}(P \parallel Q)$ measures the extra penalty caused specifically by the mismatch between $Q$ and $P$:

$$H(P, Q) = H(P) + D_{\text{KL}}(P \parallel Q)$$

For our one-hot docs target, $P = [0, 1, 0]$ has zero entropy ($H(P) = 0$): the supervised answer is completely deterministic. Every fraction of a nat in our $2.170$ loss comes from model error ($D_{\text{KL}} = 2.170$).

When targets are soft (such as in distillation or human annotations with disagreement, like $P = [0.5, 0.5, 0]$), the reference distribution has inherent entropy $H(P) = \ln(2) \approx 0.693$ nats. Even an ideal model ($Q = P$) will report a loss of $0.693$ nats, with $D_{\text{KL}} = 0$. Don't mistake inherent target uncertainty for model failure.

A non-zero loss doesn't mean the model is broken if the labels themselves carry irreducible noise. The decomposition separates target entropy from model prediction error.

The combined gradient reveals a clean error signal

Loss tells us how poorly the model performed. To update weights, backpropagation needs the gradient: the partial derivative of the loss with respect to each logit, $\frac{\partial L}{\partial z_k}$.

Softmax and cross-entropy create one of the cleanest gradient simplifications in machine learning.^{[1]Reference 1Deep Learning.https://www.deeplearningbook.org/}

Recall the loss in terms of raw logits:

$$L = \log\left(\sum_{j=1}^K e^{z_j}\right) - z_c$$

Differentiate with respect to any logit $z_k$:

$$\frac{\partial L}{\partial z_k} = \frac{\partial}{\partial z_k} \log\left(\sum_{j=1}^K e^{z_j}\right) - \frac{\partial z_c}{\partial z_k}$$

By the chain rule, the derivative of $\log(u)$ is $\frac{1}{u} \frac{\partial u}{\partial z_k}$. With $u = \sum_j e^{z_j}$:

$$\frac{\partial}{\partial z_k} \log\left(\sum_{j=1}^K e^{z_j}\right) = \frac{1}{\sum_{j=1}^K e^{z_j}} \cdot \frac{\partial}{\partial z_k} e^{z_k} = \frac{e^{z_k}}{\sum_{j=1}^K e^{z_j}} = p_k$$

The derivative of the log-sum-exp term is softmax itself.

Meanwhile, the second term $\frac{\partial z_c}{\partial z_k}$ equals $1$ if $k = c$ (the target class) and $0$ otherwise. That is exactly the target weight $y_k$:

$$\frac{\partial L}{\partial z_k} = p_k - y_k$$

Vectorized across all classes:

$$\nabla_z L = p - y$$

Predicted probability minus true target. Look at the numbers for our triage example:

The zero-sum conservation property

Sum all elements of the gradient vector:

$$\sum_{k=1}^K \frac{\partial L}{\partial z_k} = \sum_{k=1}^K (p_k - y_k) = \sum_{k=1}^K p_k - \sum_{k=1}^K y_k = 1.0 - 1.0 = 0$$

The gradients sum to zero. The upward pull on the correct class ($-0.886$) is exactly balanced by the downward push on competing classes ($+0.844 + 0.042 = +0.886$). Softmax distributes a fixed unit of probability mass; learning to favor one class inevitably pulls mass away from others.

In a complete neural network, backpropagation channels this $p - y$ residual back to shared weights via the chain rule. Testing on logits directly makes the local mechanic transparent.

Confirm the mathematical slope against numerical finite differences: nudge each logit by $\epsilon = 10^{-5}$ and check the symmetric difference quotient.

Why squared probability error saturates and stalls

Why not use Mean Squared Error (MSE) on predicted probabilities: $L_{\text{MSE}} = \frac{1}{2} \sum_k (p_k - y_k)^2$?

To see why MSE fails for classification, examine how gradients flow through the softmax Jacobian matrix:

$$\frac{\partial p_i}{\partial z_j} = p_i (\delta_{ij} - p_j)$$

When differentiating cross-entropy $L_{\text{CE}} = -\sum_i y_i \log(p_i)$ via the multivariate chain rule:

$$\frac{\partial L_{\text{CE}}}{\partial z_j} = \sum_i \frac{\partial L}{\partial p_i} \frac{\partial p_i}{\partial z_j} = \sum_i \left(-\frac{y_i}{p_i}\right) \cdot p_i (\delta_{ij} - p_j) = -\sum_i y_i (\delta_{ij} - p_j) = p_j - y_j$$

The factor $-\frac{y_i}{p_i}$ has $p_i$ in the denominator, which cancels the $p_i$ in the softmax Jacobian numerator.

Under MSE, the derivative with respect to probability is $\frac{\partial L_{\text{MSE}}}{\partial p_i} = (p_i - y_i)$. There is no $p_i$ denominator to cancel the softmax Jacobian:

$$\frac{\partial L_{\text{MSE}}}{\partial z_j} = \sum_i (p_i - y_i) p_i (\delta_{ij} - p_j)$$

When the model is confidently wrong (for instance, $p_{\text{bug}} = 0.9993$ and $p_{\text{docs}} = 0.0003$), the softmax derivative term $p_i(1 - p_i) \approx 0.0003$ vanishes. The gradient flattens out, shrinking the correction by three orders of magnitude. The network enters a saturated plateau and stops learning.

With cross-entropy, the gradient on the target logit is $p_{\text{docs}} - 1 = 0.0003 - 1.0 = -0.9997$, delivering a massive, full-strength corrective kick.

The numerical gap is striking: cross-entropy provides an aggressive gradient around $-1.0$, while squared probability error stalls with a feeble signal around $-0.0007$.

PyTorch expects unnormalized logits directly

In PyTorch, nn.CrossEntropyLoss expects unnormalized logits, not probabilities.^{[2]Reference 2CrossEntropyLosshttps://docs.pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html}

Internally, PyTorch fuses log_softmax with negative log-likelihood (NLLLoss). Fusing them performs the stable log-sum-exp trick in a single optimized GPU kernel and avoids materializing intermediate probabilities in memory.

If you mistakenly insert an nn.Softmax() activation before nn.CrossEntropyLoss, PyTorch treats those already-normalized probabilities as logits. It applies log-softmax a second time: $\log(\text{softmax}(p)) - p_c$. The code won't crash because tensor shapes match, but the objective function is corrupted and gradients shrink drastically.

Now reproduce the common bug: softmax first, then CrossEntropyLoss. The shapes match, so Python raises no warning. Compare the target gradients.

The distorted gradient reveals the defect: the corrective gradient for docs shrank from $-0.886$ to $-0.127$. Leave your final classification layer linear during training, and feed raw logits directly into nn.CrossEntropyLoss. Apply softmax only during inference when a client or caller explicitly requests probabilities.

Temperature scaling shapes inference entropy

During training, we optimize the network at temperature $T = 1.0$. During inference or knowledge distillation, scaling logits by a positive temperature parameter $T > 0$ alters distribution entropy:

$$p_i(T) = \frac{e^{z_i / T}}{\sum_{j=1}^K e^{z_j / T}}$$

Dividing logits by $T$ changes how sharply probabilities concentrate:

• When $T < 1.0$, logit differences expand. The leading class absorbs more probability mass. As $T \to 0$, softmax approaches an argmax one-hot distribution (greedy selection).
• When $T > 1.0$, logit differences shrink. Mass spreads more evenly across competitors. As $T \to \infty$, the distribution approaches uniform uncertainty ($1/K$).

Temperature doesn't alter class ranking: the highest logit before scaling remains the highest logit after scaling. Temperature is a decoding and calibration control, not a parameter update.

In knowledge distillation, Hinton et al. use high temperature ($T = 3.0$ to $5.0$) to reveal soft relationships between non-target classes (so-called "dark knowledge").^{[3]Reference 3Distilling the Knowledge in a Neural Network.https://arxiv.org/abs/1503.02531} In large language model serving, temperature modulates sampling diversity.

If docs is correct but bug is currently the highest logit, does lowering temperature repair the model?

Answer

No. Lowering temperature sharpens the current ranking, strengthening the mistaken bug preference to 98%. Training changes logits by following loss gradients; temperature only rescales probabilities generated from current logits.

Label smoothing curbs overconfident logits

Training models on hard one-hot targets introduces a structural hazard: overconfidence.

With a hard target $y_{\text{docs}} = 1.0$, minimizing cross-entropy loss $L = -\log(p_{\text{docs}})$ requires driving $p_{\text{docs}} \to 1.0$. In softmax, $p_{\text{docs}} \to 1.0$ only occurs when:

$$z_{\text{docs}} - z_k \to +\infty \quad \text{for all } k \ne \text{docs}$$

Because logits are linear combinations of weights and activations, the optimizer continuously drives weight norms to grow larger and larger. The network becomes dogmatically overconfident, degrades its calibration, and memorizes noise in the training labels.

Label smoothing resolves this by softening the hard target distribution.^{[4]Reference 4Attention Is All You Need.https://arxiv.org/abs/1706.03762} Given smoothing parameter $\alpha \in (0, 1)$ (typically $\alpha = 0.1$) and $K$ classes:

$$\tilde{y}_k = (1 - \alpha) y_k + \frac{\alpha}{K}$$

For our 3-class triage problem with target docs ($c=1$), $\alpha = 0.1$, and $K = 3$:

• Target class: $\tilde{y}_{\text{docs}} = (1 - 0.1) \cdot 1.0 + \frac{0.1}{3} = 0.900 + 0.0333 = 0.9333$
• Competitor classes: $\tilde{y}_{\text{bug}} = (1 - 0.1) \cdot 0.0 + \frac{0.1}{3} = 0.0333$
• Competitor classes: $\tilde{y}_{\text{security}} = 0.0333$

What logit gap does the model need now to achieve zero gradient? Instead of an infinite score difference, the target probability reaches equilibrium when:

$$\frac{p_{\text{docs}}}{p_{\text{competitor}}} = \frac{0.9333}{0.0333} = 28 \implies z_{\text{docs}} - z_{\text{competitor}} = \ln(28) \approx 3.332$$

The model only needs a modest, finite gap of $\approx 3.33$ logits. If weights push $z_{\text{docs}}$ higher so that $p_{\text{docs}} > 0.933$, the gradient $p_k - \tilde{y}_k$ becomes positive, pushing the target logit back down. Label smoothing acts as an automatic brake on runaway weights.

PyTorch provides built-in label smoothing via nn.CrossEntropyLoss(label_smoothing=0.1).

The output verifies the mathematical correspondence: PyTorch's label-smoothed loss matches manual NLL against $\tilde{y}$ exactly, and the gradient equals $p - \tilde{y}$.

Next-token prediction extends the loss across sequences

Our triage classifier evaluated one decision. Autoregressive language models evaluate the exact same categorical decision at every sequence position.

Given input context $x_1, x_2, \dots, x_t$, a language model emits hidden states $h_t \in \mathbb{R}^d$. The language model head projects $h_t$ against an embedding matrix to produce vocabulary logits $z_t \in \mathbb{R}^{|V|}$, where $|V|$ is vocabulary size. The observed next token $x_{t+1}$ provides the supervised target.

For a sequence of length $S$, the sequence loss is the mean negative log-likelihood across all predicted positions:

$$L_{\text{seq}} = -\frac{1}{S} \sum_{t=1}^S \log P(x_{t+1} \mid x_1, \dots, x_t)$$

Reuse our running numbers with a tiny vocabulary: ["bug", "docs", "today"].

• Position 1 predicts docs ($c=1$) from logits [3.0, 1.0, 0.0]. Target probability is $0.114$, producing $L_1 = 2.170$ nats.
• Position 2 predicts today ($c=2$) from logits [0.5, 0.0, 2.0]. Target probability is $0.736$, producing $L_2 = 0.306$ nats.
• Sequence mean loss: $L_{\text{seq}} = \frac{2.170 + 0.306}{2} = 1.238$ nats.

Notice the proportion: position 1 accounts for $2.170 / 2.476 \approx 87.6\%$ of the sequence loss.

Tensor contracts and vocabulary memory scaling

In PyTorch, a language model batch emits logits shaped [B, S, V] (batch size, sequence length, vocabulary size) and targets shaped [B, S]. PyTorch's multi-dimensional cross-entropy expects classes on dimension 1: [B, V, S]. Alternatively, reshape both into 2D tensors: [B * S, V] with targets [B * S].

In modern language models with vocabularies between $32\text{k}$ and $128\text{k}+$ tokens (such as Llama 3 or GPT-4o), materializing full 3D logit tensors [B, S, V] creates severe VRAM bottlenecks. At batch size $B=16$, sequence length $S=4096$, and $|V|=128256$, a float32 logit tensor requires over 33 GB of GPU memory just for the scores. Production pipelines use chunked cross-entropy or fused GPU kernels that stream hidden states in small tiles, computing log-sum-exp and gradients in SRAM without ever writing the full vocabulary matrix to global GPU memory.

The same reduction can be implemented without PyTorch using standard library loops:

Diagnostic failure modes in production pipelines

When classification or language modeling runs misbehave, trace errors systematically from arithmetic through interface contracts:

Diagnose issues in strict dependency order: numerical stability first, interface contract second, gradient sign third, decoding controls fourth, and sequence reduction last.

Try a second target

Experiment with the running example: switch the correct label to security ($c = 2$) while keeping logits [3.0, 1.0, 0.0].

Before running the code, calculate the expected loss and gradient by hand:

• Target probability: $p_{\text{security}} \approx 0.042$.
• Loss: $-\log(0.0420) \approx 3.170$ nats.
• Gradient $p - y$:
  • bug: $0.844 - 0 = +0.844$ (pushed down)
  • docs: $0.114 - 0 = +0.114$ (pushed down)
  • security: $0.042 - 1.0 = -0.958$ (pulled up strongly) Notice the zero-sum balance: $+0.844 + 0.114 - 0.958 = 0.000$. Subtracting $\eta(p - y)$ raises the security logit while suppressing both competitors.

What values should the security-target experiment produce before the update?

Answer

The probabilities remain [0.844, 0.114, 0.042], because changing labels doesn't alter the forward pass. The loss jumps to approximately 3.170 nats. The gradient becomes [0.844, 0.114, -0.958]. Subtracting that gradient raises security rather than docs.