Softmax and Cross-Entropy
Deriving the softmax Jacobian and the cross-entropy gradient, and why their composition collapses into prediction minus target.
Softmax Function and Its Derivative
The softmax function is defined as follows:
\[\operatorname{softmax}(\mathbf{z})_i = \sigma(\mathbf{z})_i = \frac{e^{z_i}}{\sum_{k=1}^{K} e^{z_k}} \qquad \text{for } i = 1, 2, \ldots, K\]To compute its derivative with respect to $z_j$, we apply the quotient rule. Throughout, $\sum_k$ denotes $\sum_{k=1}^{K}$, and $\delta_{ij}$ is the Kronecker delta ($1$ if $i = j$, $0$ otherwise). A useful property we will rely on is sifting: $\sum_k a_k \delta_{kj} = a_j$.
\[\begin{aligned} \frac{\partial \sigma(\mathbf{z})_i}{\partial z_j} &= \frac{\left(\dfrac{\partial}{\partial z_j}e^{z_i}\right)\sum_k e^{z_k} - e^{z_i}\left(\dfrac{\partial}{\partial z_j}\sum_k e^{z_k}\right)} {\left(\sum_k e^{z_k}\right)^2} \\[8pt] &= \frac{e^{z_i}\,\delta_{ij}\sum_k e^{z_k} - e^{z_i}e^{z_j}} {\left(\sum_k e^{z_k}\right)^2} \\[8pt] &= \frac{e^{z_i}}{\sum_k e^{z_k}} \left(\delta_{ij} - \frac{e^{z_j}}{\sum_k e^{z_k}}\right) \\[8pt] &= \sigma(\mathbf{z})_i \left(\delta_{ij} - \sigma(\mathbf{z})_j\right) \\[8pt] &= \begin{cases} \sigma(\mathbf{z})_i \left(1 - \sigma(\mathbf{z})_i\right) & \text{if } i = j \\ -\,\sigma(\mathbf{z})_i \,\sigma(\mathbf{z})_j & \text{if } i \neq j \end{cases} \end{aligned}\]Now we have the full Jacobian matrix of the softmax function:
\[J_{\sigma}(\mathbf{z}) = \begin{bmatrix} \sigma(\mathbf{z})_1 (1 - \sigma(\mathbf{z})_1) & -\sigma(\mathbf{z})_1 \sigma(\mathbf{z})_2 & \cdots & -\sigma(\mathbf{z})_1 \sigma(\mathbf{z})_K \\ -\sigma(\mathbf{z})_2 \sigma(\mathbf{z})_1 & \sigma(\mathbf{z})_2 (1 - \sigma(\mathbf{z})_2) & \cdots & -\sigma(\mathbf{z})_2 \sigma(\mathbf{z})_K \\ \vdots & \vdots & \ddots & \vdots \\ -\sigma(\mathbf{z})_K \sigma(\mathbf{z})_1 & -\sigma(\mathbf{z})_K \sigma(\mathbf{z})_2 & \cdots & \sigma(\mathbf{z})_K (1 - \sigma(\mathbf{z})_K) \end{bmatrix}\]or, more compactly,
\[J_{\sigma}(\mathbf{z}) = \operatorname{diag}\left(\sigma(\mathbf{z})\right) - \sigma(\mathbf{z})\,\sigma(\mathbf{z})^{\top},\]where the $\delta_{ij}$ term corresponds to the diagonal and the $-\sigma_i\sigma_j$ term to the outer product.
But do we really need the full Jacobian? In practice, we only need the derivative of the loss with respect to the input of the softmax. That is, we are not interested in $\frac{\partial \sigma(\mathbf{z})_i}{\partial z_j}$ for all $i$ and $j$, but rather in
\[\frac{\partial L}{\partial z_j} = \sum_{i=1}^{K} \frac{\partial L}{\partial \sigma(\mathbf{z})_i} \frac{\partial \sigma(\mathbf{z})_i}{\partial z_j}.\]Cross-Entropy Loss and Its Derivative
In our setting, the predicted probabilities are the output of the softmax: $\hat{\mathbf{y}} = \sigma(\mathbf{z})$. We assume the target $\mathbf{y}$ is a probability distribution, i.e., $\sum_{i} y_i = 1$ (a one-hot label being the most common case). The cross-entropy loss for a single example is defined as:
\[L(\mathbf{y}, \hat{\mathbf{y}}) = -\sum_{i=1}^{K} y_i \log(\hat{y}_i)\]The derivative with respect to the predicted probabilities is:
\[\frac{\partial L}{\partial \hat{y}_i} = -\frac{y_i}{\hat{y}_i}\]Combining Softmax and Cross-Entropy
Substituting both derivatives into the chain rule, something remarkable happens:
\[\begin{aligned} \frac{\partial L}{\partial z_j} &= \sum_{i=1}^{K} \frac{\partial L}{\partial \hat{y}_i} \frac{\partial \hat{y}_i}{\partial z_j} \\[4pt] &= \sum_{i=1}^{K} \left(-\frac{y_i}{\hat{y}_i}\right) \hat{y}_i \left(\delta_{ij} - \hat{y}_j\right) && (\hat{y}_i = \sigma(\mathbf{z})_i) \\[4pt] &= \sum_{i=1}^{K} (-y_i)\left(\delta_{ij} - \hat{y}_j\right) && \left(\tfrac{1}{\hat{y}_i} \cdot \hat{y}_i = 1\right) \\[4pt] &= -\sum_{i=1}^{K} y_i\,\delta_{ij} + \hat{y}_j \sum_{i=1}^{K} y_i \\[4pt] &= -y_j + \hat{y}_j \underbrace{\sum_{i=1}^{K} y_i}_{=\,1} && (\text{sifting}) \\[4pt] &= \hat{y}_j - y_j \end{aligned}\]The $K \times K$ Jacobian has collapsed into a simple subtraction: prediction minus target. When $\mathbf{y}$ is one-hot, this is exactly the probs - onehot computation found in every from-scratch implementation, and it is why frameworks such as PyTorch fuse the two operations into a single F.cross_entropy — computing the full Jacobian would be wasteful when the combined gradient is this simple. Note that the derivation only used $\sum_i y_i = 1$, so the result holds for any distributional target, including smoothed labels.