Softmax and Cross-Entropy

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.

Alternative: A Shortcut for One-Hot Targets

The general derivation above works for any distributional target, but if we commit to a one-hot $\mathbf{y}$ from the start, there is a shorter path that avoids the Jacobian entirely.

Let $c$ denote the true class, so that $y_i = \delta_{ic}$. The cross-entropy loss collapses to a single term:

\[L = -\sum_{i=1}^{K} y_i \log(\hat{y}_i) = -\log(\hat{y}_c)\]

Now substitute the softmax directly and split the logarithm:

\[L = -\log\left(\frac{e^{z_c}}{\sum_k e^{z_k}}\right) = -z_c + \log\sum_k e^{z_k}\]

The logarithm has undone the exponential in the numerator, leaving $L$ as an explicit function of the logits — no intermediate $\hat{\mathbf{y}}$ to chain through. We can differentiate term by term:

\[\frac{\partial L}{\partial z_j} = \underbrace{-\,\delta_{jc}}_{\partial(-z_c)/\partial z_j} + \underbrace{\frac{e^{z_j}}{\sum_k e^{z_k}}}_{\partial \log\sum_k e^{z_k} / \partial z_j} = \hat{y}_j - \delta_{jc}\]

The first term is $-1$ when $j = c$ and $0$ otherwise. The second term is worth pausing on: the derivative of log-sum-exp is the softmax itself. In hindsight, this is the real reason the cancellation in the previous section had to happen — softmax was the gradient of $\log\sum_k e^{z_k}$ all along, and composing it with a logarithm simply recovers that structure.

Since $\delta_{jc}$ is exactly the $j$-th component of the one-hot target, $\hat{y}j - \delta{jc}$ agrees with the general result $\hat{y}_j - y_j$, as it must.

So why bother with the longer derivation? The shortcut is quicker, but the one-hot assumption is baked in from the first line. The general path costs a Jacobian and a few index manipulations, and in exchange the result holds for any target with $\sum_i y_i = 1$ — smoothed labels included.