Softmax Cross-Entropy and Its Collapsed Gradient

A classifier ends with a vector of scores, one per class, and produces a probability distribution over those classes. Written as two separate operations, the last two stages of such a network are an exponential normalizing function followed by a loss that takes the negative logarithm of the probability assigned to the correct class. Implemented that way they work, and they cost more arithmetic and lose more precision than necessary, because differentiating the composition cancels almost everything. The cancellation is not an optimization that happens to be available; it is the reason the pair is standard.

The normalizing function and its Jacobian

Given unnormalized scores $z \in \mathbb{R}^{K}$, the softmax output is

$$p_i = \frac{\exp(z_i)}{\sum_{j=1}^{K} \exp(z_j)}$$

Differentiating with respect to the scores gives a $K \times K$ matrix,

$$\frac{\partial p_i}{\partial z_j} = p_i,(\delta_{ij} - p_j)$$

which is dense: every output probability depends on every score, because the denominator depends on all of them. The matrix is singular, with rank at most $K-1$, and the all-ones vector is in its null space, which is the algebraic statement of the fact that adding a constant to every score leaves the output unchanged. That invariance is what makes the numerically stable implementation possible.

The collapse: the gradient is the prediction minus the target

The loss for a single example with correct class $y$ is $L = -\log p_y$, equivalently $L = -\sum_i \tilde{y}_i \log p_i$ with $\tilde{y}$ a one-hot vector. Composing the two derivatives, the derivative of the loss with respect to the probability has entries $-\tilde{y}_i / p_i$, and multiplying by the Jacobian above collapses the sum:

$$\frac{\partial L}{\partial z_k} = \sum_i \left(-\frac{\tilde{y}i}{p_i}\right) p_i (\delta{ik} - p_k) = p_k - \tilde{y}_k$$

The division by $p_i$ that made the intermediate derivative unbounded is cancelled by the factor $p_i$ in the Jacobian, and what remains is bounded in $[-1, 1]$ regardless of how confident or how wrong the prediction is. For a uniform distribution over 1000 classes, $p_i = 0.001$, so the derivative with respect to the correct logit is $0.001 - 1 = -0.999$ while the other 999 entries are $+0.001$: the step is dominated by lifting the correct logit relative to everything else. The same collapse occurs for a single logistic output unit paired with binary cross-entropy, where the derivative with respect to the logit is $a - y$ with $a$ the sigmoid output. In both cases the loss of a squashing function’s derivative from the gradient path means there is no saturation factor to suppress the signal when the prediction is confidently wrong, the failure mode described in saturating activations.

The arithmetic consequence follows directly. Evaluating softmax and then the loss separately and differentiating through both builds the $K \times K$ Jacobian and performs $K$ multiply-adds per output element, so $O(K^2)$ work per example. The collapsed form needs one pass to compute $p$, one subtraction, and one division by the number of examples, so $O(K)$ work per example, and a $K$-dimensional gradient vector instead of a $K \times K$ matrix retired for each example.

Two reduction conventions set the scale of what comes out of that expression. The derivative above is for a single example, and the loss is normally averaged over the examples in a batch, so the batch gradient is the sum of the per-example gradients divided by the number of examples in it. Replacing that mean with a sum multiplies every gradient entry by the batch size, a factor of a thousand at batch size 1000, which is the shape of an instability that appears only when the batch size changes. The scale of the loss itself is a useful sanity check on a fresh network, since a uniform prediction over $K$ classes has loss $\log K$: about 0.693 for two classes, 2.303 for ten, and 6.908 for a thousand, in natural units. A loss below that value on a balanced problem means the model is doing better than chance, and a loss far below it on a large batch means the predictions have become confident, which is where the logit magnitudes that threaten the exponential come from.

Why the naive implementation overflows and what to subtract

The exponential is the numerically fragile part. In 32-bit floating point the largest finite value is about $3.4 \times 10^{38}$, and $\exp(z)$ exceeds it for $z$ beyond roughly 88.7; in 64-bit arithmetic the corresponding threshold is roughly 709.8. Scores of that magnitude are not exotic: a dense layer at the end of a network can produce them after enough steps of gradient descent if nothing constrains the scale of its output. Once a single exponential overflows to infinity, the ratio becomes a division of infinities and the whole gradient vector is not-a-number.

The shift invariance of the normalizing function removes the problem entirely, since $\exp(z_i - c) / \sum_j \exp(z_j - c)$ is the same number for any $c$. Subtracting the largest score makes the largest exponent $\exp(0) = 1$ and every other exponent at most 1, so the numerator cannot overflow and the denominator is at least 1 and at most $K$, which cannot overflow either. Terms whose shifted score is very negative underflow to zero, which is the correct rounding for probabilities below the representable range, and the denominator stays strictly positive because at least one term is exactly 1.

For the loss itself, the same idea applied in logarithmic form is more accurate than exponentiating and taking a logarithm. The quantity needed is

$$\log p_y = z_y - \log \sum_j \exp(z_j)$$

and the sum of exponentials is evaluated by factoring out the maximum, which is the log-sum-exp pattern. Forming $p_y$ first and taking its logarithm loses the relative precision of the loss when $p_y$ is small, and loses the value entirely if $p_y$ rounds to zero, at which point the logarithm is negative infinity instead of a large finite number. The error analysis of these two expressions is the subject of the numerical-analysis literature on log-sum-exp and softmax: the shifted form has a small relative error, while the unshifted form does not.

import numpy as np

def loss_and_grad(logits, labels):
    shift = logits - np.max(logits, axis=-1, keepdims=True)
    lse = np.log(np.sum(np.exp(shift), axis=-1, keepdims=True))
    logp = shift - lse                              # log-softmax, no rounding to zero
    loss = -logp[np.arange(len(labels)), labels].mean()
    p = np.exp(logp)
    p[np.arange(len(labels)), labels] -= 1.0        # p - onehot: the whole gradient
    return loss, p / len(labels)

Target manipulation: smoothing, weighting, and masking

The gradient $p - \tilde{y}$ is zero only when the prediction matches the target exactly, and with a one-hot target that is a fixed point the network can approach arbitrarily closely. That is often undesirable, because it requires logits to grow without bound. Replacing the one-hot target with $(1-\epsilon)\tilde{y} + \epsilon/K$ adds a floor of $\epsilon/K$ to every target coordinate, so the gradient can never vanish and the logits stay finite; a common choice for the smoothing weight is 0.1. Per-class weights multiply the per-example loss before reduction and therefore rescale that example’s entire gradient vector, which changes the relative size of the steps taken on different classes rather than merely the loss value.

Masking follows from the same algebra. Setting a score to negative infinity forces the corresponding probability to exactly zero and its gradient entry to exactly zero, because $\exp(-\infty) = 0$ and the collapsed gradient for that class is zero regardless of the target. This is how positions that must be excluded from a comparison are removed without adding a branch to the distance computation.

Degenerate cases

Several inputs are well-formed and useless. With $K = 1$ the single output probability is identically 1 and the loss is identically 0, so the gradient is identically 0 and no amount of training changes the parameters. Equal scores produce a uniform distribution, and with $K = 1000$ classes that is a loss of $\log 1000 \approx 6.908$ in natural units, the expected starting point for a balanced problem and a useful sanity check on a fresh network. A single non-finite score poisons the entire row, since the maximum propagates the non-finite value and every output depends on every input through the denominator. Finally, applying the normalizing function before passing values to a loss that performs the normalization again double-applies it, squashing the distribution toward uniform and shrinking the gradient by a factor that depends on the scores; the correct interface for such a loss takes raw logits, which is why the fused form accepts them rather than probabilities.

The choice of whether the nonlinearity sits inside the loss or outside it also determines what the layer below stores. When the last operation is a fused loss on logits, the backward pass never needs the probability tensor, and the last dense layer hands its raw output directly to the loss, which is the arrangement that the collapsed gradient makes cheap.