Build a tiny GPT from scratch

The big idea

Every lesson in this module described one part of a language model. Here you build the whole thing, small enough to read in one sitting and train in about a second on a laptop: a character-levelA language model whose tokens are single characters, so the vocabulary is tiny and no tokenizer training is needed. Common for teaching. GPT with one transformer block, written in plain NumPy. No deep learning framework, so nothing is hidden: the forward pass, every gradient and the optimizer are all on the page.

The model is tiny and the training text is a few nursery-rhyme sentences, so it will not write essays. But the parts are the real parts. Scale the same recipe up in width, depth, data and compute, swap characters for BPE tokens, and you have the architecture behind modern language models.

Think of it as building a working model engine from a kit. It will never power a car, but every piston, valve and spark plug is the real shape, and after you have assembled one, a full engine stops being mysterious.

How it works

The full program is in the Try it section. Here it is in six steps, with the key lines from it.

Step 1: data and tokenization

The vocabulary is just the set of characters in the text: 15 of them (letters, a space and a full stop). Each character maps to an integer id. Real models use BPE tokens ( Tokenization ) with vocabularies in the tens of thousands, but the idea is identical.

chars = sorted(set(text)); V = len(chars)
stoi = {c: i for i, c in enumerate(chars)}
data = np.array([stoi[c] for c in text])

Training examples are random windows of 16 characters. The target for each position is simply the next character, so one window gives 16 training examples at once:

def batch():
    ix = rng.integers(0, len(data) - T - 1, B)
    x = np.stack([data[i:i + T] for i in ix])
    return x, np.stack([data[i + 1:i + T + 1] for i in ix])

Step 2: embeddings

Each id looks up a learned 32-number vector in the embedding table E, and each position 0 to 15 looks up a learned vector in pos. Adding them gives a vector that says both "which character" and "where" ( Embeddings and positional encoding ). All weights start as small random numbers.

Step 3: causal self-attention

Each position makes a query, key and value. Scores are query-key dot products scaled by $\sqrt{d}$; the causal maskA mask that blocks each position from attending to later positions, so a model trained on next-token prediction cannot see the answer it is asked to predict. sets scores for future positions to a huge negative number so softmax gives them zero weight. Without it the model could read the very character it is supposed to predict, and training would be meaningless.

q, k, v = X @ P["Wq"], X @ P["Wk"], X @ P["Wv"]
s = q @ k.transpose(0, 2, 1) / np.sqrt(d)
s = np.where(mask[:t, :t], -1e9, s)
a = np.exp(s - s.max(-1, keepdims=True))
a /= a.sum(-1, keepdims=True)
h = a @ v
x1 = X + h @ P["Wo"]                          # residual around attention

Step 4: the MLP and the output

A two-layer MLP with a ReLU, four times wider in the middle, processes each position on its own, again with a residual connectionA connection that adds a layer's input to its output, so the layer learns a change to its input rather than a replacement. Makes deep networks trainable.. A final linear layer turns each position's vector into 15 logits, one per character.

pre = x1 @ P["W1"] + P["b1"]; m = np.maximum(pre, 0)
x2 = x1 + m @ P["W2"] + P["b2"]               # residual around the MLP
logits = x2 @ P["Wout"] + P["bout"]

Step 5: loss and gradients

Softmax over the logits gives a probability for each character; the loss is the average negative log probability of the true next character (cross-entropy, Pretraining: next-token prediction ). An untrained model guesses uniformly, so the first loss should be close to $\ln 15\approx 2.71$. That is a free sanity check: if your first loss is far from $\ln V$, something is wrong.

The gradient of cross-entropy with respect to the logits is beautifully simple: predicted probabilities minus a one-hot vector of the truth. From there, the chain rule runs backwards through every layer. The trickiest line is the softmax inside attention, whose gradient is $a\odot (da-\sum a\odot da)$:

da, dv = dh @ v.transpose(0, 2, 1), a.transpose(0, 2, 1) @ dh
ds = a * (da - (da * a).sum(-1, keepdims=True)) / np.sqrt(d)
dq, dk = ds @ k, ds.transpose(0, 2, 1) @ q

Hand-written gradients are easy to get wrong, so the program checks one against a finite differenceChecking a hand-derived gradient by nudging a weight up and down by a tiny amount and measuring how the loss changes.: nudge one weight up and down by $10^{-5}$ and measure the change in loss. Frameworks such as PyTorch compute all of this automatically with autograd, which is why nobody writes backprop by hand in practice.

Step 6: training and sampling

The optimizer is Adam with a small weight decay, the same idea as AdamW: keep running averages of each gradient and its square, and step each weight by the ratio. Then generation is the decoding loop from Decoding and sampling : feed the last 16 characters, take the logits for the final position, divide by the temperature, sample, append, repeat.

logits, _ = forward(P, np.array([idx[-T:]]))
z = logits[0, -1] / temp                  # temperature sampling
pr = np.exp(z - z.max()); pr /= pr.sum()
idx.append(rng.choice(V, p=pr))

Notice that this loop recomputes attention over the whole window at every step. A real implementation keeps a KV cache instead ( Context windows and the KV cache ).

Try it

The complete program, about 100 lines. Copy it and run it with Python and NumPy; it finishes in a couple of seconds.

# A tiny character-level GPT in pure NumPy: one block, trained by hand.
import numpy as np

text = ("the cat sat on the mat. the dog sat on the log. "
        "the cat saw the dog. the dog saw the cat on the mat. ") * 4
chars = sorted(set(text)); V = len(chars)
stoi = {c: i for i, c in enumerate(chars)}
data = np.array([stoi[c] for c in text])

rng = np.random.default_rng(0)
d, T, B, steps, lr = 32, 16, 16, 600, 3e-3
def init(*shape):
    return rng.normal(size=shape) * 0.02 if len(shape) > 1 else np.zeros(shape)
P = {"E": init(V, d), "pos": init(T, d), "Wq": init(d, d), "Wk": init(d, d),
     "Wv": init(d, d), "Wo": init(d, d), "W1": init(d, 4 * d),
     "b1": init(4 * d), "W2": init(4 * d, d), "b2": init(d),
     "Wout": init(d, V), "bout": init(V)}
mask = np.triu(np.ones((T, T)), 1).astype(bool)   # causal: no peeking ahead

def forward(P, idx):
    t = idx.shape[1]
    X = P["E"][idx] + P["pos"][:t]                # token + position embeddings
    q, k, v = X @ P["Wq"], X @ P["Wk"], X @ P["Wv"]
    s = q @ k.transpose(0, 2, 1) / np.sqrt(d)
    s = np.where(mask[:t, :t], -1e9, s)
    a = np.exp(s - s.max(-1, keepdims=True))
    a /= a.sum(-1, keepdims=True)
    h = a @ v
    x1 = X + h @ P["Wo"]                          # residual around attention
    pre = x1 @ P["W1"] + P["b1"]; m = np.maximum(pre, 0)
    x2 = x1 + m @ P["W2"] + P["b2"]               # residual around the MLP
    logits = x2 @ P["Wout"] + P["bout"]
    return logits, (X, q, k, v, a, h, x1, pre, m, x2)

def loss_and_grads(P, idx, tgt):
    logits, (X, q, k, v, a, h, x1, pre, m, x2) = forward(P, idx)
    p = np.exp(logits - logits.max(-1, keepdims=True))
    p /= p.sum(-1, keepdims=True)
    n = idx.size
    loss = -np.log(p.reshape(n, V)[np.arange(n), tgt.ravel()]).mean()
    dl = p.copy()                                 # dLoss/dlogits = p - onehot
    dl.reshape(n, V)[np.arange(n), tgt.ravel()] -= 1
    dl /= n
    G = {}
    f = lambda A: A.reshape(-1, A.shape[-1])      # flatten batch and time
    G["Wout"], G["bout"] = f(x2).T @ f(dl), dl.sum((0, 1))
    dx2 = dl @ P["Wout"].T
    G["W2"], G["b2"] = f(m).T @ f(dx2), dx2.sum((0, 1))
    dpre = (dx2 @ P["W2"].T) * (pre > 0)
    G["W1"], G["b1"] = f(x1).T @ f(dpre), dpre.sum((0, 1))
    dx1 = dx2 + dpre @ P["W1"].T
    G["Wo"] = f(h).T @ f(dx1)
    dh = dx1 @ P["Wo"].T
    da, dv = dh @ v.transpose(0, 2, 1), a.transpose(0, 2, 1) @ dh
    ds = a * (da - (da * a).sum(-1, keepdims=True)) / np.sqrt(d)
    dq, dk = ds @ k, ds.transpose(0, 2, 1) @ q
    G["Wq"], G["Wk"], G["Wv"] = (f(X).T @ f(g) for g in (dq, dk, dv))
    dX = dx1 + dq @ P["Wq"].T + dk @ P["Wk"].T + dv @ P["Wv"].T
    G["pos"] = dX.sum(0)
    G["E"] = np.zeros_like(P["E"]); np.add.at(G["E"], idx, dX)
    return loss, G

def batch():
    ix = rng.integers(0, len(data) - T - 1, B)
    x = np.stack([data[i:i + T] for i in ix])
    return x, np.stack([data[i + 1:i + T + 1] for i in ix])

# Train with AdamW-style updates (Adam + small weight decay).
mom = {k_: np.zeros_like(w) for k_, w in P.items()}
vel = {k_: np.zeros_like(w) for k_, w in P.items()}
print(f"vocab {V} chars, start loss should be near ln({V}) = {np.log(V):.2f}")
for step in range(1, steps + 1):
    xb, yb = batch()
    loss, G = loss_and_grads(P, xb, yb)
    for k_ in P:
        mom[k_] = 0.9 * mom[k_] + 0.1 * G[k_]
        vel[k_] = 0.999 * vel[k_] + 0.001 * G[k_] ** 2
        mh, vh = mom[k_] / (1 - 0.9 ** step), vel[k_] / (1 - 0.999 ** step)
        P[k_] -= lr * (mh / (np.sqrt(vh) + 1e-8) + 0.01 * P[k_])
    if step in (1, 100, 300, 600):
        print(f"step {step:4d}  loss {loss:.3f}")

# Gradient check: backprop vs a finite difference on one weight.
xb, yb = batch(); _, G = loss_and_grads(P, xb, yb); e = 1e-5
P["W1"][0, 0] += e; lp, _ = loss_and_grads(P, xb, yb)
P["W1"][0, 0] -= 2 * e; lm, _ = loss_and_grads(P, xb, yb); P["W1"][0, 0] += e
print(f"grad check: backprop {G['W1'][0, 0]:.6f}, "
      f"numeric {(lp - lm) / (2 * e):.6f}")

def generate(start, n, temp):
    idx = [stoi[c] for c in start]
    for _ in range(n):
        logits, _ = forward(P, np.array([idx[-T:]]))
        z = logits[0, -1] / temp                  # temperature sampling
        pr = np.exp(z - z.max()); pr /= pr.sum()
        idx.append(rng.choice(V, p=pr))
    return "".join(chars[i] for i in idx)

for temp in (0.5, 1.0):
    print(f"T={temp}: {generate('the ', 60, temp)!r}")

vocab 15 chars, start loss should be near ln(15) = 2.71
step    1  loss 2.707
step  100  loss 0.714
step  300  loss 0.350
step  600  loss 0.291
grad check: backprop 0.000383, numeric 0.000383
T=0.5: 'the dog. the cat on the cat on sat on the cat on the dog. the do'
T=1.0: 'the mat. the cat sat on the log. son the dog. s the cat on the c'

Read the output from top to bottom:

• The first loss is 2.707, right at $\ln 15=2.71$: the untrained model knows nothing.
• After 100 steps the loss is 0.714, after 600 it is 0.291. The model has learned which characters follow which, given up to 16 characters of context.
• The gradient check agrees to six decimal places (0.000383 both ways), so the backward pass is right.
• At temperature 0.5 the samples are made of real words and sentence patterns from the training text, recombined ("the cat on the cat on sat on the cat"). At temperature 1.0 they are more varied and include slips such as "son the dog".

Things to try on your own copy:

1. Delete the line that applies the causal mask and retrain. Predict first: what will happen to the training loss, and to the samples? (The model can now see the answer during training but not during generation.)
2. Remove the position embeddings (+ P["pos"][:t]). Attention by itself cannot tell positions apart. How does that change what the model can learn?
3. Replace the text with a paragraph of your own, raise steps, and watch how the first loss changes with the vocabulary size.

Going deeper

What this model leaves out

A production GPT has the same skeleton plus:

• Layer normalisation before attention and before the MLP ( The transformer block ), which keeps activations in a stable range in deep stacks.
• Multiple heads ( Multi-head attention ) and many blocks, dozens of them, each with its own weights.
• A subword tokenizer, so each token carries more text and the vocabulary has tens of thousands of entries.
• Dropout or other regularisation, and a held-out validation set to detect overfitting. This lesson's model trains on a tiny repetitive text and mostly memorises it; the loss above is training loss only.
• A GPU and a framework with autograd, mixed precision and efficient attention kernels.

The same block in PyTorch

For comparison, here is one pre-norm block in PyTorch, as found in typical small GPT codebases. It needs PyTorch installed and is not run here.

import torch, torch.nn as nn

class Block(nn.Module):
    def __init__(self, d, n_head):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, n_head, batch_first=True)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, x):
        T = x.size(1)
        mask = torch.triu(torch.ones(T, T, dtype=torch.bool, device=x.device), 1)
        h = self.ln1(x)
        x = x + self.attn(h, h, h, attn_mask=mask, need_weights=False)[0]
        return x + self.mlp(self.ln2(x))

Stack a few of these between an embedding layer and a final linear layer, compute F.cross_entropy on the logits, call loss.backward() and step torch.optim.AdamW. That replaces the whole hand-written backward pass above.

Debugging habits that transfer

• Check that the first loss is close to $\ln V$.
• Overfit a single small batch first: if the loss cannot go near zero on 16 examples, there is a bug.
• Gradient-check anything you wrote by hand.
• Sample during training, not just at the end; the samples show what the loss number hides.

Trade-offs

Common mistakes

• Forgetting the causal mask. Training loss drops suspiciously fast, then generation produces garbage.
• Off-by-one targets: the target sequence must be the input shifted by exactly one position.
• Not scaling attention scores by $\sqrt{d}$, which pushes softmax into saturation as $d$ grows.
• Reporting training loss as if it measured generalisation. Hold out some text and measure loss on it.
• Feeding generation a context longer than the position table. Keep only the last T tokens, as the sampling loop does.

In the wild

Check yourself

Recap