What does softmax do, and why is it used in the output layer?
Softmax turns class logits into positive values that sum to one by exponentiating and normalizing them. It is used for mutually exclusive multiclass classification because the result forms a categorical distribution and pairs naturally with cross-entropy loss, although models are usually trained directly from logits for numerical stability.
How to think about it
A classifier’s final layer might output [3.0, 1.0, 0.0]. Those are raw scores, not probabilities: they can be negative, larger than one, and they do not add up to anything useful. Softmax converts them into a probability-like distribution, such as [0.8438, 0.1142, 0.0420]. It is used when the classes are mutually exclusive, meaning exactly one class should be correct.
What softmax does
The final linear layer produces one logit, a raw score for each class. For a logit vector z, softmax computes each class probability as:
p_i = exp(z_i) / Σ_j exp(z_j)
The exponential makes every finite score positive. Dividing by the sum makes all the results add to one.
That gives us a categorical distribution: the model assigns its entire probability mass across the possible classes. If the classes are billing, technical, and refund, the model cannot assign 70 percent to billing and 60 percent to technical. The total must remain 100 percent.
There are two important consequences.
First, softmax preserves the ordering of the logits. The largest logit still produces the largest probability, so choosing the class with the largest logit gives the same answer as choosing the class with the largest softmax output.
Second, softmax compares classes against one another. Increasing the billing logit does not merely increase billing’s value. Because the probabilities must still sum to one, it reduces the relative share available to the other classes.
The logits are also shift-invariant. Adding the same number to every logit changes nothing:
softmax([3, 1, 0]) = softmax([103, 101, 100])
The common factor introduced by exponentiation cancels during normalization. This is useful because the absolute zero point of logits has no meaning. Their differences matter.
A concrete example
Suppose a support-ticket model has three output classes:
| Class | Logit | Softmax probability |
|---|---|---|
| Billing | 3.0 | 0.8438 |
| Technical | 1.0 | 0.1142 |
| Refund | 0.0 | 0.0420 |
For the billing class, exp(3) is about 20.09. For technical, exp(1) is about 2.72. For refund, exp(0) is 1. The total is about 23.81, so billing receives 20.09 / 23.81, or roughly 0.8438.
The difference between billing and technical logits is 2. Their unnormalized weights therefore differ by exp(2), about 7.39. A two-unit logit advantage means about 7.39 times as much relative probability weight, before the other classes are included.
A common mistake is to read 0.8438 as “this model is correct 84.38 percent of the time.” That interpretation is only justified if the model is well calibrated. The number first means: given this model and its current parameters, billing has the largest relative score and receives 84.38 percent of the normalized output mass.
Why it pairs with cross-entropy
For single-label classification, the training target usually identifies one correct class. Cross-entropy loss, also called negative log-likelihood in this setting, penalizes the model according to the probability it assigned to that class:
loss = -log(p_true)
For the ticket above, if the correct answer really is billing, the loss is approximately -log(0.8438) = 0.1698. That is a small penalty.
If the correct answer is refund, the same prediction gives a loss of approximately -log(0.0420) = 3.1698. The model is not merely wrong; it is confidently wrong, so the loss is much larger. This is exactly what we usually want during training.
Softmax and cross-entropy also produce a particularly clean gradient. For the combined operation, the gradient with respect to a logit is the predicted probability minus the target indicator. For the billing example, the gradient is approximately:
[-0.1562, 0.1142, 0.0420]
The negative value on billing tells gradient descent to raise the billing logit. The positive values tell it to lower the competing logits. The model learns both the correct class and the separation between classes.
This connection is why softmax is the standard final transformation for multiclass classification. It creates the distribution that cross-entropy evaluates, and the resulting gradient is simple and useful for optimization.
Why training code often omits softmax
In a mathematical diagram, people often draw:
linear layer → softmax → cross-entropy loss
In production training code, the usual pattern is:
linear layer → cross-entropy-from-logits loss
The loss function computes the equivalent of log-softmax and negative log-likelihood internally, using a numerically stable implementation. In PyTorch, for example, torch.nn.functional.cross_entropy expects raw logits, not probabilities:
import torch
import torch.nn.functional as F
logits = torch.tensor([[3.0, 1.0, 0.0]])
target = torch.tensor([0]) # billing is class 0
probs = F.softmax(logits, dim=-1)
loss = F.cross_entropy(logits, target)
print(probs)
print(loss)
The output is approximately:
tensor([[0.8438, 0.1142, 0.0420]])
tensor(0.1698)
Passing probs into cross_entropy would apply the wrong computation. It can also lose numerical precision, especially when a probability becomes extremely close to zero. Use softmax to inspect or consume probabilities. Use the logits directly with a loss function documented to accept logits.
Numerical stability: subtract the maximum
Exponentials grow very quickly. A naive implementation applied to logits [1000, 999, 998] tries to calculate exp(1000), which overflows in ordinary floating-point arithmetic and becomes infinity.
The standard fix is to subtract the largest logit first:
m = max(z)
softmax(z_i) = exp(z_i - m) / Σ_j exp(z_j - m)
For [1000, 999, 998], the shifted logits are [0, -1, -2]. The result is approximately [0.6652, 0.2447, 0.0900], exactly the same mathematical distribution as the naive formula would produce.
The reason it works is the shift-invariance described earlier. Subtracting the maximum makes the largest exponent equal to exp(0), which is one, while every other exponent is at most one.
A production symptom of getting this wrong is NaN or inf in the loss, followed by model weights becoming NaN. Use the framework’s tested softmax and cross-entropy implementations unless there is a good reason to write your own.
Temperature and confidence
A temperature, a positive scaling value applied to logits before softmax, controls how sharp the distribution is:
softmax(z / T)
For the ticket logits [3, 1, 0]:
- At
T = 1, the probabilities are approximately[0.8438, 0.1142, 0.0420]. - At
T = 2, they flatten to approximately[0.6285, 0.2312, 0.1402]. - At
T = 0.5, they sharpen to approximately[0.9796, 0.0179, 0.0024].
A lower temperature magnifies logit differences. A higher temperature reduces them. Temperature is useful when sampling from language models, where it controls randomness, and in knowledge distillation, where a teacher’s softer distribution can convey information about class similarities.
It is also used for calibration. If a classifier is systematically overconfident, a temperature learned on held-out validation data can make its probabilities better match observed accuracy. This changes the confidence values, not the ordering of the classes.
Common mistake: Softmax does not guarantee calibrated probabilities. A model can output
0.99for the wrong class, particularly on unfamiliar or out-of-distribution inputs. A high softmax value means the model strongly prefers that class relative to the alternatives; it does not prove that the input belongs to the class.
When not to use softmax
Softmax is appropriate when the alternatives compete and one label is expected to be correct. It is not the right output transformation for every prediction problem.
For binary classification, a model commonly uses one logit followed by a sigmoid. Sigmoid maps one score independently to a value between zero and one. A two-logit softmax is mathematically equivalent to a sigmoid applied to the difference between the two logits, so the choice is mainly about representation and the loss API.
For multilabel classification, several labels can be true at the same time. Consider an image that contains both a dog and a bicycle. Independent sigmoid outputs can assign high values to both labels. Softmax would force the dog and bicycle probabilities to compete and sum to one, which misrepresents the task.
For regression, where the output might be a temperature of 21.7 degrees or a price of $42,000, softmax is also inappropriate. There is no set of competing classes to normalize.
Softmax can also be unnecessary when all you need is the winning class or top-k ranking. Since the largest logit is always the largest softmax output, applying the transformation adds work without changing the decision. You need softmax when you need a distribution, probabilities for downstream logic, or sampling.
What they’ll ask next
Does softmax improve classification accuracy?
Not by itself. Softmax preserves the ranking of the logits, so applying it after a trained model does not change the predicted class. Its value is that it produces a normalized distribution and works naturally with cross-entropy during training. The training objective, data, architecture, and decision threshold determine accuracy.
Why use sigmoid instead of softmax for multilabel classification?
Sigmoid treats each label as a separate yes-or-no decision. A dog label can be high and a bicycle label can also be high. Softmax assumes the labels are alternatives and forces their probabilities to share one fixed total, so it is wrong when multiple labels may be present.
Can I use the largest softmax probability as an uncertainty estimate?
Only cautiously. Softmax confidence is relative and may be poorly calibrated. Check calibration on held-out data, inspect performance on the types of inputs seen in production, and consider an explicit abstention or out-of-distribution strategy if the cost of confident mistakes is high.
Say this in the interview
“Softmax converts a vector of class logits into positive values that sum to one, giving a categorical distribution; it is used for mutually exclusive multiclass classification because it pairs with cross-entropy, although during training we normally pass raw logits to a numerically stable combined loss.”