Numerical Stability: Floating Point, Log-Sum-Exp and Avoiding NaNs

One of the most frustrating experiences in ML is watching a loss curve suddenly turn into NaN. The mathematics was correct; the arithmetic was not. Computers do not work with real numbers — they work with floating-point approximations. Knowing where those approximations break is a core professional skill.

Floating-point numbers#

A floating-point number stores a sign, an exponent and a mantissa (significand), representing $\pm m \times 2^e$. Common formats:

Machine epsilon is the gap between 1 and the next representable number: about $2.2 \times 10^{-16}$ for float64 and $1.2 \times 10^{-7}$ for float32. Every operation may introduce a relative error of this size.

import numpy as np
print(0.1 + 0.2 == 0.3)                      # False!
print(np.float32(1) + np.float32(1e-8) == 1) # True: 1e-8 is below float32 epsilon
print(np.float16(70000))                     # inf: overflow in half precision

The main failure modes#

1. Overflow — a result exceeds the largest representable number and becomes inf. np.exp(1000) overflows in float64.
2. Underflow — a result is smaller than the smallest representable number and becomes 0. Multiplying hundreds of probabilities underflows.
3. Catastrophic cancellation — subtracting nearly equal numbers destroys significant digits. Computing variance as $\mathbb{E}[X^2] - \mathbb{E}[X]^2$ for data with a large mean can even produce negative variance.
4. Invalid operations — inf - inf, 0 * inf, 0/0 and log(0) produce NaN or -inf, which then spread through every later computation.

The log-sum-exp trick#

Softmax and many likelihoods require $\ln\sum_i e^{x_i}$. If some $x_i = 1000$, $e^{1000}$ overflows. If all $x_i = -1000$, every term underflows to 0 and the log gives $-\infty$. The fix: factor out the maximum $m = \max_i x_i$:

Now the largest exponent is $e^0 = 1$ — no overflow — and at least one term equals 1 — no total underflow.

Stable softmax#

Subtracting the max does not change the result mathematically but keeps it computable.

import numpy as np

def naive_softmax(x):
    e = np.exp(x); return e / e.sum()

def stable_softmax(x):
    z = x - x.max(); e = np.exp(z); return e / e.sum()

def log_softmax(x):
    m = x.max()
    return x - (m + np.log(np.exp(x - m).sum()))

x = np.array([1000.0, 1001.0, 1002.0])
with np.errstate(all="ignore"):
    print("naive: ", naive_softmax(x))      # [nan nan nan]
print("stable:", stable_softmax(x))         # [0.090 0.245 0.665]
print("log-softmax:", log_softmax(x))

Stable cross-entropy and sigmoid#

Computing log(sigmoid(z)) naively fails for large negative $z$ (sigmoid underflows to 0). Use the identity

and implement softplus stably as $\max(z, 0) + \ln(1 + e^{-|z|})$. This is why deep-learning frameworks provide BCEWithLogitsLoss and cross_entropy functions that take logits, not probabilities. Always pass logits to these fused losses.

Other practical techniques#

• Work in log space for products of probabilities (HMMs, Naive Bayes, sequence likelihoods).
• Add small epsilons carefully: log(p + 1e-12) or x / (norm + 1e-8) prevent log(0) and division by zero — but choose epsilon relative to the precision (1e-12 is meaningless in float16).
• Welford's algorithm computes running mean and variance stably in one pass.
• Solve, don't invert: use np.linalg.solve or Cholesky decomposition rather than explicit matrix inverses.
• Check condition numbers: np.linalg.cond(A) — if it approaches $1/\epsilon$, results are unreliable.
• Gradient clipping limits the norm of gradients to prevent a single bad batch from causing overflow.

Mixed-precision training#

Modern GPUs are much faster in 16-bit arithmetic, so large models train in mixed precision: compute in float16 or bfloat16, keep a float32 "master copy" of weights, and accumulate sums in float32.

• float16 has a narrow range: small gradients underflow to zero. Loss scaling multiplies the loss by a large factor before backpropagation (and divides the gradients afterwards) to keep gradients in range; dynamic loss scaling lowers the factor when overflow occurs.
• bfloat16 keeps float32's exponent range with fewer mantissa bits, so overflow and underflow are rare and loss scaling is usually unnecessary — one reason it became the default for training large models.