Files
gpt_from_scratch/02_makemore_mlp.ipynb
T
2026-06-10 19:45:12 +08:00

581 lines
22 KiB
Plaintext

{
"cells": [
{
"cell_type": "markdown",
"id": "b9c2fa10",
"metadata": {},
"source": [
"# makemore — the MLP character model\n",
"\n",
"Predict the next letter from the **previous few letters** with a small neural net — **concept -> code -> Your turn** each step.\n",
"\n",
"This is notebook **02** of the series. Notebook `00` built backprop by hand; the bigram in `01`\n",
"only ever looked at **one** character back. Here we take a real step up: a multi-layer\n",
"perceptron (MLP) that looks at a small **window** of previous characters, following Bengio et\n",
"al. 2003 and Karpathy's *\"makemore part 2: MLP\"* lecture.\n",
"\n",
"Dataset: `names.txt` — 32,033 real first names. Goal: **make more** name-like words."
]
},
{
"cell_type": "markdown",
"id": "69f5f6bc",
"metadata": {},
"source": [
"## Prologue — in plain English\n",
"\n",
"A bigram only remembers the last letter, so it cannot tell that `bri` should probably be\n",
"followed by `a` (Brianna) rather than a random vowel. We fix that by letting the model read a\n",
"**block** of the last few letters at once.\n",
"\n",
"The trick that makes this work is the **embedding**: instead of treating each letter as an\n",
"isolated symbol, we give every letter a short list of numbers (a vector) that the network\n",
"learns. Letters that behave similarly drift close together.\n",
"\n",
"Real-life picture: imagine each letter is a person, and the embedding is their personality\n",
"profile. The network learns these profiles so that \"similar\" letters (say, the vowels) end up\n",
"with similar profiles, and it can generalise: a context it never saw can still be handled\n",
"because it resembles ones it did see."
]
},
{
"cell_type": "markdown",
"id": "2d5dee6c",
"metadata": {},
"source": [
"### 0.1 Imports\n",
"\n",
"We import **torch first** (before matplotlib) to avoid a Windows OpenMP DLL clash, and make\n",
"plotting optional so the notebook runs even without matplotlib.\n",
"\n",
"**-> Training:** `torch` holds the weight tensors and does autograd; `F.cross_entropy` turns\n",
"logits + the true next letter into the single `loss` number we minimise."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "07c035c6",
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"os.environ.setdefault(\"KMP_DUPLICATE_LIB_OK\", \"TRUE\") # avoid duplicate-OpenMP crash on Windows\n",
"\n",
"import torch\n",
"import torch.nn.functional as F\n",
"import random\n",
"\n",
"try:\n",
" import matplotlib.pyplot as plt\n",
" HAS_PLT = True\n",
"except Exception:\n",
" HAS_PLT = False\n",
" print(\"matplotlib not installed -> plots will be skipped (everything else still works)\")\n",
"\n",
"print(\"PyTorch version:\", torch.__version__)"
]
},
{
"cell_type": "markdown",
"id": "dc58a8fc",
"metadata": {},
"source": [
"### 0.2 Load the names and build the alphabet\n",
"\n",
"Each name becomes letters plus a special **`.`** marker (id 0) that means *start* and *end*.\n",
"So `emma` is really `. e m m a .` — the dots tell the model where words begin and stop."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6b7e8198",
"metadata": {},
"outputs": [],
"source": [
"words = open(\"names.txt\").read().splitlines()\n",
"print(\"number of names:\", len(words))\n",
"print(\"examples:\", words[:8])\n",
"print(\"shortest/longest:\", min(len(w) for w in words), \"/\", max(len(w) for w in words))\n",
"\n",
"chars = sorted(list(set(\"\".join(words))))\n",
"stoi = {ch: i + 1 for i, ch in enumerate(chars)}\n",
"stoi[\".\"] = 0\n",
"itos = {i: ch for ch, i in stoi.items()}\n",
"vocab_size = len(itos)\n",
"print(\"vocab_size:\", vocab_size)\n",
"print(\"stoi:\", stoi)"
]
},
{
"cell_type": "markdown",
"id": "d7a9d8ab",
"metadata": {},
"source": [
"### 1.1 The context window (`block_size`)\n",
"\n",
"`block_size` = how many previous letters the model sees before guessing the next one. With\n",
"`block_size = 3`, predicting the letters of `emma` looks like this:\n",
"\n",
"```\n",
"... -> e (start: three dots, predict 'e')\n",
"..e -> m\n",
".em -> m\n",
"emm -> a\n",
"mma -> . (predict the end marker)\n",
"```\n",
"\n",
"Each arrow is one training example. Real-life picture: a sliding window of 3 letters moves\n",
"left to right across the word; at each stop the model must guess what comes next."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4c0464a0",
"metadata": {},
"outputs": [],
"source": [
"block_size = 3\n",
"\n",
"def build_dataset(word_list):\n",
" X, Y = [], []\n",
" for w in word_list:\n",
" context = [0] * block_size # start with '...'\n",
" for ch in w + \".\":\n",
" ix = stoi[ch]\n",
" X.append(context)\n",
" Y.append(ix)\n",
" context = context[1:] + [ix] # slide the window\n",
" return torch.tensor(X), torch.tensor(Y)\n",
"\n",
"# show the examples for one name\n",
"print(\"examples generated from 'emma':\")\n",
"Xtmp, Ytmp = build_dataset([\"emma\"])\n",
"for ctx, target in zip(Xtmp.tolist(), Ytmp.tolist()):\n",
" print(\"\".join(itos[i] for i in ctx), \"->\", itos[target])"
]
},
{
"cell_type": "markdown",
"id": "5f607e40",
"metadata": {},
"source": [
"### 1.2 Train / dev / test split\n",
"\n",
"We shuffle the names and cut three piles:\n",
"\n",
"- **train (80%)** — the model learns from these.\n",
"- **dev (10%)** — we tune choices (size, learning rate) against these.\n",
"- **test (10%)** — touched only at the very end, to get an honest score.\n",
"\n",
"Real-life picture: train = homework you study, dev = practice quiz you check yourself against,\n",
"test = the final exam you only sit once."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5c8e4f0c",
"metadata": {},
"outputs": [],
"source": [
"random.seed(42)\n",
"random.shuffle(words)\n",
"n1 = int(0.8 * len(words))\n",
"n2 = int(0.9 * len(words))\n",
"\n",
"Xtr, Ytr = build_dataset(words[:n1])\n",
"Xdev, Ydev = build_dataset(words[n1:n2])\n",
"Xte, Yte = build_dataset(words[n2:])\n",
"\n",
"print(\"train:\", tuple(Xtr.shape), \"dev:\", tuple(Xdev.shape), \"test:\", tuple(Xte.shape))"
]
},
{
"cell_type": "markdown",
"id": "b10d4c19",
"metadata": {},
"source": [
"### 2.1 The model — embedding, hidden layer, output\n",
"\n",
"Three learnable pieces:\n",
"\n",
"1. **`C`** — the embedding table: each of the 27 symbols gets an `n_embd`-dim vector.\n",
"2. **`W1, b1`** — a hidden layer with `tanh`: it mixes the 3 letter-vectors into `n_hidden`\n",
" features. Analogy: a committee that reads the 3 personalities and forms opinions.\n",
"3. **`W2, b2`** — the output layer: turns those opinions into `vocab_size` scores (logits) for\n",
" the next letter.\n",
"\n",
"**-> Training:** `parameters` is the full list of knobs; every one gets a gradient and a nudge."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4a0bcc2f",
"metadata": {},
"outputs": [],
"source": [
"g = torch.Generator().manual_seed(2147483647)\n",
"\n",
"n_embd = 10 # size of each letter's vector\n",
"n_hidden = 200 # neurons in the hidden layer\n",
"\n",
"C = torch.randn((vocab_size, n_embd), generator=g)\n",
"W1 = torch.randn((n_embd * block_size, n_hidden), generator=g) * 0.2\n",
"b1 = torch.randn(n_hidden, generator=g) * 0.01\n",
"W2 = torch.randn((n_hidden, vocab_size), generator=g) * 0.01\n",
"b2 = torch.randn(vocab_size, generator=g) * 0.0\n",
"\n",
"parameters = [C, W1, b1, W2, b2]\n",
"for p in parameters:\n",
" p.requires_grad = True\n",
"\n",
"print(\"total parameters:\", sum(p.nelement() for p in parameters))"
]
},
{
"cell_type": "markdown",
"id": "77778ae3",
"metadata": {},
"source": [
"### 2.2 One forward pass, step by step\n",
"\n",
"Watch the shapes flow. The key move is `emb.view(-1, n_embd*block_size)`: it **flattens** the\n",
"3 letter-vectors of each example into one long row so the hidden layer can read them together.\n",
"\n",
"**-> Training:** `F.cross_entropy` = softmax + negative-log-likelihood in one numerically-stable\n",
"call; lower means more probability on the true next letter."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "13123079",
"metadata": {},
"outputs": [],
"source": [
"ix = torch.randint(0, Xtr.shape[0], (4,), generator=g) # 4 random examples\n",
"Xb, Yb = Xtr[ix], Ytr[ix]\n",
"\n",
"emb = C[Xb] # (4, 3, n_embd)\n",
"flat = emb.view(-1, n_embd * block_size) # (4, 30)\n",
"h = torch.tanh(flat @ W1 + b1) # (4, n_hidden)\n",
"logits = h @ W2 + b2 # (4, vocab_size)\n",
"loss = F.cross_entropy(logits, Yb)\n",
"\n",
"print(\"emb :\", tuple(emb.shape))\n",
"print(\"flat :\", tuple(flat.shape))\n",
"print(\"h :\", tuple(h.shape))\n",
"print(\"logits:\", tuple(logits.shape))\n",
"print(\"loss :\", loss.item())"
]
},
{
"cell_type": "markdown",
"id": "b90331ba",
"metadata": {},
"source": [
"### 3.1 Finding a good learning rate\n",
"\n",
"The learning rate (how big each nudge is) matters a lot: too small and training crawls, too\n",
"big and the loss bounces or explodes. A cheap way to find a good one: sweep the rate from tiny\n",
"to large over a few hundred steps and watch where the loss drops fastest.\n",
"\n",
"Real-life picture: testing how big a step to take downhill in fog — tiny steps are safe but\n",
"slow; huge steps overshoot the valley.\n",
"\n",
"We use a **fresh** set of weights here so the search does not disturb the real model."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "acae02ad",
"metadata": {},
"outputs": [],
"source": [
"# fresh throwaway weights just for the search\n",
"gs = torch.Generator().manual_seed(2147483647)\n",
"Cs = torch.randn((vocab_size, n_embd), generator=gs)\n",
"W1s = torch.randn((n_embd * block_size, n_hidden), generator=gs) * 0.2\n",
"b1s = torch.randn(n_hidden, generator=gs) * 0.01\n",
"W2s = torch.randn((n_hidden, vocab_size), generator=gs) * 0.01\n",
"b2s = torch.randn(vocab_size, generator=gs) * 0.0\n",
"params_s = [Cs, W1s, b1s, W2s, b2s]\n",
"for p in params_s:\n",
" p.requires_grad = True\n",
"\n",
"# candidate learning rates spaced from 10^-3 up to 10^0\n",
"lre = torch.linspace(-3, 0, 1000)\n",
"lrs = 10 ** lre\n",
"lr_used, losses = [], []\n",
"\n",
"for i in range(1000):\n",
" ix = torch.randint(0, Xtr.shape[0], (32,), generator=gs)\n",
" Xb, Yb = Xtr[ix], Ytr[ix]\n",
" emb = Cs[Xb]\n",
" h = torch.tanh(emb.view(-1, n_embd * block_size) @ W1s + b1s)\n",
" loss = F.cross_entropy(h @ W2s + b2s, Yb)\n",
" for p in params_s:\n",
" p.grad = None\n",
" loss.backward()\n",
" lr = lrs[i].item()\n",
" for p in params_s:\n",
" p.data += -lr * p.grad\n",
" lr_used.append(lre[i].item())\n",
" losses.append(loss.item())\n",
"\n",
"best = lre[min(range(len(losses)), key=lambda k: losses[k])].item()\n",
"print(f\"loss dropped fastest around 10^{best:.2f} ~= {10**best:.3f}\")\n",
"print(\"good practical choice: lr = 0.1\")\n",
"\n",
"if HAS_PLT:\n",
" plt.figure(figsize=(5, 4))\n",
" plt.plot(lr_used, losses)\n",
" plt.xlabel(\"log10(learning rate)\")\n",
" plt.ylabel(\"minibatch loss\")\n",
" plt.title(\"learning-rate sweep\")\n",
" plt.show()"
]
},
{
"cell_type": "markdown",
"id": "13569252",
"metadata": {},
"source": [
"### 4.1 Train the model\n",
"\n",
"Now the real run on the actual `parameters`. Each step:\n",
"\n",
"1. grab a random **minibatch** of 32 examples (fast, noisy, good enough)\n",
"2. forward -> loss\n",
"3. `loss.backward()` -> gradients\n",
"4. nudge every parameter downhill\n",
"\n",
"We start at `lr = 0.1` and **decay** to `0.01` near the end to settle into the valley."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c5ec35d5",
"metadata": {},
"outputs": [],
"source": [
"max_steps = 20000\n",
"batch_size = 32\n",
"lossi = []\n",
"\n",
"for i in range(max_steps):\n",
" ix = torch.randint(0, Xtr.shape[0], (batch_size,), generator=g)\n",
" Xb, Yb = Xtr[ix], Ytr[ix]\n",
"\n",
" emb = C[Xb]\n",
" h = torch.tanh(emb.view(-1, n_embd * block_size) @ W1 + b1)\n",
" logits = h @ W2 + b2\n",
" loss = F.cross_entropy(logits, Yb)\n",
"\n",
" for p in parameters:\n",
" p.grad = None\n",
" loss.backward()\n",
"\n",
" lr = 0.1 if i < 15000 else 0.01\n",
" for p in parameters:\n",
" p.data += -lr * p.grad\n",
"\n",
" lossi.append(loss.log10().item())\n",
" if i % 4000 == 0 or i == max_steps - 1:\n",
" print(f\"step {i:5d}/{max_steps} minibatch loss {loss.item():.4f}\")\n",
"\n",
"if HAS_PLT:\n",
" plt.figure(figsize=(5, 4))\n",
" plt.plot(lossi)\n",
" plt.xlabel(\"step\")\n",
" plt.ylabel(\"log10(loss)\")\n",
" plt.title(\"training loss\")\n",
" plt.show()"
]
},
{
"cell_type": "markdown",
"id": "e6baab73",
"metadata": {},
"source": [
"### 4.2 Honest scores on train and dev\n",
"\n",
"Minibatch loss is noisy. Evaluate on the **whole** train and dev sets for a stable number.\n",
"If dev loss is much worse than train, the model is memorising; if they are close, it is\n",
"generalising. (A plain bigram on names scores around 2.45 — we should beat that.)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "579fe1c8",
"metadata": {},
"outputs": [],
"source": [
"@torch.no_grad()\n",
"def split_loss(split):\n",
" x, y = {\"train\": (Xtr, Ytr), \"dev\": (Xdev, Ydev), \"test\": (Xte, Yte)}[split]\n",
" emb = C[x]\n",
" h = torch.tanh(emb.view(-1, n_embd * block_size) @ W1 + b1)\n",
" logits = h @ W2 + b2\n",
" return F.cross_entropy(logits, y).item()\n",
"\n",
"print(\"train loss:\", round(split_loss(\"train\"), 4))\n",
"print(\"dev loss:\", round(split_loss(\"dev\"), 4))"
]
},
{
"cell_type": "markdown",
"id": "054682f1",
"metadata": {},
"source": [
"### 5.1 Generate new names\n",
"\n",
"Sampling, letter by letter: start from `...`, predict a distribution, **draw** one letter,\n",
"slide the window, repeat until the model emits the end marker `.`.\n",
"\n",
"These are brand-new names the model invented — not copied from the file."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "06bf121c",
"metadata": {},
"outputs": [],
"source": [
"g2 = torch.Generator().manual_seed(2147483647 + 10)\n",
"\n",
"for _ in range(20):\n",
" out = []\n",
" context = [0] * block_size\n",
" while True:\n",
" emb = C[torch.tensor([context])]\n",
" h = torch.tanh(emb.view(1, -1) @ W1 + b1)\n",
" logits = h @ W2 + b2\n",
" probs = F.softmax(logits, dim=1)\n",
" ix = torch.multinomial(probs, num_samples=1, generator=g2).item()\n",
" context = context[1:] + [ix]\n",
" if ix == 0:\n",
" break\n",
" out.append(itos[ix])\n",
" print(\"\".join(out))"
]
},
{
"cell_type": "markdown",
"id": "0a0de526",
"metadata": {},
"source": [
"**Your turn 5** — Re-run section 5.1 a few times (different random draws) and read the\n",
"names. Then try raising `n_hidden` (e.g. 300) or `block_size` (e.g. 4 or 5) back in section 2,\n",
"retrain, and see whether dev loss improves and the names look more realistic."
]
},
{
"cell_type": "markdown",
"id": "bc2003e7",
"metadata": {},
"source": [
"### 6.1 Bonus — visualise letters in 2-D\n",
"\n",
"When the embedding is 2-D we can literally plot it. We quickly train a tiny `n_embd = 2`\n",
"version and scatter the 27 symbols. Watch the **vowels** drift into their own neighbourhood —\n",
"the network discovered that they play a similar role, without ever being told what a vowel is.\n",
"\n",
"(Optional: needs matplotlib. Smaller/shorter than the main model, so the clustering is rough.)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7786f75a",
"metadata": {},
"outputs": [],
"source": [
"g3 = torch.Generator().manual_seed(2147483647)\n",
"C2 = torch.randn((vocab_size, 2), generator=g3)\n",
"W1b = torch.randn((2 * block_size, n_hidden), generator=g3) * 0.2\n",
"b1b = torch.randn(n_hidden, generator=g3) * 0.01\n",
"W2b = torch.randn((n_hidden, vocab_size), generator=g3) * 0.01\n",
"b2b = torch.randn(vocab_size, generator=g3) * 0.0\n",
"params2 = [C2, W1b, b1b, W2b, b2b]\n",
"for p in params2:\n",
" p.requires_grad = True\n",
"\n",
"for i in range(8000):\n",
" ix = torch.randint(0, Xtr.shape[0], (32,), generator=g3)\n",
" Xb, Yb = Xtr[ix], Ytr[ix]\n",
" emb = C2[Xb]\n",
" h = torch.tanh(emb.view(-1, 2 * block_size) @ W1b + b1b)\n",
" loss = F.cross_entropy(h @ W2b + b2b, Yb)\n",
" for p in params2:\n",
" p.grad = None\n",
" loss.backward()\n",
" for p in params2:\n",
" p.data += -0.1 * p.grad\n",
"\n",
"print(\"2-D model dev-ish minibatch loss:\", round(loss.item(), 4))\n",
"\n",
"if HAS_PLT:\n",
" pts = C2.detach()\n",
" plt.figure(figsize=(6, 6))\n",
" plt.scatter(pts[:, 0], pts[:, 1], s=300)\n",
" for i in range(vocab_size):\n",
" plt.text(pts[i, 0].item(), pts[i, 1].item(), itos[i],\n",
" ha=\"center\", va=\"center\", color=\"white\")\n",
" plt.title(\"learned 2-D letter embedding (look for the vowels clustering)\")\n",
" plt.grid(True)\n",
" plt.show()\n",
"else:\n",
" print(\"matplotlib missing - skipping the embedding scatter\")"
]
},
{
"cell_type": "markdown",
"id": "c5e5b7a7",
"metadata": {},
"source": [
"## What's next\n",
"\n",
"You went from \"one letter back\" (bigram) to a real neural net that reads a **window** of\n",
"letters through a learned **embedding**, and it makes noticeably better names.\n",
"\n",
"Where this leads in the series:\n",
"\n",
"- `03_batchnorm_activations.ipynb` — why deep nets are fragile, and how BatchNorm keeps the\n",
" signal healthy so we can train bigger models.\n",
"- `04_backprop_ninja.ipynb` — compute all these gradients **by hand**, cementing notebook `00`.\n",
"- `06_build_gpt_attention.ipynb` — replace the fixed window with **attention**, so the model\n",
" can look back as far as it needs.\n",
"\n",
"**Checklist**\n",
"- [ ] Names -> context windows (`block_size`) -> (X, Y)\n",
"- [ ] train / dev / test split\n",
"- [ ] embedding `C` + hidden `tanh` + output logits\n",
"- [ ] cross-entropy loss, minibatch training\n",
"- [ ] learning-rate search\n",
"- [ ] sampled new names; visualised the 2-D embedding"
]
}
],
"metadata": {
"kernelspec": {
"display_name": ".venv (3.12.10)",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python",
"version": "3.12.10"
}
},
"nbformat": 4,
"nbformat_minor": 5
}