Wiki Home
Softmax and logits
Definition
A classifier’s or language model’s last layer emits one unbounded real number per class or per vocabulary token. These are the logits. The softmax turns them into a probability distribution: exponentiate each logit, then divide by the sum of the exponentials. The result is non-negative and sums to one, which is what makes it usable as a distribution and what pairs it with the cross-entropy loss.
Two properties explain most of how it is used. Softmax is monotone, so it preserves the ranking of the logits and the arg-max is the same before and after; turning logits into probabilities changes what the model reports, never which class it picks. It is also shift-invariant: adding a constant to every logit leaves the output unchanged, so logits are meaningful only up to an offset, and implementations subtract the maximum before exponentiating to avoid overflow. Temperature divides the logits before the softmax, which is the one knob that does change the distribution’s shape.
The same operation appears twice more. Inside a transformer’s self-attention, a softmax over query-key scores produces the weights each position uses to mix the others. And at a deployed prediction API, the softmax output is what a caller gets back: the full confidence vector, a truncated top-k, or only the arg-max label. That choice is a leakage decision, since the vector carries the model’s relative scores for every class it did not pick.
Papers that use this concept
- Membership Inference Attacks From First Principles — the confidence the model assigns to a candidate’s true label is the observable, and the attack rescales it before fitting a distribution to it.
- Deep Learning with Differential Privacy — in its transfer-learning setting only the softmax layer and the fully connected layers above it are trained privately, which is part of what the reported epsilon covers.
- A Watermark for Large Language Models — the watermark is added to the logits before the softmax, which is why it needs no change to the weights.
- Secure Transformer Inference Made Non-interactive — softmax is one of the nonlinearities that homomorphic evaluation cannot take at face value.
- zkGPT: An Efficient Non-interactive Zero-knowledge Proof Framework for LLM Inference — the softmax inside attention needs comparison and division, which an arithmetic circuit expresses only at a cost.
Variants and traps
- A softmax probability is not a calibrated confidence. It is what the model’s scores normalize to, and a model can assign near-certainty to an answer it gets wrong; a claim that reads a probability as a likelihood of correctness needs separate evidence.
- “Logits” also names the log-odds transform
log(p / (1 - p))in statistics. Privacy papers use that second sense deliberately, rescaling a confidence into log-odds so its distribution across models is closer to Gaussian. - Softmax is cheap on a GPU and expensive nearly everywhere else in this course. Exponentiation, division, and the maximum it subtracts are all comparisons or transcendental operations, which homomorphic encryption, multiparty protocols, and proof circuits have to approximate or emulate.