Softmax
For logits $\mathbf{z} = (z1, \dots, zK)$:
Turns a vector of raw scores (logits) into a probability distribution — all outputs in $(0,1)$ and summing to 1. The standard output activation for multi-class classification.
Definition
For logits $\mathbf{z} = (z_1, \dots, z_K)$:
\[\text{softmax}(z_i) = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}}\]Each output is positive and they sum to 1 — a valid distribution over $K$ classes.
Numerical Stability — always subtract the max
$e^{z_i}$ overflows for large logits. Softmax is shift-invariant, so subtract $\max_j z_j$ first:
\[\text{softmax}(z_i) = \frac{e^{z_i - \max_j z_j}}{\sum_k e^{z_k - \max_j z_j}}\]This changes nothing mathematically but keeps exponents $\le 0$.
Key Properties
- Shift-invariant: $\text{softmax}(\mathbf{z} + c) = \text{softmax}(\mathbf{z})$ — only differences between logits matter.
- Not scale-invariant: multiplying logits by a temperature $T$ sharpens ($T<1$) or softens ($T>1$) the distribution: $\text{softmax}(z_i / T)$.
- Generalizes the sigmoid: for $K=2$, softmax reduces to the sigmoid.
Why it pairs with Cross-Entropy
When softmax is followed by Cross-Entropy Loss, the gradient w.r.t. the logits collapses to the beautifully simple:
\[\frac{\partial L}{\partial z_i} = \hat{y}_i - y_i \quad (\text{predicted} - \text{true})\]This clean gradient is why frameworks fuse them (e.g. PyTorch’s CrossEntropyLoss takes raw logits, not softmax outputs — applying softmax twice is a common bug).
Softmax vs Sigmoid
| Softmax | Sigmoid | |
|---|---|---|
| Output | Distribution over $K$ classes (sums to 1) | Independent prob per output |
| Use | Multi-class (one label) | Binary or multi-label |
Related
- Cross-Entropy Loss — its natural loss partner
- Logistic Regression — softmax = multi-class generalization of sigmoid
- Attention — softmax normalizes attention weights
- Neural Networks · Activation Functions and Optimizers