Numerically Stable Softmax From Scratch

Difficulty: Medium · Concept: Activation Functions — Auxiliary Functions

$\text{softmax}(\mathbf{z})_i = \frac{e^{z_i}}{\sum_j e^{z_j}}$

Unlike sigmoid, softmax operates on a whole vector at once, turning raw scores ("logits") into a probability distribution that sums to 1 — the standard output layer for multi-class classification, and exactly what turns an LLM's raw next-token scores into a probability distribution over the vocabulary.

Your task: implement softmax(z) for a list of real numbers. The naive formula above overflows for large inputs — math.exp(1000) raises OverflowError in plain Python — so your implementation needs to be numerically stable: it must return correct results even when every value in z is large.

Implement it yourself

assert abs(sum(softmax([1, 2, 3])) - 1.0) < 1e-9
assert abs(softmax([1, 2, 3])[2] - 0.6652409557748219) < 1e-9
assert abs(softmax([0, 0, 0])[0] - 1 / 3) < 1e-9
assert max(abs(p - 1 / 3) for p in softmax([1000, 1000, 1000])) < 1e-9
assert abs(softmax([1, 2])[0] - 0.2689414213699951) < 1e-9

Next: Softmax Gradient (Vector-Jacobian Product) (a harder variant)