Training a Language Model

Your program will compute the number a GPT is trained to shrink: the average cross-entropy of its next-token predictions. You will also compute the perplexity that people quote, and the gradient that starts backpropagation.

One objective

Pretraining is next-token prediction on a huge pile of text. Take a sequence of T+1 token ids. Feed the first T into the model; the output at position t is a vector of scores (logits), one per vocabulary entry, and the correct answer for position t is the token that actually came next, tokens[t+1].

tokens:    [ 12    7    40    7    3 ]          T + 1 = 5
inputs:    [ 12    7    40    7 ]               positions 0..3
targets:   [      7    40    7    3 ]           the same list shifted left by one

Because of the causal mask (lesson 8), one forward pass produces valid predictions for all T positions at once: T training examples for the price of one.

Cross-entropy for one position

Turn the logits z into probabilities with softmax, then take the negative log of the probability given to the right token c:

loss = -log( softmax(z)[c] ) = logsumexp(z) - z[c]        logsumexp(z) = log( sum_i exp(z_i) )

The two forms are equal. Use the second. exp(1000) overflows a double, but logsumexp(z) = m + log( sum_i exp(z_i - m) ) with m = max(z) never does, for the same reason the softmax in lesson 5 subtracts the maximum.

Worked example, z = [2, 1, 0]:

softmax = [0.6652, 0.2447, 0.0900]
target 0 (the favourite):   loss = -ln 0.6652 = 0.4076
target 2 (the last choice): loss = -ln 0.0900 = 2.4076

A confident correct prediction costs almost nothing; a confident wrong one costs a lot. With all logits equal, softmax is uniform and the loss is ln V: about 10.8 for a 50,000-token vocabulary, which is where a freshly initialised model starts. Training pushes it below 3.

The mean, and perplexity

The training loss is the mean over scored positions. Its exponential is the perplexity: exp(mean loss). It reads as an effective branching factor, the number of equally likely tokens the model is choosing between. A uniform guess over V tokens has perplexity V.

Ignored positions. Sequences in a batch are padded to a common length, and the padding must not count. The convention (PyTorch's ignore_index) is a target id that is skipped: no loss for that position, not counted in the mean.

The gradient

Backpropagation needs the derivative of the loss with respect to every logit. For softmax followed by cross-entropy it is unusually simple:

d loss / d z_i = softmax(z)_i - (1 if i is the target else 0)

For z = [2, 1, 0] and target 0 that is [0.6652 - 1, 0.2447, 0.0900] = [-0.3348, 0.2447, 0.0900]. It always sums to zero: the target logit is pushed up by exactly what the others are pushed down. For the mean loss over n scored positions, divide each position's row by n.

You can check any gradient numerically: nudge one logit by 0.00001, recompute the loss, and divide the change by 0.00001.

Around the loop

The rest of a training run is engineering around this number: AdamW updates the weights from the gradients; the learning rate warms up for the first thousand steps and then decays along a cosine; batches of a million tokens are spread over many GPUs; scaling laws (Kaplan et al., Chinchilla) say loss falls smoothly with parameters, data and compute, with about 20 training tokens per parameter as the compute-optimal ratio.

Common mistakes

• Using tokens[t] as the target of position t instead of tokens[t+1]. The model then learns to copy its input. The first and last tokens are the usual victims: the last logits row has no target and the first token is never a target.
• Computing -log(softmax(z)[c]) directly. It works for small logits and returns inf or nan for large ones.
• Dividing the gradient by T instead of by the number of scored positions, or counting an ignored position in the mean.
• Taking the mean of the per-position perplexities instead of the exponential of the mean loss.

Your exercise

Write the cross-entropy of one position, the gradient row, and the function that applies the shift and the ignore rule over a whole sequence. The starter parses the input and prints the report.