Softmax was one of those functions I used correctly for a long time before I actually understood why it worked. I knew it turned a list of numbers into “probabilities that sum to 1,” and that was enough to get my classifiers working. But it wasn’t until I implemented it manually — and ran straight into a numerical overflow bug — that I really understood what’s happening under the hood. Let me take you through that same journey.
What Problem Does Softmax Solve?
When I build a multi-class classifier, the final layer of my network produces raw, unbounded scores called logits — one per class. These logits could be any real number: negative, positive, huge, or tiny. But I need a probability distribution over classes: numbers between 0 and 1 that sum to exactly 1, so I can interpret them as “the model’s confidence that this example belongs to each class.” Softmax is the function that converts raw logits into exactly this kind of probability distribution.
The Mathematical Definition
Given a vector of logits $z = [z_1, z_2, \ldots, z_C]$ for $C$ classes, the softmax function computes:
$$\text{softmax}(z)i = \frac{e^{z_i}}{\sum{j=1}^{C} e^{z_j}}$$
Each output $\text{softmax}(z)_i$ is guaranteed to be in the range $(0, 1)$, and the full vector always sums to 1:
$$\sum_{i=1}^{C} \text{softmax}(z)_i = 1$$
Why Exponentials?
I used to wonder why softmax uses $e^{z_i}$ instead of some simpler normalization like $z_i / \sum z_j$. There are a few good reasons:
- Exponentials are always positive, so I never end up with negative “probabilities” even if some logits are negative.
- Exponentials amplify differences, so a logit that’s slightly larger than another produces a disproportionately larger probability — this makes the output distribution more decisive and interpretable.
- Smooth and differentiable, which is essential for backpropagation.
A Worked Numerical Example
Let’s say my model outputs logits $z = [2.0, 1.0, 0.1]$ for three classes.
$$e^{2.0} \approx 7.389, \quad e^{1.0} \approx 2.718, \quad e^{0.1} \approx 1.105$$
$$\text{sum} = 7.389 + 2.718 + 1.105 = 11.212$$
$$\text{softmax}(z) = \left[\frac{7.389}{11.212}, \frac{2.718}{11.212}, \frac{1.105}{11.212}\right] = [0.659, 0.242, 0.099]$$
Notice how the largest logit (2.0) ends up with the largest share of probability (65.9%), while the differences between the logits get amplified in the resulting probabilities.
The Numerical Stability Trick
Here’s the bug I mentioned earlier: if my logits are large (say, 1000), computing $e^{1000}$ directly overflows to infinity in floating point arithmetic. The standard fix is to subtract the maximum logit from every value before exponentiating:
$$\text{softmax}(z)i = \frac{e^{z_i – \max(z)}}{\sum{j=1}^{C} e^{z_j – \max(z)}}$$
This is mathematically identical to the original formula (the $\max(z)$ terms cancel out in the ratio), but it keeps all the exponentiated values in a numerically safe range, since the largest exponent becomes $e^0 = 1$.
The Gradient of Softmax
The derivative of softmax is a bit more involved than most activation functions because each output depends on every input, not just the corresponding one. The Jacobian matrix entry is:
$$\frac{\partial \text{softmax}(z)_i}{\partial z_j} = \text{softmax}(z)i \left(\delta{ij} – \text{softmax}(z)_j\right)$$
where $\delta_{ij}$ is 1 if $i=j$ and 0 otherwise. In practice, when softmax is paired with categorical cross-entropy loss, this messy Jacobian combines with the loss gradient to produce the beautifully simple result:
$$\frac{\partial \mathcal{L}}{\partial z_i} = \hat{y}_i – y_i$$
This is exactly why softmax and categorical cross-entropy are almost always used together — the combined gradient is clean and efficient to compute.
Visualizing Where Softmax Fits
flowchart TD
A[Input Features] --> B[Neural Network Layers]
B --> C["Raw Logits: z1, z2, ..., zC"]
C --> D["Softmax: exp(zi) / sum(exp(zj))"]
D --> E[Probability Distribution Summing to 1]
E --> F[Categorical Cross-Entropy Loss]
F --> G["Gradient: y-hat minus y"]
G --> H[Backpropagation]
Code Examples
PyTorch
import torch
import torch.nn as nn
import torch.nn.functional as F
logits = torch.tensor([2.0, 1.0, 0.1])
probs = F.softmax(logits, dim=0)
print(f"Softmax probabilities: {probs}")
# In practice, use log_softmax + NLLLoss, or CrossEntropyLoss which
# combines both internally for numerical stability
model = nn.Sequential(
nn.Linear(20, 64),
nn.ReLU(),
nn.Linear(64, 3) # raw logits — do NOT apply softmax here if using CrossEntropyLoss
)
TensorFlow / Keras
import tensorflow as tf
logits = tf.constant([2.0, 1.0, 0.1])
probs = tf.nn.softmax(logits)
print(f"Softmax probabilities: {probs.numpy()}")
model = tf.keras.Sequential([
tf.keras.layers.Dense(64, activation='relu', input_shape=(20,)),
tf.keras.layers.Dense(3, activation='softmax')
])
model.compile(optimizer='adam', loss='categorical_crossentropy')
Softmax Temperature
An important variant I use often, especially in knowledge distillation and language model sampling, introduces a temperature parameter $T$:
$$\text{softmax}(z, T)i = \frac{e^{z_i / T}}{\sum{j=1}^{C} e^{z_j / T}}$$
- Low temperature ($T < 1$) makes the distribution sharper/more confident — the model behaves more deterministically.
- High temperature ($T > 1$) makes the distribution softer/more uniform — useful for exposing more information about relative class similarities, as in knowledge distillation.
- $T = 1$ recovers standard softmax.
Softmax vs Sigmoid
This distinction trips up a lot of people building multi-label classifiers, so it’s worth being explicit:
| Aspect | Softmax | Sigmoid |
|---|---|---|
| Output sums to 1? | Yes | No |
| Classes mutually exclusive? | Yes (single-label) | No (independent per class) |
| Typical task | Multi-class classification | Binary or multi-label classification |
| Loss pairing | Categorical cross-entropy | Binary cross-entropy |
If I’m building a classifier where an image can have multiple correct labels simultaneously (e.g., “cat” AND “indoor”), I should use independent sigmoid outputs per label, not softmax — softmax would force the probabilities to compete with each other, which doesn’t make sense when multiple labels can be true at once.
Advantages of Softmax
- Produces a valid probability distribution, making outputs directly interpretable as class confidences.
- Differentiable everywhere, enabling smooth gradient-based optimization.
- Naturally pairs with cross-entropy loss to produce clean, efficient gradients.
- Temperature scaling flexibility allows adapting output sharpness for different applications (sampling, distillation).
Disadvantages and Limitations
- Only appropriate for mutually exclusive classes — misusing it for multi-label problems produces incorrect probability semantics.
- Sensitive to logit scale — very large or very small logits can produce extremely peaked or extremely flat distributions unless properly normalized or regularized.
- Computationally more expensive than simpler activations due to exponentials and normalization across all classes.
- Can be overconfident — softmax outputs are known to sometimes poorly reflect true model uncertainty, which is why techniques like temperature scaling and label smoothing are used for calibration.
Real-World Use Cases
- Image classification — final layer of CNNs classifying images into one of many categories.
- Language modeling — predicting the next word from a vocabulary, where softmax converts logits over tens of thousands of tokens into a probability distribution.
- Attention mechanisms — the attention weights in Transformer architectures are computed using softmax over similarity scores.
- Reinforcement learning — softmax policies convert action-value scores into a probability distribution over actions for stochastic policies.
- Knowledge distillation — temperature-scaled softmax outputs are used to transfer knowledge from teacher to student models.
Best Practices
- Never manually apply softmax before a cross-entropy loss function that expects raw logits — most frameworks apply softmax internally for numerical stability (
CrossEntropyLossin PyTorch,from_logits=Truein TensorFlow). - Use the max-subtraction trick if implementing softmax manually, to avoid numerical overflow.
- Use temperature scaling when you need softer or sharper distributions, such as in distillation or exploration-heavy RL policies.
- Avoid softmax for multi-label problems — use independent sigmoids per class instead.
- Consider label smoothing alongside softmax outputs to avoid overconfident predictions and improve generalization.
Summary
Softmax converts a vector of raw, unbounded logits into a valid probability distribution over mutually exclusive classes, making it the standard choice for the output layer of multi-class classifiers, attention mechanisms in Transformers, and stochastic policies in reinforcement learning. Its exponential formulation amplifies differences between logits, its gradient combines beautifully with cross-entropy loss, and its temperature parameter offers useful flexibility for tasks like distillation and sampling. Understanding its numerical stability requirements and its proper use case (mutually exclusive classes only) will save you from some very common and frustrating bugs.
References
- Bishop, C.M. (2006). “Pattern Recognition and Machine Learning.” Springer.
- Goodfellow, I., Bengio, Y., & Courville, A. (2016). “Deep Learning.” MIT Press.
- Vaswani, A., et al. (2017). “Attention Is All You Need.” arXiv:1706.03762.
- PyTorch Documentation: https://pytorch.org/docs/stable/generated/torch.nn.Softmax.html
- TensorFlow Documentation: https://www.tensorflow.org/api_docs/python/tf/nn/softmax
