Softmax

Start from brute force

The formula as written: exps = [math.exp(x) for x in logits], then divide each by sum(exps). That is O(n) time and O(n) space, and exact for modest logits, but it fails at the extremes that real logits reach. math.exp(1000) raises OverflowError in Python and is Infinity in JavaScript, where infinity over infinity gives NaN. At the other end, every exp(-1000) underflows to 0.0, and dividing by a zero total raises ZeroDivisionError. A slower variant recomputes the sum inside the loop, once per element, which is O(n²) for the same wrong answers.

The approach

Shift first, then exponentiate. The reference softmax computes m = max(logits), then exps = [math.exp(x - m) for x in logits] and total = sum(exps), and returns [e / total for e in exps]. Subtracting a constant changes nothing mathematically, because exp(x − m) = exp(x) · exp(−m) and the common factor exp(−m) cancels between every numerator and the denominator. Numerically it changes everything: every exponent is now at most 0, so the largest term is exactly exp(0) = 1 and nothing can overflow, and total is at least 1, so the division is always defined.

Very negative shifted logits can still underflow to 0.0, and that is harmless: their true probability is below the smallest positive double, about 5e-324. The output keeps the order of the logits, is unchanged when a constant is added to all of them, and sums to 1 up to rounding. The TypeScript version does the same three steps with Math.max(...logits) and reduce.

Walk through an example

For the second example, [2, 1, 0], the shift makes the largest exponent exactly 0, and the result sums to 1 up to the last bit.

Complexity

Time O(n): one pass for the max, one to exponentiate, one to sum and one to divide. Space O(n) for exps and the output; overwriting the input in place would make the extra space O(1). The unshifted version has the same O(n) time and space but overflows or underflows on large logits, and recomputing the sum per element makes it O(n²).

Edge cases to test

• [1000, 1000] → [0.5, 0.5]. Unshifted, math.exp(1000) raises OverflowError, and the JavaScript version returns [NaN, NaN].
• [-1000, -1000] → [0.5, 0.5]. Unshifted, both exponentials underflow to 0.0 and the division by a zero total fails.
• [1000, 0] → [1.0, 0.0]. exp(-1000) underflows to exactly 0.0, so a probability that is positive in theory prints as 0. Use log-softmax rather than log(softmax(x)), which would give negative infinity.
• [2, 1, 0] and [3, 2, 1] return identical lists. Adding a constant to every logit never changes the softmax; multiplying them, as temperature does, does change it.
• []: Python’s max raises ValueError, while the TypeScript version returns []. The constraints exclude it, but the two languages disagree.
• A very large array in TypeScript: Math.max(...logits) passes every element as a separate argument and can throw RangeError past the engine’s argument limit. A plain loop for the max avoids it.

Common mistakes

• Skipping the max-subtraction — exp() of a large logit overflows to inf and the result becomes nan.
• Dividing by n (the count) instead of the sum of the exponentials.
• Treating it element-wise — softmax is defined over the whole vector, so the denominator couples every entry.

Follow-ups an interviewer asks next

Add a temperature parameter.

Divide every logit by T first: softmax([x / T for x in logits]). A temperature below 1 sharpens the distribution toward the largest logit, above 1 flattens it toward uniform, and as T approaches 0 it approaches greedy argmax. Keep the max shift on the scaled values, because a small T makes them large enough to overflow.

You need log-probabilities for a loss.

Compute x - m - log(sum(exp(x_j - m))) directly, the log-sum-exp trick, instead of log(softmax(x)). The direct form never takes the log of an underflowed 0.0, so [1000, 0] gives log-probabilities of 0 and -1000 instead of 0 and negative infinity. Deep learning libraries fuse log-softmax with the negative log-likelihood loss for this reason.

The logits arrive in chunks and you cannot hold them all.

Keep a running maximum m and a running sum s of exp(x - m). When a new value exceeds m, rescale the sum first, s = s * exp(m - x_new) + 1, then make it the new maximum; otherwise add exp(x - m). One pass yields the normalizer, and this online softmax is the core of FlashAttention.

You only need to sample one token, not the whole distribution.

The Gumbel-max trick samples from softmax(x) without normalizing: add independent Gumbel noise, -log(-log(u)) with u uniform in (0, 1), to every logit and take the argmax. It is exact, needs no normalizing sum, and parallelizes well over a large vocabulary.

Where this shows up

Softmax is the output layer of essentially every classifier and language model (a distribution over the next token) and sits inside attention (a distribution over the keys). The max-subtraction trick is a classic “implement it from scratch” interview check precisely because forgetting it is the bug that silently produces nan in real training code.

Reference solution

import math

def softmax(logits):
    m = max(logits)                       # shift for numerical stability
    exps = [math.exp(x - m) for x in logits]
    total = sum(exps)
    return [e / total for e in exps]


print(softmax([2, 1, 0]))   # -> [0.665..., 0.245..., 0.090...]
print(sum(softmax([2, 1, 0])))  # -> 1.0

function softmax(logits: number[]): number[] {
  const m = Math.max(...logits);                 // shift for numerical stability
  const exps = logits.map((x) => Math.exp(x - m));
  const total = exps.reduce((a, b) => a + b, 0);
  return exps.map((e) => e / total);
}

console.log(softmax([2, 1, 0]));   // -> [0.665..., 0.245..., 0.090...]