Implement Softmax + Cross-Entropy Stably

The Numerical Overflow Problem

The Softmax Function converts a vector of $K$ real-valued logits $z \in \mathbb{R}^K$ into a probability distribution $\hat{p}$:

$$\hat{p}_i = \frac{e^{z_i}}{\sum_{j=1}^K e^{z_j}}$$

If a neural network outputs a raw logit $z_i = 1000$:

$$e^{1000} \implies \text{Overflow! Output is } \texttt{inf} \text{ or } \texttt{NaN} \text{ in 32-bit floats!}$$

Conversely, if $\hat{p}_i = 0.0$, computing Categorical Cross-Entropy Loss $-\log(\hat{p}_i)$ produces $-\log(0) \implies \texttt{-inf}$!

Math Foundations for Numerical Stability

1. Logit Max Shift (Softmax Trick)

Notice that shifting all logits by an arbitrary constant $C$ leaves the Softmax probabilities completely unchanged:

$$\frac{e^{z_i + C}}{\sum_j e^{z_j + C}} = \frac{e^C \cdot e^{z_i}}{e^C \cdot \sum_j e^{z_j}} = \frac{e^{z_i}}{\sum_j e^{z_j}}$$

Set $C = -\max(z)$. This guarantees that the largest logit becomes $0.0$, so $e^0 = 1.0$, preventing floating point overflow completely!

2. Combined Softmax + Cross-Entropy Gradient

Combining Softmax with Categorical Cross-Entropy Loss $\mathcal{L} = -\sum y_i \log(\hat{p}_i)$ simplifies the loss derivative with respect to logits $z$:

$$\frac{\partial \mathcal{L}}{\partial z_i} = \hat{p}_i - y_i$$

Where $y_i$ is the one-hot target label vector.

NumPy Implementation from Scratch

import numpy as np


class StableSoftmaxCrossEntropy:

  def softmax(self, z):
    # z shape: [batch_size, num_classes]
    # Subtract max logit per row for numerical stability
    z_max = np.max(z, axis=-1, keepdims=True)
    exp_z = np.exp(z - z_max)
    return exp_z / np.sum(exp_z, axis=-1, keepdims=True)

  def loss(self, z, y_true_indices):
    # z shape: [batch_size, num_classes]
    # y_true_indices shape: [batch_size] (integer class labels)
    batch_size = z.shape[0]

    # Compute stable Softmax probabilities
    probs = self.softmax(z)

    # Pick probability of true class for each sample
    correct_class_probs = probs[np.arange(batch_size), y_true_indices]

    # Clip probabilities to prevent log(0)
    correct_class_probs_clipped = np.clip(correct_class_probs, 1e-15, 1.0)

    # Compute Cross Entropy Loss
    loss = -np.mean(np.log(correct_class_probs_clipped))
    return loss, probs

  def gradient(self, probs, y_true_indices):
    # dz = Softmax(z) - OneHot(y)
    batch_size = probs.shape[0]
    dz = probs.copy()
    dz[np.arange(batch_size), y_true_indices] -= 1.0
    return dz / batch_size

Say this out loud

Numerically stable Softmax subtracts the maximum logit value per sample before exponentiation, shifting the highest exponent to zero and preventing floating point overflow. Combining Softmax with Cross Entropy loss yields a clean loss gradient derivative dz = softmax(z) - y.

Followups to expect

1. What is LogSumExp trick? Calculating $\log \sum e^{z_i}$ as $m + \log \sum e^{z_i - m}$ where $m = \max(z)$, preventing numerical underflow and overflow when computing log probabilities.
2. What is Label Smoothing? Replacing hard one-hot target vectors $[0, 1, 0]$ with smoothed targets $[0.05, 0.90, 0.05]$ to prevent Softmax logits from growing infinitely large during training.