Softmax is the small function that does an outsized amount of work in a transformer. It appears twice on every forward pass — once inside every attention head, turning raw compatibility scores into a set of weights that sum to one, and once at the very end, turning the model’s output logits into a probability distribution over the vocabulary. Its job is deceptively simple: take an arbitrary vector of real numbers and return a probability distribution — all entries positive, all entries summing to exactly one — while preserving the order of the inputs and being smoothly differentiable so gradients can flow. This piece derives softmax from first principles: why the exponential is not an arbitrary choice, how shift-invariance yields the max-subtraction trick that keeps floats alive, what the temperature knob does, and how the famously clean Jacobian ∂p_i/∂z_j = p_i(δ_ij − p_j) collapses the backward pass to a subtraction. We work a full numeric example and close on the CPU-SLM overflow and implementation details that decide whether the math survives contact with 32-bit floats.
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.
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 = 1The 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.
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)_iThe 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 underflowAfter 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.000Now 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_labelThe 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 valuesThe 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.
The two-class case is the sigmoid
It is worth seeing that softmax generalizes the logistic sigmoid rather than competing with it. Apply softmax to a two-element logit vector (z_1, z_2) and use shift-invariance to subtract z_2:
p_1 = exp(z_1) / (exp(z_1) + exp(z_2))
= 1 / (1 + exp(-(z_1 - z_2)))
= σ(z_1 - z_2)So a two-class softmax is exactly a sigmoid applied to the difference of the two logits. This clarifies a common design question: for a binary decision you can use either a single logit through a sigmoid or two logits through a softmax — they are the same model, merely over-parameterized in the softmax case (the extra degree of freedom is the one shift-invariance says is irrelevant). The sigmoid is the special case; softmax is the general multi-class extension. Recognizing the connection also explains why the softmax-plus-cross-entropy gradient p − y is the multi-class twin of the equally clean logistic gradient σ(z) − y — both are ‘prediction minus target,’ and both fall straight out of the exponential family.
CPU-SLM implementation and overflow
On a CPU-hosted small language model the softmax is small in FLOPs but surprisingly rich in pitfalls, because CPUs commonly run inference in float32 or even quantized paths where the exponent range is tight. The non-negotiable rules:
| Concern | What to do |
|---|---|
| Overflow | Always subtract max(z) before exp — never call exp on a raw logit |
| Log path | Use log_softmax / logsumexp, never log(softmax(z)) |
| Loss | Fuse softmax + cross-entropy so the backward pass is the single subtraction p − y |
| Masking | Add −∞ (or a large negative) to masked scores before softmax, not after |
| exp cost | exp is a transcendental; over a 128K vocab it is real time — vectorize / use a fast polynomial approximation |
Two subtleties bite specifically on CPU. First, an all-masked row (every score −∞) yields 0/0 = NaN; guard against it or ensure at least one position is always visible. Second, in low-precision paths the exp and the sum should accumulate in float32 even when weights are int8 or bf16, because the normalization is where precision is lost. The softmax is cheap in arithmetic but is the numerical linchpin of both attention and the output head — treat its precision as load-bearing.
Common pitfalls and misconceptions
A handful of misunderstandings recur. Softmax logits are not probabilities and not log-probabilities until after the transform; a logit of 10 is not ‘ten times likelier,’ only likelier by exp of the difference from its peers. Softmax is not invariant to scaling, only to shifting — multiplying all logits by a constant is exactly the temperature operation and does change the distribution, a fact people conflate with the additive shift-invariance.
Softmax never outputs a true zero or one; every probability is strictly inside (0, 1), which is a feature — it keeps gradients alive for every class — but means ‘confident’ predictions still leak a little mass everywhere. Saturated softmax kills gradients: when one logit dominates, p_i(1 − p_i) → 0, so the Jacobian vanishes and learning stalls — the reason attention scales by sqrt(d_k) and the reason very high-confidence classifiers train slowly. Finally, softmax is permutation-equivariant but not position-aware: it treats its inputs as an unordered set, which is why transformers must inject positional information separately rather than expecting the attention softmax to supply it.