Softmax and Numerical Stability

Why This Matters

The softmax function appears in virtually every classification neural network, every attention mechanism, every language model, and every reinforcement learning policy. It converts a vector of real numbers (logits) into a probability distribution. This sounds simple, but a naive implementation will produce NaN or Inf on perfectly reasonable inputs.

Understanding why softmax breaks numerically, and how the log-sum-exp trick fixes it, is the first real lesson in numerical computing for ML. If you have ever seen NaN losses during training, there is a good chance a softmax-related instability was the cause.

Mental Model

You have a vector of scores $z = (z_1, \ldots, z_K)$ called logits. You want to convert them into probabilities that sum to 1. Softmax exponentiates each score and normalizes:

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

The problem: if $z_i = 1000$, then $e^{1000}$ overflows to Inf in floating point. If $z_i = -1000$, then $e^{-1000}$ underflows to $0$. Both are fatal for the computation.

The Softmax Function

The translation invariance property is the key to numerical stability. Shifting all logits by the same constant does not change the output. We will exploit this shortly.

Why It Overflows and Underflows

In IEEE 754 double precision floating point:

• $e^{709} \approx 8.2 \times 10^{307}$ (near DBL_MAX)
• $e^{710}$ overflows to Inf
• $e^{-745}$ underflows to 0

In single precision (float32, the ML default):

• $e^{88}$ overflows to Inf
• $e^{-104}$ underflows to 0

A logit vector like $z = (100, 200, 300)$ in float32 produces $\exp(z) = (\text{Inf}, \text{Inf}, \text{Inf})$, and $\text{Inf}/\text{Inf} = \text{NaN}$.

Even with more modest logits, the ratio of a very small exponential to a very large sum gives underflow in the numerator, producing $0$ probabilities when they should be small but nonzero. This matters for cross-entropy loss, where $\log(0) = -\infty$.

The Log-Sum-Exp Trick

Stable Softmax and Log-Softmax

Using the trick, the stable softmax computation is:

$\text{softmax}(z)_i = \frac{\exp(z_i - c)}{\sum_j \exp(z_j - c)}, \qquad c = \max_j z_j$

Even more important is log-softmax, because cross-entropy loss uses $\log p_i$ directly:

$\log \text{softmax}(z)_i = z_i - c - \log \sum_j \exp(z_j - c)$

$= z_i - \text{logsumexp}(z)$

Computing $\log \text{softmax}$ directly (without first computing softmax and then taking $\log$) avoids the catastrophic loss of precision that occurs when softmax outputs a number very close to 1 and $\log$ maps it to nearly 0.

This is why PyTorch has separate F.softmax and F.log_softmax functions, and why F.cross_entropy takes raw logits rather than probabilities.

Temperature Scaling

Temperature is used in:

• Knowledge distillation: high $T$ softens teacher outputs to reveal inter-class relationships
• Language model sampling: $T < 1$ makes generation more deterministic, $T > 1$ makes it more random
• Reinforcement learning: Boltzmann exploration policies use temperature to trade off exploitation vs. exploration

Connection to Probability Distributions

Softmax arises naturally as the canonical link function for the categorical (multinoulli) distribution in the exponential family. If $X \sim \text{Categorical}(\pi_1, \ldots, \pi_K)$ with natural parameters $\eta_i = \log \pi_i$, then the mean parameters are recovered by:

$\pi_i = \frac{\exp(\eta_i)}{\sum_j \exp(\eta_j)} = \text{softmax}(\eta)_i$

In physics, this is the Boltzmann distribution (or Gibbs distribution): a system at temperature $T$ with energy levels $E_1, \ldots, E_K$ occupies state $i$ with probability:

$p_i = \frac{\exp(-E_i / T)}{\sum_j \exp(-E_j / T)} = \text{softmax}(-E / T)_i$

The partition function $Z = \sum_j \exp(-E_j / T)$ is precisely the denominator of softmax. Computing $\log Z$ is the log-sum-exp.

Gumbel-Softmax for Differentiable Sampling

A key problem: sampling from a categorical distribution is not differentiable, so you cannot backpropagate through it. The Gumbel-softmax trick provides a differentiable approximation.

This is used extensively in VAEs with discrete latent variables, neural architecture search, and any setting where you need to "differentiate through a categorical choice."

The Gumbel-max theorem guarantees that $\arg\max_i (z_i + g_i)$ is an exact sample from $\text{Categorical}(\text{softmax}(z))$. Gumbel-softmax relaxes the $\arg\max$ to a $\text{softmax}$, making it differentiable at the cost of approximation.

Common Confusions

Summary

• Softmax: $p_i = \exp(z_i) / \sum_j \exp(z_j)$. maps logits to probabilities
• Naive implementation overflows for large logits, underflows for very negative logits
• Log-sum-exp trick: subtract $\max_j z_j$ before exponentiating
• Always compute log_softmax directly, never log(softmax(z))
• Temperature $T$: high = uniform, low = peaked at argmax
• Softmax is the canonical link for the categorical distribution and the Boltzmann distribution
• Gumbel-softmax provides differentiable approximate sampling from categoricals

Exercises

References

Canonical:

• Goodfellow, Bengio, Courville, Deep Learning (2016), Section 4.1 (numerical stability)
• Bishop, Pattern Recognition and Machine Learning (2006), Section 4.3.4

Current:

• Jang, Gu, Poole, "Categorical Reparameterization with Gumbel-Softmax" (2017)

• Maddison, Mnih, Teh, "The Concrete Distribution" (2017)

• Hastie, Tibshirani, Friedman, The Elements of Statistical Learning (2009)

Next Topics

The natural next step from numerical stability:

• Conditioning and condition number: formalizing when a computation is inherently sensitive to perturbations, beyond specific tricks