Cross-Entropy Loss › Gradient Behavior and Why Cross-Entropy Works

Understanding why cross-entropy produces good gradients requires examining what happens when we differentiate through the softmax and cross-entropy combination. This is one of those places where the mathematical elegance of the formulation reveals why certain design choices are not arbitrary but deeply principled.

Let $z_k$ be the unnormalized logit for class $k$, and let $\hat{y}_k = \text{softmax}(z_k)$ be the predicted probability after normalization. The gradient of the cross-entropy loss with respect to the logit $z_j$ is:

$$\frac{\partial \mathcal{L}}{\partial z_j} = \hat{y}_j - y_j$$

where:

• $\frac{\partial \mathcal{L}}{\partial z_j}$: the gradient of the loss with respect to logit $z_j$
• $\hat{y}_j$: the model's predicted probability for class $j$
• $y_j$: the true label for class $j$ (1 for the true class, 0 otherwise)

This remarkably simple result comes from the cancellation of terms when differentiating through softmax and logarithm together. For the true class (where $y_j = 1$), the gradient is $\hat{y}_j - 1$: a negative quantity that pushes the logit upward, increasing the true class probability. For all other classes (where $y_j = 0$), the gradient is $\hat{y}_j$: a positive quantity that pushes those logits downward, reducing their probability. The gradient magnitude is proportional to the prediction error: when the model is right with high confidence ($\hat{y}_j \approx 1$), gradients are near zero. When the model is wrong with high confidence (the model assigned probability near 1 to the wrong class), gradients are large.

This gradient formula has another important property: it is bounded. No matter how wrong the model is, the gradient of cross-entropy with respect to logits is at most 1.0 in magnitude. Compare this to squared error loss, where gradients grow without bound as predictions diverge from targets. The bounded gradients of cross-entropy make training more stable and reduce the need for gradient clipping in practice. As discussed in Gradient Clipping, managing gradient magnitude is important for stable transformer training, and the naturally bounded logit gradients from cross-entropy help.

A natural question arises: why not use mean squared error (MSE) as a classification loss? MSE between predicted probabilities and one-hot targets is:

$$\mathcal{L}_{\text{MSE}} = \frac{1}{K} \sum_{k=1}^{K} (y_k - \hat{y}_k)^2$$

The gradient of MSE with respect to logits, after accounting for the softmax, involves terms like $(\hat{y}_j - y_j) \cdot \hat{y}_j(1 - \hat{y}_j)$. When the model is very wrong (say, $\hat{y}_j \approx 0$ for the true class), this gradient includes the factor $\hat{y}_j(1 - \hat{y}_j) \approx 0$, which nearly vanishes. This is the saturation problem: MSE produces tiny gradients exactly when the model is most wrong and needs correction most urgently. Cross-entropy avoids this problem because its gradient is simply $\hat{y}_j - y_j$, which remains large when the model is wrong. This is why cross-entropy consistently outperforms MSE for training classifiers, and why you will never see MSE used as the training objective for a language model.