Language-Model Loss with Masking & Label Smoothing

The problem

A language model predicts the next token at every position of every text in the batch. So instead of a table of logits it outputs a 3-D array:

• logits of shape $(B, T, V)$: $B$ texts, $T$ positions per text, $V$ scores per position (one per token of the vocabulary);
• targets of shape $(B, T)$: the right next token at every position.

The loss is the cross-entropy of the previous exercise, averaged over all positions, with two extras that real training loops use. This exercise reproduces exactly what PyTorch's torch.nn.CrossEntropyLoss does with them.

Extra 1: ignored positions

Some positions must not count. Padding tokens are added only so that every text in the batch has the same length $T$; in fine-tuning, the prompt tokens are not something the model should learn to write. Their target is set to a special value, ignore_index (default -100), and:

• they add nothing to the loss;
• the mean is taken over the $M$ positions that are not ignored.

Example. One text ($B = 1$), $T = 3$ positions, $V = 3$ tokens, targets = [[0, 2, -100]]: the last position is padding. The logits of the first two positions are [2, 1, 0] and [0, 0, 0], whose losses are $0.408$ and $1.099$ (computed in the previous exercise). So $M = 2$, and the loss is $(0.408 + 1.099) / 2 \approx 0.753$, whatever the logits of the padding position are.

Extra 2: label smoothing

Label smoothing with a parameter $\varepsilon$ (epsilon, between 0 and 1) makes the target less absolute: instead of "100% the right token", it becomes $1 - \varepsilon$ on the right token plus $\varepsilon$ spread evenly over the whole vocabulary. The loss of one position becomes:

and the loss of the batch is the mean over the kept positions:

where:

• $q$ is the softmax of the position's logits, and $-\log q_v$ is its log-softmax with the sign flipped;
• $y$ is the right token of the position, so $-\log q_y$ is the usual cross-entropy;
• the second term is $\varepsilon$ times the average of $-\log q_v$ over all $V$ tokens;
• $M$ is the number of positions that are not ignored.

With $\varepsilon = 0$ this is plain cross-entropy.

Example. Same batch, $\varepsilon = 0.1$.

• Position 0: $-\log q \approx (0.408, 1.408, 2.408)$, right token 0. $\ell = 0.9 \cdot 0.408 + 0.1 \cdot \tfrac{0.408 + 1.408 + 2.408}{3} \approx 0.367 + 0.141 = 0.508$.
• Position 1: all logits equal, so $-\log q_v = 1.099$ for every token and $\ell = 0.9 \cdot 1.099 + 0.1 \cdot 1.099 = 1.099$.
• Position 2 is ignored. Mean of the other two: $(0.508 + 1.099) / 2 \approx 0.803$.

What to implement

lm_cross_entropy(logits, targets, ignore_index=-100, label_smoothing=0.0): return the single number $\mathcal{L}$.

Requirements

• Stable for any logits: compute the log-softmax with log-sum-exp.
• ignore_index can be any integer, including -100. Never use an ignored target as an index. With a vocabulary of 100 tokens or more, logp[..., -100] does not fail: it silently reads a real token, the 100th from the end. With a smaller vocabulary it raises IndexError. Remove the ignored positions first.
• If every position is ignored, return 0.0 (there is nothing to learn from), not nan.
• No Python loops over positions.

Hints

1. Flatten first: logits to shape $(B \cdot T, V)$ and targets to $(B \cdot T,)$, with reshape. Now it is a table again, as in the previous exercise.
2. keep = targets != ignore_index is a boolean mask; logits[keep] and targets[keep] keep only the rows that count.
3. The second term of $\ell$ is label_smoothing times the mean of the log-softmax row, with the sign flipped.