Softmax: from logits to probabilities

Why we need it

A classifier with $K$ classes ends with $K$ numbers — logits: for example, $(2,\,1,\,0)$ for “cat”, “dog”, “bird”. They are just scores: they can be negative, they do not add up to 1, and they do not tell how sure the model is. Softmax turns such a vector into a probability distribution: all numbers are positive, they sum to 1, and the order is preserved — a higher score means a higher probability.

Where it appears:

• the output of a multi-class classifier — from logistic regression to ResNet and ViT: $\hat y=\operatorname{softmax}(Wx+b)$;
• language models — at every step an LLM outputs logits for every token in its vocabulary (tens or hundreds of thousands), softmax turns them into the distribution of the next token, and temperature controls how bold the generation is;
• attention in transformers — the weights $\operatorname{softmax}(QK^\top/\sqrt d)$ decide which tokens to look at;
• mixture of experts (MoE) — the router distributes a token among experts through a softmax;
• reinforcement learning — a Boltzmann policy picks action $a$ with probability proportional to $e^{Q(a)/T}$;
• statistical physics — it is the Boltzmann–Gibbs distribution $p_i\propto e^{-E_i/(kT)}$; hence the word “temperature”.

Intuition: a vote with amplification

Think of the classes as candidates and of the logits as their scores. To get shares of the “vote”, the scores must be made positive and normalized. Simply dividing by the sum does not work: scores can be negative and the sum can be zero. The exponential solves both problems: $e^{z}$ is always positive, and a difference in scores becomes a ratio: every unit of advantage makes a candidate’s share $e\approx2.718$ times its rival’s.

Drag the tops of the logit bars on the left and watch the probabilities on the right. Then press “+1 to every logit”: the bars on the left rise, while nothing changes on the right.

The formula, bone by bone

$z=(z_1,\dots,z_K)$

the vector of logits — raw outputs of the last layer, any real numbers;

$e^{z_i}$

the exponential: makes the number positive and turns logit differences into ratios;

$\sum_{j}e^{z_j}$

the normalizing sum (in physics, the partition function $Z$);

$s_i$

the probability of class $i$: $s_i\in(0,\,1)$ and $\sum_i s_i=1$.

Let us walk through the cheat sheet’s running example — the logits $z=(2,\,1,\,0)$:

Temperature

Logits are often divided by a number $T>0$, the temperature, before the softmax:

The table shows what temperature does to the same vector $(2,\,1,\,0)$:

Two limits explain the name. As $T\to0$ only the largest logit wins: softmax tends to the one-hot vector of the $\operatorname{argmax}$ (if the maximum is unique). As $T\to\infty$ the differences vanish and the distribution becomes uniform, $s_i\to1/K$. Softmax is a “soft” argmax: something between a hard choice and complete indifference.

Properties

• A distribution: $s_i>0$ and $\sum_i s_i=1$ for any logits. Exactly 0 or 1 never happens — only in the limit (and on a computer, through rounding).
• Order-preserving: $z_i>z_j\iff s_i>s_j$, so $\operatorname{argmax}s=\operatorname{argmax}z$. To merely predict the class you can skip the softmax.
• Ratios: $s_i/s_j=e^{z_i-z_j}$ — the logit difference is the logarithm of the probability ratio. Conversely, $z_i=\ln s_i+C$: logits are recovered from probabilities up to a common constant.
• Shift invariance: $\operatorname{softmax}(z+c\mathbf 1)=\operatorname{softmax}(z)$ for any number $c$. Of $K$ logits only $K-1$ differences matter.
• No scale invariance: $\operatorname{softmax}(az)$ is sharper for $a>1$ and flatter for $0<a<1$. This is the same temperature, $a=1/T$.
• Competition: raising one logit increases its probability and decreases all the others — mass is redistributed, not created.

Two classes are a sigmoid

For two classes softmax is the sigmoid of the logit difference, which is why a binary classifier needs only one output. A similar formula holds for $K$ classes if we look at one class against all the others:

Every probability is a sigmoid of how far its logit is ahead of the “smooth maximum” of the rest. To check it, divide the numerator and the denominator of $s_i$ by $e^{z_i}$.

Geometry: the probability simplex

All distributions over $K$ classes form the probability simplex:

For $K=3$ it is an equilateral triangle. The vertices are confident answers $(1,\,0,\,0)$, $(0,\,1,\,0)$, $(0,\,0,\,1)$, the centre is the uniform distribution $(\tfrac13,\,\tfrac13,\,\tfrac13)$, and a point inside is a “mixture” of the vertices with weights $s_i$ (barycentric coordinates).

Softmax maps the whole logit space $\mathbb R^K$ onto the interior of the simplex: the edges and vertices are out of reach. All points of the line $z+c\mathbf 1$ land on one point, so in essence softmax is a one-to-one correspondence between the $(K-1)$-dimensional space of logit differences and the interior of the simplex. The inverse map is the logarithm: $z=\ln s$ up to a shift.

The derivative: the Jacobian matrix

Softmax has $K$ inputs and $K$ outputs, so its derivative is not a number but a $K\times K$ matrix: the entry $J_{ij}=\partial s_i/\partial z_j$ shows how the $i$-th probability responds to the $j$-th logit.

Properties of this matrix:

• the diagonal $s_i(1-s_i)$ is exactly the derivative of the sigmoid: at most $\tfrac14$, largest at $s_i=\tfrac12$;
• off the diagonal every entry is negative: raising one logit takes probability away from the others;
• it is symmetric, and every row (and column) sums to 0: $J\mathbf 1=s-s\,(s^\top\mathbf 1)=0$. This is the derivative of $\sum_i s_i=1$, and the vector $\mathbf 1$ is the shift direction that changes nothing. Hence $J$ is singular;
• it is positive semi-definite: $v^\top Jv=\sum_i s_iv_i^2-\bigl(\sum_i s_iv_i\bigr)^2$ is the variance of $v$ under the distribution $s$, and a variance is never negative;
• it fades with confidence: when $s$ is close to a vertex of the simplex, every entry of $J$ is close to zero — saturation, as with the sigmoid.

A backward pass without the matrix

Backprop needs not $J$ itself but the product $J^\top g$, where $g=\partial L/\partial s$ comes from above. Thanks to the structure $\operatorname{diag}(s)-ss^\top$ it costs $O(K)$ operations instead of $O(K^2)$:

For an LLM vocabulary of $K\approx10^5$ tokens this is the difference between $10^5$ and $10^{10}$ numbers per position. Paired with cross-entropy it gets simpler still: $\partial L/\partial z=s-y$ (derived on the cross-entropy page).

The 3D picture: softmax over the logit plane

Fix $z_3=0$ — no loss of generality, since a shift can make any logit zero — and draw $s_1$ as a function of the two variables $z_1$ and $z_2$. The result is a “ramp”: the surface is close to 1 where $z_1$ exceeds both $z_2$ and zero, and close to 0 where the first class loses. Across the boundary it rises along a sigmoid — the one-against-all formula in action.

The top view splits the logit plane into three regions. The boundaries are the sets of equal logits: the ray $z_1=z_2\ge0$ separates classes 1 and 2, the ray $z_1=0$, $z_2\le0$ separates classes 1 and 3, and the ray $z_2=0$, $z_1\le0$ separates classes 2 and 3. All three rays start at the origin, where $s=(\tfrac13,\,\tfrac13,\,\tfrac13)$.

Log-sum-exp: a smooth maximum

Softmax has an “antiderivative” — the log-sum-exp function:

LSE is a smooth maximum: the largest exponential dominates the sum, and the logarithm returns almost its exponent. More precisely,

The left inequality holds because a sum is at least its largest term, the right one because each of the $K$ terms is at most the largest. For $(2,\,1,\,0)$: $\operatorname{LSE}=\ln11.107\approx2.408$ — between $2$ and $2+\ln3\approx3.099$. With a temperature, $\operatorname{LSE}_T(z)=T\ln\sum_j e^{z_j/T}$, the gap shrinks to $T\ln K$: as $T\to0$ LSE becomes the ordinary maximum, and its gradient $\operatorname{softmax}(z/T)$ becomes the one-hot argmax vector.

Softmax in training

The simplest model with a softmax is multinomial logistic (softmax) regression: $K$ linear functions give the logits, softmax gives the probabilities, and the loss is cross-entropy:

Here $y$ is the one-hot vector of the correct class. The whole Jacobian, paired with the logarithm, has collapsed into the difference “prediction − truth”.

Softmax inside the transformer: attention

Each row of $\operatorname{softmax}(\cdot)$ is how one token spreads its attention over the tokens of the sequence (itself included), with weights summing to 1. Dividing by $\sqrt d$ is a temperature: the dot product of $d$-dimensional vectors with independent zero-mean, unit-variance components has variance $d$. Without the scaling the attention logits are large, softmax saturates to nearly one-hot, and the Jacobian — and with it the gradients — nearly vanishes.

Temperature in text generation

A language model draws the next token from $\operatorname{softmax}(z/T)$. As $T\to0$ (greedy decoding) the most likely token is taken; with $T<1$ the text is more predictable, with $T>1$ more diverse and riskier. A temperature $T>1$ also appears in knowledge distillation (Hinton et al., 2015): the teacher’s “soft” probabilities carry more information than one-hot labels.

In code: stability, log-softmax and PyTorch

The naive formula breaks on large logits: $e^{1000}\approx2\cdot10^{434}$ fits in neither float32 nor float64 and becomes infinity, which then gives $\infty/\infty=\text{NaN}$. In float32 (and in bfloat16, which has the same exponent range) the limit is closer: $e^{x}$ overflows already for $x>88.72$ (the largest float32 is $\approx3.4\cdot10^{38}$). On very negative logits the trouble is the opposite: the exponentials underflow to zero, giving $0/0$.

The cure is shift invariance: subtract $m=\max_j z_j$. The largest exponent becomes 0, $e^0=1$, and the denominator lies between 1 and $K$ — no overflow, no division by zero:

pythonimport numpy as np
np.set_printoptions(precision=3)

def softmax(z, axis=-1, T=1.0):
    """Stable softmax: subtract the maximum, the result does not change."""
    z = np.asarray(z, dtype=float) / T
    z = z - z.max(axis=axis, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=axis, keepdims=True)

def log_softmax(z, axis=-1):
    # log s_i = z_i − logsumexp(z): no log(0) even for tiny probabilities
    z = np.asarray(z, dtype=float)
    m = z.max(axis=axis, keepdims=True)
    return z - m - np.log(np.exp(z - m).sum(axis=axis, keepdims=True))

def softmax_backward(s, g):
    # J^T g = s * (g − <g, s>): O(K) instead of a K×K matrix
    return s * (g - (g * s).sum(axis=-1, keepdims=True))

z = np.array([2.0, 1.0, 0.0])
print(softmax(z))            # [0.665 0.245 0.09 ]
print(softmax(z + 1000))     # [0.665 0.245 0.09 ] — the shift does not matter
print(softmax(z, T=0.5))     # [0.867 0.117 0.016] — sharper
print(log_softmax([0.0, -800.0]))  # [   0. -800.], while np.log(softmax(...)) gives -inf

pythonimport torch
import torch.nn.functional as F

z = torch.tensor([2.0, 1.0, 0.0], requires_grad=True)
s = F.softmax(z, dim=-1)          # tensor([0.6652, 0.2447, 0.0900])
s[0].backward()                   # the first row of the Jacobian
print(z.grad)                     # tensor([ 0.2227, -0.1628, -0.0599])

# For training you do not need a separate softmax:
logits = torch.randn(4, 3)        # a batch: 4 examples, 3 classes
target = torch.tensor([0, 2, 1, 0])
loss = F.cross_entropy(logits, target)   # log_softmax + NLL inside
log_p = F.log_softmax(logits, dim=-1)    # if you need log-probabilities

Common mistakes

1. A softmax before nn.CrossEntropyLoss. This loss expects logits and applies log-softmax itself. A double softmax squeezes the probabilities: for three classes the second softmax never exceeds $e/(e+2)\approx0.576$, the loss cannot drop below $0.551$, and the gradients are distorted.
2. The wrong axis. For a batch of shape $(N,K)$ normalize over classes (dim=-1), not over examples (dim=0). No exception is raised — the model simply learns nonsense.
3. A naive exponential without subtracting the maximum — NaN on large logits, especially in mixed precision: in float16 $e^{x}$ overflows already for $x>11.09$.
4. Softmax for multi-label. If an object can carry several labels at once (a cat and a dog in a photo), the classes do not compete — use a sigmoid on every output and BCE.
5. A softmax probability as a guarantee. $s_{\max}=0.99$ does not mean “right 99% of the time”: deep networks are often overconfident. Calibration is a separate task; the simplest method is temperature scaling, fitting a single $T$ on validation data.
6. Comparing logits of different examples. Logits are defined up to a shift, so “a logit of 5 for one example is more than 3 for another” means nothing. Compare probabilities, or logit differences within one example.

Links to other formulas

• Sigmoid: $\operatorname{softmax}(z_1,z_2)_1=\sigma(z_1-z_2)$, and the Jacobian’s diagonal $s_i(1-s_i)$ is the derivative of the sigmoid.
• Cross-entropy: softmax’s natural partner; together they give the gradient $s-y$, and $\mathrm{CE}=\operatorname{LSE}(z)-z_y$.
• BCE: the two-class special case of cross-entropy, where softmax reduces to the sigmoid.
• Focal Loss: multiplies $-\ln p_t$ by $(1-p_t)^\gamma$, where $p_t$ is the softmax probability of the correct class.
• ReLU: softplus $\ln(1+e^{x})=\operatorname{LSE}(x,0)$ is a smooth version of $\max(x,0)$.

Summary

Check yourself