5. Training a bigram baseline
Prove the complete data, loss, optimisation, checkpoint, and generation pipeline before adding attention.
Before building attention, establish that the surrounding system can already learn something. A bigram model predicts the next character using only the current character.
Its entire parameter set is a 65 × 65 table:
self.token_transition_logits = nn.Embedding(vocabulary_size, vocabulary_size)
Looking up the current token ID returns 65 next-character logits. Cross-entropy compares those logits with the shifted target, and AdamW updates the table.
Why this baseline matters
The bigram exercises the real tokenizer, batch builder, loss, optimiser, validation split, checkpoint format, reload path, and sampling loop. If it cannot reduce loss, adding a Transformer only makes the fault harder to locate.
.venv/bin/python lessons/04-bigram-baseline/train.py
The retained model has 4,225 parameters. Over 1,000 updates, validation loss fell from 4.7263 to 2.4902.
ROMEO:
Fida rdly le swen's weasin d th heas tin and my f poit en's hisod pee'tonore
Sththerirot t ma suchiswswher sootrley oou, the,
The output has plausible adjacent characters but no stable words or sentence structure. That is the expected ceiling: after the current character, the model has forgotten everything else.
Reload the checkpoint independently:
.venv/bin/python lessons/04-bigram-baseline/generate.py \
--checkpoint output/bigram/checkpoint.pt \
--prompt 'ROMEO:'
The Transformer must do more than reduce loss. It must beat this baseline by using a longer context. The final model reaches validation loss 2.1298, showing that attention and deeper computation provide useful information beyond the current character.