Optimization
Assemble a tiny causal language-model experiment with a clear data split, baseline, training loop and evaluation.
THE IDEA
Optimisation: start small and check
Before a long run, verify the pipeline on a case you can do by hand. At the start, small random or zero logits mean a nearly uniform prediction, so the loss should be about ln V. Check that the analytic gradient (softmax − onehot, averaged over positions) matches a finite-difference estimate, and that a single gradient step lowers the loss. Then record the initialisation, optimiser settings, data order and loss curve so runs can be compared.
One training example with three possible next tokens; the true one is index 1. Start with all logits 0 (loss ln 3 = 1.099) and take gradient steps.
| step | logits | P(true) | loss | |
|---|---|---|---|---|
| 0 | [0, 0, 0] | 0.333 | 1.099 | |
| 1 | [-0.33, 0.67, -0.33] | 0.576 | 0.551 | |
| 2 | [-0.55, 1.09, -0.55] | 0.72 | 0.329 | |
| 3 | [-0.69, 1.37, -0.69] | 0.796 | 0.228 | |
| 4 | [-0.79, 1.57, -0.79] | 0.841 | 0.173 | |
| 5 | [-0.87, 1.73, -0.87] | 0.871 | 0.139 | |
| 6 | [-0.93, 1.86, -0.93] | 0.891 | 0.115 | |
| 7 | [-0.99, 1.97, -0.99] | 0.906 | 0.099 | |
| 8 | [-1.03, 2.07, -1.03] | 0.917 | 0.086 |
gradient = softmax − onehot, so the true logit goes up and the others go down. A falling training loss shows the optimiser works on this data; it does not by itself show the model will do well on text it has not seen. Always begin with a tiny case like this and check that the initial loss equals ln V.
Predict first. Does a falling training curve alone show generalisation?
Chapter context
The capstone joins the earlier pieces into one inspectable experiment. The goal is to understand the model and its evidence, not to claim that a toy trained on a few paragraphs competes with a production assistant.
TRY IT YOURSELF
Part 1. Do it hands-on
Compute a gradient step by hand, then write the loss and gradient, check it numerically and train a bigram logit table.
One training example with three possible next tokens; the true one is index 1. Start with all logits 0 (loss ln 3 = 1.099) and take gradient steps.
| step | logits | P(true) | loss | |
|---|---|---|---|---|
| 0 | [0, 0, 0] | 0.333 | 1.099 | |
| 1 | [-0.33, 0.67, -0.33] | 0.576 | 0.551 | |
| 2 | [-0.55, 1.09, -0.55] | 0.72 | 0.329 | |
| 3 | [-0.69, 1.37, -0.69] | 0.796 | 0.228 | |
| 4 | [-0.79, 1.57, -0.79] | 0.841 | 0.173 | |
| 5 | [-0.87, 1.73, -0.87] | 0.871 | 0.139 | |
| 6 | [-0.93, 1.86, -0.93] | 0.891 | 0.115 | |
| 7 | [-0.99, 1.97, -0.99] | 0.906 | 0.099 | |
| 8 | [-1.03, 2.07, -1.03] | 0.917 | 0.086 |
gradient = softmax − onehot, so the true logit goes up and the others go down. A falling training loss shows the optimiser works on this data; it does not by itself show the model will do well on text it has not seen. Always begin with a tiny case like this and check that the initial loss equals ln V.
- To do: Three classes, all logits 0 (so every probability is 1/3), the true class is class 2 (index 1). The gradient of the loss w.r.t. the logits is softmax − onehot. Type its entry for the TRUE class (3 decimals)
- To do: One gradient step with learning rate 1 subtracts the gradient: logits become [−1/3, 2/3, −1/3]. The probability of the true class now (3 decimals)
- To do: The loss on this example after the step (3 decimals; it was ln 3 = 1.099)
Part 2. Do it in code
- softmax_rows(z): row-wise stable softmax of a 2-D array.
- loss_and_grad(logits, targets): logits (n, V), targets length n. Return (mean cross-entropy loss, gradient w.r.t. the logits) where the gradient is (softmax − onehot) / n.
- numeric_grad_check(logits, targets, eps=1e-6): the largest absolute difference between your analytic gradient and a central finite-difference estimate over all entries.
- train_bigram_logits(tokens, V, lr, steps): a (V, V) table of logits starting at zeros. Each step: logits for every current token, loss_and_grad against the next tokens, then table[cur] row updates with lr (use np.add.at to accumulate gradients of repeated current tokens). Return (table, list of losses before each step).
Stuck, or something not working? You can come back to it any time.
CHECK YOUR UNDERSTANDING
Does a falling training curve alone show generalization?
Chapter reference: math, worked example & code
This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.
Math and notation
For a token sequence of length n and context length T, stride-one windows containing T inputs and their shifted targets require T+1 original tokens. The number of complete windows is max(0,n-T). Keep this counting convention distinct from nonoverlapping batching.
Worked example
A 400-token text is split into 320 training tokens and 80 validation tokens before windowing. With T=8, the training part contains 320−8=312 complete shifted windows, while validation contains 72. Constructing windows before the split would risk overlap across the boundary.
Reference code
The following example uses the Python standard library and does not call a remote service.
text='abcd'*100
train,validation=text[:320],text[320:]
context=8
windows=[train[i:i+context+1] for i in range(len(train)-context)]
assert len(text) == 400
assert len(windows) == 312
assert all(len(w)==9 for w in windows)
assert len(validation)-context == 72
Python examples are provided for study and local execution; this page does not execute Python.
Common mistakes
Toy output quality is not a benchmark claim. Splitting overlapping windows after construction leaks shared text. Save the tokenizer and architecture configuration, not only raw weights.