The Softmax Cross-Entropy Gradient

Two facts about softmax and cross-entropy get stated often without derivation: that the gradient reaching the logits is simply predicted - true, and that production frameworks always fuse the two operations into one kernel instead of computing softmax, then log, then the loss, as three separate steps. Both facts have the same root cause, and working through the derivation once makes both obvious.

Setting up the derivative

Let z be the logit vector, p = softmax(z) (see softmax-logits), and y the one-hot target with y_c = 1 for the true class c. Cross-entropy loss is:

L = - sum_i y_i * log(p_i) = - log(p_c)

since y is zero everywhere except at c. To get the gradient with respect to the logits, dL/dz_j for every j, apply the chain rule through p_c, which itself depends on every logit through the shared softmax denominator. Softmax's own Jacobian is:

dp_i / dz_j = p_i * (delta_ij - p_j)

where delta_ij is 1 if i == j and 0 otherwise. This is the key fact that makes softmax's derivative depend on every logit, not just z_i, because of the shared normaliser.

The collapse

Substitute into the chain rule for dL/dz_j = -(1/p_c) * dp_c/dz_j:

dL/dz_j = -(1/p_c) * p_c * (delta_cj - p_j)
        = -(delta_cj - p_j)
        = p_j - delta_cj
        = p_j - y_j

Every p_c factor cancels exactly, leaving a result with no division, no log, and no reference to which class happened to be correct except through y_j. The gradient flowing back into the logits is, position by position, just the predicted probability minus the true (one-hot) probability. This is the same result quoted, without derivation, in next-token-prediction-cross-entropy; the point of deriving it here is to see why it is exact rather than an approximation, and why it survives the introduction of label smoothing largely unchanged: with a smoothed target y_smooth in place of y (see label-smoothing), the identical derivation gives p_j - y_smooth_j, still a single subtraction.

Why frameworks fuse softmax and cross-entropy

Computing softmax and cross-entropy as two separate operations means materialising the full probability vector p in memory, then taking log(p_c) of it. This has two costs. First, precision: p_c for a confident correct prediction can be extremely close to 1, or for a confidently wrong prediction extremely close to 0, and taking log of an already-rounded floating point probability compounds rounding error that a direct computation avoids. The numerically stable path computes log(p_c) via log-sum-exp directly on the logits, log(p_c) = z_c - max(z) - log(sum_j exp(z_j - max(z))) (the max-subtraction trick from softmax-logits), never forming p_c as an intermediate value at all. Second, memory: the backward pass for the fused operation only ever needs to produce p - y, so a fused CrossEntropyLoss(logits, labels) kernel can compute the forward loss via log-sum-exp and the backward gradient via the same softmax it needed anyway, without storing a separate full-precision probability tensor for the backward pass to consume.

Chunked cross-entropy at large vocabulary

The fusion argument gets sharper as vocabulary size V grows. A batch of B sequences of length T, scored against a vocabulary of V (commonly 100,000 to 250,000 for a modern LLM tokeniser, see tokenisation-bpe), produces a logits tensor of shape B x T x V. For long sequences and large batches this tensor can be larger than the rest of the model's activations combined, purely because V is so large. Because the exact gradient into the logits is known in closed form (p - y), production training kernels avoid ever materialising the full B x T x V logits tensor in global memory: they compute the loss and its gradient in chunks along the vocabulary dimension, accumulating the log-sum-exp normaliser and the loss incrementally, and only ever holding one vocabulary chunk of logits at a time. None of this is possible without first knowing, from the derivation above, exactly what gradient the chunked computation needs to reproduce.

When it falls down

• The per-logit gradient is bounded even when the loss is not. |p_j - y_j| <= 1 always, since both are probabilities, but the loss value -log(p_c) grows without bound as p_c -> 0. A very wrong, very confident prediction produces a large loss value but not an unusually large logit gradient magnitude; the emphasis cross-entropy places on confident wrongness lives in the loss curve and its effect through further computation, not in an outsized per-position logit gradient.
• Adding z-loss breaks the clean cancellation. With a z-loss term added to the objective (see z-loss-logit-regularisation), the total gradient into z_j gains an extra term proportional to p_j * log(Z), so the pure p - y result only describes the cross-entropy component, not the full training gradient in a model that also regularises logit scale.
• The derivation assumes a single correct class per position. Anywhere the target is genuinely a distribution rather than one-hot or smoothed-one-hot, such as training against soft labels from a teacher model, the same chain rule applies but y is no longer sparse, and the intuition of "one class gets penalised" no longer holds.
• Fused kernels trade flexibility for speed. A fused cross-entropy op that never materialises full logits makes it harder to inspect or modify the per-token probability distribution mid-training (for debugging, for custom per-token loss weighting), which is a real cost when you need that visibility.

Further reading

• Goodfellow, Bengio and Courville, 2016, Deep Learning, Chapter 6 - the maximum-likelihood derivation of the softmax and cross-entropy pairing.
• Milakov and Gimelshein, 2018, Online normalizer calculation for softmax - the streaming log-sum-exp computation that chunked cross-entropy kernels build on.