Lesson 5 of 6 · The Shakespeare project
Train your transformer on Shakespeare
One update improved the examples used to make that update. Now we want to know whether the model learns patterns that help on other text. We will repeat the same operation on many Shakespeare windows and check predictions on a separate part of the file.
Keep an answer sheet the model cannot train on
Before extracting any windows, we split the 1,115,394-character file into 1,003,854 training characters and 111,540 validation characters. The validation portion is the last 10% of the stream.
last 10%
Training text supplies parameter updates. Validation text supplies a check. We calculate the same prediction penalty on validation examples, but never use their gradients to update the model. This tests predictions on text excluded from those updates.
Split first, then draw windows entirely within each part. Otherwise overlapping windows or windows crossing the boundary could put held-out characters into training. The two parts may also contain different scenes, so a gap between their losses alone does not diagnose memorization.
Repeat the update on batches
A batch is several examples processed together. Here each update uses 16 windows of 64 input characters: 1,024 next-character targets. Average their penalties, calculate gradients, and apply one update. Then draw another batch.
for step in range(4000):
x, y = get_batch(train, 16, 64, rng)
logits = model(x)
loss = prediction_loss(logits, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()The real script also measures progress and saves snapshots. It starts fresh from the same random initialization as lesson 3; the one-update demonstration does not secretly give this run a head start. Windows are sampled repeatedly, so some text is seen more than once.
Where the batch dimension goes
The arrays are now [16, 64] input IDs, [16, 64, 64] representations, and [16, 64, 65] scores. In that middle array, one 64 means positions and the other means features per position; they happen to be equal in our settings. Every score row is paired with its own next-character target.
Check improvement on held-out text
Recorded results from the downloadable Python project. These controls replay saved measurements; they do not train or run a model in your browser.
Across 4,000 updates, the measured training penalty fell from 4.198 to 1.592. Validation penalty fell from 4.192 to 1.802. The second change is evidence of better prediction on held-out text; the first alone would not establish that.
Training 1.592 · validation 1.802
ROMEO: OF QOEXEN EDWARD: 'Tis and a last fast depass yoursess! even his deple's a seld-tightly, and a such I sweet soon As it their had belp's drongers; and of it. Ma
Every sample uses prompt ⏎ROMEO:⏎, seed 2026, and 160 new characters with unmodified softmax shares (temperature 1, explained next). Notice the change from random characters toward words and dialogue formatting. The final sample still contains invented words and awkward sentences. A lower penalty does not mean that the model writes coherent plays.
See the measurements and evaluation procedure
| Updates | Training penalty | Validation penalty |
|---|---|---|
| 0 | 4.198 | 4.192 |
| 250 | 2.457 | 2.484 |
| 500 | 2.295 | 2.334 |
| 750 | 2.143 | 2.204 |
| 1000 | 2.044 | 2.131 |
| 1250 | 1.954 | 2.070 |
| 1500 | 1.884 | 2.006 |
| 1750 | 1.834 | 2.002 |
| 2000 | 1.783 | 1.955 |
| 2250 | 1.741 | 1.932 |
| 2500 | 1.718 | 1.901 |
| 2750 | 1.690 | 1.887 |
| 3000 | 1.658 | 1.862 |
| 3250 | 1.638 | 1.842 |
| 3500 | 1.620 | 1.811 |
| 3750 | 1.610 | 1.822 |
| 4000 | 1.592 | 1.802 |
At each measurement, we evaluate the same 128 sampled windows from each split, scoring every position. Reusing those windows makes the curve comparable. They are a sample of each split, not a score for every character in the file.
A model can get better at its homework and worse at the check
We also ran a deliberate counterexample: restart the model, then repeat only the first training window 200 times. We measure the same held-out windows as before.
| Updates | Penalty on that window | Held-out penalty |
|---|---|---|
| 0 | 4.199 | 4.192 |
| 20 | 2.051 | 3.739 |
| 200 | 0.028 | 5.263 |
By the end, it predicts that one window extremely well while the held-out penalty is worse than at initialization. That is overfitting: improvements on the examples used for updates stop transferring to the check. This is why we keep validation separate. Better held-out scores still do not prove that a model never memorizes any training passages.
Save something you can use again
A checkpoint is a saved model snapshot. Ours stores the learned numbers, the model dimensions, and the character-to-ID mapping. The same number must still identify the same character when you reload it.
The recorded 4,000-update CPU run took 31 seconds on the machine used to make this lesson, using two CPU threads and about 303 MiB peak process memory. Your runtime can differ. A short --steps 50 run checks setup; it will not reproduce the trained sample.
Generation checkpoints versus resuming an interrupted training run
These checkpoints are sufficient for generation and inspection. They do not save the optimizer’s running statistics or the training random-generator state, so they are not an exact training-resume feature. The included trainer starts a new run.
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 train.py --steps 4000 --out my-run
python generate.py --checkpoint my-run/step-4000.ptThe 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.
Quick check
Training penalty falls, but held-out penalty rises. What should you conclude?
You now have a checkpoint that can be reloaded without training again. Next, change its prompt and sampling settings, then inspect the attention patterns it actually computes.
Your turn
Explain it in your own words.
In the repeated-window experiment, the training-window penalty became tiny while the held-out penalty got worse. What does this tell you about the improvement? Why must validation windows stay out of the parameter updates?
Answer the question in your own words. A short explanation is enough.
Your feedback
Work through a hint
Hint 1
Which examples were used to change the model’s numbers?
Hint 2
Would the validation text still check transfer if its errors also supplied parameter updates?
A worked explanation
It became very good at the repeated window, but worse on text excluded from training, so that improvement did not transfer. Training on the validation windows would make the check part of its practice material.
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.