Building a tiny language model from scratch

You do not need a GPU cluster or a billion parameters to understand how a language model works. You can build the important parts yourself, train the result on a modest laptop, and watch it progress from random characters to recognisable, if still nonsensical, language.

This series is aimed at developers who are comfortable reading code but are new to language-model concepts. We will begin with the problem each component solves, build the smallest working version, and then look under the hood at how the code produces the behaviour we want.

You can treat it as both a guided build and an experiment of your own. The complete code is public, deliberately small, and designed to be forked. Change the corpus, remove a component, alter the model size, or break an assumption and see what happens.

We will implement a decoder-only Transformer from ordinary PyTorch tensor operations without using a pretrained model, nn.Transformer, nn.MultiheadAttention, Hugging Face Trainer, or a hidden attention implementation.

PyTorch still provides tensors, automatic differentiation, optimisation, and device access. “From scratch” means building the language model rather than rebuilding the numerical computing stack beneath it.

What you will build

The finished character-level model is intentionally small enough to inspect:

Every part is introduced separately: tokenisation, shifted training targets, embeddings, positional information, scaled dot-product attention, causal masking, multiple heads, feed-forward layers, residual paths, LayerNorm, training, checkpoints, and autoregressive generation.

Here is the complete mental model we are working towards:

text
  -> character token IDs
  -> token embeddings + position information
  -> [causal attention -> feed-forward network] x 4
  -> one score for every possible next character
  -> sample one character, append it, and repeat

What the trained model actually produces

The retained run began from seeded random weights and processed 6,144,000 sampled characters in 1,500 optimiser updates. Held-out loss fell from 4.6709 to 2.1298. The measured loop took 84.46 seconds on an Apple M1 through PyTorch’s Metal backend; no discrete GPU or hosted training service was required.

I think that is genuinely exciting. From roughly 1.1 MB of Shakespeare and a model small enough to understand, recognisable dramatic structure emerges on an ordinary laptop. The result is not a useful writer, but it makes the underlying idea feel tangible rather than remote.

Here is a generation performed after loading the saved checkpoint in a new CPU process. The distinction between input and output is deliberate: the prompt is supplied by the reader; everything after it is sampled by the model.

KING:

Fid trely levew his weas bud thin as the as
And for it en cong on preite hee
Sthes ill the mavence I swher soot
And his, the ago I fore mor neve and.

LUCERIO:
Gord you ur labll ooneack my his con fopy, sour,
Dar and fand athe chind sthin ath some f
Their nlime e my foreer or sithe bean.

LUSENGOLA

This is not coherent English and it is not a chatbot. It is nevertheless real learned behaviour: the untrained model emitted mostly repetitive arbitrary characters, while the trained model learned spaces, punctuation, speaker labels, line breaks, word-like fragments, and longer-range formatting.

Get the code and make it yours

The complete experiment lives in the Tiny Transformer Course repository. You can fork it on GitHub, clone it with Git, or download the source as a ZIP file.

git clone https://github.com/steve-obrien/tiny-transformer-course.git
cd tiny-transformer-course

Each article links to the matching source and gives you something small to change. I would encourage you to run the original example first, confirm that you can reproduce its behaviour, and only then start experimenting. That makes surprising results far easier to reason about.

Follow the complete build

Each part of the implementation has its own focused article. Read them in order if you want to build the model alongside the code, or jump directly to the component you want to understand.

1. What a language model predicts
2. Preparing a small, reproducible corpus
3. Turning text into character tokens
4. Constructing inputs, targets, and batches
5. Training a bigram baseline
6. Reading Attention Is All You Need as an implementer
7. Token embeddings and positional encodings
8. Does king − man + woman equal queen?
9. Implementing scaled dot-product attention
10. Causal masking: preventing future leakage
11. Combining several attention heads
12. The position-wise feed-forward network
13. Residual connections, LayerNorm, and dropout
14. Assembling the complete decoder
15. Training from random weights
16. Reading the loss curve and generated samples
17. Reloading a checkpoint and generating text
18. Testing the architecture with an ablation
19. Trying other corpora and scaling the model

The shortest runnable path

After cloning the repository, install the locked environment and prepare the pinned corpus:

uv sync
.venv/bin/python lessons/02-corpus-and-tokenisation/prepare.py

Run the tests and train the complete model:

.venv/bin/pytest -p no:cacheprovider
.venv/bin/python lessons/11-language-model-training/train.py

Then load the checkpoint in a separate process and supply your own prompt:

.venv/bin/python lessons/12-generation/generate.py \
	--checkpoint output/transformer/checkpoint.pt \
	--prompt 'KING:'

The repository retains the exact corpus checksum, model configuration, seeds, loss measurements, environment, checkpoint hash, and unedited outputs from the published run. The model is deliberately not useful; the value is that every important operation remains visible enough to question, test, and change.