machine learning · math

Softmax: turning opinions into odds

Jul 18, 2017 · 7 min read

A classifier's final layer is blunt to the point of rudeness. Ask it whether an image is a cat, a dog, or a canary and it doesn't answer with percentages — it answers with three numbers like 2.0, 1.0, 0.1, called logits, and it offers no apology for the fact that they're unbounded, can go negative, and don't add up to anything in particular. They encode a ranking and a rough sense of how much more one option is favoured than another, and nothing else. To turn that into something you can actually act on — a probability you can threshold, sample from, or feed to a loss function — you need a translator. Softmax is that translator, and it does the job so cleanly that it has become the default last step of essentially every classification network on Earth.

From logits to a probability distribution

We have a vector of logits z = (z_1, \ldots, z_K) and we want a vector of probabilities: K numbers that are each non-negative and together sum to one. Softmax gets there in two moves. First exponentiate, which forces every entry positive no matter how negative the logit was; then divide by the total, which drags the whole lot down onto the simplex where they sum to one:

\text{softmax}(z)_i = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}}.

That's the entire definition, and every property worth knowing falls out of staring at it. The exponential is monotone, so the biggest logit stays the biggest probability — softmax never reorders your preferences, it only rescales them. And because e^x grows so fast, gaps between logits get amplified into the probabilities: a logit that's a couple of units ahead of the pack doesn't win by a couple of units, it wins by a landslide. Run our (2.0, 1.0, 0.1) through it and you get roughly (0.66, 0.24, 0.10) — the leader took two-thirds of the mass off a one-unit lead. The name is a good one, once you read it right: this is a soft version of the hard \max, which would have handed all the probability to the winner and nothing to anyone else. Softmax keeps the runners-up in the conversation.

Shift invariance: the constant that vanishes

Here is a small fact that turns out to carry a lot of weight. Add the same constant c to every logit and the output doesn't budge:

\text{softmax}(z + c)_i = \frac{e^{z_i + c}}{\sum_j e^{z_j + c}} = \frac{e^{c}\, e^{z_i}}{e^{c} \sum_j e^{z_j}} = \frac{e^{z_i}}{\sum_j e^{z_j}} = \text{softmax}(z)_i.

The e^c factors out of every term and cancels top against bottom. So softmax cares only about the differences between logits, never their absolute level — the vectors (2, 1, 0.1) and (102, 101, 100.1) produce identical probabilities. One immediate consequence: softmax is not injective. A whole line of logit vectors, all differing by a constant, maps to the same distribution, which is why you can never recover logits from probabilities uniquely (there's always a free additive constant). But the reason this fact earns its own section is that it hands us a free pass to subtract whatever constant we like before exponentiating — and choosing that constant wisely is exactly what stops softmax from blowing up. We'll cash that in below.

Temperature: one knob from sharp to uniform

Now give the logits a dial. Divide them all by a positive temperature T before the softmax:

\text{softmax}(z/T)_i = \frac{e^{z_i / T}}{\sum_j e^{z_j / T}}.

The name is borrowed from statistical physics, where this is the Boltzmann distribution and T is literally temperature, but you don't need the physics to feel what the knob does. Turn T down toward zero and you're dividing by a tiny number, so the logit differences get magnified enormously; the exponential of the largest term dwarfs everything else, and the distribution collapses onto a single spike — a one-hot vector on the top logit. That's \text{argmax}. Turn T up toward infinity and z/T \to 0, so every e^{z_j/T} \to 1, every entry becomes 1/K, and the distribution flattens into the uniform — the model shrugs and says all classes are equally likely.

same logits z = (2.0, 1.2, 0.5, 0.1) → softmax(z / T) low T sharp → argmax T = 1 the plain softmax high T flat → uniform
The same three logits, three temperatures. Low T sharpens the distribution toward a one-hot spike on the winner; high T flattens it toward uniform. The logits never changed — only how loudly the softmax listens to their differences.

So T is a confidence knob that never touches the underlying scores. Cold means decisive, hot means hedging. This is why sampling from a language model at "low temperature" produces safe, repetitive text (the distribution is peaked, so the likeliest token nearly always wins) and "high temperature" produces surprising, sometimes unhinged text (the distribution is flat, so long shots get a real chance). Same logits, same ranking — only the softness changes.

The two-class case is the sigmoid in disguise

Suppose K = 2. Softmax gives the probability of class 1 as

\frac{e^{z_1}}{e^{z_1} + e^{z_2}}.

Divide top and bottom by e^{z_1} — legal, because that's just the shift invariance from before, factoring a constant out of the exponentials:

\frac{1}{1 + e^{z_2 - z_1}} = \frac{1}{1 + e^{-(z_1 - z_2)}} = \sigma(z_1 - z_2).

That is precisely the logistic sigmoid \sigma(u) = 1/(1 + e^{-u}), evaluated at the difference of the two logits. Which is the whole point in a nutshell: with two classes, only the gap between the logits matters (shift invariance again), so a genuinely two-dimensional problem was always secretly one-dimensional. Binary logistic regression and two-class softmax are not cousins — they are the same model, and the sigmoid you meet in logistic regression is just softmax that hasn't noticed it only has two options. Softmax is the honest generalization of the sigmoid to K classes.

Numerical stability: subtract the max before you exponentiate

Now to collect on the promise from the shift-invariance section. The formula as written is a trap for a computer. Suppose a logit is 1000 — not exotic; deep networks produce large pre-softmax activations all the time. Then e^{1000} overflows a 64-bit float to Infinity, the denominator becomes Infinity, and \infty/\infty evaluates to NaN. Your beautiful probability distribution is now a vector of garbage, and the poison spreads through every gradient downstream.

Shift invariance is the cure, and it's free. Because we can subtract any constant from every logit without changing the answer, subtract the largest one, m = \max_j z_j:

\text{softmax}(z)_i = \frac{e^{z_i - m}}{\sum_j e^{z_j - m}}.

Now the largest exponent is exactly z_m - m = 0, so the biggest term you ever exponentiate is e^0 = 1 — no overflow is possible. Every other exponent is negative, so those terms land safely in (0, 1]; the worst that can happen is a very negative logit underflowing to 0, which is harmless (it should be nearly zero). This is the log-sum-exp trick, and it's why a stable implementation of \log \sum_j e^{z_j} is written as m + \log \sum_j e^{z_j - m}. Same mathematics, subtracted constant and all, but one version returns a probability distribution and the other returns NaN. The playground has a toggle that runs the naive and stabilized versions side by side on deliberately huge logits — flip it and watch the naive column light up with Infinity and NaN while the stabilized column stays perfectly finite.

The Jacobian, and the gradient that collapses

The last thing worth doing by hand is differentiating softmax, because the answer is prettier than it has any right to be. Write p_i = \text{softmax}(z)_i. Differentiating p_i with respect to logit z_j splits into two cases — the diagonal, where you differentiate a term with respect to its own logit, and the off-diagonal, where the logit only appears in the shared denominator — and both fold into one line using the Kronecker delta \delta_{ij} (which is 1 when i = j and 0 otherwise):

\frac{\partial p_i}{\partial z_j} = p_i (\delta_{ij} - p_j).

That's the Jacobian. It's neat, but on its own it's a full K \times K matrix that you'd have to multiply through on every backward pass — not obviously a bargain. The magic happens when softmax is followed by the loss it was born to work with: cross-entropy. For a one-hot target y, the loss is L = -\sum_i y_i \log p_i. Push the gradient of L back through the Jacobian, let the sums telescope (using \sum_i y_i = 1), and nearly everything cancels:

\frac{\partial L}{\partial z_j} = p_j - y_j.

Prediction minus target. That's the whole gradient. No Jacobian matrix survives into the final expression, no exponentials, no division — just the vector of probabilities minus the vector of truth. It is the same disarmingly simple form you get from linear regression with squared error, and it is not a coincidence: softmax-with-cross-entropy and linear-regression-with-squared-error are both instances of the same exponential-family story, where this clean residual gradient is guaranteed rather than lucky. It's also exactly why these two are always taught, implemented, and fused together as a single "softmax cross-entropy" op: separately their gradients are a mess of exponentials and reciprocals, but composed, the mess annihilates and leaves you with a subtraction. The diplomat and its favourite loss function were made for each other.


Thoughts? Find me on LinkedIn.