Why softmax subtracts the row maximum first

L1ScreenCode reading6 minCodingScreen

RolesLLM Product EngineeringInference EngineeringAI Data SystemsAI Evaluation and Reliability

The question

A candidate writes this softmax for a batch of logits and the tests pass on small inputs but return NaN on a real attention row:

python
def softmax(x):
    e = np.exp(x)
    return e / e.sum(axis=-1, keepdims=True)

What is wrong, how do you fix it, and why does the fix not change the result?

Answer it first, aloud or on paper. Then open the answer.

The answer

Show short answer

A large logit overflows the exponential to infinity, and infinity divided by infinity is NaN. Subtract the maximum of each row before the exponential. The softmax of a row is unchanged when the same constant is subtracted from every entry, because the constant cancels between the numerator and the denominator.

Show full answer

The failure

Attention scores and language model logits can reach values in the tens or hundreds. The exponential of such a value does not fit in a float, so the row becomes infinite and the division gives NaN. Small test inputs never show it.

The fix

python
def softmax(x):
    shifted = x - np.max(x, axis=-1, keepdims=True)
    e = np.exp(shifted)
    return e / e.sum(axis=-1, keepdims=True)

After the shift the largest entry of each row is zero, so the largest exponential is one and nothing overflows. Entries far below the maximum underflow to zero, which is the correct limit.

Why the result is the same

Softmax of a row divides each exponential by the sum of all exponentials in that row. Subtracting a constant multiplies every exponential by the same factor, and that factor cancels in the ratio. The shift must be per row, over the last axis. A single maximum taken across the entire batch changes nothing about correctness but lets one hot row push every other row into underflow.

Two more details that bite

  • Keep the dimension with keepdims=True so the subtraction broadcasts along the row and not along the batch.
  • Do not subtract in place when the caller still needs the logits.

Follow-ups

  1. L1ScreenWhere does this same shift appear inside attention?

    Show expected answer

    in the softmax over the scaled scores, which is why the attention lesson writes the shifted version first.

  2. L2WorkingWhy is a global maximum over the whole batch a bug?

    Show expected answer

    one row with a very large logit shifts every other row far below zero, so those rows underflow to all zeros and divide by zero.

  3. L2WorkingThe practice task also asks for the gradient of the mean cross-entropy. What is it?

    Show expected answer

    the softmax minus the one-hot target, divided by the batch size.

Wrong answers

Show the wrong answers that interviewers hear

  • “Clip the logits to a safe range.”

    Clipping changes the distribution; the shift does not.

  • “Use float64 and the problem goes away.”

    It moves the threshold, and attention scores still overflow in long rows.

  • “Subtract the mean instead of the maximum.”

    The mean does not bound the largest entry, so a row can still overflow.

Taught in

The self-attention lesson builds the shifted softmax as its first step, and the softmax and cross-entropy practice checks it with large logits.