For a softmax over K classes with cross-entropy loss and one-hot label y, derive the gradient of the loss with respect to the pre-softmax logits z. Explain why the result is simpler than the chain rule through the softmax Jacobian suggests.

For a softmax over K classes with cross-entropy loss and one-hot label y, derive the gradient of the loss with respect to the pre-softmax logits z. Explain why the result is simpler than the chain rule through the softmax Jacobian suggests.

Approach: Write the softmax Jacobian as diag(p) - p p^T, contract it with the derivative of cross-entropy, and use the fact that the label vector sums to one.

dL/dz = p - y, the predicted probability vector minus the one-hot label. The softmax Jacobian is dp_k/dz_j = p_k(delta_{kj} - p_j), and L = -sum_k y_k log p_k has dL/dp_k = -y_k/p_k. Contracting, dL/dz_j = -sum_k (y_k/p_k) p_k (delta_{kj} - p_j) = -y_j + p_j sum_k y_k, and sum_k y_k = 1 for a one-hot label, so dL/dz_j = p_j - y_j. The cancellation happens because the 1/p_k from the log derivative exactly kills the p_k factor in the Jacobian, which is why cross-entropy is paired with softmax rather than squared error. That pairing also removes a saturation problem: with squared error the gradient carries an extra p(1 - p) factor and vanishes when the model is confidently wrong, while p - y stays order one. Implementations fuse the two operations and compute the logits through a log-sum-exp with the maximum subtracted for numerical stability.

Follow-up: Softmax is invariant to adding a constant to every logit. What does that imply about the rank of the Hessian, and how does weight decay interact with it?

Key concepts: softmax, cross-entropy, Jacobian, logits.