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.