Small language model

For Applied Machine Learning at Columbia, I implemented a small decoder-only Transformer¹ in PyTorch and trained it on Tiny Shakespeare² (about 1.1M characters). The attention, RMSNorm, rotary embeddings and blocks are written out from nn.Linear and nn.Embedding rather than PyTorch’s built-in Transformer layers.

Here’s what it wrote when I prompted it with ROMEO:, next to a better-known Romeo.

ROMEO:
How now,
More than the unplease of all which looks:
But for you, with much women, 'tis thus:
 your house remedy for her gossiping than
I written laid to be a sent to you.
You have been follow'd in his turning a liberty friar
butt, but on that seems shall be ready:
O, and you again, my daughter's breath!

Shepherd:
Whose it did not d

It gets the format right (speaker names, line breaks, punctuation), but the sentences don’t make sense.

Model

The model reads text as tokens, chunks a few characters long. Common chunks get a token of their own, and rarer words are built from pieces, the way you read a familiar word at a glance but sound out a new name syllable by syllable. GPT models do the same with far bigger vocabularies, and OpenAI’s tokenizer tool shows how they split any text you paste in. Each token then becomes a list of 128 numbers, a bit like a color described by its red, green and blue values, except with 128 channels whose meaning comes from training.

• BPE tokenizer³ with a 500-token vocabulary, trained with Hugging Face tokenizers (about 517k tokens on the corpus)
• 2 pre-norm blocks, 4 heads, 128-dim embeddings, GELU feed-forward 128 → 256 → 128
• RMSNorm⁴, rotary positional embeddings⁵, tied input and output embeddings
• 50-token context, 327,552 parameters
• AdamW (lr 3e-4, weight decay 0.01), cosine schedule, batch 64, dropout 0.1, gradient clipping at 1.0, 20 epochs

class TransformerBlock(nn.Module):
    def __init__(self, embed_dim, num_heads, ff_dim, dropout=0.1):
        super().__init__()
        self.norm1 = RMSNorm(embed_dim)
        self.attn  = CausalSelfAttention(embed_dim, num_heads, dropout)
        self.norm2 = RMSNorm(embed_dim)
        self.ff    = FeedForward(embed_dim, ff_dim, dropout)

    def forward(self, x):
        # pre-norm residual style used in llama
        x = x + self.attn(self.norm1(x))
        x = x + self.ff(self.norm2(x))
        return x

Next-token predictions

This is the trained model running in your browser, starting from ROMEO: and How are. Click a token to append it, or sample one at temperature 0.8. Underlines mark token boundaries.

After How are, it puts 41% on you. All six of its top guesses fit grammatically, and one of them, thy, is Early Modern English.

Training

Training loss fell from 3.36 to 2.68 over 20 epochs. Validation loss, measured on the last 20% of the text, reached 3.39 by epoch 2 and bottomed out at 3.375 at epoch 7. After that it stayed between 3.381 and 3.391, so I generated from the epoch-7 checkpoint.

Its validation perplexity is 29.22. Perplexity is the standard score for language models (OpenAI scored GPT-2 this way⁶), and 29 means the model is about as unsure of each next token as if it were rolling a 29-sided die. A model that only knows how common each token is, without reading the text before it, scores 246.37.

Validation loss stopped improving long before training loss did, meaning the model was fitting the training text without getting better on new text. A bigger model or a longer context would probably help more than more epochs. OpenAI’s scaling-laws work⁷ found that larger language models get more out of the same data. The 50-token context covers only about 100 characters, two or three lines of verse, while Claude’s context window holds up to a million tokens.⁸

Attention

Attention is how each position decides which earlier tokens matter for predicting the next one. If you’re filling in “I grabbed my keys and my ___”, you lean on “grabbed” and “keys” and mostly ignore “I” and “and”. Each row of the map below is one token doing that: its weights over the tokens before it, summing to 1. Later tokens are masked so it can’t see the answer. I saved these maps for all 8 heads at epochs 1, 10 and 20, on the same 20 validation tokens.

In all eight heads, most of the change happens between epochs 1 and 10, and the epoch 10 and 20 maps look alike. That matches validation loss flattening early.

Layer 2, head 3 (shown first) puts most of its weight on the current or previous token at every epoch, but the later t tokens change. At epoch 1, the t of “tender” puts 0.63 on the to right before it, and the t of the second “twenty” puts 0.74 on itself. By epoch 20, both put the most weight on an earlier t: 0.57 on the one in the first “twenty” and 0.67 on the one in “tender”.

Anthropic found that two-layer models like this one can learn induction heads⁹, which find an earlier copy of the current token and attend to the token right after it. This head attends to the copy itself, not the token after it.

Notebook on GitHub

References