Lesson 9. Cross-Entropy Loss Over the Vocabulary

Mission link: Stage 4 opens the training loop: the full forward pass (stages 1 to 3) produces logits, and this lesson turns those logits into the single scalar number training actually improves the model against.
Primary source: Docs: "torch.nn.functional.cross_entropy", PyTorch
Prerequisites: Lesson 8, Scaled dot-product attention

Warm-up

1. ▢ Why can the input embedding and output layers share a single weight matrix (weight tying), rather than needing two separate ones?

2. ▢ How does causal masking restrict what a query position can attend to?

Know this

From logits to a probability, from a probability to a loss

Lesson 8's output layer produces a vocab_size-length logit vector at every position, one raw score per possible next token. Softmax turns those logits into a proper probability distribution. Training then needs a single number measuring how wrong that distribution was, given what the actual next token in the training data turned out to be: cross-entropy loss is that number, computed as the negative log of the probability the model assigned to the actual correct token: loss = -log(p_correct).

Why negative log probability specifically

A perfect prediction, probability 1 for the correct token, gives a loss of exactly 0. As the assigned probability for the correct token shrinks toward 0, -log(p) grows without bound, an unboundedly increasing penalty for a confidently wrong prediction. Compare p = 0.8 for the correct token, giving -log(0.8) ≈ 0.223, against p = 0.1, giving -log(0.1) ≈ 2.303: the loss grows sharply, not just linearly, as the model's confidence moves further from the truth. A simpler-looking alternative like 1 - p doesn't have this property: it saturates at 1 no matter how confidently wrong the model is, giving no extra signal to distinguish "somewhat wrong" from "catastrophically, confidently wrong." Cross-entropy also has a clean gradient with respect to the logits, softmax_output - one_hot_true_label, a simple, numerically well-behaved expression that backpropagation (next lesson) can use directly, which is part of why softmax and cross-entropy are almost always paired together and often implemented as one fused, numerically stable operation taking raw logits directly, rather than a separate softmax step followed by a log.

Teacher forcing: every position's loss, computed in one pass

During training, the model predicts the token at position i+1 using only positions 0 through i, exactly what causal masking (lesson 3) already enforces. Because causal masking lets every position's prediction be computed simultaneously in a single forward pass, the loss at every position in the sequence, comparing that position's predicted distribution against the actual next token from the training data, can also be computed in that same single pass. This is teacher forcing: training always feeds the true, ground-truth previous tokens as context, never the model's own possibly-wrong past predictions, which is different from generation, where the model necessarily conditions on its own previously generated tokens instead.

One scalar for the whole batch

Each position produces its own cross-entropy loss value; these are averaged across every position in the sequence and across every sequence in the batch, collapsing down to one single scalar number. A single scalar is what the backward pass (next lesson) needs: gradient descent updates weights with respect to one number to minimize, not a whole tensor of separate loss values, so the averaging step is what turns "how wrong was every prediction across an entire batch" into the one quantity training actually optimizes.

Practice

1. ▢ The model assigns probability 0.8 to the correct next token in one case, and 0.1 in another. Compute the cross-entropy loss for each, and describe how the loss changes as the assigned probability for the correct token drops.

2. ▢ Why does cross-entropy use -log(p) as its penalty rather than something like 1 - p?

3. ▢ How does causal masking (lesson 3) make it possible to compute the training loss for every position in a sequence using a single forward pass, rather than one pass per position?

4. ▢ Why does training need the per-position, per-batch losses averaged down to a single scalar, rather than working with the full tensor of individual loss values directly?

5. ▢ Which claim is true of teacher forcing during training?

   • a) The model conditions each prediction on its own previously generated tokens, the same way it does during inference
   • b) The model conditions each position's prediction on the true, ground-truth previous tokens from the training data, not on its own predictions
   • c) Teacher forcing requires a separate forward pass for each position in the sequence
   • d) Teacher forcing and causal masking address unrelated problems with no connection to each other

Real-world reps

• [ ] Implement cross-entropy loss from raw operations (softmax, then negative log of the correct token's probability) for a small example, and confirm it matches PyTorch's F.cross_entropy on the same logits and target.
• [ ] For a training sequence, compute the per-position loss for each position by hand from a small set of made-up logits, then average them into a single scalar.
• [ ] Tomorrow: read the PyTorch cross_entropy docs and note why it takes raw logits rather than requiring you to call softmax yourself first.

Going further

• Docs: "torch.nn.functional.cross_entropy", PyTorch
• Repo: nanoGPT, Karpathy
• Resources

Not landing? Reread the primary source at the top, since this lesson compresses it and compression is where understanding leaks. Check the glossary for any term that felt slippery.

If the lesson itself is unclear rather than the material, that is a defect: open an issue.