The problem softmax solves

A neural network layer outputs a vector of unbounded real numbers — call them logits, z = (z_1, …, z_K). They can be negative, positive, large, or tiny; nothing constrains them. But a great many tasks need a probability distribution as output: the next-token probabilities over a vocabulary, the attention weight each key receives, the class posteriors of a classifier. A distribution has two hard requirements: every entry must be non-negative, and the entries must sum to one.

So the real question is: what is the right map from an unconstrained vector z ∈ ℝ^K to the probability simplex? We want more than just any map. We want one that (1) is order-preserving — a larger logit should get a larger probability; (2) is smooth and differentiable everywhere, so it can sit in the middle of a network trained by gradient descent; (3) never assigns a hard zero, so no option is ever ruled out and gradients never vanish for an unpicked class; and (4) has a single tunable notion of ‘how peaked’ the result is. Softmax is the function that satisfies all four, and the exponential is what makes it work.

Advertisement

The definition

Softmax is defined component-wise. For a logit vector z of length K:

softmax(z)_i  =  exp(z_i) / Σ_j exp(z_j)     for i = 1 … K

shapes:  z : [K]   →   p : [K]     with   p_i > 0   and   Σ_i p_i = 1

The numerator, exp(z_i), is a positive number for any real z_i — that alone guarantees non-negativity. The denominator, Σ_j exp(z_j), is the sum of all the numerators; dividing by it forces the outputs to sum to one. That is the whole trick: exponentiate to force positivity, then normalize to force summation. The name is a portmanteau of ‘soft’ and ‘argmax’: it is a smooth, differentiable stand-in for the hard argmax that would put all mass on the single largest logit. As the inputs spread apart, softmax approaches that hard one-hot; as they bunch together, it approaches the uniform distribution. Between those extremes it gives a graded, differentiable answer — exactly what a network being trained by backpropagation needs.

Advertisement

Why the exponential, not something else

Positivity could be had many ways — squaring, taking absolute values, or applying any non-negative function — so why exp specifically? Three properties single it out. First, monotonicity: exp is strictly increasing, so the ordering of logits is preserved exactly and the largest logit always receives the largest probability. Squaring fails this — it would map −3 above +2.

Second, the ratio structure. The relative odds of option i over option j are p_i/p_j = exp(z_i − z_j) — they depend only on the difference of the logits. This log-linear form is exactly the maximum-entropy distribution consistent with the logits as sufficient statistics: among all distributions matching the given constraints, the exponential family is the one that assumes the least (highest entropy). Softmax is not a hack; it is the principled answer to ‘least committal distribution given these scores.’ Third, a clean derivative: because d/dx exp(x) = exp(x), the gradient of softmax folds back into softmax itself, giving the tidy Jacobian we derive below. No other choice yields all three at once.

Shift-invariance: the property everything hangs on

Add the same constant c to every logit and the distribution does not move:

softmax(z + c)_i = exp(z_i + c) / Σ_j exp(z_j + c)
                 = exp(c)·exp(z_i) / ( exp(c)·Σ_j exp(z_j) )
                 = exp(z_i) / Σ_j exp(z_j)  =  softmax(z)_i

The exp(c) factors out of numerator and denominator and cancels. Softmax therefore depends only on the differences between logits, never on their absolute level — which is why a network can freely shift its logits up or down without changing its predictions. This is not just an elegant identity; it is the single most useful property in practice. It means we are allowed to subtract any constant we like from the logits before exponentiating, and the mathematics is guaranteed to give the identical answer. The obvious choice of constant — the maximum logit — turns a formula that overflows into one that cannot. Shift-invariance is the bridge between the clean math and the numerically safe implementation, and it is worth internalizing before anything else.

The numerical-stability trick: subtract the max

Logits in a real model routinely reach values like 30, 50, or more. exp(50) is about 5×10^21; exp(89) already overflows a 32-bit float (whose ceiling is about 3.4×10^38). Once any single exp(z_j) becomes inf, the sum is inf, and the result is inf/inf = NaN — the whole distribution is poisoned. Shift-invariance rescues us. Subtract the maximum logit m = max(z) from every entry first:

m      = max(z)
z_safe = z - m            # every entry ≤ 0, the largest is exactly 0
e      = exp(z_safe)      # every value in (0, 1]; no overflow possible
p      = e / sum(e)       # sum ≥ 1, so no divide-by-tiny underflow

After the shift, the largest exponent is exp(0) = 1 and every other is between 0 and 1, so overflow is impossible. Underflow of a small exp(z_safe) to zero is harmless — that entry genuinely deserves near-zero probability — and because at least one term equals 1, the denominator is always ≥ 1. Every production softmax, on GPU or CPU, performs this subtraction. The output is bit-for-bit the distribution the naive formula intends; only the intermediate floats are kept in range.

Temperature: one knob for peakedness

Softmax gains a control parameter when we divide the logits by a temperature T > 0 before exponentiating:

softmax(z / T)_i = exp(z_i / T) / Σ_j exp(z_j / T)

The name borrows from statistical physics, where this is the Boltzmann distribution and T is literal temperature. The effect is intuitive. Low temperature (T → 0) magnifies the differences z_i / T, so the distribution sharpens toward a one-hot on the largest logit — softmax becomes hard argmax, greedy and deterministic. High temperature (T → ∞) shrinks every difference toward zero, so the distribution flattens toward uniform — maximally random. T = 1 is the plain softmax. In LLM sampling, temperature is the primary creativity dial: T < 1 makes generations more focused and repetitive, T > 1 more diverse and risky. Note that temperature is a special case of shift-and-scale that does not cancel — unlike an additive constant, dividing by T changes the differences and therefore genuinely changes the distribution.

A worked numeric example

Take three logits z = (2.0, 1.0, 0.1). Exponentiate: exp(2.0) ≈ 7.389, exp(1.0) ≈ 2.718, exp(0.1) ≈ 1.105. The sum is 11.212. Divide:

p = (7.389, 2.718, 1.105) / 11.212
  = (0.659, 0.242, 0.099)          # sums to 1.000

Now verify shift-invariance by subtracting the max m = 2.0: z_safe = (0, −1.0, −1.9), exp = (1.000, 0.368, 0.150), sum = 1.518, and p = (0.659, 0.242, 0.099) — identical, as promised, but with no value above 1. Next apply temperature. At T = 0.5 the logits become (4.0, 2.0, 0.2) and p ≈ (0.864, 0.117, 0.019) — much sharper. At T = 2.0 they become (1.0, 0.5, 0.05) and p ≈ (0.502, 0.304, 0.194) — noticeably flatter. Same logits, same ordering, three very different distributions. This single vector is worth keeping in mind: it makes every later claim — the Jacobian, log-softmax, the cross-entropy gradient — concrete to check by hand.

The Jacobian: differentiating softmax

Softmax maps K inputs to K outputs, so its derivative is a K×K Jacobian. Write p_i = exp(z_i)/S with S = Σ_k exp(z_k), and differentiate p_i with respect to z_j using the quotient rule. Two cases appear, depending on whether i = j:

i = j:  ∂p_i/∂z_i = (exp(z_i)·S - exp(z_i)·exp(z_i)) / S^2
                       = p_i - p_i^2  =  p_i (1 - p_i)

i ≠ j:  ∂p_i/∂z_j = (0·S - exp(z_i)·exp(z_j)) / S^2
                       = - p_i p_j

both cases:  ∂p_i/∂z_j = p_i (δ_ij - p_j)

The Kronecker delta δ_ij (1 if i = j, else 0) unifies the two lines into one compact expression: ∂p_i/∂z_j = p_i(δ_ij − p_j). The diagonal terms are positive — raising a logit raises its own probability — and the off-diagonal terms are negative — raising one logit steals mass from the others, exactly as a normalized distribution must. From our worked vector: ∂p_1/∂z_1 = 0.659×0.341 ≈ 0.225 and ∂p_1/∂z_2 = −0.659×0.242 ≈ −0.160.

Why that Jacobian makes training cheap

Softmax almost never appears alone; it is followed by a cross-entropy loss against a target class t: L = −log p_t. The magic is what happens when you compose the two derivatives. Using the chain rule and the Jacobian above, the gradient of the loss with respect to the logits collapses to something startlingly simple:

∂L/∂z_j = Σ_i (∂L/∂p_i)(∂p_i/∂z_j)
             = p_j - y_j            # where y is the one-hot target vector

in words:    gradient = softmax_output - one_hot_label

The entire backward pass through softmax-plus-cross-entropy is a single subtraction: the predicted distribution minus the true one-hot label. If the model predicts p = (0.659, 0.242, 0.099) and the true class is the first, the gradient on the logits is (−0.341, 0.242, 0.099) — push the correct logit up, push the rest down, in proportion to how wrong each was. This is why softmax and cross-entropy are always implemented as a fused operation rather than two separate steps: fusing them avoids ever forming the full K×K Jacobian and is both faster and more numerically stable.

Softmax in attention

Inside every attention head, softmax is what turns raw similarities into a convex combination. For a query q and keys K, the head computes compatibility scores, scales them, and softmaxes each query’s row:

scores = Q K^T / sqrt(d_k)          # [N, N], one row per query
weights = softmax(scores, axis=-1)  # each row sums to 1
output  = weights · V             # [N, d_v] weighted average of values

The softmax is taken along the key axis, so each query produces a probability distribution over all keys — the ‘attention weights.’ Because they are non-negative and sum to one, the output is a genuine weighted average of the value vectors: attention never extrapolates beyond the convex hull of the values, it interpolates within it. The division by sqrt(d_k) exists precisely to protect this softmax: dot products of d_k-dimensional vectors grow with d_k, and without the scaling the scores would be large enough to push softmax into its saturated, near-one-hot regime where gradients vanish. Scaling keeps the scores in a range where softmax stays soft and trainable. Causal masking is applied by setting future positions’ scores to −∞ before the softmax, so exp(−∞) = 0 zeroes their weight exactly.

Softmax in the output layer

At the top of the network, the final hidden state is projected to a logit vector of length K = |vocabulary| — typically 32K to 130K for modern LLMs — and softmax turns those logits into the next-token distribution P(token | context). During training this feeds cross-entropy against the actual next token, using the fused gradient = p − y shortcut from above. During inference the distribution is what sampling operates on: greedy decoding takes its argmax, temperature reshapes its peakedness, and top-k / top-p truncate its tail before drawing.

This output softmax is also the single most expensive softmax in the model, because K is enormous. A 128K-way softmax over the vocabulary is a 128K-element exponentiate-and-normalize per generated token, and the logit projection that produces it (the ‘LM head’) is often the largest single matrix in a small model. On CPU-hosted SLMs this matters: the vocabulary softmax and its projection can dominate per-token latency, which is part of why smaller vocabularies and tied embeddings are attractive when squeezing a model onto a CPU.

Log-softmax and log-sum-exp

Training does not want p; it wants log p, because cross-entropy is −log p_t. Computing softmax and then taking its logarithm is wasteful and unstable — a probability that underflowed to 0 gives log(0) = −∞. Compute the log directly instead:

log_softmax(z)_i = z_i - log(Σ_j exp(z_j))
                 = z_i - logsumexp(z)

logsumexp(z)     = m + log(Σ_j exp(z_j - m))     where m = max(z)

The log-sum-exp identity is the same shift-invariance trick wearing a different hat: factor out exp(m), and the log(exp(m)) = m comes back outside the sum, leaving an inner sum of terms all ≤ 1. From our example, m = 2.0, the inner sum is 1.518, so logsumexp = 2.0 + log(1.518) = 2.417, and log_softmax = (−0.417, −1.417, −2.317) — and indeed exp(−0.417) = 0.659, matching p_1. Log-softmax turns a product of probabilities into a numerically safe sum of log-probabilities, which is why perplexity, beam-search scores, and the training loss are all computed in log space.