Problem 103 · Tokenization & Sampling

Numerically Stable Softmax with Temperature

EasyAI engineering

01Problem

Write `softmax(logits, temperature=1.0)` that returns a probability distribution. It must not overflow for logits like `[1000, 1001]`, must treat `-inf` logits as masked (probability 0), must return a one-hot on the first argmax when `temperature == 0` (greedy), and must raise `ValueError` on empty input, negative temperature, or when every logit is masked.

02Key insight

Subtract the max before exponentiating: softmax is shift-invariant, and the shift keeps every exponent at or below zero so nothing overflows. Temperature divides logits, so T -> 0 is a limit, not a value; special-case it as argmax rather than dividing by zero. All-masked rows are a real production bug (a grammar or tool constraint that leaves no legal token) and must fail loudly instead of returning NaNs that silently sample token 0.

03Solution

solution.pypython18 lines

import math


def softmax(logits: list[float], temperature: float = 1.0) -> list[float]:
    if not logits:
        raise ValueError("empty logits")
    if temperature < 0:
        raise ValueError("temperature must be >= 0")
    m = max(logits)
    if m == -math.inf:
        raise ValueError("every token is masked")
    if temperature == 0:
        idx = logits.index(m)
        return [1.0 if i == idx else 0.0 for i in range(len(logits))]
    scaled_max = m / temperature
    exps = [math.exp(x / temperature - scaled_max) for x in logits]
    total = sum(exps)
    return [e / total for e in exps]

Try it first.

TimeO(V)

SpaceO(V)

tests.py · 7 passing

tests.pypython20 lines

p = softmax([1.0, 2.0, 3.0])
assert abs(sum(p) - 1) < 1e-12 and p[2] > p[1] > p[0]
assert abs(softmax([1000.0, 1001.0])[1] - 1 / (1 + math.exp(-1))) < 1e-12
assert softmax([3.0, 5.0, 5.0], 0) == [0.0, 1.0, 0.0]
cold = softmax([1.0, 2.0], 0.1)
hot = softmax([1.0, 2.0], 100.0)
assert cold[1] > 0.9999 and abs(hot[0] - 0.5) < 0.01
masked = softmax([0.0, -math.inf, 0.0])
assert masked == [0.5, 0.0, 0.5]
for bad in ([], [-math.inf, -math.inf]):
    try:
        softmax(bad)
        assert False
    except ValueError:
        pass
try:
    softmax([1.0], -1)
    assert False
except ValueError:
    pass

04Staff-level follow-up

On GPU the same idea becomes the online softmax used by FlashAttention: keep a running max and running sum, rescaling the sum when the max changes, so you never materialise the full row. Also ask what 'temperature 0' means in a given serving stack: batched inference with non-deterministic kernels can still diverge between runs even at argmax.