Gradients of Loss Functions

1. Gradient of Mean Squared Error Loss

Let's begin with the simpler case: regression. In a typical regression problem, a network produces a vector of predictions, which we can call $\hat{\mathbf{y}}$. For a regression task, these outputs are often the raw outputs of the final linear layer, which we've been calling the logits. The Mean Squared Error (MSE) loss measures the squared Euclidean distance between these predictions and the true target values $\mathbf{y}$.

Let $\hat{\mathbf{y}},\mathbf{y}\in {\mathbb{R}}^K$. The MSE loss is defined as:

For simplicity in differentiation, we often work with the sum of squared errors, omitting the constant $1/K$ and sometimes adding a factor of $\frac{1}{2}$ to cancel the exponent after differentiation. Let's use $L=\frac{1}{2}{\sum }_{i=1}^K({\hat{y}}_i-y_i{)}^2$. Our goal is to find the gradient of this scalar loss $L$ with respect to the prediction vector $\hat{\mathbf{y}}$, which is ${\nabla }_{\hat{\mathbf{y}}}L$.

We can compute this by finding the partial derivative with respect to each component ${\hat{y}}_j$:

Due to the linearity of the derivative, we can move it inside the sum. The derivative $\frac{\partial }{\partial {\hat{y}}_j}$ will be zero for all terms where $i\ne j$. The only term that survives is the one where $i=j$:

Since this holds for every component $j$, the full gradient vector is:

The result is simply the error vector—the difference between the predicted and true values. This is a wonderfully intuitive outcome. The gradient points in the direction of the error, telling each output unit how far and in which direction it needs to move to reduce the loss. This vector $\hat{\mathbf{y}}-\mathbf{y}$ is the initial upstream gradient that is fed into the final layer of the network for backpropagation.

2. Gradient of Categorical Cross-Entropy Loss

Now for the more involved, and perhaps more satisfying, case: multi-class classification. Here, the standard setup involves a final linear layer producing logits $\mathbf{z}$, a softmax activation function to convert logits into probabilities $\hat{\mathbf{y}}$, and a Categorical Cross-Entropy loss function to measure the discrepancy between the predicted probabilities $\hat{\mathbf{y}}$ and the one-hot encoded true labels $\mathbf{y}$.

The pipeline is:

1. Logits: $\mathbf{z}\in {\mathbb{R}}^K$
2. Softmax: $\hat{\mathbf{y}}=\text{softmax}(\mathbf{z})$
3. Loss: $L=\text{CE}(\mathbf{y},\hat{\mathbf{y}})=-{\sum }_{i=1}^K{y}_i\log ({\hat{y}}_i)$

Our goal is to compute the gradient of the loss with respect to the logits, $\frac{\partial L}{\partial \mathbf{z}}$. Using the chain rule, we can decompose this as:

This involves the product of the gradient of the loss with respect to the softmax outputs (a row vector) and the Jacobian of the softmax function. Let's derive each part separately.

The Derivative of Softmax and Cross-Entropy

The following article provides a clear, step-by-step derivation for both the softmax Jacobian and the gradient of the cross-entropy loss. We'll use it to build our argument.

Let's summarize the key results from the reading:

1. Softmax Jacobian ($\frac{\partial \hat{\mathbf{y}}}{\partial \mathbf{z}}$): The partial derivative of the $i$-th softmax output with respect to the $j$-th logit is:

   $$\frac{\partial {\hat{y}}_i}{\partial z_j}=\{\begin{array}{ll}{\hat{y}}_i(1-{\hat{y}}_i)& \text{if }i=j\\ -{\hat{y}}_i{\hat{y}}_j& \text{if }i\ne j\end{array}$$

   This can be concisely written using the Kronecker delta, ${\delta }_{ij}$, as ${\hat{y}}_i({\delta }_{ij}-{\hat{y}}_j)$.

2. Cross-Entropy Gradient (${\nabla }_{\hat{\mathbf{y}}}L$): For a one-hot true label vector $\mathbf{y}$ where $y_c=1$ for the correct class $c$, the loss simplifies to $L=-\log ({\hat{y}}_c)$. The partial derivative with respect to an arbitrary output probability ${\hat{y}}_j$ is:

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

   This is because $y_j$ is only non-zero (and equal to 1) for the correct class, so the derivative is non-zero only for that component.

Combining the Derivatives

Now we combine these two parts using the chain rule. The $j$-th component of our target gradient ${\nabla }_{\mathbf{z}}L$ is:

Let's plug in our results:

This sum looks intimidating, but a beautiful cancellation occurs. Let's expand it by considering the two cases for the softmax Jacobian.

We split the sum into the term where $i=j$ and all terms where $i\ne j$:

Substitute the corresponding Jacobian elements:

Now, simplify each part:

• First part: $-y_j(1-{\hat{y}}_j)=-y_j+y_j{\hat{y}}_j$
• Second part: ${\sum }_{i\ne j}{y}_i{\hat{y}}_j={\hat{y}}_j{\sum }_{i\ne j}{y}_i$

Combine them:

The term in the parentheses is ${\sum }_{i=1}^K{y}_i$. Since $\mathbf{y}$ is a one-hot vector, the sum of its elements is exactly 1.

This stunningly simple result holds for every component $j$. Therefore, the gradient vector is:

This is a cornerstone result in deep learning. The complex derivatives involving exponentials and logarithms completely cancel out, leaving the same intuitive form as the MSE gradient: prediction minus true value.

The following video provides an alternative walk-through of this exact derivation. Watching it can help solidify the algebraic steps.

The fact that softmax and cross-entropy are such a compatible pair is no accident. This pairing provides a stable and meaningful gradient for training classification models. The error signal $\hat{\mathbf{y}}-\mathbf{y}$ directly tells the model how to adjust its logits: if the predicted probability for the correct class is too low, the gradient will be positive for that logit, pushing it higher.

Conclusion

In this lesson, we applied the chain rule to derive the initial gradients for the backpropagation algorithm. We found that for two of the most common setups in deep learning, the result is remarkably consistent and intuitive.

Key Takeaways:

• The gradient of the Mean Squared Error loss with respect to the network's predictions $\hat{\mathbf{y}}$ is simply the error vector: ${\nabla }_{\hat{\mathbf{y}}}L=\hat{\mathbf{y}}-\mathbf{y}$.
• The gradient of the Categorical Cross-Entropy loss combined with a Softmax activation, taken with respect to the pre-activation logits $\mathbf{z}$, also simplifies to the error vector: ${\nabla }_{\mathbf{z}}L=\hat{\mathbf{y}}-\mathbf{y}$.
• This "gradient collapse" is a key reason for the prevalence of the Softmax/Cross-Entropy pair in classification, as it provides a numerically stable and computationally efficient starting point for backpropagation.

We've now seen how to compute the Jacobian for intermediate layers and the initial gradient from the loss function. In the next lesson, we will tie these concepts together by deriving the Jacobian of the softmax function itself, solidifying the derivation we used today and preparing you to trace gradients through any part of a network.