{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "# Step 3 — Autodiff, and a Neural $n$-gram Model\n\n**Solution notebook**, and the densest of the eight. Part 1 builds\nreverse-mode automatic differentiation from nothing — about 60 lines\nimplementing Lecture 3's Theorem 3.1 — and validates it three ways,\nthe third being Step 2's diamond network differentiated by hand, by the\nengine, by finite differences, and finally by PyTorch. Part 2 uses\nPyTorch's autograd to train the Bengio (2003) model, and ends with a\nmeasured result that motivates the entire second half of the course.\n\nRuntime: the engine is instant; the $k$-sweep is ~20 s per value of\n$k$ single-threaded.",
   "id": "cell-00"
  },
  {
   "cell_type": "markdown",
   "id": "cell-01",
   "metadata": {},
   "source": "# Part 1 — A scalar autograd engine\n\nEach `Value` holds a number, a gradient slot, its parents, and a\nclosure saying how to push its own gradient to those parents. That\nclosure *is* the local partial $\\partial\\varphi_i/\\partial v_j$ of\nTheorem 3.1; `backward()` is the reverse topological sweep.\n\nTwo details carry the whole proof, and both are where bugs live:\n**reverse topological order** (a node must not be processed before all\nits children) and **`+=` not `=`** (a node used twice receives a\ncontribution along each path)."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "class Value:\n    def __init__(self, data, _parents=()):\n        self.data = float(data)\n        self.grad = 0.0\n        self._parents = _parents\n        self._backward = lambda: None\n\n    def __repr__(self):\n        return f'Value({self.data:.4f}, grad={self.grad:.4f})'\n\n    def _wrap(self, other):\n        return other if isinstance(other, Value) else Value(other)\n\n    def __add__(self, other):\n        other = self._wrap(other)\n        out = Value(self.data + other.data, (self, other))\n        def _backward():\n            self.grad  += out.grad          # d(a+b)/da = 1\n            other.grad += out.grad\n        out._backward = _backward\n        return out\n\n    def __mul__(self, other):\n        other = self._wrap(other)\n        out = Value(self.data * other.data, (self, other))\n        def _backward():\n            self.grad  += other.data * out.grad   # d(ab)/da = b\n            other.grad += self.data  * out.grad\n        out._backward = _backward\n        return out\n\n    def __pow__(self, k):\n        assert isinstance(k, (int, float))\n        out = Value(self.data ** k, (self,))\n        def _backward():\n            self.grad += k * self.data ** (k - 1) * out.grad\n        out._backward = _backward\n        return out\n\n    def tanh(self):\n        import math\n        t = math.tanh(self.data)\n        out = Value(t, (self,))\n        def _backward():\n            self.grad += (1 - t * t) * out.grad   # d tanh = 1 - tanh^2\n        out._backward = _backward\n        return out\n\n    def exp(self):\n        import math\n        e = math.exp(self.data)\n        out = Value(e, (self,))\n        def _backward():\n            self.grad += e * out.grad\n        out._backward = _backward\n        return out\n\n    def log(self):\n        import math\n        out = Value(math.log(self.data), (self,))\n        def _backward():\n            self.grad += (1.0 / self.data) * out.grad\n        out._backward = _backward\n        return out\n\n    def relu(self):\n        out = Value(max(self.data, 0.0), (self,))\n        def _backward():\n            self.grad += (1.0 if self.data > 0 else 0.0) * out.grad   # subgradient 0 at the corner\n        out._backward = _backward\n        return out\n\n    # conveniences\n    def __neg__(self):       return self * -1\n    def __sub__(self, o):    return self + (-self._wrap(o))\n    def __radd__(self, o):   return self + o\n    def __rsub__(self, o):   return (-self) + o\n    def __rmul__(self, o):   return self * o\n    def __truediv__(self, o): return self * (self._wrap(o) ** -1)\n\n    def backward(self):\n        order, visited = [], set()\n        def visit(v):\n            if id(v) not in visited:\n                visited.add(id(v))\n                for p in v._parents:\n                    visit(p)\n                order.append(v)      # parents appended before children\n        visit(self)\n        self.grad = 1.0              # dL/dL = 1\n        for v in reversed(order):    # reverse topological order\n            v._backward()",
   "id": "cell-02"
  },
  {
   "cell_type": "markdown",
   "id": "cell-03",
   "metadata": {},
   "source": "### Validation 1 — fan-out, against hand computation\n\n$\\mathcal L = (ab + a)\\tanh b$ at $a=2$, $b=-1$. Both variables are\nused twice, so this is precisely the test that `=` instead of `+=`\nfails. By hand, with $t = \\tanh(-1)$:\n$\\partial\\mathcal L/\\partial a = (b+1)t$ and\n$\\partial\\mathcal L/\\partial b = a\\,t + (ab+a)(1-t^2)$."
  },
  {
   "cell_type": "code",
   "id": "cell-04",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "import math\na, b = Value(2.0), Value(-1.0)\nL = (a * b + a) * b.tanh()\nL.backward()\n\nt = math.tanh(-1.0)\nda_hand = ((-1.0) + 1) * t\ndb_hand = 2.0 * t + (2.0 * -1.0 + 2.0) * (1 - t * t)\nprint(f'engine: dL/da = {a.grad:.10f}   hand: {da_hand:.10f}')\nprint(f'engine: dL/db = {b.grad:.10f}   hand: {db_hand:.10f}')\nassert abs(a.grad - da_hand) < 1e-12 and abs(b.grad - db_hand) < 1e-12\nprint('fan-out test passed')"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "### Validation 2 — finite differences\n\nIn double precision (Python floats are float64), so the resolution\nproblem that bites float32 finite differences does not arise here. The\nReLU test is placed away from the corners: at a corner the derivative\ndoesn't exist, and a central difference straddling it returns $\\tfrac12$\nwhile the engine's subgradient returns $0$ or $1$ — both defensible,\nneither \"wrong\".",
   "id": "cell-05"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "def fd_check(f, xs, h=1e-6):\n    vals = [Value(x) for x in xs]\n    out = f(vals)\n    out.backward()\n    errs = []\n    for i in range(len(xs)):\n        up = list(xs); up[i] += h\n        dn = list(xs); dn[i] -= h\n        num = (f([Value(v) for v in up]).data\n               - f([Value(v) for v in dn]).data) / (2 * h)\n        errs.append(abs(num - vals[i].grad) / max(abs(vals[i].grad), 1e-12))\n    return max(errs)\n\ntests = {\n    '(ab+a)tanh(b)': (lambda v: (v[0]*v[1] + v[0]) * v[1].tanh(), [2.0, -1.0]),\n    'exp(ab)/(1+a^2)': (lambda v: (v[0]*v[1]).exp() / (1 + v[0]**2), [0.7, -0.4]),\n    'log(a^2+b^2+1)': (lambda v: (v[0]**2 + v[1]**2 + 1).log(), [1.3, 0.5]),\n    'relu(a)+relu(-b)': (lambda v: v[0].relu() + (-v[1]).relu(), [0.9, -0.6]),   # away from the corners\n}\nfor name, (f, xs) in tests.items():\n    e = fd_check(f, xs)\n    print(f'{name:<18} max relative error {e:.2e}')\n    assert e < 1e-6",
   "id": "cell-06"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "### Validation 3 — the full circle: Step 2's diamond network\n\nRebuild one forward pass of the diamond classifier from scalar `Value`s\nat $\\mathbf x=(0.5,0.25)^\\top$, $\\gamma=4$, with target *inside*, and\ndifferentiate the loss with respect to the **input** coordinates. Both\ncoordinates are positive there, so only two of the four ReLUs are\nactive and $r=x_1+x_2$; the loss is\n\n$$\n\\mathcal L=-\\log p(\\mathrm{in}\\mid\\mathbf x)\n=\\log\\bigl(1+e^{z_{\\mathrm{out}}-z_{\\mathrm{in}}}\\bigr)\n=\\log\\bigl(1+\\exp(2\\gamma(x_1+x_2-1))\\bigr),\n$$\n\nwhose derivative in either coordinate is\n$2\\gamma\\,\\sigma(2\\gamma(x_1+x_2-1))=2\\gamma\\,p(\\mathrm{out}\\mid\\mathbf x)\n=8\\,\\sigma(-2)\\approx0.9536$. Three independent routes to that number:\nthe formula, the engine, central differences.",
   "id": "cell-07"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "def diamond_loss(x1, x2, gamma=4.0):\n    # forward pass of Step 2's network on scalars (Value objects), target = inside\n    h = [x1.relu(), (-x1).relu(), x2.relu(), (-x2).relu()]      # the four hidden units\n    r = h[0] + h[1] + h[2] + h[3]                                # |x1| + |x2|\n    z_out, z_in = (r - 1) * gamma, (1 - r) * gamma              # W2 h + b2\n    Z = z_out.exp() + z_in.exp()                                 # softmax denominator\n    return Z.log() - z_in                                        # -log softmax(z)[in]\n\nx1, x2 = Value(0.5), Value(0.25)\nL = diamond_loss(x1, x2)\nL.backward()\n\np_out = 1 / (1 + math.exp(-2 * 4.0 * (0.5 + 0.25 - 1)))          # sigma(2γ(x1+x2-1))\nanalytic = 2 * 4.0 * p_out\nh = 1e-6\nfd = lambda f, a, b: (f(Value(a + h), Value(b)).data - f(Value(a - h), Value(b)).data) / (2 * h)\nprint(f'loss                 : {L.data:.10f}   (= log(1+e^-2) = {math.log(1 + math.exp(-2)):.10f})')\nprint(f'analytic  dL/dx1     : {analytic:.10f}')\nprint(f'engine    dL/dx1, dx2: {x1.grad:.10f}, {x2.grad:.10f}')\nprint(f'finite-difference    : {fd(diamond_loss, 0.5, 0.25):.10f}')\nassert abs(x1.grad - analytic) < 1e-6 and abs(x2.grad - analytic) < 1e-6\nassert abs(fd(diamond_loss, 0.5, 0.25) - analytic) < 1e-6\nprint('three routes, one gradient: 0.9536')",
   "id": "cell-08"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 3.2 Switch to PyTorch autograd\n\nThe same check once more, with PyTorch's tensor-valued engine and Step\n2's `diamond_net` reproduced verbatim. `requires_grad=True` marks the\npoint as a leaf whose gradient we want; `loss.backward()` is our\n`Value.backward()` — the topological sort and reverse sweep — filling\n`point.grad`. PyTorch is your `Value` class with three upgrades:\ntensors instead of scalars, a compiled backend, and a library of\nprimitives. Conceptually nothing else. After this cell you may call\n`.backward()` for the rest of your life with a clear conscience.",
   "id": "cell-09"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "import torch\n\ndef diamond_net(X, gamma=4.0):                     # Step 2's solution, unchanged\n    W1 = torch.tensor([[1., 0.], [-1., 0.], [0., 1.], [0., -1.]], dtype=X.dtype)\n    b1 = torch.zeros(4, dtype=X.dtype)\n    W2 = gamma * torch.tensor([[1., 1., 1., 1.], [-1., -1., -1., -1.]], dtype=X.dtype)\n    b2 = gamma * torch.tensor([-1.0, 1.0], dtype=X.dtype)\n    a1 = X @ W1.T + b1\n    h = torch.relu(a1)\n    logits = h @ W2.T + b2\n    shifted = logits - logits.max(dim=1, keepdim=True).values\n    weights = shifted.exp()\n    probs = weights / weights.sum(dim=1, keepdim=True)\n    return a1, h, logits, probs\n\npoint = torch.tensor([[0.5, 0.25]], dtype=torch.float64,\n                     requires_grad=True)\n_, _, _, probs = diamond_net(point, gamma=4.0)\nloss = -probs[0, 1].log()             # target = inside\nloss.backward()\n\nexpected = 2 * 4.0 * probs.detach()[0, 0]\nprint(point.grad)                      # both entries ≈ 0.9536\nprint(expected)\nassert torch.allclose(point.grad, torch.full((1, 2), analytic, dtype=torch.float64), atol=1e-6)\nassert abs(point.grad[0, 0].item() - x1.grad) < 1e-6\nprint('PyTorch agrees with the engine, the formula, and finite differences')",
   "id": "cell-10"
  },
  {
   "cell_type": "markdown",
   "id": "cell-11",
   "metadata": {},
   "source": "# Part 2 — The Bengio (2003) model\n\n$$x = (C[a_1],\\dots,C[a_k]) \\in \\mathbb R^{kd},\\qquad\n  z = W_2\\tanh(W_1x + b_1) + b_2,\\qquad p = \\operatorname{softmax}(z).$$\n\nThe embedding table $C$ is the new idea: each character becomes a\n*learned vector*, and because $C$ is shared across positions and\ncontexts, evidence about one character informs every context containing\nit. That sharing is what a count table cannot do (Lecture 2 §2).\n\nWe write minibatch SGD by hand — no `torch.optim` until Step 6."
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## Setup: data and tokenizer (Step 1's solution, reproduced)\n\nEvery notebook in this series is self-contained: it re-creates what it\nneeds from earlier steps in one compact cell, so you can run it top to\nbottom without opening the others. On Colab, uncomment the download.",
   "id": "cell-12"
  },
  {
   "cell_type": "code",
   "id": "cell-13",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "# !wget -q https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt\nimport torch\n\nwith open('input.txt') as f:\n    text = f.read()\nvocab = sorted(set(text))\nV = len(vocab)\nstoi = {ch: i for i, ch in enumerate(vocab)}\nitos = {i: ch for i, ch in enumerate(vocab)}\nencode = lambda s: [stoi[c] for c in s]\ndecode = lambda ids: ''.join(itos[i] for i in ids)\nids = torch.tensor(encode(text), dtype=torch.long)\n\nassert len(text) == 1_115_394 and V == 65\nprint(f'{len(text):,} characters, vocab {V}')"
  },
  {
   "cell_type": "code",
   "id": "cell-14",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "import torch.nn.functional as F\ntorch.set_num_threads(1)\n\nn_split = int(0.9 * len(ids))\ntr, va = ids[:n_split], ids[n_split:]\n\ndef windows(data, k):\n    return data.unfold(0, k, 1)[:-1], data[k:]   # (N,k) contexts, targets\n\ndef run(k, d=16, hid=128, steps=20_000, B=64, seed=0, verbose=False):\n    torch.manual_seed(seed)\n    Xtr, Ytr = windows(tr, k)\n    Xva, Yva = windows(va, k)\n    C  = (torch.randn(V, d) * 0.1).requires_grad_()\n    W1 = (torch.randn(k*d, hid) * (5/3) / (k*d)**0.5).requires_grad_()\n    b1 = torch.zeros(hid, requires_grad=True)\n    W2 = (torch.randn(hid, V) * 0.01).requires_grad_()\n    b2 = torch.zeros(V, requires_grad=True)\n    params = [C, W1, b1, W2, b2]\n    fwd = lambda Z: torch.tanh(C[Z].view(Z.shape[0], -1) @ W1 + b1) @ W2 + b2\n\n    for s in range(steps):\n        ix = torch.randint(0, Xtr.shape[0], (B,))\n        loss = F.cross_entropy(fwd(Xtr[ix]), Ytr[ix])\n        for p in params:\n            p.grad = None\n        loss.backward()\n        lr = 0.1 if s < steps * 0.75 else 0.01     # crude two-stage decay\n        with torch.no_grad():\n            for p in params:\n                p -= lr * p.grad\n\n    @torch.no_grad()\n    def ev(Xd, Yd, iters=30, BB=4096):\n        tot = 0.0\n        for _ in range(iters):\n            ix = torch.randint(0, Xd.shape[0], (BB,))\n            tot += F.cross_entropy(fwd(Xd[ix]), Yd[ix]).item()\n        return tot / iters\n\n    return ev(Xtr, Ytr), ev(Xva, Yva), sum(p.numel() for p in params), (C, fwd)"
  },
  {
   "cell_type": "markdown",
   "id": "cell-15",
   "metadata": {},
   "source": "## The context-length sweep — and a surprise\n\nPredict before running: does longer context help?"
  },
  {
   "cell_type": "code",
   "id": "cell-16",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "import math\nprint(f\"{'k':>3} {'params':>8} {'train':>8} {'val':>8} {'bits/ch':>9} {'table would need':>18}\")\nresults = {}\nfor k in (1, 3, 5, 8):\n    trl, val, npar, _ = run(k)\n    results[k] = val\n    print(f'{k:>3} {npar:>8,} {trl:>8.4f} {val:>8.4f} {val/math.log(2):>9.4f}'\n          f' {float(V**(k+1)):>18.1e}', flush=True)"
  },
  {
   "cell_type": "markdown",
   "id": "cell-17",
   "metadata": {},
   "source": "Reference: **2.4805, 1.9821, 2.0018, 2.0329** for $k = 1,3,5,8$.\nRead that carefully — it says three things.\n\n1. **$k{=}1$ reproduces the bigram** (2.48 val against Step 1's 2.45\n   train-set CE), as it must: the $k{=}1$ model's hypothesis class\n   *contains* the bigram table.\n2. **$k{=}3$ buys 0.5 nats for 4,000 extra parameters** — against the\n   $1.8\\times10^{7}$-entry table a trigram counting model would need\n   for the same context. That ratio is the whole argument of Lecture 2.\n3. **$k{=}5$ and $k{=}8$ are *worse* than $k{=}3$**, and this is not a\n   bug. Train loss barely improves either, so it is underfitting, not\n   overfitting: flattening a longer context into one $kd$-vector\n   spreads fixed capacity thinner and hands SGD a harder problem.\n\nPoint 3 is the hinge of the course. The fix is not more parameters. It\nis a way of reading context that does not concatenate everything into a\nflat vector, and that can decide **per input** which earlier positions\nmatter. That is attention."
  },
  {
   "cell_type": "markdown",
   "id": "cell-18",
   "metadata": {},
   "source": "## Samples, and the diagnostic failure\n\nReference output from the $k{=}3$ model (seeded; the cell below\nreproduces it) contains real words in runs and — strikingly — the\n`SPEAKER:` line format with its preceding blank line. But the speakers\nit emits, `ENCHAR:` and `DONLY:`, occur **zero** times in the corpus.\nThree characters of context cannot span a name, so the model has\nlearned the *shape* of a speaker line with no ability to remember any\nparticular one. Hold that thought until Step 6, where the transformer\nemits `WARWICK:` and `Montague` correctly spelled."
  },
  {
   "cell_type": "code",
   "id": "cell-19",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "torch.manual_seed(0)\n_, _, _, (C3, fwd3) = run(3, steps=20_000)\nctx, out = [stoi['\\n']] * 3, []\nwith torch.no_grad():\n    for _ in range(300):\n        p = F.softmax(fwd3(torch.tensor([ctx]))[0], dim=-1)\n        i = int(torch.multinomial(p, 1))\n        out.append(itos[i]); ctx = ctx[1:] + [i]\nprint(''.join(out))"
  },
  {
   "cell_type": "markdown",
   "id": "cell-20",
   "metadata": {},
   "source": "## The embedding table, visualized\n\nTwo principal components of $C$. Vowels cluster; digits and punctuation\nseparate from letters; capitals separate from lowercase. **Nobody told\nthe model these categories exist** — they fall out of predicting the\nnext character. Geometry learned from prediction alone is the thesis of\nthe whole enterprise, and here it is at 11,601 parameters."
  },
  {
   "cell_type": "code",
   "id": "cell-21",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "import matplotlib.pyplot as plt\nCd = C3.detach()\nCc = Cd - Cd.mean(0)\nU, S, Vh = torch.linalg.svd(Cc, full_matrices=False)\nxy = (Cc @ Vh[:2].T).numpy()\nplt.figure(figsize=(9, 7))\nfor i, (x, y) in enumerate(xy):\n    ch = itos[i]\n    label = {'\\n': '\\\\n', ' ': '␣'}.get(ch, ch)\n    color = ('tab:red' if ch in 'aeiouAEIOU' else\n             'tab:blue' if ch.isupper() else\n             'tab:green' if ch.islower() else 'tab:gray')\n    plt.text(x, y, label, color=color, fontsize=13, ha='center')\nplt.xlim(xy[:,0].min()*1.15, xy[:,0].max()*1.15)\nplt.ylim(xy[:,1].min()*1.15, xy[:,1].max()*1.15)\nplt.title('character embeddings, first two principal components\\n'\n          'red = vowels, blue = capitals, green = lowercase, grey = other')\nplt.tight_layout()"
  },
  {
   "cell_type": "markdown",
   "id": "cell-22",
   "metadata": {},
   "source": "→ Continue with [Step 4](../project/step-4): one attention head, and\nthe proofs from Lecture 4 executed as assertions."
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}