Softmax

Softmax converts logits into a probability distribution over mutually exclusive choices.

$$p_i=\mathrm{softmax}(z{)}_i=\frac{\exp (z_i)}{{\sum }_{j=1}^K\exp (z_j)}$$

Here $z_i$ is a logit and $p_i$ is the normalized probability for class or option $i$.

Softmax is invariant to adding the same constant to all logits:

$$\mathrm{softmax}(z{)}_i=\mathrm{softmax}(z-c1{)}_i$$

This is used for numerical stability:

$$p_i=\frac{\exp (z_i-{\max }_j{z}_j)}{{\sum }_{j=1}^K\exp (z_j-{\max }_j{z}_j)}$$

Log-Softmax

Many losses use log probabilities directly:

$$\log p_i=z_i-\log \sum _{j=1}^K\exp (z_j)$$

The stable form uses log-sum-exp:

$$\log p_i=z_i-c-\log \sum _j\exp (z_j-c),c=\underset{j}{\max }{z}_j$$

This avoids overflow and is why implementations often combine log-softmax with negative log-likelihood or cross-entropy.

Temperature

A temperature rescales logits:

$$p_i(\tau )=\frac{\exp (z_i/\tau )}{{\sum }_{j=1}^K\exp (z_j/\tau )}$$

Small $\tau$ makes the distribution sharper. Large $\tau$ makes it flatter. Temperature is used in decoding, calibration, distillation, and contrastive losses.

Cross-Entropy Link

With one-hot target $y$, softmax cross-entropy is:

$$\mathcal{L}=-\sum _{i=1}^K{y}_i\log p_i$$

For logits $z$ and probabilities $p=\mathrm{softmax}(z)$:

$$\frac{\partial \mathcal{L}}{\partial z_i}=p_i-y_i$$

This simple gradient is one reason softmax and cross-entropy are usually implemented together.

Masked Softmax

For attention or constrained choices, invalid positions are masked before softmax:

$$p_i=\frac{\exp (z_i+m_i)}{{\sum }_j\exp (z_j+m_j)}$$

where $m_i=0$ for valid positions and $m_i=-\infty$ for invalid positions.

Attention Interpretation

In attention, softmax normalizes scores over keys for each query:

$${\alpha }_{ij}=\frac{\exp (s_{ij})}{{\sum }_k\exp (s_{ik})}$$ $$o_i=\sum _j{\alpha }_{ij}{v}_j$$

The axis matters: normalizing over keys answers “where should this query attend?” Normalizing over queries would mean a different operation.

Where It Appears

• Attention weights in Attention.
• Classification heads.
• Autoregressive next-token distributions.
• Routing probabilities in Mixture of experts.

Checks

• Softmax is sensitive to logit scale.
• Masks usually add $-\infty$ to impossible positions before softmax.
• For multi-label prediction, sigmoid is usually more appropriate than softmax.
• Probabilities can still be poorly calibrated; see Calibration.
• Check the axis: class dimension, token dimension, expert dimension, or candidate list.
• Check whether logits are raw scores, normalized similarities, or temperature-scaled scores.
• Use log-softmax or fused cross-entropy for training losses when possible.
• Confirm that masked positions are impossible after normalization, not merely small.

Related

• Probability distribution
• Entropy and KL divergence
• Attention
• Cross-entropy loss
• Probabilistic prediction
• Autoregressive model
• Calibration