Numerically Stable Softmax

Implement the softmax function over the rows of a 2-D array, using the max-subtraction trick so it does not overflow on large inputs. Softmax turns a row of raw scores (logits) into a probability distribution, and it sits at the output of nearly every classifier. Interviewers ask for the stable version because the naive formula silently breaks on large logits.

Definition

For a row of logits z = [z_1, ..., z_k], softmax is:

softmax(z)_i = exp(z_i) / sum_j exp(z_j)

The naive form overflows: exp(1000) is inf in float64, and inf / inf is nan. Because softmax is shift-invariant, subtracting the row maximum m = max(z) from every entry gives the identical result with no overflow:

softmax(z)_i = exp(z_i - m) / sum_j exp(z_j - m)

Now the largest exponent is exp(0) = 1, and every other term is in (0, 1].

Task

Complete softmax(logits) so it returns an array of the same shape where each row is a probability distribution over the k columns. Subtract the per-row maximum before exponentiating. Operate along the last axis (axis=1). Do not call scipy.special.softmax or any library softmax.

Example

logits = np.array([[1.0, 2.0, 3.0],
                   [1000.0, 1000.0, 1000.0]])
softmax(logits)
# -> array([[0.09003057, 0.24472847, 0.66524096],
#           [0.33333333, 0.33333333, 0.33333333]])

The first row is the ordinary softmax of [1, 2, 3]. The second row would overflow to nan without the max-subtraction, but stably resolves to a uniform distribution.

Constraints

• logits is a 2-D array of shape (n, k) with 1 <= n, k.
• Values may be large in magnitude (e.g. ±1000); the output must never contain inf or nan.
• Use vectorized NumPy with broadcasting, not Python loops over rows.

Approach

Softmax is shift-invariant: subtracting a constant from every entry of a row leaves the output unchanged, because the shared factor cancels in the numerator and denominator. Subtract the per-row maximum so the largest logit becomes 0, its exponential becomes 1, and no exp can overflow. Everything is done with keepdims=True reductions along axis=1 so the max and the sum broadcast back over each row.

Solution

import numpy as np

def softmax(logits: np.ndarray) -> np.ndarray:
    row_max = np.max(logits, axis=1, keepdims=True)   # shape (n, 1)
    shifted = logits - row_max                         # largest entry per row is 0
    exps = np.exp(shifted)                             # all in (0, 1], no overflow
    return exps / np.sum(exps, axis=1, keepdims=True)  # normalize each row to sum 1

Why subtract the max

exp grows extremely fast: exp(710) already exceeds the float64 range and returns inf. On a row like [1000, 1000, 1000], the naive exp(z_i) / sum_j exp(z_j) computes inf / inf, which is nan — the result is destroyed even though the true answer is a clean uniform distribution.

Subtracting the row max m fixes this without changing the math. For any constant c:

exp(z_i - c) / sum_j exp(z_j - c) = (exp(-c) * exp(z_i)) / (exp(-c) * sum_j exp(z_j))
                                  = exp(z_i) / sum_j exp(z_j)

The exp(-c) factor cancels, so the output is identical for every choice of c. Choosing c = m = max(z) is the safe choice: the largest shifted logit is z_max - m = 0, so exp(0) = 1 is the biggest term and every other term lies in (0, 1]. There is no overflow, and the denominator is at least 1, so there is no division blow-up either. (Underflow of very negative shifted logits to 0 is harmless — those classes genuinely have negligible probability.)

Walkthrough

On the example:

logits = [[   1,    2,    3],
          [1000, 1000, 1000]]

• row_max = [[3], [1000]]
• shifted = [[-2, -1, 0], [0, 0, 0]]
• exps = [[0.1353, 0.3679, 1.0], [1, 1, 1]]
• row sums = [[1.5032], [3]]
• result = [[0.0900, 0.2447, 0.6652], [0.3333, 0.3333, 0.3333]]

The first row matches plain softmax of [1, 2, 3]; the second row is uniform, computed with no nan.

Complexity & notes

• Time O(n·k), space O(n·k) — one max reduction, one elementwise exp, and one sum reduction over the (n, k) array.
• keepdims=True is what makes the broadcast work: row_max and the row sums keep shape (n, 1) so they subtract/divide across each row of shape (n, k). Without it the shapes would be (n,) and broadcasting would fail or align on the wrong axis.
• For training you usually want log_softmax(z) = z - m - log(sum(exp(z - m))) instead, which avoids re-exponentiating and is what cross-entropy loss uses under the hood.
• The trick generalizes to any axis; here axis=1 is the last axis of a 2-D batch of logits.