Why does the gradient of softmax cross-entropy with respect to the logits reduce to p - y?

## Setting it up A classification head produces `K` unbounded real scores called **logits**, `z_1 ... z_K`. The softmax turns them into probabilities: ``` p_k = exp(z_k) / sum_j exp(z_j) ``` The categorical cross-entropy loss for one example whose true class is `c` is ``` L = -log p_c ``` The question is what `dL/dz_m` looks like - the sensitivity of the loss to each raw score, which is what training actually needs. ## The softmax Jacobian Softmax maps a vector to a vector, so its derivative is a `K x K` matrix. Two cases: - **Same index** (`k = m`): `dp_k/dz_k = p_k * (1 - p_k)`. - **Different index** (`k != m`): `dp_k/dz_m = -p_k * p_m`. Both collapse into one expression using the indicator `delta_km`, which is 1 when `k = m` and 0 otherwise: ``` dp_k/dz_m = p_k * (delta_km - p_m) ``` The off-diagonal term is negative and it is the whole reason the classes compete: raising any logit pushes every other probability down. ## The cancellation The loss touches the logits only through `p_c`, so the chain rule has a single term: ``` dL/dz_m = dL/dp_c * dp_c/dz_m ``` The first factor is `d(-log p_c)/dp_c = -1/p_c`. The second is the Jacobian entry `p_c * (delta_cm - p_m)`. Multiply: ``` dL/dz_m = -(1/p_c) * p_c * (delta_cm - p_m) = -(delta_cm - p_m) = p_m - delta_cm ``` The `p_c` in the numerator of the Jacobian and the `1/p_c` from the logarithm annihilate each other. Since `delta_cm` is exactly the one-hot target vector `y`, the whole gradient is ``` dL/dz = p - y ``` That cancellation is the point of the whole derivation. The Jacobian factor `p_c` is the quantity that becomes vanishingly small when the model is badly wrong; the logarithm's `1/p_c` is exactly its reciprocal. Compose them and the dangerous factor disappears. Compose the softmax with almost any other loss and it does not. ## Reading the result For a three-class example with logits `[2.0, 1.0, 0.1]` and true class 0, the probabilities are about `[0.659, 0.242, 0.099]`, the loss is about `0.417`, and the gradient is `[-0.341, 0.242, 0.099]`. Three things to notice: 1. **The components sum to zero.** `sum_k (p_k - y_k) = 1 - 1 = 0`. This is the gradient counterpart of the fact that adding a constant to every logit does not change the probabilities - the loss has a flat direction along the all-ones vector, and the gradient never points along it. 2. **Every component lies in `[-1, 1]`.** The correct class gets `p_c - 1`, which is negative (push that logit up); every wrong class gets `+p_m` (push it down). The loss itself is unbounded, but its gradient with respect to the logits is not. 3. **The magnitude scales with the error.** Confidently correct (`p_c` near 1): the gradient is near zero, so the example stops contributing. Confidently wrong (`p_c` near 0): the correct-class component approaches `-1` and the winning wrong class approaches `+1`, the strongest signal available. That monotone, never-saturating relationship between wrongness and gradient magnitude is why this loss trains well. ## Soft targets Nothing in the derivation required `y` to be one-hot beyond the step where a single `p_c` appeared. Redo it with a general target distribution `q` and loss `-sum_k q_k * log p_k`, and you get ``` dL/dz = p - q ``` for any target vector whose entries sum to one - an averaged annotator distribution, a distillation target, a soft assignment. The form is identical; the one-hot case is just `q = y`. ## The batch, and the sign convention For a batch of `N` examples the reported loss is normally the mean, so each example's `p - y` is divided by `N` before it flows backward. Training then subtracts a multiple of that gradient. It is worth keeping the sign straight in an interview: `p - y` is the derivative of the loss you **minimise**, so the update moves the correct logit **up** and the others **down**. ## The mistake to avoid When this gradient is implemented as a fused unit that takes logits and a target and returns `p - y`, the softmax derivative is already inside it. Multiplying the softmax Jacobian in again downstream double-counts and produces an update that is wrong by a factor of `p * (1 - p)` - a bug that does not crash, merely trains badly, which is why interviewers like the derivation. The second common error is quoting `p * (1 - p)` as the answer: that is the derivative of a scalar sigmoid, not the gradient of this composed loss with respect to its logits.