Softmax Function

Softmax turns a vector of scores into a probability distribution:

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

Every output lands in \((0,1)\) and they sum to 1. Logits of \([2, 1, 0]\) become \([0.66, 0.24, 0.09]\), so the first class takes about two thirds of the mass. It is the standard final layer for multiclass classification, and with two classes it reduces exactly to the sigmoid.

Temperature, gaps, and where confidence comes from

import numpy as np

def softmax(z, temperature=1.0):
    scaled = np.asarray(z, dtype=float) / temperature
    e = np.exp(scaled - scaled.max())
    return e / e.sum()

logits = np.array([2.0, 1.0, 0.0])
print(f"  logits {logits.tolist()}, read at different temperatures")
for t in (0.25, 0.5, 1.0, 2.0, 10.0):
    print(f"    T = {t:>5.2f}   {np.array2string(softmax(logits, t), precision=4, floatmode='fixed')}")

print("\n  with all other logits fixed, increasing the top logit raises its probability")
for gap in (0.5, 1.0, 2.0, 5.0, 10.0):
    p = softmax(np.array([gap, 0.0, 0.0]))
    print(f"    top logit ahead by {gap:>4.1f}   top probability {p[0]:.4f}")

print("\n  adding classes at low logits causes a small, nonzero decrease in top probability")
for extra in (0, 3, 10, 100):
    z = np.concatenate([logits, np.full(extra, -10.0)])
    print(f"    {extra:>3} extra classes at logit -10   top probability {softmax(z)[0]:.6f}")
#   logits [2.0, 1.0, 0.0], read at different temperatures
#     T =  0.25   [9.8169e-01 1.7980e-02 3.2932e-04]
#     T =  0.50   [0.8668 0.1173 0.0159]
#     T =  1.00   [0.6652 0.2447 0.0900]
#     T =  2.00   [0.5065 0.3072 0.1863]
#     T = 10.00   [0.3672 0.3322 0.3006]
#
#   with all other logits fixed, increasing the top logit raises its probability
#     top logit ahead by  0.5   top probability 0.4519
#     top logit ahead by  1.0   top probability 0.5761
#     top logit ahead by  2.0   top probability 0.7870
#     top logit ahead by  5.0   top probability 0.9867
#     top logit ahead by 10.0   top probability 0.9999
#
#   adding classes at low logits causes a small, nonzero decrease in top probability
#       0 extra classes at logit -10   top probability 0.665241
#       3 extra classes at logit -10   top probability 0.665233
#      10 extra classes at logit -10   top probability 0.665214
#     100 extra classes at logit -10   top probability 0.664969

Temperature rescales the logits before exponentiating. The same three numbers give \([0.9817, 0.0180, 0.0003]\) at \(T = 0.25\) and \([0.3672, 0.3322, 0.3006]\) at \(T = 10\). Low temperature sharpens toward a hard maximum; high temperature flattens toward uniform. Nothing about the model changed between those rows.

That single parameter does a lot of work in practice. It is the knob in temperature scaling, which Guo and colleagues showed is a simple and effective way to fix the overconfidence of modern networks. It controls how adventurous sampling from a language model is. And it is what makes knowledge distillation work — a high temperature exposes the teacher’s relative beliefs about the wrong classes, which is the part worth transferring.

Logit gaps determine pairwise probability ratios. In this block the other logits are fixed, so increasing the top gap from 1 to 5 raises the top probability from 0.5761 to 0.9867. Across arbitrary vectors, the same top-two gap need not give the same confidence: [2,1,-10] and [2,1,1] both have gap 1 but top probabilities about 0.7311 and 0.5761. Every term in the denominator matters.

So a softmax output near 0.99 says the top logit is a few units clear of the rest. It does not say the model is right 99% of the time, and on a trained network those two are usually quite different — see calibration.

It always names a winner

The last block adds classes the model has effectively ruled out, at a logit of −10. A hundred of them move the top probability from 0.665241 to 0.664969 — four decimal places in.

That is the intended behavior, and it has an awkward corollary: because the outputs are forced to sum to one, softmax has no way to express “none of these.” Shown something outside its training distribution, it will still produce a distribution, and the shape of that distribution is not evidence that the input belonged to any of the classes.

Where that matters, the usual answers are an explicit “other” class, a threshold on the maximum probability with an abstain option, or a separate out-of-distribution detector. A multi-label head of independent sigmoids also has the property softmax lacks — it can say no to everything.

Implementation notes

• Subtract the maximum before exponentiating. \(e^{1000}\) overflows; the shifted form is mathematically identical, as logits shows.
• Do not apply softmax before the loss. PyTorch’s CrossEntropyLoss takes logits and fuses the log_softmax internally; applying softmax first is a common bug that silently degrades training.
• Cost grows with the number of classes. At a vocabulary of a hundred thousand, the denominator is the expensive part of a language model’s output layer, which is why sampled and hierarchical approximations exist.

References

• PyTorch documentation, CrossEntropyLoss
• Guo, Pleiss, Sun & Weinberger, On Calibration of Modern Neural Networks, ICML 2017