Stable Softmax in NumPy: Interview Problem and Solution

Given a list of real numbers x (logits), return the list of softmax probabilities p, where p_i = e^x_i / Σ_j e^x_j. Your function must return finite, correct values for any finite input, including logits in the thousands.

Examples

Input x = [1, 2, 3]

Output [0.09, 0.2447, 0.6652]

Each step of 1 in the logit multiplies the probability by e.

Input x = [1000, 1001, 1002]

Output [0.09, 0.2447, 0.6652]

The same answer: adding 999 to every logit changes nothing.

Input x = [-10000, 0, 10000]

Output [0, 0, 1]

The first two are below the smallest float64. Zero is the correct answer.

Constraints

1 ≤ len(x) ≤ 10⁴
−10⁵ ≤ x[i] ≤ 10⁵
Answers are checked to a relative tolerance of 10⁻⁹

Hints

Hint 1

Try np.exp(1000.0). Then ask what you could do to the logits that leaves the answer unchanged.

Hint 2

Adding the same constant c to every logit multiplies numerator and denominator by eᶜ. Which c guarantees no exponent is positive?

Hint 3

Subtract max(x). The largest term becomes e⁰ = 1, so the sum is at least 1 and nothing overflows or divides by zero.

Where candidates go wrong

Subtracting the mean. With x = [0, 2000] the mean is 1000 and e¹⁰⁰⁰ still overflows.
Clipping the logits to a range such as [−500, 500]. It avoids the overflow by changing the answer.
Returning a NumPy array of dtype int when the input is a list of integers: convert to float first.

The idea

Softmax is invariant to a common shift: replacing every x_i by x_i − c multiplies the numerator and the denominator by the same factor e^−c. Taking c = max(x) makes every exponent ≤ 0, so every term lies in (0, 1], the largest term is exactly 1, and the denominator lies in [1, n]. There is no overflow, and no division by anything smaller than 1.

Entries far below the maximum underflow to 0, which is the correctly rounded answer. The relative error of the result is a few units in the last place, the best float64 can do.

Time O(n) · Space O(n)

Reference solution

import numpy as np

def softmax(x):
    x = np.asarray(x, dtype=float)
    z = x - x.max()          # every exponent is now <= 0; the largest term is exactly 1
    e = np.exp(z)
    return e / e.sum()       # the denominator is between 1 and len(x)

Console⌘/Ctrl + Enter runs

Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.

Why interviewers ask it

This is the opening question in a large share of ML coding screens, because it separates candidates who have used softmax from candidates who understand it. The naive version passes every small test and fails on the first realistic one.

The theory behind it is in the lesson Numerically stable softmax and log-sum-exp.

Keep going

Easy8 min

Softmax that never overflows