Transformer from scratch

A small GPT-style language model built from scratch in PyTorch. The transformer internals (multi-head self-attention, causal masking, residual blocks and positional embeddings) and a byte-pair tokenizer are written by hand with basic operations rather than torch.nn.Transformer or an external tokenizer library, so the mechanics are visible, testable and easy to ablate.

It trains to predict the next token. The default corpus is the tiny Shakespeare dataset, and the model can use either a character vocabulary or a byte-pair one.

How the model works

1. Each token id becomes a vector through a token embedding, and its position is added through a separate position embedding.
2. The sequence passes through several transformer blocks. Each block applies, with a residual connection and pre-layer-normalisation, a self-attention layer and a small feed-forward network.
3. Attention is computed directly: one linear layer produces the query, key and value projections, which are reshaped into heads. The scaled dot-product scores are masked so a position cannot attend to later positions (this lower-triangular mask is what makes the model autoregressive), turned into weights with a softmax, and used to take a weighted sum of the values.
4. A final layer norm and a linear head produce a score for every token in the vocabulary. Training minimises the cross-entropy against the actual next token.

Tokenizers

Two tokenizers share the same interface:

• A character tokenizer: the vocabulary is the set of distinct characters.
• A byte-pair tokenizer trained from scratch: starting from the raw bytes, it repeatedly merges the most frequent adjacent pair into a new token until it reaches the target vocabulary size. This gives a smaller, subword-like vocabulary and shorter sequences.

Run it

python -m venv .venv && source .venv/bin/activate
pip install torch --index-url https://download.pytorch.org/whl/cpu
pip install -e ".[dev]"

python scripts/download_data.py              # download tiny Shakespeare
python -m chartransformer.train              # train and save a checkpoint
python -m chartransformer.generate           # sample text
python -m chartransformer.evaluate           # validation perplexity
python scripts/ablation.py                   # run the ablation study
pytest

Results

The baseline is a character-level model with 818K parameters (4 layers, 4 heads, embedding size 128, context length 64), trained for 2000 steps on CPU. It reaches a validation cross-entropy of 1.93, a perplexity of 6.89 (random guessing over 65 characters would be 65). A sample:

KING RICHUS:
Wh this to is raise stown the so the like:
Whome viste come in lecter all benistel of bereates this for him
To not man and of Eur man you that to beake.

LAM:
The grice thy to in as I will full mustion,
But thou cieful hour to ear tace a the cruch

The text is not real English at this size and budget, but the model clearly learns the structure of the plays: capitalised speaker names followed by a colon, line breaks, punctuation, and many common short words.

Ablations

All runs share the same architecture and 2000-step budget so they can be compared.

variant	val loss	val perplexity
character baseline	1.93	6.89
no positional embeddings	2.08	8.04
single attention head	1.91	6.76

• Positional embeddings matter. Removing them raises perplexity from 6.89 to 8.04: without positions, attention cannot tell tokens apart by order, so the model loses track of word order.
• At this size, multiple heads barely help. The single-head model is about the same as the four-head one (6.76 vs 6.89). Multiple heads tend to help more with a larger model and a longer context; here the capacity is so small that splitting it across heads makes little difference.
• Byte-pair tokenization shortens the sequences. The byte-pair tokenizer (vocab 512) turns the 1,115,394-character corpus into 568,210 tokens, about 1.96 characters per token, so each attention step covers roughly twice as much text. Its perplexity is not comparable to the character models, because a byte-pair token carries more information than a single character; bits-per-character would be the fair cross-tokenizer metric.

Tests

The unit tests cover the forward-pass shapes, that attention is genuinely causal (changing a later token does not change the logits at earlier positions), that the model runs with and without positional embeddings, and the character and byte-pair tokenizers (round-trip, token-count reduction, unicode, config save and load).

Notes and next steps

• A larger model and a longer context window.
• Bits-per-character evaluation to compare the tokenizers fairly.
• Learning-rate warmup and cosine decay.