Part III · How LLMs work · Chapter 20

Project: train your own tiny language model

Build two language models that write text one character at a time. The first is a bigram model in pure Python that counts which character follows which. The second is a tiny transformer in PyTorch with real attention. Both train on the same short text, so you can see exactly what attention adds.

What you need

The files are in labs/tiny-language-model/: the two scripts, requirements.txt and the training text data/savings.txt, about 6,600 characters of plain notes on saving money written for this book. Every command below runs from the repository root.

# from the repository root
python3 labs/tiny-language-model/bigram.py        # no install needed
python3 -m venv .venv && . .venv/bin/activate
pip install -r labs/tiny-language-model/requirements.txt   # PyTorch, for the transformer
python labs/tiny-language-model/tiny_transformer.py --steps 500

Part 1: the bigram model

A bigram model predicts the next character from the current one only. Training is counting: how often does each character follow each other one? With 50 distinct characters in the text, the model is a 50 × 50 table. The probability of b after a is

Why smooth? With α = 0, a pair never seen in training gets probability 0, and a single such pair in the validation text makes the loss −ln 0 = infinity. The script shows exactly that.

$ python3 labs/tiny-language-model/bigram.py --after t
text: 6586 characters, 50 distinct; training on 5927, validating on 659
uniform guess: loss ln 50 = 3.9120 nats, perplexity 50
bigram (alpha=0.1): training loss 2.3002, validation loss 2.4832, validation perplexity 11.98
most common pairs: 'e␣' 182, 's␣' 138, '␣t' 133, 'th' 129, 't␣' 120, '␣a' 119, 'he' 87, 'n␣' 84
after 't': h 0.318, ␣ 0.296, e 0.086, o 0.059, s 0.040, . 0.037, i 0.035, a 0.025

(␣ stands for a space.) Read it against earlier chapters. A uniform guess over 50 characters has loss ln 50 = 3.91 (chapter 9). Counting brings the validation loss to 2.48, a perplexity of about 12: as unsure as choosing among 12 characters instead of 50. After "t" the model expects "h" (as in "the") or a space. The training loss is lower than the validation loss, the usual gap (chapter 11).

Smoothing is regularisation

The run here gave these validation losses for different α:

More smoothing always raises the training loss and first lowers, then raises the validation loss: the U shape of chapter 11. α is a regularisation strength chosen on validation data, which is why the script's default is 0.1.

Sampling

Generation starts from "T", draws the next character from its row of probabilities, and repeats. The temperature works as in chapter 16. From the run here:

temperature 1:   Tis, merangsallimpust. the isttholor s th ls ft n yon eet morrewirendaladded ststrther...
temperature 0.5: The chen moun thatout the andithe th t th ns fond yor fes pleres s s r and marouroun t thisthe...

Word-like fragments, common pairs such as "th" and "he", spaces in sensible places, but no words beyond the shortest. A model that sees one character of context cannot do better. Lower temperatures repeat "the" and "th" more; higher ones produce capital letters and punctuation in odd places.

The same table, learned by gradient descent

With --learn the script also trains a 50 × 50 table of logits, W, by gradient descent: the next-character probabilities are softmax(W[current character]). This is a neural network with one layer and a one-hot input. Its gradient for each row is "predicted minus actual", summed over the pairs starting with that character (chapter 9).

$ python3 labs/tiny-language-model/bigram.py --learn
learning a 50 x 50 table of logits by gradient descent (lr 50, 300 steps):
  step     1  training loss 3.9120
  step    30  training loss 2.4735
  step   150  training loss 2.3267
  step   300  training loss 2.2999
  the counting model with alpha=0 has training loss 2.2733: the floor to aim for

(Some lines are left out.) It starts at the uniform guess and falls towards 2.2733, the loss of plain counting, which is the best any table can do on the training text. Counting and gradient descent reach the same model; gradient descent is the method that still works when the model is not a simple table.

Part 2: the tiny transformer

tiny_transformer.py keeps the task and the data, but each prediction may look back at up to 64 characters through attention. Its parts map one to one onto chapter 18:

The default model has width 64, 4 heads, 2 blocks and a 64-character context. Counting its weights as in chapter 18 gives 107,392 parameters (49,984 per block, plus the tables and the final layer norm). Training for 2,000 steps of 32 windows of 64 characters is 4,096,000 character positions, so by chapter 19 about 6 × 107,392 × 4,096,000 ≈ 2.6 × 10¹² operations. At 3 to 10 billion useful operations per second, a fair range for PyTorch on a laptop CPU, that is roughly 4 to 15 minutes. The script prints its own count when it starts, and the time at each evaluation.

What to look for

• The loss starts near ln 50 = 3.91 and should fall well below the bigram model's 2.48 validation loss, because the model can use up to 64 characters of context instead of one.
• The text is short, so a model this size can start to memorise it. If the validation loss stops falling while the training loss keeps dropping, that is overfitting (chapter 11). The script reports the step with the lowest validation loss.
• Samples should contain whole words and phrases from the text. With enough training, long stretches may be copied verbatim: memorisation, the same thing that makes large models repeat their training data.

Experiments

1. Context length. Run the transformer with --block 1. With one character of context, it can do no better than the bigram model. Then try 8, 32 and 64.
2. Bigger text. Pass any plain text you have the right to use with --text, ideally hundreds of thousands of characters. The gap between training and validation loss should shrink.
3. Bigger model. Try --layers 4 --width 128 on the short text and watch overfitting set in sooner. Then raise --dropout to 0.2.
4. No mask. Comment out the masked_fill line in CausalSelfAttention.forward. The training loss drops far lower, because every position except the last can now see the character it should predict; generation becomes nonsense. This is why the mask exists.
5. No scaling. Remove the division by math.sqrt(self.head_size) and compare the first few hundred steps.
6. Temperature and top-k. Sample with --temperature 0.5, 1.2 and --top-k 5.