27 September 2026 · 7 min read

Softmax explained: probabilities, model confidence, and the stable softmax trick

Softmax turns a classifier's raw scores into probabilities that add up to 1. This post explains the formula with a worked example, how to read its output as the model's confidence and when not to trust it, and why softmax code subtracts the largest score first.

Softmax is a function that turns a list of raw scores, one per class, into probabilities that are each between 0 and 1 and add up to 1. It sits at the output of most classifiers that pick one class out of several, such as an image model choosing between cat, dog, and bird, or a language model choosing the next word out of its vocabulary.

The problem softmax solves

A classifier with $K$ classes ends with a layer that outputs $K$ numbers, one per class, called logits. For a photo and the classes cat, dog, and bird, the logits might be 2.3, 0.6, and -1.2. A higher logit means the model favors that class more, but logits are not probabilities: they can be negative or larger than 1, and they do not add up to any fixed total. Training with cross-entropy loss, and reading the output as how likely each class is, both need a probability distribution: $K$ numbers, each between 0 and 1, that sum to 1.

The softmax formula

Softmax makes every logit positive by exponentiating it, then divides each result by the total. For logits $z_1, z_2, \ldots, z_K$, the probability for class $k$ is

$$\text{softmax}(z)k = \frac{e^{z_k}}{\sum{j=1}^{K} e^{z_j}}$$

The numerator $e^{z_k}$ is positive for any $z_k$, and the denominator adds up the exponentials of all $K$ logits, so the $K$ outputs sum to exactly 1.

With cat 2.3, dog 0.6, and bird -1.2:

$$e^{2.3} \approx 9.97, \quad e^{0.6} \approx 1.82, \quad e^{-1.2} \approx 0.30$$

The sum is about 12.10, so the probabilities are cat $9.97 / 12.10 \approx 0.824$, dog about 0.151, and bird about 0.025. Three properties follow from the formula:

Every output is strictly between 0 and 1, so every class gets some probability, even bird.
The outputs sum to 1.
The order is kept, so the largest logit gets the largest probability.

The exponential also stretches gaps. The cat logit is 1.7 above the dog logit, and the cat probability is $e^{1.7} \approx 5.5$ times the dog probability, because the ratio of two softmax outputs depends only on the difference between their logits. For the same reason, adding the same number to every logit leaves the output unchanged, which the stable version below relies on.

The name comes from argmax, which puts all of the probability on the largest logit: cat 1, dog 0, bird 0. Softmax gives most of it to the largest logit and a share to every other class.

Softmax with two classes is sigmoid

With $K = 2$ and logits $z_1$ and $z_2$, dividing the top and bottom of the formula by $e^{z_1}$ gives

$$\text{softmax}(z)_1 = \frac{1}{1 + e^{-(z_1 - z_2)}} = \sigma(z_1 - z_2)$$

where $\sigma$ is the sigmoid function, $\sigma(z) = 1/(1 + e^{-z})$. A two-class softmax and a one-output sigmoid classifier such as logistic regression compute the same probabilities, with two logits in one and their difference in the other.

Softmax in PyTorch

For training, the classifier returns logits and nn.CrossEntropyLoss applies softmax inside the loss:

1import torch
2from torch import nn
3import torch.nn.functional as F
4 
5torch.manual_seed(0)
6batch_size, d, K = 4, 8, 3
7 
8model = nn.Linear(in_features=d, out_features=K)
9loss_fn = nn.CrossEntropyLoss()            # applies log-softmax inside
10 
11x = torch.randn(batch_size, d)
12y = torch.tensor([0, 2, 1, 0])             # class indices, shape (batch_size,)
13 
14logits = model(x)                          # shape (batch_size, K)
15loss = loss_fn(logits, y)
16loss.backward()
17 
18probs = F.softmax(logits.detach(), dim=-1) # shape (batch_size, K)
19print(probs.sum(dim=-1))                   # 1.0 for every example

dim=-1 tells softmax to normalize across the last axis, the $K$ classes, so each example's row sums to 1. The labels are class indices such as 0, 1, and 2. The cross-entropy for one example is $-\log$ of the softmax probability of its correct class, so the cat photo above, with 0.824 on cat, costs $-\log 0.824 \approx 0.19$. The cross-entropy post explains where that loss comes from.

Reading softmax as model confidence

The class with the largest softmax probability is the model's prediction, and that probability is often read as the model's confidence in it:

1probs = F.softmax(logits, dim=-1)
2confidence, prediction = probs.max(dim=-1)   # both shape (batch_size,)

For a binary classifier, torch.sigmoid(logit) gives the same kind of number, the probability of class 1.

Confidence helps decide what to do with a prediction. A system can act on predictions above 0.95 automatically and send the ones between 0.5 and 0.95 to a person to check. A wrong prediction made at 95% confidence also does more damage than one made at 51%, because nothing downstream is likely to question it.

Confidence is not the same as being correct. A model is calibrated when its confidence matches its accuracy: of all the predictions it makes at 70% confidence, about 70% turn out correct. Many neural networks come out of training overconfident, with predictions at 95% confidence that are right much less than 95% of the time. Calibration is checked on held-out data by grouping predictions by confidence and comparing each group's average confidence with its accuracy.

Model confidence on the x axis and actual accuracy on the y axis. A dashed diagonal marks perfect calibration, a teal calibrated model follows it, saying 70% sure and right 70% of the time, and a red overconfident model sits far below it, 95% confident but often wrong.

A common fix is temperature scaling: divide every logit by one number $T$ before softmax, with $T$ picked on validation data. $T > 1$ spreads the probabilities out and lowers confidence, and $T < 1$ sharpens them. With $T = 2$ the cat logits become 1.15, 0.3, and -0.6, and the cat probability drops from 0.824 to about 0.62. Dividing every logit by the same positive $T$ keeps their order, so the predicted class does not change.

The overflow problem and stable softmax

The exponential grows fast. In 32-bit floats, $e^{z}$ overflows to infinity for any $z$ above about 88.7. Logits that large are rare in a model that trains normally but show up when training goes wrong, for example with a learning rate that is too high. With logits $[1000, 998, 1002]$, every exponential is inf and inf / inf is nan, so every output is nan, and a nan in the loss spreads to every weight on the next update.

The fix is to subtract the largest logit from every logit before exponentiating. With $c = \max(z)$, the largest logit:

$$\text{softmax}(z)_k = \frac{e^{z_k - c}}{\sum_j e^{z_j - c}}$$

This gives exactly the same result. $e^{z_k - c} = e^{z_k} e^{-c}$, and the $e^{-c}$ factor appears in the numerator and in every term of the denominator, so it cancels:

$$\frac{e^{z_k} e^{-c}}{\sum_j e^{z_j} e^{-c}} = \frac{e^{z_k}}{\sum_j e^{z_j}}$$

After the shift the largest logit is 0 and the rest are negative, so the largest exponential is $e^{0} = 1$ and nothing can overflow. For $[1000, 998, 1002]$, $c = 1002$, the shifted logits are $[-2, -4, 0]$, the exponentials are 0.135, 0.018, and 1, and the probabilities come out about 0.117, 0.016, and 0.867.

Before the shift, z equals 1000, 998, 1002 exponentiates to inf, inf, inf, and inf over inf is nan. After subtracting the max, 1002, the logits are minus 2, minus 4, 0, the exponentials are 0.135, 0.018 and 1, and the result is valid probabilities with no overflow. The two versions are mathematically identical because the max term cancels.

Written from scratch for a batch of logits:

1import torch
2 
3def stable_softmax(z: torch.Tensor) -> torch.Tensor:
4    # z has shape (batch_size, K); subtract each row's own max
5    z = z - z.max(dim=-1, keepdim=True).values
6    exp_z = torch.exp(z)
7    return exp_z / exp_z.sum(dim=-1, keepdim=True)
8 
9z = torch.tensor([[1000.0, 998.0, 1002.0]])
10naive = torch.exp(z) / torch.exp(z).sum(dim=-1, keepdim=True)
11print(naive)               # tensor([[nan, nan, nan]])
12print(stable_softmax(z))   # tensor([[0.1173, 0.0159, 0.8668]])

keepdim=True keeps the max with shape (batch_size, 1), so it broadcasts across the $K$ columns and each row is shifted by its own max. F.softmax, F.log_softmax, and nn.CrossEntropyLoss do this shift internally, so the hand-written version is only needed when softmax is written from scratch.

When to use softmax, and common mistakes

Softmax fits when each input belongs to exactly one of $K$ classes. When an input can belong to several classes at once, such as a photo tagged both outdoor and dog, one sigmoid per class fits better, because softmax forces the classes to share a total of 1. The multi-class vs multi-label post covers that case.

Common mistakes:

Softmax before nn.CrossEntropyLoss. The loss applies its own softmax, so passing probabilities in applies it twice. Training runs, but the twice-squashed outputs stay close to uniform and the gradients shrink. Pass raw logits.
The wrong dim. On a (batch_size, K) tensor, F.softmax(logits, dim=0) normalizes across the batch, so each column sums to 1 instead of each row, and no error is raised.
Treating confidence as accuracy. A 0.99 softmax output can still be wrong, and softmax has no output for "none of these classes", so an input from outside the training classes still gets its probability spread across them.
Hand-written softmax without the max shift. It works on small test logits and returns nan once a logit passes about 88 in 32-bit floats.

Sigmoid is the two-class case of softmax, covered with logistic regression. Log-softmax, $\log \text{softmax}(z)$, is computed directly in a stable way and is what nn.CrossEntropyLoss uses internally. Attention layers in transformers also use softmax, on each row of similarity scores between tokens, to turn the scores into weights that add up to 1.

QuiddityML teaches softmax, model confidence, and stable softmax as three concepts in the ML Foundation track, and the exercises include implementing softmax over a batch, finding the buggy line in a stable softmax, and following the shape of a (32, 10) batch of logits through it.