ML Katas

One of the most satisfying derivations in deep learning is the gradient of the combined Softmax and Cross-Entropy loss. For a multi-class classification problem with $K$ classes, given true labels $𝐲$ (one-hot encoded) and predicted logits $𝐳$, the softmax output is ${\hat{y}}_k=\frac{e^{z_k}}{{\sum }_j{e}^{z_j}}$. The categorical cross-entropy loss is $L=-{\sum }_k{y}_k\log ({\hat{y}}_k)$.

1. Partial Derivative of Softmax: First, derive $\frac{\partial {\hat{y}}_k}{\partial z_i}$. You'll need to consider two cases:
   • When $k=i$ (the derivative of a specific output with respect to its own logit).
   • When $k\ne i$ (the derivative of a specific output with respect to a different logit). Show that $\frac{\partial {\hat{y}}_k}{\partial z_i}={\hat{y}}_k(1-{\hat{y}}_i)$ if $k=i$, and $\frac{\partial {\hat{y}}_k}{\partial z_i}=-{\hat{y}}_k{\hat{y}}_i$ if $k\ne i$. This can be compactly written as ${\hat{y}}_k({\delta }_{ki}-{\hat{y}}_i)$, where ${\delta }_{ki}$ is the Kronecker delta.
2. Chain Rule Application: Now, use the chain rule to find the gradient of the loss $L$ with respect to a logit $z_i$: $\frac{\partial L}{\partial z_i}={\sum }_k\frac{\partial L}{\partial {\hat{y}}_k}\frac{\partial {\hat{y}}_k}{\partial z_i}$.
   • Recall that $\frac{\partial L}{\partial {\hat{y}}_k}=-\frac{y_k}{{\hat{y}}_k}$.
3. The Result: Show that the final gradient simplifies to a remarkably clean form: $\frac{\partial L}{\partial z_i}={\hat{y}}_i-y_i$. This means the gradient is simply the difference between the predicted probability and the true one for that class.
4. Intuition: Why is this derivative so elegant? What does it imply about how the network updates its logits based on the error? How does this simplicity contribute to the efficiency of backpropagation in classification tasks?
5. Verification: You can conceptually verify this by considering a simple case (e.g., 2 classes, one-hot labels) and thinking about what would happen if ${\hat{y}}_i$ is much larger or much smaller than $y_i$. Numerically, you could perform a small forward pass, calculate the loss, and then compute the analytical gradient and compare it to a numerical gradient approximation (e.g., using finite differences) for a single logit change.

import numpy as np

def softmax(z):
    exp_z = np.exp(z - np.max(z)) # for numerical stability
    return exp_z / np.sum(exp_z)

def cross_entropy_loss(y_true, y_pred_probs):
    # y_true is one-hot encoded
    # y_pred_probs are softmax outputs
    # Ensure to handle log(0) cases, usually by clipping predictions
    epsilon = 1e-10
    y_pred_probs = np.clip(y_pred_probs, epsilon, 1 - epsilon)
    return -np.sum(y_true * np.log(y_pred_probs))

# Example logits and true label
# logits = np.array([2.0, 1.0, 0.1])
# y_true = np.array() # Class 0 is true

# y_pred = softmax(logits)
# print(f"Predicted probabilities: {y_pred}")

# Analytical gradient for logits:
# grad_analytical = y_pred - y_true
# print(f"Analytical gradient: {grad_analytical}")

# For numerical verification (optional, more advanced):
# You'd perturb each logit slightly and observe change in loss.
# e.g., for z: (loss(z + epsilon_vec) - loss(z - epsilon_vec)) / (2 * epsilon)