Numerically Stable Softmax - Deep Learning Basics

Numerically Stable Softmax

What you'll build

The softmax function for a 1-D vector of scores (logits), implemented so it never overflows: softmax(x) returns a probability vector of the same length.

Background

Classifiers and language models output unnormalised scores. Softmax turns them into a probability distribution: every output is positive, they sum to 1, and larger scores get larger probabilities. Done naively, np.exp(1000) overflows to inf and the result is nan. The fix is a one-line identity that every interviewer expects you to know.

Formula

softmax(x)_i = exp(x_i) / sum_j exp(x_j)

Shift invariance: softmax(x - c) = softmax(x) for any constant c, because exp(-c) cancels.
Choose c = max(x): the largest exponent becomes exp(0) = 1, so nothing overflows,
and the denominator is at least 1, so nothing divides by zero.

Inputs

• x: 1-D NumPy array of length n >= 1 (ints or floats, possibly very large or very negative).

Expected return format

A 1-D float array of length n with positive entries summing to 1.

Edge cases and conventions

• Equal inputs give the uniform distribution 1/n.
• softmax([100, 101, 102]) equals softmax([1, 2, 3]) (shift invariance), and softmax([1000, 1001, 1002]) must not produce nan.
• Very negative entries may get probabilities that are tiny but still positive.

Constraints

NumPy only; vectorised, no loops.

Follow-up

• Generalise to a batch: softmax(X, axis=-1), using keepdims=True for the max and the sum.
• The Jacobian is diag(p) - p pᵀ; combined with cross-entropy, the gradient with respect to the logits is simply p - y (used in MLP Forward & Backward Pass).
• Log-softmax: x - max(x) - log(sum(exp(x - max(x)))) is more stable than log(softmax(x)) (used in Masked Cross-Entropy with Label Smoothing).

Input: softmax(np.array([1.0, 2.0, 3.0]))
Output: [0.09, 0.2447, 0.6652] (rounded to 4 decimals)

Input: softmax(np.array([1000.0, 1001.0, 1002.0]))
Output: [0.09, 0.2447, 0.6652] (rounded to 4 decimals)

import numpy as np

def softmax(x: np.ndarray) -> np.ndarray:
    """
    Numerically stable softmax of a 1-D vector.

    Args:
        x: Scores (logits), 1-D array of length n

    Returns:
        Probabilities, 1-D float array of length n (positive, sums to 1)
    """
    # Your code here
    pass