Optional Foundation — Build a Small Language Model
The forward pass and the loss
Training needs a single number that says how wrong the model currently is. This chapter builds that number: run the forward pass to get a probability for the true next token, then turn 'how much probability did we put on the right answer?' into the cross-entropy loss that everything else minimises.
To improve the model you first need to measure how wrong it is — as one number, so you can push it down. That number is the loss. Getting to it takes two steps you already half-know: run the forward pass to turn inputs into predicted probabilities, then compare those probabilities to the true next tokens with cross-entropy. This chapter connects them, and it doubles as the cleanest possible intuition for what "the model is learning" really means.
What you will understand by the end
- What the forward pass produces: a probability distribution over the vocabulary at every position.
- Why the loss is built from the probability the model gave the true next token.
- What cross-entropy / negative-log-likelihood is, in plain terms, and why the log.
- How per-token losses become one number per batch to minimise.
The forward pass, in one line
Feed a batch of input sequences through the model — embeddings, the stacked blocks, the LM head — and out come logits: an unnormalised score for every vocabulary token at every position. A softmax turns each position's logits into a probability distribution over the vocabulary. That is the model's answer: at each position, "here is how likely I think each possible next token is."
input tokens ─▶ [ model forward pass ] ─▶ logits ─softmax▶ P(next token) at each position
e.g. position 5 → { "the": .30, "a": .12, ... }
At the start of training these distributions are near-uniform noise; the whole job is to make them put high probability on the token that actually comes next.
The loss: how much probability did we put on the truth?
We have, from the data chapter, the true next token at every position (the shifted target). And we have, from the forward pass, the model's probability for that true token. A good model gives the true token high probability; a bad one gives it low probability. So the loss is built directly from that number:
For each position, look up the probability the model assigned to the token that actually came next. If it's near 1, the model was confident and right — tiny loss. If it's near 0, the model was confidently wrong — huge loss. The loss is just a way of averaging "how much probability did you put on the truth?" across every position, turned into something you can minimise.
Why the logarithm
Instead of using the probability directly, cross-entropy uses its negative logarithm,
−log p(true token). Three reasons this is the right shape:
- It is 0 when the model is perfect (
p = 1 → −log 1 = 0) and grows without bound aspapproaches 0 — confidently-wrong predictions are punished very hard, which is what you want. - It turns products into sums. The probability of a whole sequence is the product of per-token probabilities; logs turn that into a sum, which is numerically stable and easy to average.
- It matches probability space, where the meaningful difference between 0.01 and 0.001 is a factor of ten, not a tiny 0.009 gap — the log captures that.
This is cross-entropy loss, equivalently the negative log-likelihood of the true tokens under the model.
From many positions to one number
Each position gives one −log p value. Average them across all positions and all sequences in
the batch, and you get a single scalar — the batch loss. That one number is what training
minimises. Lower loss means the model is, on average, putting more probability on the right
next token.
Loss is the training-time cousin of everything in the evaluation half of this book. It is a proxy: low loss means good next-token prediction on the training distribution, which is not the same as being correct or useful on your task — exactly the "valid is not correct" gap seen from the training side. You minimise loss to train; you build evals to find out what the model is actually good at.
Mental model
Forward pass → a probability for every possible next token. Loss → the negative log of the probability you gave the true next token, averaged over the batch. Perfect prediction is loss 0; confidently wrong is huge loss. One scalar, and the whole of training is pushing it down.
Common mistakes
- Thinking the loss reads the predicted token. It reads the probability assigned to the true token — a confident wrong guess and a hesitant right guess score very differently.
- Forgetting the log's job. It makes perfect = 0, punishes confident errors hard, and turns sequence products into stable sums.
- Confusing loss with accuracy. Loss is a smooth probability-based signal you can optimise; accuracy is a blunt count. Training needs the smooth one.
- Reading low training loss as "good model." It only means good next-token prediction on the training data — quality is a separate, evaluated question.
Practical guidance
- Track loss in nats/token (or perplexity, its exponential) so the number is comparable across runs and context lengths.
- Watch both training and validation loss — the gap between them is your first read on overfitting (next chapter).
- Remember loss is a proxy: use it to drive training, but never let it stand in for a real evaluation of task quality.
- If the loss is flat at the very start, sanity-check the pipeline (targets shifted correctly, mask in place) before blaming the model.
Summary
- The forward pass turns inputs into a probability distribution over the vocabulary at every position.
- The loss is
−logof the probability the model gave the true next token — 0 when perfect, large when confidently wrong. - This is cross-entropy / negative log-likelihood; the log makes perfect = 0, punishes confident errors, and turns products into sums.
- Averaging over the batch yields one scalar to minimise — and it is a training proxy, not a substitute for evaluation.
Knowledge check
The model predicts the wrong token but was only 51% confident in it; on another example it predicts the right token with 99% confidence. Which contributes more to the loss, and why?
The loss looks at the probability given to the true token, not the argmax. On the first example the true token got low probability (well under 51%), so −log p is large. On the second the true token got 0.99, so −log 0.99 ≈ 0.01 — tiny. The first example dominates the loss. That is exactly why cross-entropy trains well: it responds to how much probability mass landed on the truth, not just to right/wrong.
Why does training minimise cross-entropy rather than something like squared error on the probabilities?
Cross-entropy is 0 exactly when the true token gets probability 1 and rises without bound as that probability approaches 0, punishing confident mistakes hard; its log turns the sequence-probability product into a stable sum; and it operates in the right (multiplicative) probability space. Squared error on probabilities is flat and weak near 0, so it gives poor gradients for the confidently-wrong predictions you most need to fix.
Related chapters
- From corpus to training examples — where the "true next token" the loss compares against comes from
- How the weights actually move — using this loss to update the parameters
- The model is only one component — why low loss is not the same as a correct system