Training a small GPT from zero

Transformers, lesson 17

Train and sample a tiny GPT in the browser, with every line in view.

This lesson puts the whole module into one small model. You need "Norm, attention, MLP, residual" and "Autoregressive generation". "Autograd you can step through", in Neural Networks, helps with the training step.

Part 1. One block, every part you've met

Here is a whole GPT, small enough to train inside this page. It reads names one character at a time and learns to write new ones. Its vocabulary is 27 characters: a to z, plus a full stop that marks where a name starts and ends. So "emma" goes in as .emma, and at every position the model is asked for the character one place ahead: emma..

Example 1.1

One forward pass over ".emma"

.emma
0513131
Five character IDs go in. The targets are the same text shifted one place: e, m, m, a, and the full stop that ends the name.
token row+position row=5 vectors of 32
Each ID picks a row of a 27 × 32 token table, and each slot adds a row of a 12 × 32 position table ("Looking up a token's vector", "Encoding word order and RoPE").
norm→2 heads of 16, masked→+ residual
Each position attends to itself and the positions before it. Two heads of 16 numbers each, scores divided by √16 = 4 ("Attention as a soft lookup" to "Heads in parallel").
norm→MLP: 32 → 128 → 32→+ residual
The MLP works on each position alone, and its change is added back ("Norm, attention, MLP, residual").
norm→32 × 27 grid→5 × 27 logits
27 logits at every position, and softmax turns each row into probabilities ("Logits and softmax over the vocabulary"). Five positions, five guesses to score.

Count the weights: the token table holds 27 × 32 = 864, the position table 12 × 32 = 384, attention 4 × 32 × 32 = 4,096, the MLP 2 × 32 × 128 = 8,192 and the output grid 32 × 27 = 864. That's 14,400 in all, and as in GPT-2 small, the MLP holds the most.

Here's the attention step on its own: one head reading .emm, with vectors of two numbers so you can follow every score. The first number says "vowel", the second "consonant".

One causal attention head, in Python

import math

def softmax(xs):
    e = [math.exp(x - max(xs)) for x in xs]
    return [v / sum(e) for v in e]

def times(x, W):  # row vector x times grid W
    return [sum(a * w for a, w in zip(x, col)) for col in zip(*W)]

tok = {".": [0.0, 0.0], "e": [1.0, 0.0], "m": [0.0, 1.0]}  # (vowel, consonant)
pos = [[0.0, 0.1], [0.1, 0.0], [0.0, -0.1], [-0.1, 0.0]]
Wq = [[0.0, 0.0], [2.0, 0.0]]  # a consonant asks for a vowel
Wk = [[1.0, 0.0], [0.0, 1.0]]
Wv = [[1.0, 0.0], [0.0, 1.0]]

text = ".emm"
x = [[a + b for a, b in zip(tok[c], pos[t])] for t, c in enumerate(text)]
q, k, v = [times(r, Wq) for r in x], [times(r, Wk) for r in x], [times(r, Wv) for r in x]
for t in range(len(text)):
    scores = [sum(a * b for a, b in zip(q[t], k[u])) / math.sqrt(2) for u in range(t + 1)]
    w = softmax(scores)  # only keys 0..t: the causal mask
    out = [sum(wu * v[u][j] for u, wu in enumerate(w)) for j in range(2)]
    print(text[t], [round(a, 2) for a in w], [round(a, 2) for a in out])
# . [1.0] [0.0, 0.1]
# e [0.5, 0.5] [0.55, 0.05]
# m [0.17, 0.67, 0.17] [0.74, 0.17]
# m [0.13, 0.62, 0.13, 0.11] [0.67, 0.25]

Both m's put most of their weight on the e. The query grid turns "I'm a consonant" into "find me a vowel". Nobody writes grids like that in the lab: training finds its own.

Your turn

Take attention out of the block and keep everything else. What can the model still learn?

Without attention, nothing moves information between positions. Each position sees its own character plus its slot's position vector, so the best it can do is a counting model with a sense of place, like the one in "A bigram language model". Attention is what lets the second m in .emm know there was an e.

Show the answer

B. Without attention, nothing moves information between positions. Each position sees its own character plus its slot's position vector, so the best it can do is a counting model with a sense of place, like the one in "A bigram language model". Attention is what lets the second m in .emm know there was an e.

Part 2. What the loss is measured against

Keep going with Pro

You've done the first part. Pro opens the rest of this lesson, and every other one.

Still to come

What the loss is measured against
Training it, one step at a time
Same recipe, more of everything

See Pro plans I have Pro: sign in

The last lesson in TransformersBack to the overviewAttention, blocks and generation Back to Transformers

Transformers overviewSpotted a mistake? Tell us.

17 · Training a small GPT from zero