softmax function

Softmax converts a vector of raw scores into a probability distribution by exponentiating each score and dividing by the sum of all the exponentials, guaranteeing every output is positive and the whole set sums to exactly 1. The exponential form traces back to the Boltzmann distribution in 19th-century statistical mechanics, and John Bridle gave the function its modern name and neural-network framing in his 1990 paper "Probabilistic Interpretation of Feedforward Classification Network Outputs." Every GPT-style language model ends its forward pass with exactly this operation: GPT-2's output layer applies softmax across a vocabulary of 50,257 BPE tokens to turn the model's raw scores for "what comes next" into an actual probability distribution the sampler can draw from.

Subtracting the maximum logit before exponentiating isn't a style preference — it's the difference between working code and silent numerical failure. A logit of 1,000 fed directly into an exponential produces a number so astronomically large it overflows standard 32-bit floating point (whose maximum representable value is around 3.4 times 10 to the 38th power) and returns infinity instead of a usable number; softmax is mathematically shift-invariant, meaning subtracting any constant from every logit before exponentiating produces the exact same final probabilities, so subtracting the largest logit specifically guarantees every remaining exponent is e to the power of zero or smaller, which can never overflow. This single trick — free mathematically, mandatory numerically — is implemented in essentially every production softmax kernel, including the fused kernels inside FlashAttention.

Temperature is the single number that controls how sharply softmax commits to its top choice, and it's exactly the same parameter whether you're sampling text from a chatbot or distilling a smaller model from a larger one. Dividing every logit by a temperature below 1 before applying softmax sharpens the distribution toward the single highest-scoring option — at a low enough temperature it behaves almost like picking the argmax outright; dividing by a temperature above 1 flattens the distribution toward uniform, making every option more equally likely to be sampled. Hinton, Vinyals, and Dean's 2015 paper "Distilling the Knowledge in a Neural Network" used exactly this lever in the opposite direction from text generation: raising the temperature during distillation reveals a teacher model's "dark knowledge" — the relative confidence it places on wrong answers — information a hard, low-temperature prediction would otherwise completely hide from the smaller student model being trained to imitate it.

Softmax's one real limitation — it can never assign a class exactly zero probability, no matter how irrelevant that class is — is precisely what sparsemax and entmax were built to fix. Because softmax exponentiates every logit, even an extremely low score still produces a small positive number after exponentiation, spreading some nonzero probability across every single class regardless of relevance; Martins and Astudillo's 2016 paper "From Softmax to Sparsemax" introduced an alternative projection that can output exact zeros for irrelevant classes, which matters specifically in settings like sparse attention or large structured output spaces where forcing probability mass onto obviously wrong options wastes both interpretability and, at large enough vocabulary sizes, real computation.

Every one of softmax's properties traces back to the same exponential-then-normalize structure: shift-invariance comes from how exponentials handle addition, the numerical-stability trick exploits that exact invariance, temperature works by rescaling the input before the same exponential is applied, and the inability to reach exact zero is a direct consequence of the fact that no real-valued exponential ever actually equals zero — which is also exactly the property sparsemax discards on purpose to get a genuinely sparse distribution instead.