Capstone: A Tiny Attention-Based Language Model

Predict the next character of Shakespeare from the 64 before it. Self-attention lets every character choose which earlier ones to read, and two small transformer blocks, built on your own engine, learn to write passable fake Shakespeare.

55 min Prerequisites: Chapters 6, 9, 10, 12 and 15 (tensor.py and optim.py)

Inspect Architecture: Transformer

1. The idea

Every network so far has mapped a fixed input to an answer. Language comes as a sequence, and the classic task is to predict what comes next. This chapter does it one character at a time, on 1.1 million characters of Shakespeare's plays. There are 65 different characters, so guessing at random scores a loss of ln 65 = 4.174. A table that counts which character follows which does much better, 2.482 on the last tenth of the text, which is held back for validation. But a table can only look one character back. Finishing the word in "What is the time o' the da" needs the letters just before; knowing a speech has begun needs a name and a colon from a line ago.

Self-attention lets each position choose what to read. Every character computes three vectors from its own embedding: a query (what am I looking for?), a key (what do I contain?) and a value (what do I pass on?). Each query is compared with the keys of all earlier characters by a dot product, a softmax turns the scores into weights that sum to 1, and the character receives the weighted average of their values. A mask hides the later characters, so the model can't cheat by reading the answer. As in Chapter 15, the same weights are used at every position. Four heads do this in parallel, each with its own weights. Wrap it with a small MLP, residual connections and Chapter 12's layer norm and you have a transformer block, the same design as the GPT models, only tiny.

The engine needs four more operations: reshape, transpose, a batched matrix product (one product per sequence and head) and softmax as a single operation with its own backward rule. A gradient check on a small masked attention calculation agrees to eight digits. Everything else is already there: embeddings are one-hot vectors times a table, the loss is Chapter 9's cross-entropy, and the training loop is Chapter 15's, gc.collect() included.

The model has two blocks and 112,193 knobs. After 3,000 steps of Adam, about 13 minutes here, its validation loss is 1.761: on average it is about as unsure as a choice between 6 characters, where the counting table is unsure between 12. It writes speaker names in capitals followed by a colon, breaks lines like verse and spells most short words correctly, though it understands nothing. Looking inside shows how it works: one head always reads the previous character, another the one two back, and a head in the second block finds where the current word began. Large language models use this same recipe, scaled up to more data, subword tokens, longer contexts and many more blocks.

Two triangular heatmaps for the line "What is the time o' the day?", one row and one column per character. In the first, from block 1, head 3, each row has a single dark cell just left of the diagonal: every character looks at the character before it. In the second, from block 2, head 2, the dark cells form short vertical bars: the letters of each word mostly look back at the space before that word. — Attention weights of the trained model on a line from the validation text. Nobody told the heads what to look for: the first learned a fixed offset, the second where the current word began.

2. The math

The task: the next character

loss = average of −ln p(next character | the ≤ 64 before it)
guess uniformly:            ln 65 = 4.174
count pairs (bigram table): 2.482 on validation

Every position in a window is a training example: one window of 64 characters gives 64 predictions, each seeing only the characters before it. The bigram table is about the best a model can do if it sees only the current character.

Embeddings: characters and positions

x_t = tok[c_t] + pos[t]          (64 numbers each)
tok[c] = one-hot(c) · tok         (65 × 64 table)
pos: one learned vector for each of the 64 positions

Looking up a row is a matrix product with a one-hot vector, so the gradient reaches the table with no new rule. Position vectors are needed because attention by itself ignores order: shuffling the earlier characters wouldn't change their weighted average.

One attention head

q_i = W_q x_i,  k_j = W_k x_j,  v_j = W_v x_j   (16 each)
score_ij = q_i · k_j / √16
weights_i = softmax_j(score_ij),  out_i = Σ_j weights_ij v_j

Character i scores every character j by how well its query matches j's key, then takes a weighted average of the values. W_q, W_k and W_v are the same at every position. For a whole batch this is two batched matrix products: Q Kᵀ for the scores, then weights · V.

The causal mask

score_ij = −∞  for j > i     (later characters)
exp(−∞) = 0:  their weights are exactly 0
the remaining weights still sum to 1

With the mask, one forward pass trains all 64 predictions in a window at once, and none of them can see its own answer. The first character can only look at itself. The code prints a 5 × 5 mask and the weights it allows.

Why divide by √16

q · k = q₁k₁ + … + q₁₆k₁₆,   entries ~ mean 0, variance 1
Var(q · k) = 16,   spread 4
score = q · k / 4  →  spread 1

Without the scaling the scores have a spread of 4, so differences of 8 or more between positions are common. The softmax then puts nearly all the weight on one position, where its slopes are almost 0 (Chapter 9). Derived in the second derivation.

Four heads, and what they cost

4 heads × 16 numbers = 64:  outputs side by side, then W_o
scores per layer: 32 windows × 4 heads × 64 × 64 = 524,288
context of 128 instead of 64:  4 times as many

Each head can learn a different pattern. The score array grows with the square of the context length, which is why long contexts are expensive and why research on cheaper attention never stops.

A transformer block, and the whole model

x ← x + attention(layer_norm(x))
x ← x + W₂ relu(W₁ layer_norm(x))     (64 → 256 → 64)
knobs: 4,160 + 4,096 + 2 × 49,792 + 4,353 = 112,193

Residual connections add each sublayer's output instead of replacing x, so the signal and its gradient have a direct path through the network. The terms are the character table, the position table, two blocks and the final layer norm with the output layer.

Training

Adam η = 0.003, 32 windows of 64 per step (about 0.27 s here)
validation loss:  4.644 at the start (random knobs)
2.090, 1.938, 1.863, 1.817, 1.785, 1.761  after 500 … 3,000 steps

The start is worse than guessing uniformly, as in Chapter 12: random knobs are confidently wrong. After 500 steps the model already beats the bigram table, because it can read more than one character back.

What it predicts and writes

after 'To be, or not to b':  'e' 65%,  'r' 8%,  'l' 8%
perplexity e^loss:  65 (uniform),  11.97 (bigram),  5.82 (model)
sample at temperature 0.8, after 'JULIET:':
    In your matter net him would will state
    To his send comentle I to entreath.

Perplexity turns a loss back into 'as unsure as choosing uniformly among this many characters'. Sampling divides the scores by a temperature before the softmax: below 1 it plays safer, above 1 it takes more risks.

What the heads learned

block 1, head 3:  1 character back at 100% of positions
block 1, head 1:  2 back at 94%
block 1, heads 2 and 4:  spread over the last 3
block 2:  wherever the content says, e.g. the start of the word

The first block learned fixed offsets, which together give each position a view of its last few characters, like a counting table with a longer memory. The second block's heads choose positions by content. The figure shows one of each.

• Let s = softmax(x), so s_i = e^(x_i) / Σ_k e^(x_k). Differentiate s_i with respect to x_j. If j = i, the quotient rule gives s_i − s_i² = s_i(1 − s_i). If j ≠ i, only the denominator depends on x_j, giving −s_i s_j. Together: ∂s_i/∂x_j = s_i(δ_ij − s_j), where δ_ij is 1 if i = j and 0 otherwise.

• The backward pass receives g = ∂L/∂s and needs ∂L/∂x_j = Σ_i g_i ∂s_i/∂x_j = Σ_i g_i s_i(δ_ij − s_j) = s_j g_j − s_j Σ_i g_i s_i.

• So ∂L/∂x = s · (g − Σ_i s_i g_i), entry by entry, with one sum per row. That's the single line in softmax's _backward. ∎

• Chapter 9 combined softmax with cross-entropy and got the shortcut p − y. Attention has no cross-entropy right after its softmax, so it needs the general rule.

• Masked positions have s_j = 0, so their gradient is 0 too: the loss never pushes on scores the model wasn't allowed to use.

• Suppose the entries of q and k are independent with mean 0 and variance 1, roughly what layer norm and a sensible initialization give. Each product q_m k_m then has mean 0 and variance E[q_m²] E[k_m²] = 1.

• A dot product of d = 16 such products is a sum of 16 independent terms, and variances of independent terms add: Var(q · k) = 16, so the spread is √16 = 4.

• Dividing by √d gives the scores variance 16 / 16 = 1, whatever d is. ∎

• Why it matters: softmax of scores with spread 4 is often almost one-hot, with weights like 0.98 and 0.001. Its slopes s_i(δ_ij − s_j) are then nearly 0 (first derivation), so the attention weights barely learn.

• Larger models use wider heads (64 or 128 numbers), where the unscaled spread would be 8 or 11. The √d in 'scaled dot-product attention' is there for exactly this reason.

3. How it works

Turn text into numbers

List the distinct characters, give each a number, and split off the end of the text for validation. Measure the uniform and bigram baselines first, so you know what 'good' means.

Embed characters and positions

Look up a learned vector for each character and add a learned vector for its position. Without positions, attention can't tell 'tops' from 'pots'.

Attend, causally

Compute queries, keys and values; score with Q Kᵀ / √d; set the scores of later positions to −∞; softmax; average the values. Use several heads and combine them with one more matrix.

Stack blocks

Each block is attention then an MLP, each with layer norm in front and a residual connection around it. Finish with a layer norm and a dense layer to one score per character.

Train on every position, then sample

Shift the window by one character to get the targets and average the cross-entropy over all positions. To generate, feed the last 64 characters, sample the next one from the softmax, append it and repeat.

4. The code (python)

core_ch16.py

# Chapter 16 (capstone): a tiny attention-based language model that writes Shakespeare-like text, one character at a time.
# Uses tensor.py (Chapter 6) and optim.py (Chapter 10's Adam). Downloads about 1.1 MB of text.
import gc
import hashlib
import os
import time
import urllib.request
import numpy as np
from tensor import Tensor
from optim import Adam

# === 1. The text, as numbers ========================================================
URL = ("https://raw.githubusercontent.com/karpathy/char-rnn/370cbcd448eb7daf32f21a6be560b70e0b33c4e3/"
       "data/tinyshakespeare/input.txt")
if not os.path.exists("shakespeare.txt"):
    urllib.request.urlretrieve(URL, "shakespeare.txt")
with open("shakespeare.txt", "rb") as f:
    raw = f.read()
assert hashlib.sha256(raw).hexdigest() == "86c4e6aa9db7c042ec79f339dcb96d42b0075e16b8fc2e86bf0ca57e2dc565ed"
text = raw.decode("utf-8")

chars = sorted(set(text))                          # the vocabulary: every character that occurs
V = len(chars)
to_id = {c: i for i, c in enumerate(chars)}
encode = lambda s: np.array([to_id[c] for c in s])
decode = lambda ids: "".join(chars[i] for i in ids)
data = encode(text)
split = int(0.9 * len(data))
train, val = data[:split], data[split:]            # the last 10% of the plays is never trained on
print(f"{len(text):,} characters, {V} different ones: {''.join(chars)!r}")
print(f"'First Citizen' as numbers: {encode('First Citizen')}")

# Two baselines. Guessing uniformly costs ln V; counting which character follows which does better.
counts = np.ones((V, V))                           # +1 so no pair has probability 0
np.add.at(counts, (train[:-1], train[1:]), 1)
probs = counts / counts.sum(axis=1, keepdims=True)
print(f"\nloss of guessing uniformly: ln {V} = {np.log(V):.3f}")
print(f"loss of the bigram table (next character from the current one only): "
      f"{-np.mean(np.log(probs[val[:-1], val[1:]])):.3f} on validation")

# === 2. Four new operations, each with its own backward rule ======================
def reshape(x, shape):
    out = Tensor(x.data.reshape(shape), (x,), "reshape")

    def _backward():
        x.grad += out.grad.reshape(x.shape)
    out._backward = _backward
    return out


def transpose(x, axes):
    out = Tensor(x.data.transpose(axes), (x,), "transpose")

    def _backward():                                   # undo the reordering
        x.grad += out.grad.transpose(np.argsort(axes))
    out._backward = _backward
    return out


def bmm(a, b):
    """A matrix product for every (sequence, head) pair at once: (..., n, k) @ (..., k, m)."""
    out = Tensor(a.data @ b.data, (a, b), "bmm")

    def _backward():                                   # Chapter 6's rule, one matrix pair at a time
        a.grad += out.grad @ b.data.swapaxes(-1, -2)
        b.grad += a.data.swapaxes(-1, -2) @ out.grad
    out._backward = _backward
    return out


def softmax(x):
    """Softmax along the last axis (Chapter 9), as one operation."""
    e = np.exp(x.data - x.data.max(axis=-1, keepdims=True))
    s = e / e.sum(axis=-1, keepdims=True)
    out = Tensor(s, (x,), "softmax")

    def _backward():                                   # ds/dx applied to g:  s · (g − Σ s·g)
        g = out.grad
        x.grad += s * (g - (s * g).sum(axis=-1, keepdims=True))
    out._backward = _backward
    return out


def layer_norm(h, gain, shift, eps=1e-5):            # Chapter 12, one row per character
    average = Tensor(np.ones((h.shape[1], 1)) / h.shape[1])
    centred = h - h @ average
    var = (centred * centred) @ average
    return centred * (var + eps) ** -0.5 * gain + shift

# Gradient check (Chapter 4's habit) on a small attention calculation
r = np.random.default_rng(0)
q, k, v = (Tensor(r.normal(size=(2, 3, 5, 4))) for _ in range(3))
mask = np.triu(np.full((5, 5), -np.inf), 1)
probe = Tensor(r.normal(size=(2, 5, 12)))
f = lambda: (reshape(transpose(bmm(softmax(bmm(q, transpose(k, (0, 1, 3, 2))) / 2 + mask), v),
                               (0, 2, 1, 3)), (2, 5, 12)) * probe).sum()
f().backward()
worst = 0.0
for t, idx in [(q, (0, 1, 2, 3)), (q, (1, 2, 4, 0)), (k, (0, 0, 1, 2)), (k, (1, 2, 3, 1)), (v, (1, 1, 0, 3))]:
    old = t.data[idx]
    t.data[idx] = old + 1e-6
    up = float(f().data)
    t.data[idx] = old - 1e-6
    down = float(f().data)
    t.data[idx] = old
    numeric = (up - down) / 2e-6
    worst = max(worst, abs(numeric - t.grad[idx]) / (abs(numeric) + abs(t.grad[idx])))
print(f"\ngradient check of transpose, bmm, masked softmax and reshape: worst relative error {worst:.1e}")

# === 3. Self-attention: every character looks back at the ones before it ===========
T, D, H = 64, 64, 4                                   # context of 64 characters, 64 numbers per character, 4 heads
DH = D // H                                           # 16 numbers per head
MASK = np.triu(np.full((T, T), -np.inf), 1)            # −∞ above the diagonal: no peeking at later characters

print(f"\na causal mask for 5 characters (0 = allowed, -inf = hidden):\n{MASK[:5, :5]}")
w = softmax(Tensor(r.normal(size=(5, 5)) + MASK[:5, :5])).data
print(f"attention weights from random scores (each row sums to 1):\n{w.round(2)}")


def attention(x, p, B, n):
    """x is (B·n, D) for B sequences of n ≤ T characters. Each head scores every earlier position,
    then takes a weighted average of their values."""
    def heads(t):                                     # (B·n, D) → (B, H, n, DH)
        return transpose(reshape(t, (B, n, H, DH)), (0, 2, 1, 3))
    q, k, v = heads(x @ p["Wq"].T), heads(x @ p["Wk"].T), heads(x @ p["Wv"].T)
    scores = bmm(q, transpose(k, (0, 1, 3, 2))) / np.sqrt(DH) + MASK[:n, :n]   # (B, H, n, n)
    weights = softmax(scores)
    mixed = reshape(transpose(bmm(weights, v), (0, 2, 1, 3)), (B * n, D))  # heads side by side again
    return mixed @ p["Wo"].T + p["bo"], weights

# === 4. The model: embeddings, two transformer blocks, and a softmax over 65 characters
rng = np.random.default_rng(0)
LAYERS = 2


def dense(n_out, n_in, scale=1.0):
    return Tensor(rng.normal(0, scale * np.sqrt(1 / n_in), (n_out, n_in)))

params = {"tok": Tensor(rng.normal(0, 1, (V, D))),    # one learned vector per character
          "pos": Tensor(rng.normal(0, 1, (T, D)))}    # and one per position: attention alone ignores order
for l in range(LAYERS):
    params |= {f"{l}.{name}": t for name, t in {
        "g1": Tensor(np.ones(D)), "s1": Tensor(np.zeros(D)),
        "Wq": dense(D, D), "Wk": dense(D, D), "Wv": dense(D, D),
        "Wo": dense(D, D, 1 / np.sqrt(2 * LAYERS)), "bo": Tensor(np.zeros(D)),
        "g2": Tensor(np.ones(D)), "s2": Tensor(np.zeros(D)),
        "W1": dense(4 * D, D, np.sqrt(2)), "b1": Tensor(np.zeros(4 * D)),
        "W2": dense(D, 4 * D, 1 / np.sqrt(2 * LAYERS)), "b2": Tensor(np.zeros(D))}.items()}
params |= {"gf": Tensor(np.ones(D)), "sf": Tensor(np.zeros(D)), "Wout": dense(V, D), "bout": Tensor(np.zeros(V))}
print(f"\nmodel: {LAYERS} blocks of 4-head attention + MLP, {sum(p.data.size for p in params.values()):,} knobs")


def forward(ids):
    """ids is (B, n) character numbers; returns (B·n, V) scores for the next character, and the attention weights."""
    B, n = ids.shape
    x = Tensor(np.eye(V)[ids.reshape(-1)]) @ params["tok"]           # look up each character's vector
    pos = Tensor(np.eye(T)[:n]) @ params["pos"]                        # the first n position vectors
    x = reshape(reshape(x, (B, n, D)) + pos, (B * n, D))               # added to the characters' vectors
    maps = []
    for l in range(LAYERS):
        p = {name.split(".")[1]: t for name, t in params.items() if name.startswith(f"{l}.")}
        a, w = attention(layer_norm(x, p["g1"], p["s1"]), p, B, n)
        x = x + a                                                       # residual: add, don't replace
        h = (layer_norm(x, p["g2"], p["s2"]) @ p["W1"].T + p["b1"]).relu()
        x = x + h @ p["W2"].T + p["b2"]
        maps.append(w.data)
    return layer_norm(x, params["gf"], params["sf"]) @ params["Wout"].T + params["bout"], maps


def cross_entropy(Z, y):                              # Chapter 9, for V classes
    E = (Z + (-Z.data.max(axis=1, keepdims=True))).exp()
    P = E / (E @ Tensor(np.ones((V, 1))))
    return -(Tensor(np.eye(V)[y]) * P.log()).sum() / Z.shape[0]


def batch(source, B, gen):
    starts = gen.integers(0, len(source) - T - 1, B)
    x = np.stack([source[s:s + T] for s in starts])
    return x, np.stack([source[s + 1:s + T + 1] for s in starts])   # targets: the same text, one character on

val_gen = np.random.default_rng(1)
val_batches = [batch(val, 64, val_gen) for _ in range(20)]


def val_loss():
    total = 0.0
    for x, y in val_batches:
        total += float(cross_entropy(forward(x)[0], y.reshape(-1)).data)
        gc.collect()
    return total / len(val_batches)

print(f"validation loss before training: {val_loss():.3f}")

# === 5. Training: predict every next character in 32 windows of 64 =================
opt = Adam(list(params.values()), 0.003)
train_gen = np.random.default_rng(0)
STEPS = 3000
t0 = time.time()
for step in range(1, STEPS + 1):
    x, y = batch(train, 32, train_gen)
    loss = cross_entropy(forward(x)[0], y.reshape(-1))
    for p in params.values():
        p.grad = np.zeros_like(p.data)
    loss.backward()
    opt.step()
    del loss
    gc.collect()                                      # free each step's graph now (Chapter 15)
    if step % 500 == 0:
        print(f"  step {step}: validation loss {val_loss():.3f}  ({time.time() - t0:.0f} s)")
np.savez("shakespeare_model.npz", **{name: t.data for name, t in params.items()})   # Chapter 14

# === 6. What the model learned ======================================================
def next_probs(context, temperature=1.0):
    ids = encode(context[-T:])[None, :]               # at most the last 64 characters
    Z = forward(ids)[0].data[-1] / temperature        # the scores after the last character
    e = np.exp(Z - Z.max())
    return e / e.sum()

for prompt in ["To be, or not to b", "ROMEO:\nO, she doth teach the torches to burn bri"]:
    p = next_probs(prompt)
    top = np.argsort(p)[::-1][:3]
    print(f"\nafter {prompt[-24:]!r}: " + ",  ".join(f"{chars[i]!r} {p[i]:.0%}" for i in top))


def generate(prompt, n, temperature, seed):
    gen = np.random.default_rng(seed)
    text_out = prompt
    for i in range(n):
        text_out += chars[gen.choice(V, p=next_probs(text_out, temperature))]
        if i % 50 == 0:
            gc.collect()
    return text_out

print("\n--- 400 characters, temperature 0.8 ---")
print(generate("JULIET:\n", 400, 0.8, seed=1))

# Where does each head look? Over a validation batch: how far back its heaviest weight usually is.
x, _ = val_batches[0]
_, maps = forward(x)
print()
for l in range(LAYERS):
    for h in range(H):
        w = maps[l][:, h, 8:]                         # skip the first 8 positions, which have little to look at
        back = np.arange(8, T) - w.argmax(axis=-1)
        common = np.bincount(back.ravel()).argmax()
        print(f"  block {l + 1}, head {h + 1}: heaviest weight {common} back in {np.mean(back == common):4.0%} "
              f"of positions; average heaviest weight {w.max(axis=-1).mean():.2f}")

5. Practice

Work these out on paper (or in Python) and type the number. Answers are checked with a small tolerance for rounding.

P1 The text has 65 different characters. What loss does guessing uniformly give? (3 decimal places.)

P2 How many predictions does one window of 64 characters provide during training?

P3 q and k each have 16 entries with mean 0 and variance 1. What is the spread (standard deviation) of q · k?

P4 One head's scores before the softmax are [2, 0, −∞]. What weight goes to the first position? (3 decimal places.)

P5 How many numbers are in one layer's score array, for 32 windows, 4 heads and a context of 64?

P6 How many knobs do W_q, W_k, W_v, W_o and the output bias b_o have together, with 64 numbers per character?

P7 The model's validation loss is 1.761. What is its perplexity, e^loss? (2 decimal places.)

P8 If the context grew from 64 to 256 characters, by what factor would the score array grow?

6. Go further

Sample at temperatures 0.5, 1.0 and 1.5, 1,000 characters each. Describe what changes, then measure it: what fraction of the generated words (split on spaces) appear somewhere in the training text?
Delete the mask and train for 500 steps. What happens to the training loss, and to the generated text? Explain both. Then put the mask back, remove the position vectors instead (use zeros) and compare validation loss after 1,000 steps with the full model.
Make the model bigger: try 4 blocks, or 128 numbers per character, or a context of 128. Before running, work out the number of knobs and the size of the score arrays, and estimate the time per step from the 0.27 s here. Which change helps most for the same training time?

7. Check yourself

Answer all 5 questions correctly to complete the chapter · 0 / 5 done

Q1/5 What do the query, key and value of a character do in self-attention?

Q2/5 Why does the model set the scores of later positions to −∞?

Q3/5 Why does the model need position vectors?

Q4/5 The model's validation loss is 1.761 and the bigram table's is 2.482. What does that mean?

Q5/5 How does this model differ from a large language model?

In the catalog

Finished this chapter?

Learn machine learning from scratch