Softmax Function

The softmax function converts a vector of raw scores into a valid probability distribution — every value between 0 and 1, summing to exactly 1. It's the standard output-layer activation for any multi-class classification network.

Formula

For a vector of \(K\) raw scores (logits) \(\mathbf{z}\), each output is the exponential of that score, divided by the sum of exponentials across all \(K\) scores — this normalization is exactly what guarantees the outputs sum to 1, satisfying the probability axioms from Probability Basics.

Numerical Example

Logits \(\mathbf{z} = [2.0, 1.0, 0.1]\):

Notice the outputs sum to 1.0 (allowing for rounding), and the largest logit produces by far the largest probability — softmax preserves relative ordering while converting to a valid distribution.

Why Exponentiate at All?

Exponentiation guarantees every output is positive (satisfying \(P(x)\ge0\)) regardless of whether the input logits are negative, and it exaggerates differences between scores — a logit just slightly larger than another produces a noticeably larger probability, which tends to produce more confident, decisive predictions than a simpler normalization (like dividing by the sum directly) would.

Numerical Stability — The Practical Detail

Computing \(e^{z_i}\) directly can overflow for large logits. The standard fix, used internally by every deep learning framework, is subtracting the maximum logit before exponentiating:

This produces mathematically identical results (the max-subtraction cancels out in the ratio) but avoids computing extremely large exponentials.

Code

import numpy as np

def softmax(z):
    z_stable = z - np.max(z)   # numerical stability trick
    exp_z = np.exp(z_stable)
    return exp_z / np.sum(exp_z)

print(softmax(np.array([2.0, 1.0, 0.1])))   # [0.659 0.242 0.099]

import torch
import torch.nn.functional as F

logits = torch.tensor([2.0, 1.0, 0.1])
print(F.softmax(logits, dim=0))   # numerically stable, matches the manual version

Where It's Used Today

The standard output-layer activation for any multi-class classification network with mutually exclusive classes. It's also the core operation inside self-attention (covered in the Attention category), where it converts raw attention scores into a distribution over which tokens to focus on.

Common Mistakes

• Applying softmax manually before passing logits to nn.CrossEntropyLoss, which already applies log_softmax internally — as flagged repeatedly across this hub, this double-application silently corrupts gradients.
• Using softmax for multi-label classification (where an example can belong to more than one class simultaneously) — softmax's outputs are forced to sum to 1 across classes, which is wrong when classes aren't mutually exclusive; independent sigmoid outputs per class are correct there instead.

Interview Relevance

Q: "Why do deep learning frameworks subtract the maximum logit before computing softmax?" Purely for numerical stability — computing \(e^{z}\) for a large logit can overflow floating-point precision. Subtracting the maximum logit from every value before exponentiating produces mathematically identical probabilities (the shift cancels out in the ratio) while keeping the exponentials in a safe numerical range.

Practice Question

For logits \([1.0, 1.0, 1.0]\) (all equal), what will the softmax output be, without computing it directly? What does this tell you about softmax's behavior when the model is maximally uncertain?

Related DL Notes

• Sigmoid vs Softmax
• Cross-Entropy
• Scaled Dot-Product Attention