Cross-Entropy from Logits

The problem

This is the loss every language model is trained with. The model has to choose among $C$ classes (for a language model, the $C$ tokens of the vocabulary), and for each of $N$ examples it outputs one raw score per class, its logits. You get:

• logits, a table $z$ with $N$ rows and $C$ columns: $z_{n,c}$ is the score of class $c$ for example $n$;
• targets, $N$ integers: $y_n$ is the right class of example $n$.

The loss of one example is minus the log of the probability the model gives to its right class, after softmax. The loss of the batch is the average over the $N$ examples.

The loss

Computing softmax and then the log is unsafe (a tiny probability can round to 0, and $\log 0 = -\infty$). The same quantity, written with the LSE (log-sum-exp) function, stays finite:

where:

• $z_n$ is row $n$ of the logits;
• $z_{n, y_n}$ is the logit of the right class of example $n$;
• $\mathrm{LSE}(z_n) = \log \sum_c e^{z_{n,c}}$, computed with the stable trick of the NumPy section.

Example. Two examples, $C = 3$: logits [[2, 1, 0], [0, 0, 0]], targets [0, 2].

• Row 0: $\mathrm{LSE}(2, 1, 0) \approx 2.408$, the right logit is $z_{0,0} = 2$, loss $0.408$.
• Row 1: $\mathrm{LSE}(0, 0, 0) = \log 3 \approx 1.099$, the right logit is $z_{1,2} = 0$, loss $1.099$. With equal logits every class gets $1/3$, so this is $-\log(1/3)$.
• Mean: $(0.408 + 1.099) / 2 \approx 0.753$.

The gradient

The gradient with respect to the logits is the softmax probabilities minus the one-hot target, divided by $N$ because of the mean (the theory page derives it):

where:

• $q_n = \mathrm{softmax}(z_n)$ are the probabilities of row $n$;
• $[c = y_n]$ is 1 for the right class and 0 for the others.

Example. Same batch.

• Row 0: $q_0 \approx (0.665, 0.245, 0.090)$; subtract 1 at class 0: $(-0.335, 0.245, 0.090)$; divide by $N = 2$: $(-0.167, 0.122, 0.045)$.
• Row 1: $q_1 = (1/3, 1/3, 1/3)$; subtract 1 at class 2: $(0.333, 0.333, -0.667)$; divide by 2: $(0.167, 0.167, -0.333)$.

What to implement

• cross_entropy(logits, targets): the mean loss, a single number.
• cross_entropy_backward(logits, targets): the gradient, shape $(N, C)$.

Requirements

• Stable for any logits: [[1000, 0]] with target 1 has loss $1000$, not inf or nan.
• No Python loops over the batch. Read the right logits with integer-array indexing, as in the one-hot-gather exercise, not with a one-hot matrix.

Hints

• For the loss, compute the log-softmax of every row, $z_n - \mathrm{LSE}(z_n)$, then pick the right entry of each row and average.
• For the gradient, start from the softmax probabilities, subtract 1 at the right class of every row (integer-array indexing again, with -=), and divide by $N$.