Softmax, cross-entropy, and gradients in language models

Softmax: what it does and how to keep it stable

Softmax turns raw scores (logits) into probabilities, but you have to be careful to avoid numerical overflow/underflow.

Softmax definition and properties

For a logits vector $z \in \mathbb{R}^K$ (one logit per class), softmax is

$$p_i = \text{softmax}(z)_i = \frac{e^{z_i}}{\sum_{j=1}^K e^{z_j}}$$

Key properties:

• Probabilities: each $p_i \in (0,1)$
• Sum to 1: $\sum_i p_i = 1$
• Order-preserving: if $z_a > z_b$ then $p_a > p_b$
• Shift-invariance: adding the same constant $c$ to all logits doesn’t change probabilities:

$$\text{softmax}(z)_i = \text{softmax}(z + c)_i$$

Proof of shift-invariance:

$$\text{softmax}(z + c)_i = \frac{e^{z_i + c}}{\sum_j e^{z_j + c}} = \frac{e^c e^{z_i}}{e^c \sum_j e^{z_j}} = \frac{e^{z_i}}{\sum_j e^{z_j}} = \text{softmax}(z)_i$$

Why subtracting max(logits) helps numerical stability

If some logits are large (say 1000), then $e^{1000}$ overflows to infinity in float32.
If some are very negative (say -1000), $e^{-1000}$ underflows to 0.

Because of shift-invariance, we can safely do:

$$z'_i = z_i - \max_j z_j$$

Then compute

$$p_i = \frac{e^{z'_i}}{\sum_j e^{z'_j}}$$

Mechanism:

• The largest logit becomes 0: $\max_i z'_i = 0$
• So the biggest exponent is $e^0 = 1$, which is safe (no overflow)
• Others become $\le 0$, so their exponentials are in $(0,1]$, again safe
• Relative differences between logits stay the same, so probabilities are unchanged

Example that would overflow without this trick:

• Logits: z = [1000, 999, 990]
  • Directly: exp(1000) overflows in float32
• Shifted:
  • max(z) = 1000
  • z' = [0, -1, -10]
  • exp(z') ≈ [1, 0.3679, 0.0000454] — all perfectly fine