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.

Why exp and not simply dividing each score by the sum of the scores? Logits can be negative or sum to zero, so plain normalization can produce negative “probabilities” or divide by zero; exp maps every real number to a positive one. It is also strictly increasing, so a larger logit always gets a larger probability, and it turns differences into ratios: p_i / p_j = exp(x_i − x_j), so a logit gap of 1 always means a factor of e, wherever the pair sits on the number line. That ratio property is the shift invariance seen from the other side, and it is what makes the max subtraction free: only the gaps between logits carry information, never their absolute level.

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.

Training. A classifier is trained with cross-entropy, -log p[y] for the true class y, and paired with softmax its gradient with respect to the logits is simply p − onehot(y): push up the true class by how much probability it is missing, push every other class down by the probability it took. That clean gradient is a large part of why softmax plus cross-entropy is the default pairing, and why frameworks compute the two together from the logits rather than from the probabilities.

Attention. Transformer attention takes softmax(QKᵀ / √d_k) over the keys for each query. The √d_k divisor exists because dot products of longer vectors have larger magnitude, and large logits drive softmax toward a one-hot output where the gradient almost vanishes; dividing keeps the logits in a range where the distribution is still soft. A causal mask sets future positions to negative infinity before the softmax, which the max shift handles fine as long as at least one position in the row is unmasked, since exp(-inf - m) is exactly 0.

Decoding. Temperature, top-k and top-p all act around this function: temperature rescales the logits before it, top-k keeps the k largest logits and renormalizes, and top-p keeps the smallest set of tokens whose probabilities reach p. Each of those filters needs the same stable normalization you just wrote.

Check yourself

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...]