Softmax and Multiclass Classification Explained

A classifier that chooses one of several mutually exclusive classes needs a probability distribution over those classes. Softmax converts the output scores into that distribution, and categorical cross-entropy scores the probability assigned to the observed class. We will derive their combined gradient and implement the loss directly from logits to preserve accuracy when probabilities become extremely small. Run the body code blocks in order.

From one score to a vector of scores

For \(C\) classes the output layer produces a vector of \(C\) raw scores called logits: \(z=W^{[L]}a^{[L-1]}+b^{[L]}\) with \(z\in\mathbb{R}^{C}\). A logit is an unrestricted score. Under softmax, differences have a probability interpretation: \(\log(\hat y_i/\hat y_j)=z_i-z_j\). A common offset of all logits has no effect on the probabilities. The question is how to convert \(C\) real numbers into \(C\) probabilities that are all positive and sum to one.

Exponentiating makes every value positive, and dividing by the total makes them sum to one. For finite logits, the definition is:

\[\hat y_i=\text{softmax}(z)_i=\frac{e^{z_i}}{\sum_{j=1}^{C}e^{z_j}}\]

Softmax normalizes across classes. Holding the other logits fixed, raising one logit decreases the other probabilities. This suits a categorical target with one observed class. For multilabel tasks, several labels may be present together, so a common design uses a separate sigmoid probability and binary loss per label. Those probabilities need not sum to one; other models can also represent dependence among labels.

The single logistic unit these ideas build on is developed end to end in Logistic Regression: A Complete Guide to Binary Classification.

Softmax generalizes the sigmoid

With \(C=2\), the first softmax probability equals a sigmoid applied to the difference between the two logits. With affine logits, this gives the same probability family as binary logistic regression. Dividing numerator and denominator by \(e^{z_1}\) gives \(\hat y_1=1/(1+e^{-(z_1-z_2)})\), which is \(\sigma(z_1-z_2)\). Only the difference between logits matters, not their absolute size. The same property holds for any \(C\): adding a constant to every logit leaves the output unchanged, because the constant factors out of numerator and denominator.

We use this shift invariance to avoid exponentiating unnecessarily large scores.

Numerical stability: subtract the maximum

In float64, directly exponentiating 1000 overflows. Subtract each example’s maximum logit before exponentiating: the shifted values are nonpositive, and at least one is zero, so the largest exponential is 1. This prevents exponential overflow for the finite, representable shifts considered here. Very small exponential terms can still underflow to zero, and nonfinite inputs require separate handling. The code uses shape \((C,m)\): classes are rows, examples are columns, and axis=0 normalizes each example.

import numpy as np

def softmax(Z):
    """Z: (C, m) logits. Returns (C, m) probabilities, columns sum to 1."""
    assert Z.ndim == 2 and np.isfinite(Z).all()
    Z_shift = Z - np.max(Z, axis=0, keepdims=True)
    E = np.exp(Z_shift)
    return E / np.sum(E, axis=0, keepdims=True)

Z = np.array([[1.0, 1000.0],
              [2.0, 1000.0],
              [3.0,  999.0]])
A = softmax(Z)
print(np.round(A, 4))
# [[0.09   0.4223]
#  [0.2447 0.4223]
#  [0.6652 0.1554]]
print(A.sum(axis=0))   # [1. 1.]

The first column shifts to \((-2,-1,0)\), with exponentials approximately \((0.1353,0.3679,1)\). Dividing by their sum gives \((0.0900,0.2447,0.6652)\). The second shifts to \((0,0,-1)\), so its large original scores cause no exponential overflow. The maximum-probability class can be selected with A.argmax(axis=0), which gives indices [2, 0] here (NumPy uses zero-based indices and selects the first maximum in a tie). Softmax preserves logit ordering, so Z.argmax(axis=0) gives the same decisions. A normalized output is a model probability estimate; normalization alone does not establish calibration.

Categorical cross-entropy

The label for example \(i\) is a one-hot vector \(y\) with a single \(1\) at the true class \(k\). The loss for that example is

\[\mathcal{L}(\hat y,y)=-\sum_{c=1}^{C}y_c\log\hat y_c=-\log\hat y_k\]

Only the predicted probability of the correct class appears. Every other term is multiplied by zero. The loss is \(0\) at probability \(1\) for the right class (a limit for finite softmax logits with more than one class) and grows without bound as that probability approaches zero. Increasing the assigned probability of the observed class reduces its loss while leaving less total probability for the remaining classes.

The cost over \(m\) examples is the mean, \(J=-\frac{1}{m}\sum_{i=1}^{m}\log\hat y^{(i)}_{k^{(i)}}\).

Deriving the gradient from log probabilities

Write the per-example loss as \(\mathcal{L}=\log\sum_j e^{z_j}-\sum_i y_i z_i\), using \(\sum_i y_i=1\). Differentiating the first term with respect to \(z_k\) gives \(e^{z_k}/\sum_j e^{z_j}=\hat y_k\); differentiating the second gives \(-y_k\). Thus:

\[\frac{\partial\mathcal{L}}{\partial z}=\hat y-y\]

For a batch mean, divide by \(m\), giving (A-Y)/m. This formula also holds for nonnegative soft targets whose class probabilities sum to one per example, under this unweighted cross-entropy. Combining log-softmax with the loss avoids taking the logarithm of a probability that has rounded to zero. PyTorch’s CrossEntropyLoss accepts logits and performs this stable calculation; an explicit log-softmax followed by negative log-likelihood is also valid.

A separate softmax backward pass need not construct its full Jacobian either. If \(g_i=\partial J/\partial\hat y_i\) is the incoming gradient, then \(\partial J/\partial z_i=\hat y_i(g_i-\sum_j g_j\hat y_j)\). One weighted sum and elementwise operations compute this product. The combined loss formula is especially convenient because it also avoids division by tiny probabilities.

For the loss code, let \(S=Z-\max_c Z_c\) separately in each column. Then log_probs = S - log(sum(exp(S))) evaluates log probabilities directly. Log-softmax remains useful even when exponentiating a log probability would round to zero.

def log_softmax(Z):
    assert Z.ndim == 2 and np.isfinite(Z).all()
    shifted = Z - Z.max(axis=0, keepdims=True)
    return shifted - np.log(np.exp(shifted).sum(axis=0, keepdims=True))

def cross_entropy_cost(Z, Y):
    assert Z.shape == Y.shape
    assert np.all(Y >= 0) and np.allclose(Y.sum(axis=0), 1)
    return -np.sum(Y * log_softmax(Z)) / Y.shape[1]

def softmax_backward(A, Y):
    assert A.shape == Y.shape
    return (A - Y) / Y.shape[1]

Y = np.array([[0.0, 1.0], [0.0, 0.0], [1.0, 0.0]])
print(round(cross_entropy_cost(Z, Y), 4))
print(np.round(softmax_backward(A, Y), 4))
# 0.6348
# [[ 0.045  -0.2888]
#  [ 0.1224  0.2112]
#  [-0.1674  0.0777]]
Z_extreme = np.array([[1000.], [0.]])
Y_extreme = np.array([[0.], [1.]])
print("probabilities", softmax(Z_extreme).ravel())
print("loss", cross_entropy_cost(Z_extreme, Y_extreme))
# probabilities [1. 0.]
# loss 1000.0

For these one-hot labels, the true-class logit gradient is negative and the other class gradients are positive. Its magnitude is the probability residual divided by the two-example batch size. If the logits themselves were updated by gradient descent, this would raise the true-class logit and lower the others. Shared network parameters couple examples, so a parameter update need not move every example’s logits that way.

Practical details that cause bugs

• Apply softmax across the class axis, not the batch axis. With shape \((C,m)\) that is axis=0; with the more common framework layout \((m,C)\) it is axis=1 or axis=-1. The wrong axis can normalize across different examples while returning the expected shape. Check sums along the intended class axis, not just the array shape.
• Pass logits to a loss that expects logits. Applying softmax first introduces a second normalization, limiting the probability assigned to a class and imposing a positive loss floor. It preserves the argmax ordering for fixed logits, so it does not itself impose an accuracy ceiling.
• For hard labels, integer class indices avoid allocating a dense one-hot target matrix. With the column layout here, -log_probs[class_indices, np.arange(m)].mean() selects the same losses. Soft labels require a target distribution instead.
• Use a stable logits-based loss in a NumPy implementation too. Adding epsilon to probabilities changes the objective and no longer has exactly the gradient A-Y. Check that input logits and the resulting loss are finite.
• For fixed finite logits and temperature \(T>0\), \(\operatorname{softmax}(z/T)\) tends to uniform probabilities as \(T\to\infty\). As \(T\to0^+\), it concentrates on a unique largest logit, or equally among tied maxima. Positive temperature preserves class ordering. This is the same knob used for sampling from language models and for distillation.

How to choose an objective for a given problem is the subject of Loss Function Design: Choosing an Objective That Matches the Problem.

Sanity checks before training

Uniform predictions give cross-entropy \(\log C\) for any hard-label class balance: about 2.30 for ten classes and 6.91 for a thousand. Random initialization need not produce nearly equal logits. A higher initial loss can reflect large logit differences; a lower one can reflect class imbalance matched by a bias, chance, or pretrained parameters. Neither direction alone establishes a bug or label leakage. For example, predictions \((0.9,0.1)\) have loss about 0.325 on labels that are 90% class 0, below \(\log2\).

Try fitting a small batch as a debugging experiment. Failure to reach a small loss can reveal an implementation problem, but also insufficient capacity, optimization difficulties, strong regularization, or identical inputs with conflicting labels. Inspect the data and compare numerical gradients before assigning a cause.

Self-attention, masking, and the full block are derived in Transformers from Scratch: Self-Attention to Encoder–Decoder.

Exercises

1. Shift invariance. Compute the softmax of \([1,2,3]\) and of \([101,102,103]\). Then compute it for \([1,2,3]\) scaled by 10. Which pair matches exactly, and what does the third tell you about temperature?

You should get: two identical distributions and one much sharper than both.

import numpy as np
def softmax(z):
    e = np.exp(z - z.max()); return e / e.sum()
print(np.round(softmax(np.array([1.,2.,3.])), 4))       # [0.0900 0.2447 0.6652]
print(np.round(softmax(np.array([101.,102.,103.])), 4)) # [0.0900 0.2447 0.6652]
print(np.round(softmax(np.array([10.,20.,30.])), 4))    # [0. 0. 1.]

2. Verify the fused gradient. For \(C=3\) classes with logits \(z\) and a one-hot label, compute \(\partial\mathcal{L}/\partial z\) numerically and compare it against \(\hat y-y\). Then confirm that the components sum to zero and explain why.

You should get: agreement to about \(10^{-8}\), and a gradient vector whose entries sum to zero.

import numpy as np
def softmax(z):
    e = np.exp(z - z.max()); return e / e.sum()
def loss(z, y):
    shifted = z - z.max()
    return np.log(np.exp(shifted).sum()) - shifted[y]
z, y = np.array([0.5, -1.2, 2.0]), 0
h = 1e-6
num = np.array([(loss(z + h*np.eye(3)[i], y) - loss(z - h*np.eye(3)[i], y))/(2*h)
                for i in range(3)])
ana = softmax(z) - np.eye(3)[y]
print(np.round(num, 6)); print(np.round(ana, 6)); print(round(ana.sum(), 12))

# [-0.823458  0.032251  0.791207]
# [-0.823458  0.032251  0.791207]
# -0.0

3. The double-softmax bug. Compute the loss for a correct prediction two ways: once passing logits to a fused softmax cross-entropy, once applying softmax first and passing the probabilities in as if they were logits. Report both losses and explain the gap.

You should get: a small loss from the correct path and a much larger, nearly irreducible one from the bug.

import numpy as np
def softmax(z):
    e = np.exp(z - z.max()); return e / e.sum()
z, y = np.array([8.0, 0.0, 0.0]), 0        # very confident, correct
print(round(-np.log(softmax(z)[y]), 6))
# 0.000671
print(round(-np.log(softmax(softmax(z))[y]), 6))
# 0.551871