What softmax is, and the range where it breaks

Softmax maps a vector of real logits x: [N] to a probability vector: p_i = exp(x_i) / Σ_j exp(x_j), with every p_i > 0 and Σ_i p_i = 1. Mathematically it is well behaved on all of R^N. In floating point it is not, because exp leaves the representable range almost immediately.

The dangerous quantity is the largest logit. exp(x) overflows when x exceeds log(FLT_MAX): about 11.09 in fp16 (max 65504), about 88.7 in fp32 and bf16, which share an 8-bit exponent. Attention scores of 15 or 30 are routine, so a literal fp16 softmax overflows on real data — and inf / inf is NaN, not a large number but a poison value that propagates through the whole batch. Underflow is the mirror image: exp(x) flushes to zero below about -87 in fp32 and -9.7 in fp16. The two failures are not symmetric, and that asymmetry is what the standard fix exploits.

Advertisement

The max-subtraction identity, and why it is exact

Shift every logit by the same constant c and the answer does not change. The proof is one line of cancellation:

softmax(x - c)_i = exp(x_i - c) / Σ_j exp(x_j - c)
                 = [exp(-c) · exp(x_i)] / [exp(-c) · Σ_j exp(x_j)]
                 = exp(x_i) / Σ_j exp(x_j)  =  softmax(x)_i

This is exact in real arithmetic — softmax is invariant to adding any constant to all its inputs, the same statement as ‘logits are defined only up to a shift.’ Choosing c = max(x) is the choice that makes the computation safe. Every shifted logit is then ≤ 0, so every exp lands in (0, 1] and overflow becomes impossible in any format. The denominator is protected too: the argmax term contributes exactly 1.0, so the sum lies in [1, N] and can never be zero. Terms far below the max still underflow — but such a term was, by construction, a probability below the format’s smallest normal, and dropping it sits well under the rounding error you already accept.

Advertisement

A worked example: naive fp16 versus the stable form

Four logits, evaluated both ways in fp16.

x            = [ 20.0,   18.0,   17.0,   -5.0 ]

naive:  exp(20.0) = 4.85e8  →  > 65504  →  inf
        sum = inf           →  inf/inf    →  [NaN, NaN, NaN, NaN]

stable: x - max(x) = [ 0.0,  -2.0,  -3.0,  -25.0 ]
        exp        = [ 1.0,  0.13534,  0.049787,  1.389e-11 → 0 ]
        sum        = 1.185122
        p          = [ 0.84380,  0.11420,  0.04201,  ~0 ]

The naive path does not degrade gracefully; it returns nothing usable, and inside a training step that NaN reaches the optimizer and destroys the weights. The stable path loses exactly one term, the fourth, whose true value is 1.17e-11 — against fp16’s unit roundoff of 2^-11 ≈ 4.9e-4, seven orders of magnitude below the format’s noise floor. Overflow is fatal, underflow is free, which is why we shift down to the max rather than up toward it.

Why the accumulator is fp32 even when the inputs are not

Max subtraction fixes range. It does nothing for precision, and the sum is where precision dies. bf16 carries 8 significand bits, a unit roundoff of u = 2^-8 ≈ 3.9e-3, against fp32’s 2^-24 ≈ 6e-8.

The failure mode is stagnation. Sequential summation of n terms has a worst-case relative error bounded by roughly n · u (and about √n · u under a random-walk model of the rounding). With n = 4096 and bf16, n · u ≈ 16 — the bound is vacuous. Concretely: bf16 values in [256, 512) are spaced 2.0 apart, so once the running sum passes 256, adding a term of size 1 rounds straight back to 256 and the term vanishes. Since every shifted exponential is ≤ 1, a long row’s entire tail can disappear into a stalled accumulator. The fix is free: keep tensors in bf16, widen the accumulator and running max to fp32. Pairwise summation, which vectorized kernels do anyway, cuts the growth from n to log n.

Online softmax: one recurrence, one fewer pass

The textbook implementation makes three passes over the row: find the max, sum the exponentials, divide. Fine for a vector in cache, terrible for an attention row you would rather never store. The online formulation collapses the first two passes into one by carrying a running max m and running sum d together, rescaling d whenever a larger max arrives:

init:   m = -inf,   d = 0

per block B (local max m_B, local sum d_B = Σ_{i∈B} exp(x_i - m_B)):
    m' = max(m, m_B)
    d' = d · exp(m - m') + d_B · exp(m_B - m')

finish: p_i = exp(x_i - m) / d

Correctness follows from the same shift identity: d is always the sum of exp(x_j - m) over everything seen, and changing m means multiplying by exp(m_old - m_new). Both correction factors are ≤ 1, so rescaling can never overflow. The price is one extra exp per block boundary — arithmetic traded for a memory pass, the right trade on any modern machine.

From online softmax to tiled attention

The recurrence is what makes FlashAttention-style kernels possible. Standard attention computes S = QK^T / √d_k as an explicit [N, N] matrix, softmaxes it, then multiplies by V. At N = 4096 that matrix is 16.8M entries — 33.5 MB in fp16, written and read back once per head, per layer.

Tiled attention never materializes it. It walks blocks of K and V, computes each score tile on chip, and folds it into three running quantities: m, d, and an unnormalized output accumulator O, which takes the same correction factor as the sum:

O' = O · exp(m - m') + exp(S_B - m') · V_B      # then O / d at the end

The result matches a reference softmax to rounding, but memory traffic drops from O(N^2) to O(N · d). Long context stops being a memory problem and becomes a compute problem — entirely because softmax can be written as a rescalable recurrence.

Log-sum-exp: log-probabilities without probabilities

Training and evaluation want log p, not p, and the worst possible route is to compute probabilities and take their logarithm — any underflowed entry becomes log(0) = -inf. Go straight to log space with log-sum-exp:

LSE(x)            = m + log( Σ_j exp(x_j - m) ),   m = max(x)
log_softmax(x)_i  = x_i - LSE(x)
CE(x, target t)   = LSE(x) - x_t

Every step is safe: the exponentials are bounded by 1, the sum is at least 1 so the logarithm is defined and non-negative, and the final subtraction is ordinary arithmetic. A confidently wrong prediction yields a large finite loss instead of inf — the difference between a run that recovers and one that diverges. It is also why cross-entropy should be handed logits, never a softmax output, and why frameworks expose log_softmax and cross_entropy_with_logits as primitives.