Softmax & Multiclass Probabilities
Understand turning multiple logits into a probability distribution over classes.
1. Plain English Intuition
When your model has to choose between K mutually exclusive classes (like identifying dog, cat, or bird), the final linear layer outputs K unbounded real numbers called logits. The Softmax function exponentiates each logit (making them all positive) and divides each by their total sum: Softmax(z_i) = exp(z_i) / ∑ exp(z_j). The result is a valid probability distribution where every value is between 0 and 1, and the entire vector sums to exactly 1.0.
2. Where It Appears in ML & Deep Learning
- 1Final output layer of all multi-class classifiers (Vision Transformers, ResNet)
- 2Language Model next-token generation heads over vocabularies (e.g. 32,000 to 128,000 tokens)
- 3Multi-Head Self-Attention mechanisms: Attention(Q, K, V) = softmax(Q K^T / √d) V
- 4Policy gradient reinforcement learning (sampling actions from a categorical policy)
3. Inputs, Outputs, Shapes & Assumptions
4. What Can Go Wrong (Gotchas & Failures)
Numerical Overflow: If a logit is large (e.g. 100), exp(100) overflows float32 to inf, yielding NaN.
Fix: Always subtract the maximum logit before exponentiating: exp(z_i - max(z)).
Overconfidence calibration: Softmax outputs can be overconfident even on random noise.
5. Python in Practice
Engineered using: torch.softmax, scipy.special.softmax, torch.log_softmax
import numpy as np
# Numerically stable Softmax implementation
def stable_softmax(logits):
# Subtract max logit along the last axis to prevent exp overflow
shifted = logits - np.max(logits, axis=-1, keepdims=True)
exp_scores = np.exp(shifted)
return exp_scores / np.sum(exp_scores, axis=-1, keepdims=True)
# Raw output logits from a 3-class classifier: [Cat, Dog, Bird]
logits = np.array([2.5, 1.0, 0.1])
probs = stable_softmax(logits)
print(f"Probabilities: {probs}")
print(f"Sum of probabilities: {np.sum(probs):.4f}")
print(f"Predicted class index: {np.argmax(probs)}")Fast Study Loop
Complete these 3 steps to mark this topic understood:
Curated Video Lesson
Neural Networks Part 7: Cross Entropy and Softmax
Watch on YouTube