Softmax Probabilities

A classifier's last layer usually emits raw scores called logits, not probabilities. Softmax turns any vector of real numbers into a probability distribution, and it is also what converts attention scores into the mixing weights used to build a context vector.

The idea

For a vector $z$ with $K$ components, softmax assigns

Every output is positive and the outputs sum to exactly $1$, so the result is a valid distribution over the $K$ options. Three properties carry most of the intuition:

• Order preserving — a larger logit always yields a larger probability, so argmax is unchanged from the logits to the probabilities.
• Shift invariance — adding the same constant $c$ to every logit cancels top and bottom, because $e^{z_i+c} = e^c e^{z_i}$. Only differences between logits matter.
• Exponential gaps — a one-unit logit gap becomes a probability ratio of $e \approx 2.718$, so softmax amplifies small score differences into large confidence differences.

The temperature $T$ controls sharpness: $\text{softmax}(z/T)$ with small $T$ approaches a one-hot argmax, while large $T$ flattens toward the uniform distribution.

Worked example

Let $z = (1, 2, 3)$. Exponentiating gives $(2.718, 7.389, 20.086)$ with sum $30.193$, so

The top logit is only $1$ above the runner-up, yet it receives about $2.7\times$ the probability. With $T = 2$ the logits become $(0.5, 1, 1.5)$ and the distribution flattens to $(0.186,\ 0.307,\ 0.506)$.

In code

import torch

z = torch.tensor([1.0, 2.0, 3.0])
p = torch.softmax(z, dim=0)
print(p)                             # tensor([0.0900, 0.2447, 0.6652])
print(p.sum())                       # tensor(1.)
print(torch.softmax(z / 2, dim=0))   # tensor([0.1863, 0.3072, 0.5065])

# Stability: softmax subtracts the max before exponentiating.
print(torch.softmax(torch.tensor([1000.0, 1000.0]), dim=0))  # tensor([0.5, 0.5])
print(torch.exp(torch.tensor(1000.0)))                       # tensor(inf) -> don't do this

Check yourself

1. Why does adding 100 to every logit leave the softmax output unchanged?
2. If lowering the temperature makes the distribution sharper, what happens to it as T grows very large?
3. Why is cross-entropy loss computed on logits rather than on probabilities that have already been passed through softmax?

Key takeaways

• Softmax maps unnormalized logits to positive weights that sum to one.
• Only logit differences matter, and those differences are amplified exponentially.
• Temperature tunes how sharp or flat the resulting distribution is.