Softmax Function Explained: From Logits to Probabilities

The softmax function takes a list of arbitrary real numbers — called logits or scores — and converts them into a probability distribution: every output is between 0 and 1, and the outputs sum to exactly 1. Larger inputs get larger probabilities, and the gaps between inputs are amplified exponentially. Softmax is everywhere in machine learning. It’s the last step of most classifiers, it’s how a large language model turns its final layer into next-token probabilities, and it sits at the core of the attention mechanism in every transformer.

The formula

For a vector of logits z = [z₁, z₂, ..., zₙ], softmax computes each output as:

softmax(zᵢ) = exp(zᵢ) / Σⱼ exp(zⱼ)

Exponentiate every score, then divide each one by the sum of all of them. Exponentiation makes everything positive; dividing by the total makes everything sum to 1.

A worked example

Suppose a model produces three logits for three classes: [2.0, 1.0, 0.1].

1. Exponentiate: e^2.0 ≈ 7.389, e^1.0 ≈ 2.718, e^0.1 ≈ 1.105.
2. Sum: 7.389 + 2.718 + 1.105 ≈ 11.212.
3. Divide each by the sum: 0.659, 0.242, 0.099.

The first class gets about 66% of the probability mass, even though its logit is only twice the second one’s. That’s the exponential at work: a gap of 1.0 in logit space means a factor of e ≈ 2.72 in probability ratio, no matter where on the scale the logits sit.

Key properties

• Outputs form a valid distribution. All positive, summing to 1. You can sample from it, take its log, or compute cross-entropy against a target.
• Order is preserved. The largest logit always gets the largest probability. Softmax never changes the ranking, only how lopsided it is.
• Shift invariance. Adding the same constant to every logit leaves the output unchanged, because the constant factors out of both numerator and denominator. Only differences between logits matter.
• It’s differentiable everywhere. That’s why it’s called softmax: it’s a smooth, trainable approximation of argmax, which simply picks the largest value and assigns it all the weight. Gradient descent can’t learn through a hard argmax, but it can through softmax.
• Two-class softmax is the sigmoid. With two logits, softmax reduces to the logistic sigmoid applied to their difference, which is why binary classifiers usually output a single score.

Computing it safely

The naive implementation breaks on large inputs. exp(1000) overflows to infinity in standard floating-point arithmetic, and infinity divided by infinity is NaN. Very negative inputs underflow to zero, which can make the denominator zero.

Shift invariance gives a clean fix: subtract the maximum logit before exponentiating.

import numpy as np

def softmax(z):
    z = np.asarray(z, dtype=np.float64)
    shifted = z - z.max()        # largest value becomes 0
    exps = np.exp(shifted)       # every term is in (0, 1]
    return exps / exps.sum()

softmax([2.0, 1.0, 0.1])  # → [0.659, 0.242, 0.099]
softmax([1000, 999, 998]) # works; no overflow

The largest term becomes exp(0) = 1, so the sum is at least 1 and nothing overflows. Every serious library does this internally.

When you need log-probabilities — for training losses or for inspecting model confidence — use a dedicated log_softmax rather than taking log(softmax(z)). It computes zᵢ − max − log(Σ exp(zⱼ − max)) directly (the log-sum-exp trick), which avoids taking the log of a probability that has underflowed to zero.

Softmax and cross-entropy

Classifiers are almost always trained with softmax followed by cross-entropy loss, which is the negative log-probability assigned to the correct class. The pairing is popular for a mathematical reason: the gradient of the combined operation with respect to the logits is simply predicted probabilities − one-hot target. That’s cheap to compute, numerically well-behaved, and pushes the correct logit up and the others down in proportion to how wrong the model was. Frameworks typically fuse the two into one operation that takes raw logits, which is why you usually don’t apply softmax yourself before calling the loss function.

Where LLMs use softmax

Next-token prediction. At each step, a large language model produces one logit per vocabulary entry — tens of thousands of scores. Softmax turns those into a probability distribution over the next token, which a decoding strategy then samples from or searches through. The logprobs that some APIs expose are the log-softmax outputs at each position.

Temperature. Sampling temperature is implemented by dividing logits by T before the softmax. With T < 1, the differences grow and the distribution sharpens toward the top choice; with T > 1, differences shrink and the distribution flattens. As T approaches zero, softmax approaches argmax. Top-k and top-p then trim the distribution after softmax — see temperature, top-p, and top-k for how they combine. A logit bias works one step earlier, nudging specific logits before softmax runs.

Attention. Inside every transformer, self-attention scores how relevant each token is to every other token, then runs those scores through softmax to get weights:

Attention(Q, K, V) = softmax(QKᵀ / √dₖ) · V

The softmax ensures each token’s attention weights sum to 1, making the output a weighted average of value vectors. The division by √dₖ keeps the dot products from growing with dimension; without it, the scores get large, softmax saturates toward a one-hot output, and gradients vanish.

Because attention requires softmax over very long rows, it’s a performance hotspot. FlashAttention relies on an online softmax that processes a row in blocks, tracking a running maximum and running sum, so the full score matrix never has to be stored in slow GPU memory.

Limitations worth knowing

• Always nonzero. Softmax never assigns exactly zero probability, so it can’t express “impossible” without help. Masking (setting logits to negative infinity) is the standard workaround, used for padding tokens and causal attention.
• Overconfidence. Softmax outputs look like calibrated probabilities but often aren’t. A model can assign 99% to a wrong answer. Calibration techniques exist precisely because the raw numbers shouldn’t be taken at face value.
• Cost scales with the number of classes. The denominator sums over every option, which is expensive for very large vocabularies. Approximations like sampled softmax are sometimes used in training for that reason.

The takeaway

Softmax exponentiates a set of scores and normalizes them so they sum to 1, turning arbitrary logits into a smooth, differentiable probability distribution that preserves ranking and amplifies gaps. It powers classification heads, next-token sampling, and transformer attention. In practice, always subtract the maximum before exponentiating, use log_softmax when you need logs, let your framework fuse softmax with cross-entropy, and remember that temperature is just a divisor applied to the logits beforehand.