Lesson 4 of 6 · The Shakespeare project
Teach the model with one update
The model can now produce a next-character guess. Training gives it feedback: the next character is already in the text, so we can measure how much probability the model assigned to it. Then we adjust the model’s numbers to reduce a prediction penalty.
Use the text as the answer key
The corpus begins First Citizen:. After F, the observed target is i; after Fi, it is r. These are the context-and-target examples from lesson 1, now computed by different rows of one transformer pass.
| Position | Context this position may use | Observed next character |
|---|---|---|
| 1 | F | i |
| 2 | Fi | r |
| 3 | Fir | s |
| 4 | Firs | t |
| 5 | First | ␠ space |
| 6 | First␠ | C |
| 7 | First␠C | i |
| 8 | First␠Ci | t |
Take 65 consecutive characters. The first 64 are inputs, and the last 64 are targets, shifted one place. The causal rule from lesson 2 matters (the code enforces it with a mask): although a later character exists elsewhere in the input array, an earlier position cannot read it to discover its own answer.
x = stream[start : start + 64]
y = stream[start + 1 : start + 65]Turn “that was an unlikely guess” into a number
We want a small penalty when the observed target gets high probability, and a large penalty when it gets very little. These examples show how the scoring rule behaves:
| Probability given to the observed target | Penalty |
|---|---|
| 100% | 0.000 |
| 50% | 0.693 |
| 5% | 2.996 |
This penalty is called cross-entropy loss. For our window, we average it over all 64 targets. The target is the character recorded in this example, not a claim that every other continuation is impossible.
An even 1/65 probability share would give a penalty of about 4.17. Our untrained model starts close to that, because it has not learned which characters to favor in each context.
Our prediction_loss helper reshapes the score rows and targets, then calls PyTorch’s cross_entropy. The PyTorch function takes raw scores and the target IDs. It performs the needed normalization internally; do not apply softmax before passing scores to cross_entropy.
Watch one update change the model
Recorded results from the downloadable Python project. These controls replay saved measurements; they do not train or run a model in your browser.
This is one update on the first 64 targets in the corpus, starting from the same random initialization as lesson 3. The smaller penalty below is a measured result for this update, not a guarantee about every future step.
At the final position, the input ends in Al and the observed next character is l. Its probability rises from 1.61% to 2.01%. The optimizer used the average penalty over all 64 targets, not just this last one.
- Predict and scoreCompute the scores and their penalty against the observed characters.
- Calculate gradientsWork out how small changes to the trainable numbers would affect the penalty. The numbers are still unchanged.
- Apply an updateThe optimizer uses those gradients to change the model’s trainable numbers.
PyTorch tracks the calculations in the prediction, so it can compute these gradients for us. backward() calculates them; step() applies an update. An optimizer is the update rule, and its learning rate controls the scale of updates.
logits = model(x)
loss = prediction_loss(logits, y) # wrapper around PyTorch cross_entropy
optimizer.zero_grad() # clear gradients left from a previous update
loss.backward() # calculate new gradients; weights have not moved
optimizer.step() # change the trainable numbersOur run uses the AdamW optimizer with learning rate 0.001. A larger learning rate is not automatically better: an overly large update can make predictions worse. This update affects the token embeddings and the network that uses them together.
Source, corpus, tests, recorded results, and reference checkpoints. You can follow every example on this page without installing Python.
Run this step on your computer
Unzip the project. Open a terminal in its shakespeare-lab folder. Use Python 3.11–3.13 and create an environment once:
python3 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements.txtOn Windows, see the README for the equivalent commands. For a smaller Linux CPU-only installation, use the command in the README. This lab uses CPU and needs no API key.
python one_update.pyThe bundle includes checkpoints from our run so you can inspect the trained model immediately. Your own run can produce different text and slightly different numbers across machines.
Where do the penalty numbers come from?
For one observed target, the penalty is −log(p), using the natural logarithm and the probability assigned to that target. Thus −log(0.5) is about 0.693. Averaging these penalties over positions is the loss used in our update.
Check what actually changed
Run python -m unittest -v in the project folder. One test snapshots the parameters before backward(), verifies they are identical afterward, then checks that step() changes them. The recorded one-update result is also in results.json.
Quick check
You call backward() but never call step(). What has happened?
We have improved the penalty on the window we just used for training. That does not tell us whether the model predicts other Shakespeare text better. Next we will repeat the update across many windows and keep some text out of training to check.
Your turn
Explain it in your own words.
A learner computes a penalty using the next characters from the training text and calls backward(), but never calls optimizer.step(). What did the observed characters let us measure, and what has—and has not—changed in the model?
Answer the question in your own words. A short explanation is enough.
Your feedback
Work through a hint
Hint 1
Where did the target characters come from, and what is the model being scored against?
Hint 2
Does backward calculate an update direction or apply the optimizer’s update?
A worked explanation
The observed next characters let us score the predictions. Backward computed how changing the trainable numbers would affect that penalty, but those numbers have not been updated; the optimizer step is missing.
Code, data, and further reading
This project uses Tiny Shakespeare, public-domain Shakespeare excerpts. Its character-level setup follows the teaching approach in nanoGPT. See PyTorch’s training-loop explanation and its loss API. Our project README documents the exact model, corpus hash, commands and deliberate simplifications.