📘

Softmax Cross-Entropy Gradient

Implement a function to compute the gradient of the cross-entropy loss with softmax activation, a crucial component in backpropagation for training deep neural networks. This task involves understanding the mathematical foundations of softmax and cross-entropy loss.

The softmax function is used for multi-class classification problems, where it maps a vector of real numbers to a vector of probabilities, ensuring that each element is in the range (0, 1) and the elements sum up to 1. The cross-entropy loss measures the difference between the predicted probabilities and the true distribution, typically represented as a one-hot encoded vector.

To compute the gradient, we follow these steps:

Compute the softmax of the input vector $z$ using the formula $p_i = \frac{e^{z_i}}{\sum_j e^{z_j}}$ .
Calculate the cross-entropy loss using $L = -\sum_i y_i \log(p_i)$ , where $y$ is the one-hot encoded target vector. The key formula for the gradient is:

\frac{\partial L}{\partial z_i} = p_i - y_i

This technique is widely used in image classification tasks.

Example:

Input:

logits = [[1.0, 2.0, 3.0]]  # batch=1, classes=3
targets = [2]  # True class is index 2

Output:

[[0.09, 0.24, -0.67]]

Reasoning:

Softmax of [1,2,3]: [0.09, 0.24, 0.67] One-hot target: [0, 0, 1]

Gradient = softmax - one_hot = [0.09, 0.24, 0.67-1] = [0.09, 0.24, -0.33]

Constraints:

logits: Raw network outputs (batch_size, num_classes)
targets: Ground truth class indices (batch_size,)
Return: Gradient tensor same shape as logits

Background Knowledge

Softmax is used in multi-class classification to turn raw scores (logits) $z_i$ into a probability distribution over classes:

p_i = \frac{e^{z_i}}{\sum_j e^{z_j}}

Each $p_i$ is in $(0,1)$ and $\sum_i p_i = 1$ . The logits $z_i$ are typically the last linear layer outputs before the activation.

Cross-entropy measures how “far” the predicted probability distribution $p$ is from the true distribution $y$ :

L = -\sum_i y_i \log(p_i)

When labels are one-hot, $y_k = 1$ for the true class $k$ and 0 otherwise, so the loss simplifies to $L = -\log(p_k)$ . In backpropagation, we need $\frac{\partial L}{\partial z_i$ } $, i.e., how changing each logit$ z_i$ affects the loss.

The key “beautiful” result is that, when you combine softmax and cross-entropy, the gradient w.r.t. logits simplifies to:

\frac{\partial L}{\partial z_i} = p_i - y_i

This is much simpler than separately differentiating softmax and cross-entropy and is what is implemented in deep learning libraries internally.

Algorithm / Approach

For this type of problem (gradient with softmax + cross-entropy):

Use the chain rule: logits $z \rightarrow$ softmax $p(z) \rightarrow$ loss $L(p,y)$ .
First compute the softmax output $p$ from logits $z$ .
Write the loss $L$ in terms of $p$ .
Compute $\frac{\partial L}{\partial p_i}$ and $\frac{\partial p_i}{\partial z_j}$ .
Combine them using:

\frac{\partial L}{\partial z_j} = \sum_i \frac{\partial L}{\partial p_i}\frac{\partial p_i}{\partial z_j}

Simplify the algebra to obtain an expression in terms of $p$ and $y$ .

The pattern is: derive component-wise, then use vector form (here, gradient = $p - y$ ).

Step-by-Step Strategy

To understand/derive the result:

Write softmax explicitly

p_i = \frac{e^{z_i}}{\sum_k e^{z_k}}

Write cross-entropy in terms of $p$

L = -\sum_i y_i \log(p_i)

Differentiate loss w.r.t. $p_i$

\frac{\partial L}{\partial p_i} = -\frac{y_i}{p_i}

👑

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

🔒

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test Results

0/0

Run code to see test results.