Numerically Stable Softmax and Cross-Entropy

Every neural network classifier ends with the same sequence:

logits → softmax → probabilities → cross-entropy loss

Although the equations look straightforward, a naive implementation quickly breaks due to floating-point overflow and underflow.

In this post, we’ll implement:

Naive softmax
Numerically stable softmax
Log-softmax
Stable cross-entropy

and understand why every deep learning library uses the stable formulation.

Softmax

Given a vector of logits

\[\mathbf{x} = [x_1, x_2, \ldots, x_n],\]

the softmax function is

\[\operatorname{softmax}(\mathbf{x})_i = \frac{e^{x_i}} {\sum_{j=1}^{n} e^{x_j}}, \qquad i=1,\ldots,n.\]

It converts arbitrary real-valued scores into a probability distribution satisfying

\[0 \le \operatorname{softmax}(\mathbf{x})_i \le 1,\]

and

\[\sum_{i=1}^{n} \operatorname{softmax}(\mathbf{x})_i = 1.\]

The output can therefore be interpreted as the model’s confidence for each class.

Why use the exponential?

The exponential function has several desirable properties:

Always produces positive values
Preserves the ordering of logits
Converts arbitrary scores into probabilities
Amplifies confidence differences
Is differentiable everywhere
Leads to simple gradients during backpropagation
Naturally appears in maximum entropy and Boltzmann distributions

Naive Softmax

A direct implementation follows the mathematical definition.

def softmax(x: np.ndarray) -> np.ndarray:
    exp_x = np.exp(x)
    return exp_x / np.sum(exp_x)

Example:

x = np.array([2.2, 1.5, -10, -100, 0, 50], dtype=np.float32)

softmax(x)
>> array([1.7406994e-21, 8.6440566e-22, 8.7565109e-27, 0.0000000e+00,
       1.9287499e-22, 1.0000000e+00], dtype=float32)

This works for moderate values.

Unfortunately, it completely fails when the logits become large.

For example,

x = np.array([2.2, 1.5, -10, -100, 0, 50], dtype=np.float32) * 10000

softmax(x)
>> array([nan, nan,  0.,  0.,  0., nan], dtype=float32)

produces

RuntimeWarning: overflow encountered in exp
RuntimeWarning: invalid value encountered in divide return exp_x / np.sum(exp_x)

because

\[e^{500000}\]

cannot be represented using floating-point numbers.

Stable Softmax

Fortunately, softmax is invariant to adding (or subtracting) the same constant from every logit.

\[\operatorname{softmax}(\mathbf{x}) = \operatorname{softmax}(\mathbf{x}-c).\]

Choosing

\[c = \max(\mathbf{x})\]

ensures the largest exponent becomes zero.

\[e^{x_i-\max(\mathbf{x})}\]

is therefore always less than or equal to 1, eliminating overflow.

The stable implementation becomes

def softmax(x: np.ndarray) -> np.ndarray:
    x = x - np.max(x)
    exp_x = np.exp(x)
    return exp_x / np.sum(exp_x)

Now even extremely large logits are handled correctly.

x = np.array([2.2, 1.5, -10, -100, 0, 50], dtype=np.float32)

softmax(x)
>> array([1.7407006e-21, 8.6440576e-22, 8.7565109e-27, 0.0000000e+00,
       1.9287499e-22, 1.0000000e+00], dtype=float32)

Large values:

x = x * 10000

softmax(x)
>> array([0., 0., 0., 0., 0., 1.], dtype=float32)

Tiny values:

x = x * 1e-9

softmax(x)
>> array([0.16668595, 0.16668479, 0.16666563, 0.1665157 , 0.16668229,
       0.16676566], dtype=float32)

Both remain numerically stable.

Cross-Entropy

Cross-entropy measures the difference between two probability distributions.

\[H(p,q) = -\sum_{i=1}^{n} p_i \log q_i.\]

In classification,

$q$ is the predicted probability distribution.
$p$ is the one-hot encoded ground-truth label.

Since only one entry of a one-hot vector equals 1,

\[\begin{aligned} H(p,q) &= -\sum_{i=1}^{n} p_i \log q_i \\ &= -\left(1 \cdot \log q_y + \sum_{i \ne y} 0 \cdot \log q_i\right) \\ &= -\log q_y \\ &= -\log\left(\operatorname{softmax}(\hat{\mathbf{y}})_y\right), \end{aligned}\]

where $y$ is the correct class.

Naive Cross-Entropy

A straightforward implementation is

def cross_entropy(y_hat: np.ndarray, y_true: int) -> float:
    return -np.log(softmax(y_hat)[y_true])

Example

cross_entropy(
    y_hat=np.random.normal(size=(10)),
    y_true=3,
)
>> np.float64(1.195837602754121)

However,

cross_entropy(
    y_hat=np.array([-1000, 1000]),
    y_true=0,
)
>> RuntimeWarning: divide by zero encountered in log return -np.log(softmax(y_hat)[y_true])

fails because

softmax([-1000, 1000])

becomes

[0, 1]

and therefore

\[\log(0) = -\infty.\]

Log-Softmax

Instead of computing

\[\log\left(\operatorname{softmax}(x)_i\right),\]

we simplify it algebraically.

\[\begin{aligned} \log\left(\operatorname{softmax}(x)_i\right) &= \log\left( \frac{e^{x_i-\max(x)}} {\sum_j e^{x_j-\max(x)}} \right) \\ &= \log\left(e^{x_i-\max(x)}\right) - \log\left(\sum_j e^{x_j-\max(x)}\right) \\ &= x_i-\max(x) - \log\left(\sum_j e^{x_j-\max(x)}\right). \end{aligned}\]

Notice that we never compute

\[\log(0),\]

which removes the numerical instability.

Stable Log-Softmax

def log_softmax(x: np.ndarray) -> np.ndarray:
    x_max = np.max(x)
    return x - x_max - np.log(np.sum(np.exp(x - x_max)))

Using this, cross-entropy becomes

def cross_entropy(y_hat: np.ndarray, y_true: int) -> float:
    return -log_softmax(y_hat)[y_true]

Example

cross_entropy(
    y_hat=np.random.normal(size=(10)),
    y_true=3,
)
>> np.float64(4.735892103199971)

Even the extreme case now works correctly.

cross_entropy(
    y_hat=np.array([-1000, 1000]),
    y_true=0,
)
>> np.float64(2000.0)

returns

2000.0

instead of producing an infinite loss.

Takeaways

The mathematical definition of softmax is simple, but a direct implementation is numerically unstable.

The key ideas are:

Shift logits by their maximum before exponentiating.
Never compute log(softmax(x)) directly.
Use log_softmax when computing cross-entropy.

These seemingly small tricks are why implementations in PyTorch, TensorFlow, and JAX remain stable even when logits become extremely large.