Offensive Security Researcher & Quantum Cryptography Analyst

Explaining the Concept of the Softmax Activation Function

Abdul Wahab Junaid
Explain the concept of the softmax activation function

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:

  1. Exponentials are always positive, so I never end up with negative “probabilities” even if some logits are negative.
  2. 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.
  3. 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}}$$

Softmax vs Sigmoid

This distinction trips up a lot of people building multi-label classifiers, so it’s worth being explicit:

AspectSoftmaxSigmoid
Output sums to 1?YesNo
Classes mutually exclusive?Yes (single-label)No (independent per class)
Typical taskMulti-class classificationBinary or multi-label classification
Loss pairingCategorical cross-entropyBinary 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

Disadvantages and Limitations

Real-World Use Cases

  1. Image classification — final layer of CNNs classifying images into one of many categories.
  2. Language modeling — predicting the next word from a vocabulary, where softmax converts logits over tens of thousands of tokens into a probability distribution.
  3. Attention mechanisms — the attention weights in Transformer architectures are computed using softmax over similarity scores.
  4. Reinforcement learning — softmax policies convert action-value scores into a probability distribution over actions for stochastic policies.
  5. Knowledge distillation — temperature-scaled softmax outputs are used to transfer knowledge from teacher to student models.

Best Practices

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

Exit mobile version