Softmax, Logits & Numerical Stability

What Are Logits?

In deep learning, Logits are the raw unnormalized numerical outputs produced by the final layer of a neural network before any activation function is applied.

Logits can be any real number: positive, negative, zero, small, or extremely large (for example, [ 2.5, -1.2, 5.8 ]).

  Final Network Layer ──► Raw Logits z = [ 2.5, -1.2, 5.8 ] ──► [ SOFTMAX ] ──► Probabilities = [ 0.03, 0.00, 0.97 ]

The Softmax Function

Softmax converts a vector of raw logits $z = [z_1, z_2, \dots, z_K]$ into a valid Probability Distribution:

$$\text{Softmax}(z_i) = \frac{e^{z_i}}{\sum_{j=1}^K e^{z_j}}$$

1. Exponentiation ($e^{z_i}$): Turns all numbers into positive values ($e^{\text{negative}} > 0$).
2. Normalization ($\sum e^{z_j}$): Divides each value by the sum of all exponentiated values, ensuring all outputs sum to exactly 1.0.

Numerical Stability: The Max Subtraction Trick

In 32 bit floating point arithmetic, $e^{x}$ overflows to infinity (NaN) for any input $x > 88.7$.

If your network outputs raw logits like $z = [1000, 1001, 999]$, calling $\text{Softmax}(z)$ directly calculates $e^{1000}$, causing an immediate system crash.

Numerically Stable Softmax: Subtract the maximum logit value $m = \max(z)$ from all inputs:

$$\text{Softmax}(z_i) = \frac{e^{z_i - m}}{\sum_{j=1}^K e^{z_j - m}}$$

Because $\frac{e^{z_i - m}}{\sum e^{z_j - m}} = \frac{e^{z_i} e^{-m}}{\sum e^{z_j} e^{-m}} = \frac{e^{z_i}}{\sum e^{z_j}}$, this math produces the exact same result while guaranteeing that the largest exponent is $e^0 = 1.0$, preventing floating point overflow.

Softmax with Temperature

Temperature $T > 0$ controls probability sharpness during sampling:

$$\text{Softmax}(z_i, T) = \frac{e^{z_i / T}}{\sum_{j=1}^K e^{z_j / T}}$$

1. High Temperature ($T > 1.0$): Softens differences, creating a more random and creative distribution.
2. Low Temperature ($T < 1.0$): Sharpens differences, making the model confident and deterministic.
3. Temperature Near Zero ($T \to 0$): Becomes identical to One Hot Argmax (picking the single top prediction 100 percent of the time).

Say this out loud

Logits are raw unnormalized output scores from a neural network. Softmax exponentiates logits and divides by their sum to create a valid probability distribution summing to 1. To prevent floating point overflow crashes, implementations subtract the maximum logit value before exponentiating. Temperature scaling smooths or sharpens probabilities for language model generation.

Followups to expect

1. Why does PyTorch CrossEntropyLoss accept raw logits instead of probabilities? PyTorch combines LogSoftmax and NLLLoss internally into a single fused GPU kernel for numerical stability, avoiding log of zero errors.
2. What is Sparsemax? An alternative to Softmax that maps low logit values to exact zeros, producing sparse probability distributions.