Cross-Entropy Loss › Numerical Stability in Practice

When implementing cross-entropy in real training systems, numerical stability requires careful attention. The naive formulation involves taking logarithms of very small probabilities, which can produce infinities or cause gradient explosions. Modern deep learning frameworks address this through a combination of mathematical reformulations and careful floating-point management.

The standard cross-entropy computation involves first computing softmax probabilities, then taking logarithms. This two-step process is numerically unstable because softmax can produce probabilities that underflow to zero or overflow to infinity for logits with large magnitudes. The solution is the log-sum-exp trick, which computes $\log(\text{softmax}(z))$ in a single numerically stable step:

$$\log \hat{y}_k = z_k - \log \sum_{j=1}^{K} \exp(z_j)$$

where:

• $\log \hat{y}_k$: the log-probability for class $k$ (what we need for cross-entropy)
• $z_k$: the raw logit for class $k$
• $\log \sum_{j=1}^{K} \exp(z_j)$: the log-partition function, computed stably as $m + \log \sum_{j} \exp(z_j - m)$ where $m = \max_j z_j$

By subtracting the maximum logit $m$ before exponentiating, we ensure that at least one term equals $\exp(0) = 1$, and all other terms are at most 1. This prevents overflow while preserving the relative differences between logits that determine the probability distribution. PyTorch's F.cross_entropy and F.log_softmax implement this automatically.

Modern language model training uses mixed precision (float16 or bfloat16) to reduce memory usage and increase throughput. However, cross-entropy loss computation is often kept in float32 precision even when the forward pass uses float16. This is because float16 has limited range (approximately $\pm 65,000$) and limited precision (about 3 decimal digits), which can cause loss values to become NaN or Inf during computation. Frameworks like PyTorch's automatic mixed precision (AMP) automatically identify loss computation as a sensitivity point and upcasts to float32 for that portion of the computation, maintaining stability without sacrificing efficiency elsewhere.