Implement a numerically stable softmax and cross-entropy loss from scratch.

TL;DR: Naive softmax overflows because exp(large logit) is inf. Subtract the row max before exponentiating: the result is identical mathematically but numerically safe. For cross-entropy, do not compute log of softmax separately; use the log-sum-exp form so you never take log(0). The shift cancels exactly because softmax is invariant to adding a constant to all logits.

How to approach it. State the failure of the naive version (overflow/underflow) and the fix (subtract the max), and prove it is exact, not an approximation. Then write both functions vectorized over a batch, with the cross-entropy fused via log-sum-exp.

A strong answer. Softmax is invariant to adding a constant c to every logit: exp(x_i - c) / Σ exp(x_j - c) = exp(x_i)/Σ exp(x_j). Choosing c = max(x) makes the largest exponent exp(0)=1, so nothing overflows and the dominant term never underflows to zero.

import numpy as np

def softmax(x):                       # x: (batch, classes)
    x = x - x.max(axis=-1, keepdims=True)     # shift: exact, prevents overflow
    e = np.exp(x)
    return e / e.sum(axis=-1, keepdims=True)

def cross_entropy(logits, y):         # logits: (B, C), y: (B,) int labels
    z = logits - logits.max(axis=-1, keepdims=True)
    logsumexp = np.log(np.exp(z).sum(axis=-1))   # stable normalizer
    # log_softmax = z - logsumexp ; pick the true-class log-prob
    log_probs = z[np.arange(len(y)), y] - logsumexp
    return -log_probs.mean()

Two things signal experience: computing cross-entropy via log-softmax (z - logsumexp) rather than log(softmax(...)) avoids ever evaluating log of a number that underflowed to 0, which would give -inf; and the max-subtraction is applied in both functions. This is exactly why frameworks expose a fused log_softmax and a cross_entropy that takes raw logits, not probabilities.

Key takeaways

• Subtract the row max before exp: exact (softmax is shift-invariant) and it caps the largest exponent at 1.
• Compute cross-entropy as z - logsumexp to dodge log(0) = -inf, never log(softmax(...)).
• Pass raw logits to a fused loss; the softmax-CE gradient collapses to softmax(logits) - one_hot(y).
• The overflow risk grows in fp16, so accumulate the loss in fp32 even when activations are half precision.

What interviewers probe next.

• "Why subtract the max specifically, not any constant?" Any constant keeps it exact, but the max guarantees the largest exponent is 1, bounding everything in (0,1] so neither overflow nor catastrophic underflow occurs.
• "Gradient of softmax-cross-entropy?" It simplifies beautifully to softmax(logits) - one_hot(y), which is why the two are fused in practice.
• "fp16 implications?" The overflow risk is worse in low precision; the shift plus computing the loss in fp32 is standard.
• "Temperature?" Divide logits by T before softmax; higher T flattens the distribution, lower sharpens it.

Common mistakes.

• Naive exp(x)/sum(exp(x)) that overflows on large logits.
• Computing log(softmax(x)) and hitting log(0) = -inf instead of using log-sum-exp.
• Forgetting keepdims=True, so broadcasting silently does the wrong thing.
• Taking softmax then feeding probabilities into a separate log step in training, losing precision and speed versus a fused loss on logits.