Softmax Cross-Entropy Loss and Gradient in NumPy

You are given a batch of logits logits of shape (N, C) and integer class labels labels of length N, each in 0…C−1. Return a list [loss, grad] where loss is the mean over the batch of −log softmax(logits_n)_{labels_n}, and grad is the N × C matrix ∂loss/∂logits.

The function must stay finite for logits of any size.

Examples

Input logits = [[2, 1, 0]], labels = [0]

Output [0.4076, [[-0.3348, 0.2447, 0.09]]]

p = [0.665, 0.245, 0.090]; the gradient is p minus the one-hot label.

Input logits = [[0, 0], [0, 0]], labels = [0, 1]

Output [0.6931, [[-0.25, 0.25], [0.25, -0.25]]]

Uniform predictions: the loss is log 2, and the 1/N shows in the gradient.

Input logits = [[1000, 0]], labels = [1]

Output [1000, [[1, -1]]]

Confidently wrong. The loss is large and finite.

Constraints

1 ≤ N ≤ 512, 2 ≤ C ≤ 1000
|logits[n][c]| ≤ 10⁴
Checked to a relative tolerance of 10⁻⁹

Hints

Hint 1

Compute log-probabilities, not probabilities. log softmax(z)ᵢ = zᵢ − logsumexp(z).

Hint 2

Use fancy indexing to pick each row's correct class: logp[np.arange(N), labels].

Hint 3

The gradient of the mean loss with respect to the logits is (softmax − one_hot) / N. You already have softmax as exp(logp).

Where candidates go wrong

Taking np.log of the softmax output. For a confidently wrong prediction the probability underflows to 0 and the loss becomes inf.
Forgetting the 1/N in the gradient while using the mean in the loss.
Building the one-hot matrix with a Python loop. Subtract 1 at [np.arange(N), labels] instead.
Modifying the caller's array in place when shifting by the row maximum.

The idea

For one example with logits z and label y, the loss is −log softmax(z)_y = logsumexp(z) − z_y. Computing logsumexp with the row maximum subtracted keeps every exponent ≤ 0, so the loss is exact and finite even when the model is confidently wrong.

The gradient is the reason this pairing is used everywhere. Differentiating logsumexp gives softmax, and differentiating −z_y gives minus the one-hot vector, so ∂ℓ/∂z = p − e_y. The softmax's own Jacobian, which looks intimidating, never needs to be formed. Averaging over the batch divides by N.

Vectorised: shift each row by its maximum, compute logp = z − log Σ exp(z), gather logp[n, y_n] for the loss, and take exp(logp) − one_hot for the gradient.

Time O(NC) · Space O(NC)

Reference solution

import numpy as np

def cross_entropy(logits, labels):
    Z = np.asarray(logits, dtype=float)
    y = np.asarray(labels)
    n = Z.shape[0]
    Z = Z - Z.max(axis=1, keepdims=True)                       # shift each row by its own maximum
    logp = Z - np.log(np.exp(Z).sum(axis=1, keepdims=True))    # log softmax, never log(0)
    loss = -logp[np.arange(n), y].mean()
    grad = np.exp(logp)                                         # softmax
    grad[np.arange(n), y] -= 1.0                                # minus the one-hot target
    return [float(loss), grad / n]

Console⌘/Ctrl + Enter runs

Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.

Why interviewers ask it

Every classifier is trained on this loss. The question tests three things at once: whether you work in log space, whether you can vectorise over a batch without a loop, and whether you know the gradient by heart. Interviewers often ask for the gradient first and the code second.

The theory behind it is in the lesson Numerically stable softmax and log-sum-exp.

Keep going

Easy8 min

Cross-entropy loss and its gradient from logits