beastgitea2026/gpt_from_scratch
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

c

Browse Source
This commit is contained in:
betterbeast2023
2026-06-10 19:45:12 +08:00
parent 1183d10849
commit a3ebb28ce0
12 changed files with 35469 additions and 491 deletions
Download Patch File Download Diff File Expand all files Collapse all files
03_batchnorm_activations.ipynb
+496
View File
@@ -0,0 +1,496 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "aedd6523",
"metadata": {},
"source": [
"# makemore — activations, gradients, and BatchNorm\n",
"\n",
"Make the MLP **trainable and healthy** — **concept -> code -> Your turn** each step.\n",
"\n",
"This is notebook **03**. In `02` the MLP worked, but we got a bit lucky with how the weights\n",
"started. Here we look *inside* the network at the numbers flowing through it, see two common\n",
"ways training goes wrong, and fix them — first by hand (good initialisation), then with the\n",
"robust tool everyone uses: **Batch Normalization**. Follows Karpathy's *\"makemore part 3\"*."
]
},
{
"cell_type": "markdown",
"id": "2980029a",
"metadata": {},
"source": [
"## Prologue — in plain English\n",
"\n",
"A neural net is a stack of layers. At each layer numbers flow forward (activations) and\n",
"gradients flow backward. If those numbers get **too big** or **too small**, training stalls.\n",
"\n",
"Two failures we will see and fix:\n",
"\n",
"1. **Confidently wrong at the start.** If the output layer starts with large random weights,\n",
" the model makes wild over-confident guesses, so the first loss is huge and the first many\n",
" steps are wasted just calming it down.\n",
"2. **Saturated `tanh`.** If the hidden numbers are too spread out, `tanh` slams them to -1 or\n",
" +1. There the curve is flat, so its gradient is ~0 and those neurons stop learning — they\n",
" are \"dead.\"\n",
"\n",
"Real-life picture: think of each layer as an audio stage. Too quiet and you hear nothing; too\n",
"loud and everything **clips** into distortion. We want every stage at a healthy volume.\n",
"**BatchNorm** is the auto-leveler that keeps the volume right at every stage."
]
},
{
"cell_type": "markdown",
"id": "18a296f3",
"metadata": {},
"source": [
"### 0.1 Imports\n",
"\n",
"`torch` first (avoids a Windows OpenMP DLL clash), matplotlib optional for the diagnostics."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ac69da81",
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"os.environ.setdefault(\"KMP_DUPLICATE_LIB_OK\", \"TRUE\")\n",
"\n",
"import math\n",
"import random\n",
"import torch\n",
"import torch.nn.functional as F\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": "92143322",
"metadata": {},
"source": [
"### 0.2 Data — same names, same context windows as notebook 02"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d99c17df",
"metadata": {},
"outputs": [],
"source": [
"words = open(\"names.txt\").read().splitlines()\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",
"block_size = 3\n",
"\n",
"def build_dataset(word_list):\n",
" X, Y = [], []\n",
" for w in word_list:\n",
" context = [0] * block_size\n",
" for ch in w + \".\":\n",
" ix = stoi[ch]\n",
" X.append(context)\n",
" Y.append(ix)\n",
" context = context[1:] + [ix]\n",
" return torch.tensor(X), torch.tensor(Y)\n",
"\n",
"random.seed(42)\n",
"random.shuffle(words)\n",
"n1, n2 = int(0.8 * len(words)), int(0.9 * len(words))\n",
"Xtr, Ytr = build_dataset(words[:n1])\n",
"Xdev, Ydev = build_dataset(words[n1:n2])\n",
"Xte, Yte = build_dataset(words[n2:])\n",
"print(\"train:\", tuple(Xtr.shape), \" vocab:\", vocab_size)"
]
},
{
"cell_type": "markdown",
"id": "36b836bc",
"metadata": {},
"source": [
"### 1.1 Problem 1 — confidently wrong at the start\n",
"\n",
"At step 0 the model knows nothing, so the *best* it can do is spread probability evenly over\n",
"all 27 symbols. That ideal first loss is `-ln(1/27) = ln(27) ~= 3.30`.\n",
"\n",
"But with large random output weights the logits are all over the place, so the first loss is\n",
"much higher. Let's measure it."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1219a887",
"metadata": {},
"outputs": [],
"source": [
"g0 = torch.Generator().manual_seed(2147483647)\n",
"n_embd, n_hidden = 10, 200\n",
"\n",
"# deliberately naive (large) initialisation\n",
"Cn = torch.randn((vocab_size, n_embd), generator=g0)\n",
"W1n = torch.randn((n_embd * block_size, n_hidden), generator=g0)\n",
"b1n = torch.randn(n_hidden, generator=g0)\n",
"W2n = torch.randn((n_hidden, vocab_size), generator=g0)\n",
"b2n = torch.randn(vocab_size, generator=g0)\n",
"\n",
"emb = Cn[Xtr[:32]]\n",
"h = torch.tanh(emb.view(-1, n_embd * block_size) @ W1n + b1n)\n",
"logits = h @ W2n + b2n\n",
"print(\"naive initial loss :\", F.cross_entropy(logits, Ytr[:32]).item())\n",
"print(\"ideal (ln 27) :\", math.log(27))"
]
},
{
"cell_type": "markdown",
"id": "4c418501",
"metadata": {},
"source": [
"### 1.2 Fix 1 — make the output layer humble at the start\n",
"\n",
"We want the starting logits near zero (no opinion yet). So scale `W2` **down** and set `b2`\n",
"to zero. Now the first loss sits right near the ideal `~3.30`, and no steps are wasted."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e1b9bb4a",
"metadata": {},
"outputs": [],
"source": [
"g0 = torch.Generator().manual_seed(2147483647)\n",
"Cn = torch.randn((vocab_size, n_embd), generator=g0)\n",
"W1n = torch.randn((n_embd * block_size, n_hidden), generator=g0)\n",
"b1n = torch.randn(n_hidden, generator=g0)\n",
"W2n = torch.randn((n_hidden, vocab_size), generator=g0) * 0.01 # tiny\n",
"b2n = torch.randn(vocab_size, generator=g0) * 0.0 # zero\n",
"\n",
"emb = Cn[Xtr[:32]]\n",
"h = torch.tanh(emb.view(-1, n_embd * block_size) @ W1n + b1n)\n",
"logits = h @ W2n + b2n\n",
"print(\"fixed initial loss :\", F.cross_entropy(logits, Ytr[:32]).item())\n",
"print(\"ideal (ln 27) :\", math.log(27))"
]
},
{
"cell_type": "markdown",
"id": "2ab522f7",
"metadata": {},
"source": [
"### 2.1 Problem 2 — saturated `tanh` (dead neurons)\n",
"\n",
"Now look at the hidden activations `h = tanh(...)`. With large `W1`, the inputs to `tanh` are\n",
"very spread out, so most outputs are pinned at -1 or +1. That is the flat part of `tanh`,\n",
"where the gradient is ~0 — those neurons learn nothing.\n",
"\n",
"We measure the fraction of \"saturated\" activations (|h| > 0.99) and, if matplotlib is present,\n",
"show the tell-tale picture: bright cells are saturated; a fully bright column = a dead neuron."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8fd25db6",
"metadata": {},
"outputs": [],
"source": [
"emb = Cn[Xtr[:1000]]\n",
"hpre = emb.view(-1, n_embd * block_size) @ W1n + b1n # pre-activation\n",
"h = torch.tanh(hpre)\n",
"sat = (h.abs() > 0.99).float().mean().item()\n",
"print(f\"fraction of saturated activations: {sat:.1%} (high = bad)\")\n",
"\n",
"if HAS_PLT:\n",
" plt.figure(figsize=(10, 3))\n",
" plt.imshow(h.abs() > 0.99, cmap=\"gray\", aspect=\"auto\")\n",
" plt.title(\"white = saturated (|tanh| > 0.99); a full white column = a dead neuron\")\n",
" plt.xlabel(\"neuron\")\n",
" plt.ylabel(\"example\")\n",
" plt.show()"
]
},
{
"cell_type": "markdown",
"id": "e11aee42",
"metadata": {},
"source": [
"### 2.2 Fix 2 — scale the hidden layer (Kaiming init)\n",
"\n",
"The cure is to make `W1` smaller so the pre-activations have a sensible spread (roughly unit\n",
"standard deviation). The principled recipe is **Kaiming init**: multiply random weights by\n",
"`gain / sqrt(fan_in)`, where `fan_in` is the number of inputs and the `gain` for `tanh` is\n",
"`5/3`.\n",
"\n",
"Real-life picture: setting the input volume so the signal fills the range without clipping."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "71f808d3",
"metadata": {},
"outputs": [],
"source": [
"fan_in = n_embd * block_size\n",
"W1k = torch.randn((fan_in, n_hidden), generator=torch.Generator().manual_seed(1)) * ((5/3) / fan_in**0.5)\n",
"b1k = torch.zeros(n_hidden)\n",
"\n",
"emb = Cn[Xtr[:1000]]\n",
"h = torch.tanh(emb.view(-1, fan_in) @ W1k + b1k)\n",
"sat = (h.abs() > 0.99).float().mean().item()\n",
"print(f\"saturated now: {sat:.1%} (much lower)\")\n",
"print(f\"activation std: {h.std().item():.3f} (we want it near ~0.6-0.7 for tanh)\")\n",
"\n",
"if HAS_PLT:\n",
" plt.figure(figsize=(5, 3))\n",
" plt.hist(h.view(-1).tolist(), bins=50)\n",
" plt.title(\"healthy tanh activations (not all piled at +/-1)\")\n",
" plt.show()"
]
},
{
"cell_type": "markdown",
"id": "50d6f683",
"metadata": {},
"source": [
"### 3.1 BatchNorm — the auto-leveler\n",
"\n",
"Hand-tuning `W1` works, but in deep nets it is fiddly. **Batch Normalization** does it\n",
"automatically: before `tanh`, it **normalises** the pre-activations so that, across the\n",
"current batch, each neuron has mean 0 and std 1. Then it lets the network rescale/shift them\n",
"with two learnable knobs, `gamma` (gain) and `beta` (bias), in case 0/1 is not ideal.\n",
"\n",
"Real-life picture: an audio compressor that auto-levels every stage to a healthy volume, with\n",
"a manual trim (`gamma`, `beta`) if the engineer wants to nudge it.\n",
"\n",
"One catch: at test time we have no batch to average over, so during training we also keep a\n",
"**running** mean/std to use later for single predictions."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5aeb780e",
"metadata": {},
"outputs": [],
"source": [
"g = torch.Generator().manual_seed(2147483647)\n",
"n_embd, n_hidden = 10, 200\n",
"fan_in = n_embd * block_size\n",
"\n",
"C = torch.randn((vocab_size, n_embd), generator=g)\n",
"W1 = torch.randn((fan_in, n_hidden), generator=g) * ((5/3) / fan_in**0.5)\n",
"# no b1: BatchNorm's beta plays that role\n",
"W2 = torch.randn((n_hidden, vocab_size), generator=g) * 0.01\n",
"b2 = torch.randn(vocab_size, generator=g) * 0.0\n",
"\n",
"bngain = torch.ones((1, n_hidden))\n",
"bnbias = torch.zeros((1, n_hidden))\n",
"bnmean_running = torch.zeros((1, n_hidden))\n",
"bnstd_running = torch.ones((1, n_hidden))\n",
"\n",
"parameters = [C, W1, W2, b2, bngain, bnbias]\n",
"for p in parameters:\n",
" p.requires_grad = True\n",
"print(\"total parameters:\", sum(p.nelement() for p in parameters))"
]
},
{
"cell_type": "markdown",
"id": "2b07d16a",
"metadata": {},
"source": [
"### 3.2 Train with BatchNorm\n",
"\n",
"Each step normalises the pre-activations using the **current batch**, updates the running\n",
"stats (for later), then continues as usual. Watch the loss fall, starting near the ideal 3.3."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2410e7c4",
"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",
" hpreact = emb.view(-1, fan_in) @ W1\n",
"\n",
" # --- BatchNorm ---\n",
" bnmeani = hpreact.mean(0, keepdim=True)\n",
" bnstdi = hpreact.std(0, keepdim=True)\n",
" hpreact = bngain * (hpreact - bnmeani) / bnstdi + bnbias\n",
" with torch.no_grad():\n",
" bnmean_running = 0.999 * bnmean_running + 0.001 * bnmeani\n",
" bnstd_running = 0.999 * bnstd_running + 0.001 * bnstdi\n",
" # -----------------\n",
"\n",
" h = torch.tanh(hpreact)\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",
" 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} loss {loss.item():.4f}\")\n",
"\n",
"if HAS_PLT:\n",
" plt.figure(figsize=(5, 4))\n",
" plt.plot(lossi)\n",
" plt.xlabel(\"step\"); plt.ylabel(\"log10(loss)\"); plt.title(\"training loss (with BatchNorm)\")\n",
" plt.show()"
]
},
{
"cell_type": "markdown",
"id": "d80dccf6",
"metadata": {},
"source": [
"### 3.3 Evaluate — using the running stats\n",
"\n",
"At evaluation we have no batch, so we normalise with the **running** mean/std collected during\n",
"training. Train and dev loss should be close (good generalisation) and beat the bigram (~2.45)."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "30ae8013",
"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",
" hpreact = emb.view(-1, fan_in) @ W1\n",
" hpreact = bngain * (hpreact - bnmean_running) / bnstd_running + bnbias\n",
" h = torch.tanh(hpreact)\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": "269c8943",
"metadata": {},
"source": [
"### 4.1 Generate names with the BatchNorm model\n",
"\n",
"Same letter-by-letter sampling as before, but normalising with the running stats."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "63711cd1",
"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",
" hpreact = emb.view(1, -1) @ W1\n",
" hpreact = bngain * (hpreact - bnmean_running) / bnstd_running + bnbias\n",
" h = torch.tanh(hpreact)\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": "365be7ef",
"metadata": {},
"source": [
"**Your turn 4** — In section 2.1, change `W1n`'s scale (e.g. `* 0.1` vs `* 3.0`) and re-run\n",
"the saturation check. Bigger scale -> more dead neurons. Then confirm BatchNorm (section 3.2)\n",
"trains fine even from a less careful init, because it re-levels the signal every step."
]
},
{
"cell_type": "markdown",
"id": "65b60eec",
"metadata": {},
"source": [
"## What's next\n",
"\n",
"You learned to look **inside** the network and keep it healthy:\n",
"\n",
"- start the output layer humble (good first loss)\n",
"- keep `tanh` out of saturation (Kaiming init)\n",
"- or just use **BatchNorm** to auto-level every layer\n",
"\n",
"Where this leads in the series:\n",
"\n",
"- `04_backprop_ninja.ipynb` — compute every gradient here **by hand**, including through\n",
" BatchNorm, to truly own backprop.\n",
"- `05_wavenet.ipynb` — stack these healthy layers into a deeper, hierarchical model.\n",
"- `06_build_gpt_attention.ipynb` — the transformer, which uses LayerNorm (a cousin of\n",
" BatchNorm) for the same \"keep the signal healthy\" reason.\n",
"\n",
"**Checklist**\n",
"- [ ] Ideal initial loss = ln(vocab_size)\n",
"- [ ] Humble output init (small W2, zero b2)\n",
"- [ ] Saturation / dead neurons diagnosed\n",
"- [ ] Kaiming init (gain / sqrt(fan_in))\n",
"- [ ] BatchNorm: per-batch normalise + gamma/beta + running stats\n",
"- [ ] Trained, evaluated, sampled names"
]
}
],
"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
}