Softmax

Softmax converts a vector of raw scores (logits) $\mathbf{z} \in \mathbb{R}^K$ into a probability distribution:

\text{softmax}(z_i) = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}}

Properties:

Output sums to 1, all entries in $(0, 1)$ → valid probability distribution
Monotonic: larger logits get larger probabilities
Translation invariant: $\text{softmax}(\mathbf{z} + c) = \text{softmax}(\mathbf{z})$ — in practice, subtract $\max(\mathbf{z})$ for numerical stability
Temperature scaling: $\text{softmax}(z_i / \tau)$ — low $\tau$ → sharper (approaches argmax), high $\tau$ → more uniform

Gradient: $\frac{\partial}{\partial z_j}\text{softmax}(z_i) = \text{softmax}(z_i)(\delta_{ij} - \text{softmax}(z_j))$

Numerical stability: $e^{z_i}$ overflows for large $z_i$ . Exploit translation invariance and subtract the max so the largest exponent is $e^0 = 1$ :

\text{softmax}(z)_i = \frac{e^{z_i - z_{\max}}}{\sum_j e^{z_j - z_{\max}}}

For log-softmax, use $\log\text{softmax}(z)_i = z_i - \operatorname{logsumexp}(z)$ where $\operatorname{logsumexp}(z) = z_{\max} + \log\sum_j e^{z_j - z_{\max}}$ . This avoids both overflow (largest term is 1) and $\log(0)$ underflow (the sum is $\geq 1$ ).

Online softmax trick: fuse the max-finding and denominator into a single pass by keeping a running max $m_k$ and running denominator $d_k$ , rescaling when the max changes:

d_{k+1} = \underbrace{d_k \cdot e^{m_k - m_{k+1}}}_{\text{rescale old}} + \underbrace{e^{x_{k+1} - m_{k+1}}}_{\text{new term}}

When the max is unchanged, the rescale factor is 1. The same recurrence applied to a running weighted sum $o_k = \sum_i e^{x_i - m_k} v_i$ is the core of FlashAttention, where it lets attention be computed in tiles without ever materializing the full score matrix.

Where it appears:

Classification output layer — paired with cross-entropy loss
Self-Attention — softmax over scaled dot-product scores turns them into attention weights
Reinforcement learning — softmax policy converts action values to action probabilities
Contrastive learning — InfoNCE / softmax over similarity scores
Sampling and Decoding Strategies — temperature-scaled softmax controls randomness of generation

Softmax vs sigmoid: sigmoid is the 2-class special case. For multi-label problems (multiple independent binary decisions), use sigmoid per class, not softmax.

Softmax

Linked from