Softmax & Cross-Entropy

A Multiple-Choice Exam, Translated to Math

A multiple-choice question gives you four options and asks you to pick one. You might feel most confident about option B (60%), less confident about C (25%), and almost rule out A and D (10% each). The teacher then reveals the correct answer was C — you were confidently wrong, and your “loss” for this question is higher than if you had been confidently right or even uniformly unsure.

That informal reasoning is exactly the math behind every classification network. Softmax converts a vector of raw scores into a probability distribution. Cross-entropy measures how badly that distribution differs from the truth. Together they are the loss for the health-classification head we will attach to the backbone in Chapter 11 — one of the two heads of every dual-task model in this book.

Softmax: From Logits to Probabilities

Given a vector of $C$ raw scores (called logits) $\mathbf{z} \in \mathbb{R}^{C}$, softmax produces a probability distribution $\mathbf{p} \in [0,1]^{C}$:

$p_c \;=\; \frac{e^{z_c}}{\sum_{k=1}^{C} e^{z_k}}.$

Three properties make it the right choice. (1) Non-negative: $e^{z_c} > 0$ for every $c$. (2) Sums to 1: divisor is the sum of numerators. (3) Smooth and differentiable: gradients flow back to every $z_c$. The C-class health head in Chapter 11 has $C = 3$ — normal, degrading, critical.

Interactive: Drag Logits, Watch Probabilities

The visualization below lets you drag four token logits and see the softmax probabilities update live. Toggle the temperatureslider: high temperature flattens the distribution toward uniform, low temperature sharpens it toward one-hot. The same temperature knob shows up in language-model sampling where it controls creativity vs determinism.

Cross-Entropy: Distance Between Distributions

With $\mathbf{p}$ the predicted distribution and $\mathbf{q}$ the target distribution (one-hot for hard labels), the cross-entropy is

$H(\mathbf{q}, \mathbf{p}) \;=\; -\sum_{c=1}^{C} q_c \log p_c.$

For one-hot $\mathbf{q}$ with the correct class at index $y$, this collapses to

$H = -\log p_y.$

Three intuitions. (a) If $p_y = 1$ (perfect): loss is $-\log 1 = 0$. (b) If $p_y = 0.5$: loss is $\log 2 \approx 0.693$. (c) If $p_y \to 0$: loss $\to \infty$. Confidence on the wrong answer is punished without bound.

Interactive: The Cross-Entropy Loss

Watch the loss spike as the model becomes confidently wrong. The gradient flows back through log and softmax to update the logits.

Numerical Stability: The Max-Subtract Trick

Naive softmax is a numerical landmine. $e^{1000}$ overflows to inf; $e^{-1000}$ underflows to 0. The fix is to exploit a softmax invariant:

$\text{softmax}(\mathbf{z}) = \text{softmax}(\mathbf{z} - c) \quad \text{for any constant } c.$

Choose $c = \max_k z_k$. After subtracting, every $z_c - c \le 0$, so $e^{z_c - c} \in (0, 1]$ — safe. The biggest logit becomes 0; $e^{0} = 1$; the denominator is at least 1. No overflow, no underflow, identical result.

Python: Softmax and CE From Scratch

Twenty lines of NumPy on a four-engine, three-class C-MAPSS-style batch. The numerical values match what PyTorch's F.cross_entropy returns to four decimal places.

1import numpy as np
2
3# ----- Numerically-stable softmax -----
4def softmax(logits: np.ndarray) -> np.ndarray:
5    """Row-wise softmax. Subtracts max for numerical stability."""
6    shifted = logits - logits.max(axis=-1, keepdims=True)
7    expd    = np.exp(shifted)
8    return expd / expd.sum(axis=-1, keepdims=True)
9
10
11# ----- Cross-entropy for one-hot targets -----
12def cross_entropy(logits: np.ndarray, targets: np.ndarray) -> float:
13    """logits: (B, C). targets: (B,) integer class IDs.
14    Returns the mean -log(p_correct) across the batch.
15    """
16    probs = softmax(logits)
17    # Pick the probability of the correct class for each row
18    p_correct = probs[np.arange(len(targets)), targets]
19    return float(-np.log(p_correct + 1e-12).mean())
20
21
22# ----- Run on a tiny C-MAPSS-style health-state batch -----
23# Three classes: 0 = normal, 1 = degrading, 2 = critical
24logits = np.array([
25    [2.0, 1.0, 0.5],       # leans normal
26    [0.1, 3.0, 0.2],       # strong degrading
27    [0.5, 0.5, 4.0],       # very critical
28    [1.0, 1.5, 1.2],       # uncertain - leaning degrading
29], dtype=np.float32)
30targets = np.array([0, 1, 2, 1])     # ground-truth health states
31
32print("probs[0]   :", softmax(logits)[0].round(4).tolist())
33# probs[0]: [0.6285, 0.2312, 0.1402]
34print("CE per row :", [
35    -float(np.log(softmax(logits)[i, targets[i]]))
36    for i in range(4)
37])
38print("avg CE     :", round(cross_entropy(logits, targets), 4))
39# avg CE     : 0.3715

Reading the per-row losses

Row 2 (logits $[0.5, 0.5, 4.0]$, target 2) produces softmax probability 0.96 on the correct class — loss 0.041. Row 3 (logits $[1.0, 1.5, 1.2]$, target 1) is much less confident — probability 0.40, loss 0.91. The average (0.37) is what the optimiser drives toward 0 during training.

PyTorch: F.cross_entropy in One Line

Production code never writes the softmax + log + NLL pipeline by hand — F.cross_entropy fuses them in a single, numerically-stable, GPU-friendly call. The CRITICAL convention is that you pass raw logits, not pre-softmaxed probabilities.

1import torch
2import torch.nn.functional as F
3
4# ----- F.cross_entropy: softmax + log + NLL fused -----
5torch.manual_seed(0)
6
7logits = torch.tensor([
8    [2.0, 1.0, 0.5],
9    [0.1, 3.0, 0.2],
10    [0.5, 0.5, 4.0],
11    [1.0, 1.5, 1.2],
12])
13targets = torch.tensor([0, 1, 2, 1])
14
15loss = F.cross_entropy(logits, targets)
16probs = F.softmax(logits, dim=-1)
17
18print("logits.shape :", tuple(logits.shape))      # (4, 3)
19print("probs[0]     :", probs[0].tolist())
20print("loss         :", round(loss.item(), 4))
21# probs[0]     : [0.6285, 0.2312, 0.1402]
22# loss         : 0.3715

Cross-Entropy Beyond RUL

The mathematics never changes. What changes is $C$ — the number of classes — and the dataset.

The Three Pitfalls

The point. Softmax is the differentiable bridge from real-valued logits to a probability distribution. Cross-entropy is the loss that punishes confident mistakes. Together they form the classification head used in every dual-task model from Chapter 11 onward.

Takeaway

• Softmax = soft argmax. Converts logits to a probability distribution that sums to 1.
• Cross-entropy = surprise. $-\log p_y$: zero when the model nails the right class, blows up when it is confidently wrong.
• Always max-subtract before exp. Otherwise $e^{1000}$ overflows. The result is mathematically identical because softmax is shift-invariant.
• F.cross_entropy expects raw logits. Never pre-softmax the input; the function fuses log_softmax + NLL internally for stability.
• Targets are integer class IDs. Not one-hot. PyTorch convention.