Why can softmax overflow, and how is it fixed in practice?

Softmax exponentiates every logit, and exp() overflows to Infinity once its input gets large enough (around 709.8 in standard double-precision floating point), which then produces NaN when Infinity is divided by Infinity. The standard fix is the numerically stable softmax: subtract the maximum logit from every logit before exponentiating. This never changes the final probabilities, because subtracting the same constant from every logit does not change the ratios between them, but it keeps every exponent at or below zero.

What loss function actually trains a language model to predict the next token?

Cross-entropy loss, applied to the softmax probability assigned to the actual correct next token: loss = -log(p_true). Its gradient with respect to each pre-softmax logit has a clean closed form, p_i minus 1 for the true token and p_i for every other token, which is the exact signal used to adjust the model’s weights at every training step.

What is perplexity and how does it relate to cross-entropy loss?

Perplexity is exp() applied to the average cross-entropy loss per token across a sequence. It restates the same information on a more interpretable scale: a perplexity of 8 means the model was, on average, about as uncertain as if it were choosing uniformly among 8 equally likely tokens at each step. A perfect model that always assigns probability 1 to the correct token has a perplexity of exactly 1.

← Back to the Intermediate level

MATHS FOR AI/ML — ADVANCED

Softmax in Production: Numerical Stability, Cross-Entropy & Perplexity

The Intermediate level covered the real temperature formula, entropy, and where top-k and top-p disagree. This level covers what actually matters once softmax runs on real hardware at real scale: why it can silently break, the exact loss function and gradient that train a model on this math at every single step, and how nucleus sampling’s adaptive set size plays out on genuinely different kinds of context.

The Quick Answer

Plain softmax has a real failure mode: exponentiating a large logit can overflow to infinity in floating point, silently producing NaN probabilities. Production implementations use a numerically stableversion that subtracts the maximum logit first — mathematically identical output, no overflow risk. On the training side, a model learns entirely through cross-entropy losson softmax’s output for the true next token, and that loss has a remarkably clean gradient: $\frac{\partial \text{loss}}{\partial z_i} = p_i - y_i$ . Perplexityrestates the same loss on an interpretable scale, and nucleus (top-p) sampling’s real advantage over a fixed top-k is that its selected set size adapts to how confident the distribution actually is at each step.

Why Naive Softmax Can Overflow

In standard double-precision floating point, Math.exp(x) overflows to Infinityonce x exceeds roughly 709.8. Logits this large aren’t exotic edge cases — unnormalized outputs from real network layers can reach into the hundreds or thousands before any scaling is applied. Once one exponentiated value is Infinity, the softmax division becomes Infinity ÷ Infinity, which evaluates to NaN— a silently broken output that can propagate through an entire generation.

The fix is the numerically stable softmax formula, used in essentially every production ML framework:

P(x_i) = \frac{e^{z_i - \max(z)}}{\sum_{j} e^{z_j - \max(z)}}

Subtracting the maximum logit from every logit doesn’t change any of the ratios between them — the same algebraic property covered in the Intermediate level’s logit-gap-to-ratio relationship — so the final probabilities come out identical. What it does change is the exponent range: every shifted logit is now 0 or negative, so every exponentiated value is capped at 1, and overflow becomes structurally impossible. See both versions computed side by side, including a case that actually overflows:

The Loss Function That Trains an LLM

Every one of a language model’s billions of parameters gets adjusted using the same underlying signal at every training step: cross-entropy losson softmax’s output for whichever token actually came next in the real training text.

\text{loss} = -\log(p_{\text{true}})

When the model already assigns high probability to the correct token, this loss is small; when it assigns low probability to the correct token, the loss grows sharply — unboundedly, in fact, as p_true approaches 0. Try adjusting the true token’s logit and watch both the probability and the resulting loss respond:

The Gradient: Why Training Reduces to One Subtraction

Combining softmax with cross-entropy loss produces an unusually clean derivative. Let y_i be 1 for the true token and 0 for every other token (a one-hot vector). The gradient of the loss with respect to each pre-softmax logit z_i is simply:

\frac{\partial \text{loss}}{\partial z_i} = p_i - y_i

No chain-rule bookkeeping needed at training time beyond this single subtraction — the true token’s gradient is always its own probability minus 1 (pushing its logit up), and every other token’s gradient is just its own probability (pushing its logit down). Explore this directly:

Perplexity: The Same Loss, Rescaled

Cross-entropy loss is mathematically sound but not very interpretable on its own. Perplexityrestates it on a scale that maps to an intuitive quantity — roughly, “how many equally likely choices does this feel like”:

\text{perplexity} = e^{\,\text{average loss}}

A perplexity of 1 is a perfect model (always fully confident and correct); higher values mean more average uncertainty across the sequence. See it computed from real per-token losses:

Nucleus Sampling’s Real Advantage: Adaptive Set Size

The Intermediate level showed top-k and top-p disagreeing on set size for one fixed distribution. The deeper point is whythat adaptivity matters: a model’s confidence genuinely varies token to token, and a fixed top-k can’t tell the difference between a nearly-certain next token and a wide-open one. Compare a confident context against an open-ended one at the same p value:

This is the practical argument made in Holtzman et al.’s paper on neural text degeneration: a constant k is either too restrictive on genuinely open-ended continuations or too permissive on nearly-deterministic ones, while a cumulative probability cutoff adapts itself to the actual shape of the distribution at each individual generation step.

Source: Holtzman et al., “The Curious Case of Neural Text Degeneration,” arXiv:1904.09751

Fun Fact

The numerically-stable softmax trick — subtracting the max logit before exponentiating — is often introduced in course materials as the “log-sum-exp trick,” because the same shift-by-the-max identity is used to compute log(sum(exp(x))) without overflow, which shows up constantly in machine learning wherever a softmax or a log-likelihood needs to be computed at scale.

Test Yourself

Why does naive softmax risk producing NaN on logits that are otherwise mathematically reasonable?

Why does subtracting the maximum logit from every logit before exponentiating produce the exact same softmax probabilities as the naive version?

What is the gradient of cross-entropy loss with respect to the true token’s pre-softmax logit?

A model has a perplexity of 1 on a test sequence. What does that mean?

Brainstorm

These don’t have a single correct answer — they’re here to make you pause and think before checking one way to reason about each one.

The numerically stable softmax subtracts max(z) from every logit before exponentiating. Could you instead subtract any other constant — say, the mean of the logits, or just 100 — and still get correct probabilities? Why or why not?

Cross-entropy loss is -log(p_true). What happens to this loss as p_true approaches 0, and why might that be a deliberately harsh design choice rather than an accident?

If a language model reaches a perplexity of 2 on English text, roughly how many "real" next-word choices does that suggest the model is choosing between at each step, on average? Does that number feel high or low to you for natural language?

The gradient p_i − y_i is the same whether the model’s current probability for the true token is 51% or 99%. What does this tell you about how training naturally slows down as a model gets more confident and correct?

Nucleus sampling’s selected set size shrinks on confident distributions and grows on uncertain ones. Can you think of a real generation scenario where you’d actually want the opposite — a fixed top-k regardless of the model’s confidence?