Softmax Cross-Entropy Loss and Gradient

Implement the fused softmax cross-entropy used at the output of essentially every classifier. Given a batch of logits and the integer labels, return both the scalar loss and the gradient with respect to the logits. Interviewers ask for the two together because the gradient collapses to a strikingly simple form once you push the softmax and the cross-entropy through the chain rule, and knowing that form is the difference between a slow, unstable backward pass and the one real frameworks ship.

Definition

Let p = softmax(logits) be the row-wise softmax, so p[i] is a probability distribution over the k classes for example i. Cross-entropy for one example with true label y[i] is:

L_i = -log(p[i, y[i]])

The batch loss is the mean over the n examples:

L = (1/n) * sum_i -log(p[i, y[i]])

The gradient of this mean loss with respect to the logits has a famously clean closed form:

dL/d(logits) = (p - onehot(y)) / n

where onehot(y) is the (n, k) matrix that is 1 at column y[i] in row i and 0 elsewhere. In other words, subtract 1 from the softmax probability of the correct class in each row, then divide the whole matrix by n.

Task

Complete softmax_cross_entropy(logits, y) so it returns the tuple (loss, grad):

• loss is a Python float: the mean cross-entropy over the batch.
• grad is an np.ndarray of shape (n, k): the gradient of loss with respect to logits, equal to (softmax(logits) - onehot(y)) / n.

Use a numerically stable softmax (subtract the per-row maximum before exponentiating). Do not call scipy.special.softmax, any library softmax, or any cross-entropy helper.

Example

logits = np.array([[1.0, 2.0, 3.0],
                   [1.0, 2.0, 3.0]])
y = np.array([2, 0])

loss, grad = softmax_cross_entropy(logits, y)

loss
# -> 1.4076059644443804

grad
# -> array([[ 0.04501529,  0.12236424, -0.16737952],
#           [-0.45498471,  0.12236424,  0.33262048]])

Both rows have softmax [0.0900, 0.2447, 0.6652]. Row 0’s label is class 2 (the confident, correct one) so its loss term -log(0.6652) = 0.4076 is small; row 1’s label is class 0 (the unlikely one) so -log(0.0900) = 2.4076 is large. The mean is 1.4076. Each gradient row is the softmax with 1 subtracted at the true-class column, divided by n = 2.

Constraints

• logits is a 2-D array of shape (n, k) with 1 <= n and 2 <= k.
• y is a 1-D integer array of shape (n,) with every entry in [0, k-1].
• Logits may be large in magnitude; the loss and gradient must never contain inf or nan.
• Use vectorized NumPy with broadcasting and integer indexing, not Python loops over rows.

Approach

Compute the stable softmax once and reuse it for both outputs. For the loss, pull out the probability of the correct class in each row with integer indexing (probs[np.arange(n), y]), take -log, and average. For the gradient, start from the full softmax matrix and subtract 1 at the true-class position of each row (that is softmax - onehot(y)), then divide by n. The whole backward pass is one copy of the softmax plus a scatter-subtract, because the exp and log cancel analytically.

Solution

import numpy as np
from typing import Tuple

def softmax_cross_entropy(logits: np.ndarray, y: np.ndarray) -> Tuple[float, np.ndarray]:
    n = logits.shape[0]

    # Numerically stable row-wise softmax.
    shifted = logits - np.max(logits, axis=1, keepdims=True)   # largest per row is 0
    exps = np.exp(shifted)                                      # all in (0, 1], no overflow
    probs = exps / np.sum(exps, axis=1, keepdims=True)         # (n, k), rows sum to 1

    # Mean cross-entropy: -log of the correct-class probability, averaged.
    correct = probs[np.arange(n), y]                           # (n,), p[i, y[i]]
    loss = float(np.mean(-np.log(correct)))

    # Gradient wrt logits: (softmax - onehot(y)) / n.
    grad = probs.copy()
    grad[np.arange(n), y] -= 1.0                               # subtract 1 at true class
    grad /= n

    return loss, grad

Why the gradient is just softmax minus one-hot

For one example, the loss is L = -log(p_t) where t = y[i] is the true class and p = softmax(z). Differentiating through the softmax:

dL/dz_j = p_j - 1[j == t]

The two cases combine into the vector p - onehot(t). The 1 at the true class comes from the -log(p_t) term wanting to push that logit up; the p_j at every class comes from the normalization coupling all logits together. Because the batch loss is the mean of the per-example losses, the batch gradient carries the extra 1/n factor, giving (p - onehot(y)) / n. No exp or log survives in the backward pass — that cancellation is the entire reason softmax is paired with cross-entropy.

Walkthrough

On the example:

logits = [[1, 2, 3],
          [1, 2, 3]]
y = [2, 0]

• shifted = [[-2, -1, 0], [-2, -1, 0]], so each row’s softmax is probs = [[0.0900, 0.2447, 0.6652], [0.0900, 0.2447, 0.6652]].
• correct = probs[[0,1], [2,0]] = [0.6652, 0.0900] — row 0 at class 2, row 1 at class 0.
• -log(correct) = [0.4076, 2.4076], mean = 1.4076. That is the returned loss.
• grad starts as a copy of probs. Subtract 1 at (0, 2) and (1, 0): [[0.0900, 0.2447, -0.3348], [-0.9100, 0.2447, 0.6652]].
• Divide by n = 2: [[0.0450, 0.1224, -0.1674], [-0.4550, 0.1224, 0.3326]], matching the expected gradient.

Each gradient row sums to 0: subtracting a full one-hot (which sums to 1) from a probability row (which sums to 1) leaves 0, and dividing by n preserves that. That zero-sum is a useful sanity check.

Complexity & notes

• Time O(n·k), space O(n·k) — one max reduction, one exp, one sum reduction, and one scatter-subtract, all over the (n, k) array.
• The stable softmax matters here too: without the max-subtraction, large logits make exp overflow to inf and the loss becomes nan. Subtracting the per-row max keeps the largest exponent at exp(0) = 1.
• grad = probs.copy() before the in-place -= 1.0 avoids mutating probs; forgetting the copy silently corrupts the probabilities if they are reused, and doing grad = probs aliases the same buffer.
• For the loss it is cheaper and steadier to index the correct-class probabilities than to build the full onehot matrix and multiply; the one-hot only conceptually appears in the gradient, where the scatter-subtract realizes it in place.
• Production implementations fuse this further with log_softmax (z - m - log(sum(exp(z - m)))) to skip re-exponentiating inside the log, but the gradient form (softmax - onehot)/n is identical.