Softmax Cross-Entropy Gradient

mediumPROBackpropagationDeep LearningCV: Deep Learning

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.33]]

Reasoning

Inputs: logits is a $1 \times 3$ batch — one sample with raw scores $z = (1, 2, 3)$ for classes 0, 1, 2. targets is [2]: the true class of that sample is index 2, so its one-hot vector is $y = (0, 0, 1)$ .
Softmax: exponentiate each score: $e^1 \approx 2.7183$ , $e^2 \approx 7.3891$ , $e^3 \approx 20.0855$ , which sum to $\approx 30.1929$ . Dividing gives $p \approx (0.0900, 0.2447, 0.6652)$ — positive and summing to 1.
Gradient: apply $\frac{\partial L}{\partial z_i} = p_i - y_i$ . Classes 0 and 1 have $y_i = 0$ , so they keep $0.0900$ and $0.2447$ ; class 2 gets $0.6652 - 1 = -0.3348$ .
Output: a $1 \times 3$ grid of floats, same shape as logits, shown rounded as [[0.09, 0.24, -0.33]].
Notice: the only negative entry sits at the true class (gradient descent pushes its score up), and each row sums to $1 - 1 = 0$ . Every sample in a batch is handled independently, with no averaging over the batch.

Constraints

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

Learn this in

Deep LearningCV Study Plan ClassificationML Study Plan

More from CV: Deep Learning

Implement Max Pooling with Indiceshard Depthwise Separable Convolutionmedium Batch Normalization Forward Passhard Flatten Feature Mapeasy Max Poolingeasy

Concept primerPREVIEW

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}

PRO

Continue the full explanation

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

Example

Constraints

Learn this in

More from CV: Deep Learning

Continue the full explanation

Editor locked