← AI Terminology
Softmax
Softmax is a function that converts a vector of raw scores (logits) into a probability distribution — each output is positive and all outputs sum to 1.
It is the standard final layer for multi-class classification and the core of attention weight computation in transformers.
It is the standard final layer for multi-class classification and the core of attention weight computation in transformers.
Why It Matters in AI
Neural networks produce unbounded logits; softmax turns them into interpretable probabilities for classification, sampling, and attention. Without it, transformers could not distribute attention across tokens or LLMs sample coherent next tokens. Almost every modern generative model ends a prediction step with softmax (or a numerically stable variant).
Key Points
| Aspect | Description |
|---|---|
| Formula | softmax(z)_i = exp(z_i) / Σ_j exp(z_j) — exponential emphasises larger logits |
| Temperature | Dividing logits by T before softmax sharpens (T→0) or flattens (T>1) |
| Alternatives | Sparsemax, entmax, Gumbel-softmax for discrete sampling |
| Role in LLMs | Next-token logits → softmax → multinomial sampling (or argmax) |
| Numerical care | Subtract max(z) before exp to avoid overflow (log-sum-exp) |
| Role in attention | Query–key scores → softmax → attention weights over values |
Simple Analogy
Turning exam raw scores into percentages that add up to 100%: the highest score gets most of the mass, but weaker options still get a share unless the gap is huge.
Common Usage Examples
F.softmax(logits, dim=-1)— PyTorchtf.nn.softmax(logits)— TensorFlow- Attention:
attn = softmax(Q @ K.T / sqrt(d_k)) @ V - LLM decode: sample from
softmax(logits / temperature)
Summary
In short: Softmax turns raw scores into a probability distribution — the bridge from logits to attention weights and next-token sampling.