Almost every probability distribution an LLM produces is the output of a softmax core operation A function that turns a vector of arbitrary real numbers (logits) into a probability distribution: softmax(z)_i = exp(z_i) / Σ_j exp(z_j). Outputs are non-negative and sum to 1. Differentiable everywhere, smooth, and 'soft' — values close in input space map to values close in output space, unlike a hard argmax. The de facto choice for normalising over a finite set in neural networks. . Cross-entropy loss takes a softmax output and a one-hot target. Attention scores are softmax-weighted (Ch.13). Sampling for generation reads from a softmax distribution. The operation is conceptually a one-liner — turn logits into probabilities by exponentiating and normalising — but the float-precision details get you in trouble within minutes if you don’t know them. This section covers (1) the definition and three properties that make softmax the right choice, (2) the max-subtraction trick that every production implementation uses, and (3) the temperature parameter that controls how “sharp” or “diffuse” the resulting distribution is.

The definition

For a vector of real-valued logits z = (z_1, …, z_C):

The shift invariance is the one operational property to commit to memory; it’s what enables the max-subtraction trick below.

Three structural reasons softmax is the default for “turn this real vector into a probability distribution”:

1. Smooth. Unlike argmax contrast A function that returns the index of the maximum element of a vector: argmax(z)_i = 1 if z_i is the largest, else 0. The 'hard' version of softmax. Non-differentiable (discontinuous when two entries are close to equal), making it impossible to use in gradient-based training. Softmax is the smooth replacement. (which is discontinuous when two logits tie), softmax is smooth and differentiable everywhere. Critical for gradient-based training.
2. Probabilistic interpretation. The output is a valid probability distribution. You can sample from it; you can take cross-entropy with a label; you can compute KL divergence against another distribution. The softmax output behaves “as if” it came from a probabilistic model.
3. Canonical link for multinomial classification. The softmax + cross-entropy pair gives the clean ∂L/∂z = p − y gradient (Ch.9 §1). This is the canonical-link-function identity that recurs anywhere “predict a category” matters.

The numerical landmines

A naïve implementation of softmax — literally compute the exponentials and normalise — fails on real inputs in two ways:

Production training routinely sees logits in the ±50 range; long-context inference can spike to ±100. Naïve softmax in float32 fails on real workloads.

The max-subtraction trick

The fix is one line, exploiting the shift invariance:

Max-subtraction standard practice The numerical stability transform for softmax: subtract the maximum logit from all logits before exponentiating. Result is mathematically identical to naïve softmax (by shift-invariance) but every exp(z_i − max) is ≤ 1 — no overflow possible. Standard in every production softmax kernel. PyTorch's torch.softmax does this; F.cross_entropy fuses it with the log + nll. Also called 'log-sum-exp trick' when applied to log-softmax. (also called the log-sum-exp trick related identity A related identity for computing log(Σ exp(z_i)) numerically stably: log(Σ exp(z_i)) = max(z) + log(Σ exp(z_i − max(z))). Used in log-softmax and in any context where the softmax denominator's log is needed (e.g., the partition function in language modelling). Mathematically equivalent to max-subtraction in softmax; just used from the log perspective. when applied in log space) is what every production softmax kernel uses. PyTorch, JAX, TensorFlow, llama.cpp — all do this. The implementation is ~3 lines:

void softmax_stable(const float* z, float* out, int N) {
    float m = z[0];
    for (int i = 1; i < N; i++) if (z[i] > m) m = z[i];   // find max
    float sum = 0;
    for (int i = 0; i < N; i++) {
        out[i] = expf(z[i] - m);                          // shift then exp
        sum += out[i];
    }
    for (int i = 0; i < N; i++) out[i] /= sum;            // normalise
}

Three passes over the input: find max, compute exponentials, divide by sum. O(N) work, O(1) extra state.

Temperature — controlling the sharpness

The softmax has a tuning knob hidden inside it that’s exposed at inference time. Temperature generation parameter A positive scalar τ that divides the logits before softmax: softmax(z / τ). τ = 1 is the standard softmax; τ < 1 'sharpens' the distribution (concentrates mass on the top logits); τ > 1 'softens' it (more uniform). Used for inference-time sampling: low τ for deterministic / accurate, high τ for creative / diverse. Standard generation parameter in every LLM API. τ scales the logits before the softmax:

Temperature directly controls how much “randomness” the model uses at inference. τ = 0.7 is a common default — slightly sharper than the trained distribution, biasing toward higher-probability completions while still leaving some diversity. τ = 0 (or τ = 10⁻⁶) gives greedy decoding (always the most probable token). τ = 1.5 gives creative writing with more diversity but more error.

END OF CH.12 §1 — Softmax + numerical stability.
Three recall items: easy (definition + shift invariance), medium (max-subtraction trick under a production NaN scenario), hard (temperature’s entropy interpretation and its operational consequence for sampling).
Coming next: §12.2 — Cross-entropy as KL divergence; the canonical-link argument made information-theoretic.