Softmax cross-entropy gradient for a batch
Medium15 min10 hidden tests
A support inbox routes each message to one of k teams. The classifier outputs logits: one raw score per team for each message. Its loss is the mean cross-entropy over the batch, where a message's cross-entropy is minus the natural log of the softmax probability of its true team. Write softmax_ce_grad(logits, labels) that returns the gradient of that mean loss with respect to every logit.
Change the logits, pick the true team
gradient = (softmax − one-hot) / nSoftmax of [2, 1, 0.1] = [0.659, 0.242, 0.099]. The gradient is p minus the one-hot target: [−0.341, 0.242, 0.099]. A descent step raises Billing’s logit and lowers the others. In a batch of n messages, divide each row by n.
Gradient −0.341, 0.242, 0.099.
Examples
The true class gets probability minus 1, the others their probability
softmax_ce_grad([[2.0, 1.0, 0.1]], [0])
returns [[-0.340999, 0.242433, 0.0985659]]
Huge logits stay finite, and the batch of 2 halves every entry
softmax_ce_grad([[1000.0, 1001.0], [0.0, 0.0]], [1, 0])
returns [[0.134471, -0.134471], [-0.25, 0.25]]
Requirements
logitsis an (n, k) array or list of lists, andlabelsholds n class indices from 0 to k - 1.- Softmax runs across each row (axis 1). Logits can be as large as 1000 or as small as -1000, so subtract each row's maximum before you call
exp. - The loss is averaged over the n messages, so the gradient carries a factor of 1 / n.
- Return an (n, k) float array with one gradient per logit.
The idea behind it
Softmax turns a row of logits into probabilities: p_j = exp(z_j) / sum_m exp(z_m). Cross-entropy for one example with true class y is -log(p_y), the natural log. Put together, the gradient with respect to logit j is p_j - 1 for the true class and p_j for every other class. In short, p - onehot(y), where the one-hot vector has a 1 at the true class and 0 elsewhere. Averaging the loss over a batch of n rows divides that gradient by n. Two details matter in code. First, softmax is unchanged when you subtract the same number from every logit in a row, and subtracting the row maximum keeps every exponent at 0 or below, so exp cannot overflow. Second, the sums run along each row. This gradient is where backprop starts for language models too: a GPT's next-token loss is exactly this, with k equal to the vocabulary size.