Softmax

The softmax function turns a vector of $K$ real numbers into a probability distribution over $K$ outcomes. It is how an image classifier’s raw class scores become class probabilities, how an attention layer turns similarity scores into weights, and how CLIP picks the caption that matches an image. The inputs are called logits; the softmax is the last step between them and anything read as a probability.

Definition

Given logits $\mathbf{z} = (z_1, \dots, z_K)$, the softmax exponentiates each one and divides by the sum of the exponentials. The results are positive, sum to one and keep the order of the logits: the largest logit gets the largest probability. A temperature $T > 0$, dividing the logits before the exponential, makes the distribution sharper ($T < 1$) or flatter ($T > 1$).

Bridle (1989) recommended the function as the output stage of a classification network, so that its outputs form a proper distribution over classes, calling it a normalised exponential: a generalisation of the logistic to several inputs and a differentiable version of picking the maximum, which is why he named it “soft max”. Goodfellow, Bengio and Courville (2016, §6.2.2.3) point out that the function softens the arg max, not the max, and that “softargmax” would be the more accurate name.

Intuition

A network that classifies a small image as horizontal stripes, vertical stripes, diagonal stripes or a checkerboard ends with four numbers, one per class. They can be negative and have no fixed scale. The exponential makes them positive and turns differences into ratios: a class whose logit is 1 higher gets $e \approx 2.7$ times the probability. Normalizing then makes the classes compete for one unit of probability, so raising one logit lowers every other class’s probability.

Formal Definition

$$\operatorname{softmax}(\mathbf{z})_k = \frac{e^{z_k}}{\sum_{j=1}^{K} e^{z_j}}, \qquad \operatorname{softmax}(\mathbf{z}/T)_k = \frac{e^{z_k / T}}{\sum_{j=1}^{K} e^{z_j / T}}.$$

The second form, with temperature $T$, is the one Hinton, Vinyals and Dean (2015) write for distillation, with $T = 1$ the ordinary softmax.

Jacobian. With $p_k = \operatorname{softmax}(\mathbf{z})_k$, differentiating the quotient gives

$$\frac{\partial p_i}{\partial z_j} = p_i \,(\delta_{ij} - p_j), \qquad J = \operatorname{diag}(\mathbf{p}) - \mathbf{p}\mathbf{p}^\top .$$

$J$ is symmetric, and its rows and columns sum to zero: the columns because the outputs always sum to one, the rows because adding a constant to every logit changes nothing. Each entry is a product of $p_i$ and $\delta_{ij} - p_j$, and when one probability approaches 1 and the others 0, one of the two factors approaches zero in every entry, so the whole Jacobian approaches zero: the softmax saturates, as Goodfellow et al. (2016, §6.2.2.3) describe. Combined with the logarithm in the cross-entropy loss, the Jacobian cancels to the gradient $\mathbf{p} - \mathbf{e}_y$; that article derives it and shows what goes wrong with losses that do not undo the exponential.

Properties

Shift invariance. Adding the same constant $c$ to every logit multiplies numerator and denominator by $e^c$, so $\operatorname{softmax}(\mathbf{z} + c) = \operatorname{softmax}(\mathbf{z})$ (Goodfellow et al., 2016, Eq. 6.32). Only differences between logits matter, and $K$ logits describe a distribution with $K - 1$ degrees of freedom.

Two classes. For $K = 2$, $\operatorname{softmax}(a, b)_1 = 1/(1 + e^{-(a - b)}) = \sigma(a - b)$, the logistic sigmoid of the logit difference. A binary classifier with one sigmoid output is a two-class softmax with the second logit fixed at zero (Goodfellow et al., 2016, §6.2.2.3).

Temperature and its limits. Dividing the logits by $T$ keeps their order, so the arg max, and with it the predicted class, does not change. As $T \to \infty$ every probability tends to $1/K$; as $T \to 0$ the distribution collapses onto the largest logit, the one-hot arg max (Guo et al., 2017, §4.2). With a tie for the largest logit the limit splits the mass equally between the tied classes. Multiplying the logits by a factor $a > 0$ is the same as a temperature of $1/a$; it follows that a network can make its outputs sharper or flatter just by scaling its last layer’s weights and biases, without changing which class wins.

Smooth maximum. The normalizer is the log-sum-exp, $\operatorname{LSE}(\mathbf{z}) = \log \sum_j e^{z_j}$. Since the largest term of the sum is $e^{\max_j z_j}$ and there are $K$ terms, $\max_j z_j \le \operatorname{LSE}(\mathbf{z}) \le \max_j z_j + \log K$: LSE is a smooth version of the max, and its gradient is the softmax.

Numerically stable computation. The exponential overflows quickly: in floating point $e^{z}$ is infinite above $z \approx 709.8$ in 64 bits, 88.7 in 32 bits and only 11.1 in 16 bits. Shift invariance gives the fix. With $m = \max_j z_j$,

$$\operatorname{softmax}(\mathbf{z}) = \operatorname{softmax}(\mathbf{z} - m), \qquad \log \operatorname{softmax}(\mathbf{z})_k = z_k - m - \log \sum_{j} e^{z_j - m},$$

(the first identity is Goodfellow et al., 2016, Eq. 6.33). After the shift every exponent is at most zero and the largest term of the sum is exactly 1, so the sum lies between 1 and $K$ and can neither overflow nor underflow to zero. The second form, log-softmax, matters because the small probabilities themselves underflow to zero long before their logarithms become unrepresentable. PyTorch’s documentation says computing log(softmax(x)) in two steps is slower and numerically unstable, and that log_softmax uses an alternative formulation instead.

Examples

The example uses four classes of $16 \times 16$ images, each a cosine pattern with a period of 4 pixels: horizontal stripes, vertical stripes, diagonal stripes and a checkerboard. The classifier is linear, with the unit-norm class patterns as its weight templates, so each logit is the dot product of the mean-subtracted image with one template. The test image is faint vertical stripes plus Gaussian noise.

import numpy as np

# Four classes of 16x16 images, each a pattern with period 4 pixels.
r, c = np.mgrid[0:16, 0:16]
patterns = {
    "horizontal": np.cos(np.pi * r / 2),
    "vertical":   np.cos(np.pi * c / 2),
    "diagonal":   np.cos(np.pi * (r + c) / 2),
    "checker":    np.cos(np.pi * r / 2) * np.cos(np.pi * c / 2),
}
names = list(patterns)
# A linear classifier whose weights are the unit-norm class patterns (templates).
W = np.stack([p.ravel() / np.linalg.norm(p) for p in patterns.values()])   # 4 x 256

rng = np.random.default_rng(0)
img = 0.5 + 0.25 * patterns["vertical"] + rng.normal(0, 0.2, (16, 16))      # pixels near [0, 1]
z = W @ (img - img.mean()).ravel()                   # logits: one dot product per class

def softmax(z, T=1.0):
    e = np.exp((z - z.max()) / T)                    # subtract the max: exponents <= 0
    return e / e.sum()

def softmax_naive(z):
    e = np.exp(z)
    return e / e.sum()

np.set_printoptions(precision=4, suppress=True)
print("classes:", names)
print("logits:            ", z)
print("softmax:           ", softmax(z))
print("softmax(logits+50):", softmax(z + 50))        # adding a constant changes nothing

# Temperature: the same logits divided by T.
for T in (0.1, 0.5, 1, 2, 10, 100):
    print(f"T = {T:5}: ", softmax(z, T))

# exp(x) is the largest finite number when x = log(max float).
for t in (np.float16, np.float32, np.float64):
    print(f"exp overflows above {np.log(float(np.finfo(t).max)):.2f} in {t.__name__}")

# Overflow: the same image with raw pixel values 0..255 instead of 0..1.
z255 = 255 * z
with np.errstate(all="ignore"):
    print("logits (0..255 pixels):", z255)
    print("naive softmax:         ", softmax_naive(z255))
print("stable softmax:        ", softmax(z255))

# Log-probabilities: log(softmax) underflows; log-sum-exp keeps them finite.
def log_softmax(z):
    m = z.max()
    return z - m - np.log(np.exp(z - m).sum())       # z - logsumexp(z)

with np.errstate(divide="ignore"):
    print("log(softmax):", np.log(softmax(z255)))
print("log_softmax: ", log_softmax(z255))

# Jacobian: diag(p) - p p^T, checked against central differences.
p = softmax(z)
J = np.diag(p) - np.outer(p, p)
h = 1e-6
J_num = np.stack([(softmax(z + h * e) - softmax(z - h * e)) / (2 * h) for e in np.eye(4)], axis=1)
print(f"Jacobian check: max abs error {np.abs(J - J_num).max():.1e}, "
      f"row sums {np.abs(J.sum(axis=1)).max():.1e}")

# Two classes: softmax of (a, b) is the logistic sigmoid of a - b.
a, b = z[1], z[0]
print(f"two classes: softmax {softmax(np.array([a, b]))[0]:.6f}, "
      f"sigmoid {1 / (1 + np.exp(-(a - b))):.6f}")

Output:

classes: ['horizontal', 'vertical', 'diagonal', 'checker']
logits:             [ 0.0517  3.0735  0.0494 -0.0411]
softmax:            [0.0427 0.8759 0.0426 0.0389]
softmax(logits+50): [0.0427 0.8759 0.0426 0.0389]
T =   0.1:  [0. 1. 0. 0.]
T =   0.5:  [0.0024 0.9933 0.0023 0.002 ]
T =     1:  [0.0427 0.8759 0.0426 0.0389]
T =     2:  [0.1336 0.6054 0.1335 0.1276]
T =    10:  [0.2302 0.3115 0.2302 0.2281]
T =   100:  [0.2482 0.2558 0.2481 0.2479]
exp overflows above 11.09 in float16
exp overflows above 88.72 in float32
exp overflows above 709.78 in float64
logits (0..255 pixels): [ 13.1917 783.75    12.5993 -10.4904]
naive softmax:          [ 0. nan  0.  0.]
stable softmax:         [0. 1. 0. 0.]
log(softmax): [-inf   0. -inf -inf]
log_softmax:  [-770.5584    0.     -771.1507 -794.2405]
Jacobian check: max abs error 5.8e-11, row sums 2.1e-17
two classes: softmax 0.953549, sigmoid 0.953549

The other three templates are orthogonal to vertical stripes, so their logits are near zero and they share what the vertical class leaves almost equally, at every temperature.

Feeding the same classifier raw pixel values from 0 to 255 instead of 0 to 1 multiplies every logit by 255, which is the same as a temperature of 1/255. The vertical logit, 783.75, is above the float64 limit: $e^{783.75}$ is infinite, the sum is infinite, and the naive formula returns 0 for the small classes and $\infty/\infty$ = NaN for the right one. The shifted version’s small probabilities, near $e^{-770}$, are below the smallest float64 number and round to 0, so taking their logarithm gives $-\infty$; log-softmax never forms them.

Common Misconceptions

• “Softmax outputs are calibrated confidences.” They are positive and sum to one, but nothing forces an 80% output to be right 80% of the time. Guo et al. (2017) found modern deep networks poorly calibrated and fitted a single temperature on validation data to fix it; the cross-entropy loss article describes their findings and temperature scaling.
• “Softmax picks out the maximum value.” It approximates the arg max, a one-hot choice of position; a smooth maximum of the values is the log-sum-exp.
• “The softmax axis does not matter.” Softmax normalizes along one dimension, and PyTorch’s softmax takes it as the dim argument. Segmentation logits of shape $(N, C, H, W)$ need the softmax over the class dimension $C$, giving one distribution per pixel; normalizing over another dimension gives numbers that sum to one across images or pixels instead.

Where It Is Used

Classification. The output layer of nearly every image classification network produces logits, and the softmax turns them into class probabilities, as in the neural network example; training uses the cross-entropy of those probabilities, computed from the logits. Semantic segmentation applies the same softmax at every pixel. Multi-label tasks, where an image can contain several classes, use one sigmoid per class instead, since the softmax forces the classes to compete.

Attention. In self-attention, each row of query-key dot products passes through a softmax to give attention weights that sum to one; dividing by $\sqrt{d_k}$ is meant to keep that softmax out of its saturated regime, and masks set scores to $-\infty$ so their weights become exactly zero.

Contrastive image-text models. CLIP computes cosine similarities between image and text embeddings and applies a softmax over the captions in a batch, and over the images, before its cross-entropy. The temperature, which Radford et al. (2021) describe as controlling the range of the logits, is learned during training as a log-parameterized scalar instead of being tuned by hand. For zero-shot classification, the arXiv version describes a softmax over the scaled similarities between the image and one text per class.

Distillation. Hinton et al. (2015) observed that the small probabilities a trained model gives to wrong classes say much about how it generalizes. They raised the temperature of a large model’s softmax to produce soft targets, trained a small model to match them at the same temperature, and used $T = 1$ in the small model after training.