{
  "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
}