Stable Softmax

Convert logits to probabilities without numeric overflow.

Start with the concrete trace below. It shows the state the algorithm must carry as it runs.

Problem trace

Stable Softmax: Subtract the row maximum, exponentiate nonpositive shifts, and divide by their shared sum.

Input and goalConvert logits to probabilities without numeric overflow.

Find the maximum logitFor logits [1000,1001,999], the row maximum is 1001. Direct exp(1001) can overflow.

10000max = 1001100119992

Recognize it

Use it when exponentials normalize scores into probabilities and large positive logits may overflow even though only relative differences should affect the result.

Keep true

Subtracting one constant from every row element leaves every exponential ratio unchanged; after max shifting, the largest exponent is exp(0)=1 and none can overflow.

Reuse it

Before exponentiating normalized scores, exploit shift invariance to center at the maximum; preserve the reduced dimension so broadcasting back across the row is explicit.

Pattern: Shift before exponentiation.

Simple idea: Softmax does not change when the same value is subtracted from every logit. Subtract the largest logit so every exponent is at most 1.

import numpy as np

def stable_softmax(logits: np.ndarray, axis: int = -1) -> np.ndarray:
   shifted = logits - np.max(logits, axis=axis, keepdims=True)
   exponentials = np.exp(shifted)
   return exponentials / np.sum(exponentials, axis=axis, keepdims=True)

Cost: $O(n)$ time and $O(n)$ output space.