{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c022",
   "metadata": {},
   "source": [
    "# Step 8 \u2014 Capstone: Track C worked example\n",
    "\n",
    "**Solution notebook.** Step 8 is open-ended \u2014 you pick one of three\n",
    "tracks \u2014 so this notebook does not 'solve' it. Instead it works Track C\n",
    "(mechanistic interpretability) end to end as a **model for the level of\n",
    "rigour expected**: a clear question, an experiment, a figure, and a\n",
    "causal confirmation.\n",
    "\n",
    "Tracks A (scale/fine-tune/LoRA) and B (alignment) are sketched at the\n",
    "end with the mathematics you need but not executed \u2014 they are yours to\n",
    "do.\n",
    "\n",
    "Requires `step6_model.pt`. Runs in about a minute."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c023",
   "metadata": {},
   "source": [
    "## Setup: data and tokenizer (Step 1's solution, reproduced)\n",
    "\n",
    "Every notebook in this series is self-contained: it re-creates what it\n",
    "needs from earlier steps in one compact cell, so you can run it top to\n",
    "bottom without opening the others. On Colab, uncomment the download."
   ]
  },
  {
   "cell_type": "code",
   "id": "c024",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# !wget -q https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt\n",
    "import torch\n",
    "\n",
    "with open('input.txt') as f:\n",
    "    text = f.read()\n",
    "vocab = sorted(set(text))\n",
    "V = len(vocab)\n",
    "stoi = {ch: i for i, ch in enumerate(vocab)}\n",
    "itos = {i: ch for i, ch in enumerate(vocab)}\n",
    "encode = lambda s: [stoi[c] for c in s]\n",
    "decode = lambda ids: ''.join(itos[i] for i in ids)\n",
    "ids = torch.tensor(encode(text), dtype=torch.long)\n",
    "\n",
    "assert len(text) == 1_115_394 and V == 65\n",
    "print(f'{len(text):,} characters, vocab {V}')"
   ]
  },
  {
   "cell_type": "code",
   "id": "c025",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import math\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class Head(nn.Module):\n",
    "    def __init__(self, d, dh, T):\n",
    "        super().__init__()\n",
    "        self.key   = nn.Linear(d, dh, bias=False)\n",
    "        self.query = nn.Linear(d, dh, bias=False)\n",
    "        self.value = nn.Linear(d, dh, bias=False)\n",
    "        self.register_buffer('tril', torch.tril(torch.ones(T, T)))\n",
    "\n",
    "    def forward(self, x):\n",
    "        B, T, _ = x.shape\n",
    "        q, k, v = self.query(x), self.key(x), self.value(x)\n",
    "        s = q @ k.transpose(-2, -1) / k.shape[-1]**0.5\n",
    "        s = s.masked_fill(self.tril[:T, :T] == 0, float('-inf'))\n",
    "        A = F.softmax(s, dim=-1)\n",
    "        self.A = A.detach()          # stashed for Step 8 interpretability\n",
    "        return A @ v\n",
    "\n",
    "class MHA(nn.Module):\n",
    "    def __init__(self, d, H, T):\n",
    "        super().__init__()\n",
    "        self.heads = nn.ModuleList(Head(d, d // H, T) for _ in range(H))\n",
    "        self.proj = nn.Linear(d, d)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.proj(torch.cat([h(x) for h in self.heads], dim=-1))\n",
    "\n",
    "class MLP(nn.Module):\n",
    "    def __init__(self, d):\n",
    "        super().__init__()\n",
    "        self.fc, self.proj = nn.Linear(d, 4 * d), nn.Linear(4 * d, d)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.proj(F.gelu(self.fc(x)))\n",
    "\n",
    "class Block(nn.Module):\n",
    "    def __init__(self, d, H, T):\n",
    "        super().__init__()\n",
    "        self.ln1, self.attn = nn.LayerNorm(d), MHA(d, H, T)\n",
    "        self.ln2, self.mlp  = nn.LayerNorm(d), MLP(d)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x + self.attn(self.ln1(x))\n",
    "        return x + self.mlp(self.ln2(x))\n",
    "\n",
    "class GPT(nn.Module):\n",
    "    def __init__(self, V, T, d, H, L):\n",
    "        super().__init__()\n",
    "        self.T = T\n",
    "        self.tok = nn.Embedding(V, d)\n",
    "        self.pos = nn.Embedding(T, d)\n",
    "        self.blocks = nn.ModuleList(Block(d, H, T) for _ in range(L))\n",
    "        self.lnf = nn.LayerNorm(d)\n",
    "        self.head = nn.Linear(d, V, bias=False)\n",
    "        for blk in self.blocks:   # GPT-2 residual scaling (Lecture 6 \u00a71)\n",
    "            nn.init.normal_(blk.attn.proj.weight, std=0.02 / math.sqrt(2 * L))\n",
    "            nn.init.normal_(blk.mlp.proj.weight,  std=0.02 / math.sqrt(2 * L))\n",
    "\n",
    "    def forward(self, idx, targets=None):\n",
    "        B, T = idx.shape\n",
    "        x = self.tok(idx) + self.pos(torch.arange(T))\n",
    "        for blk in self.blocks:\n",
    "            x = blk(x)\n",
    "        logits = self.head(self.lnf(x))\n",
    "        loss = None if targets is None else F.cross_entropy(\n",
    "            logits.view(-1, logits.size(-1)), targets.reshape(-1))\n",
    "        return logits, loss\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def generate(self, idx, n_new, temperature=1.0):\n",
    "        for _ in range(n_new):\n",
    "            logits, _ = self(idx[:, -self.T:])\n",
    "            probs = F.softmax(logits[:, -1, :] / temperature, dim=-1)\n",
    "            idx = torch.cat([idx, torch.multinomial(probs, 1)], dim=1)\n",
    "        return idx\n"
   ]
  },
  {
   "cell_type": "code",
   "id": "c026",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "torch.manual_seed(1337)\n",
    "model = GPT(V=65, T=64, d=128, H=4, L=4)\n",
    "model.load_state_dict(torch.load('step6_model.pt'))\n",
    "model.eval()\n",
    "print(f'loaded trained model, {sum(p.numel() for p in model.parameters()):,} parameters')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c027",
   "metadata": {},
   "source": [
    "**A note on what this worked example finds.** Track C is written\n",
    "honestly: on a model this small the investigation finds *one clean\n",
    "circuit and one clean absence*, and both are worth more than a forced\n",
    "success. We identify a textbook previous-token head and confirm it\n",
    "causally (C.1, C.3), and we establish rigorously that the model has\n",
    "**no** working induction head (C.2) \u2014 the expected outcome at this\n",
    "scale, and a demonstration of how to prove a negative. The two are\n",
    "connected: a previous-token head is the known *prerequisite* for\n",
    "induction, so this model has the first ingredient of the two-head\n",
    "circuit and not the second."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c028",
   "metadata": {},
   "source": [
    "## Track C.1 \u2014 The attention atlas\n",
    "\n",
    "**Question:** of this model's $L\\times H = 16$ heads, how many compute\n",
    "something a human can name?\n",
    "\n",
    "We average each head's attention pattern over many real sequences and\n",
    "classify by where the mass sits relative to the diagonal."
   ]
  },
  {
   "cell_type": "code",
   "id": "c029",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "T = 64\n",
    "@torch.no_grad()\n",
    "def head_patterns(n_seq=32, seed=0):\n",
    "    torch.manual_seed(seed)\n",
    "    acc = torch.zeros(len(model.blocks), len(model.blocks[0].attn.heads), T, T)\n",
    "    for _ in range(n_seq):\n",
    "        i = torch.randint(0, len(ids) - T - 1, (1,)).item()\n",
    "        x = model.tok(ids[i:i+T][None]) + model.pos(torch.arange(T))\n",
    "        for li, blk in enumerate(model.blocks):\n",
    "            xn = blk.ln1(x)\n",
    "            for hi, h in enumerate(blk.attn.heads):\n",
    "                h(xn)\n",
    "                acc[li, hi] += h.A[0]\n",
    "            x = blk(x)\n",
    "    return acc / n_seq\n",
    "\n",
    "pat = head_patterns()\n",
    "L, H = pat.shape[:2]\n",
    "fig, axes = plt.subplots(L, H, figsize=(2.1*H, 2.1*L))\n",
    "for li in range(L):\n",
    "    for hi in range(H):\n",
    "        ax = axes[li, hi]\n",
    "        ax.imshow(pat[li, hi], cmap='Blues')\n",
    "        ax.set_xticks([]); ax.set_yticks([])\n",
    "        if hi == 0: ax.set_ylabel(f'layer {li}', fontsize=9)\n",
    "        if li == 0: ax.set_title(f'head {hi}', fontsize=9)\n",
    "plt.suptitle('mean attention pattern per head, averaged over 32 real sequences')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "id": "c030",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Quantify: how much mass sits exactly one position back, on the diagonal,\n",
    "# or on the very first token (a common 'attention sink')?\n",
    "print(f\"{'head':<10}{'prev-token':>12}{'self':>8}{'first-tok':>11}  guess\")\n",
    "for li in range(L):\n",
    "    for hi in range(H):\n",
    "        A = pat[li, hi]\n",
    "        rows = torch.arange(2, T)\n",
    "        prev = A[rows, rows-1].mean().item()\n",
    "        self_ = A[rows, rows].mean().item()\n",
    "        first = A[rows, 0].mean().item()\n",
    "        guess = ('previous-token' if prev > 0.35 else\n",
    "                 'self/current'   if self_ > 0.35 else\n",
    "                 'first-token sink' if first > 0.35 else 'diffuse')\n",
    "        print(f'L{li}H{hi:<7}{prev:>12.3f}{self_:>8.3f}{first:>11.3f}  {guess}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c031",
   "metadata": {},
   "source": [
    "**One head stands out: L0H3 puts 0.977 of its attention exactly one\n",
    "position back** \u2014 a textbook *previous-token head*, and it is stable to\n",
    "the third decimal across random seeds (0.977, 0.977, 0.979). The rest\n",
    "are diffuse, with a mild previous-token lean in layer 1. That single\n",
    "sharp head is our object of study; the next two cells confirm it is\n",
    "real and matters."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c032",
   "metadata": {},
   "source": [
    "## Track C.2 \u2014 Ablate and confirm the previous-token head\n",
    "\n",
    "A stable pattern is a correlation. To show L0H3 *matters*, we\n",
    "**ablate** it \u2014 zero its value projection, which deletes exactly its\n",
    "OV-circuit term from the additive decomposition of Lecture 5, Prop.\n",
    "1.2 \u2014 and measure validation loss against ablating a diffuse control\n",
    "head. This is the causal step interpretability demands."
   ]
  },
  {
   "cell_type": "code",
   "id": "c033",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import copy\n",
    "\n",
    "n_split = int(0.9 * len(ids))\n",
    "val_ids = ids[n_split:]\n",
    "\n",
    "@torch.no_grad()\n",
    "def val_loss(m, iters=60, B=64, T=64, seed=0):\n",
    "    torch.manual_seed(seed)\n",
    "    tot = 0.0\n",
    "    for _ in range(iters):\n",
    "        ix = torch.randint(len(val_ids) - T - 1, (B,))\n",
    "        xb = torch.stack([val_ids[i:i+T] for i in ix])\n",
    "        yb = torch.stack([val_ids[i+1:i+T+1] for i in ix])\n",
    "        tot += m(xb, yb)[1].item()\n",
    "    return tot / iters\n",
    "\n",
    "def ablate_head(layer, head_idx):\n",
    "    m = copy.deepcopy(model)\n",
    "    with torch.no_grad():\n",
    "        m.blocks[layer].attn.heads[head_idx].value.weight.zero_()\n",
    "    return m\n",
    "\n",
    "base = val_loss(model)\n",
    "print(f'{\"intact\":<34}val {base:.4f}')\n",
    "for l, h, name in [(0, 3, 'L0H3 (previous-token head)'),\n",
    "                   (0, 1, 'L0H1 (diffuse, control)')]:\n",
    "    v = val_loss(ablate_head(l, h))\n",
    "    print(f'ablate {name:<27}val {v:.4f}   (+{v-base:.4f})')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c034",
   "metadata": {},
   "source": [
    "**A clean causal result.** Reference numbers: ablating the\n",
    "previous-token head **L0H3 costs +2.16 nats** of validation loss \u2014\n",
    "catapulting the model from 1.54 back past the bigram baseline \u2014 while\n",
    "ablating the diffuse control **L0H1 costs only +0.28**. An 8\u00d7 gap. The\n",
    "one head we could interpret is also, by a wide margin, the one the\n",
    "model most depends on. Knowing what a component *does* and knowing it\n",
    "*matters* are different claims, and this cell establishes the second."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c035",
   "metadata": {},
   "source": [
    "## Track C.3 \u2014 A rigorous negative: no induction head\n",
    "\n",
    "Now the capability the atlas did **not** turn up. An induction head\n",
    "implements *find an earlier copy of the current token, predict what\n",
    "followed it* \u2014 the mechanism behind in-context learning. The standard\n",
    "test builds a random sequence and repeats it, `[base | base]`: a model\n",
    "with induction predicts the second copy far better than the first,\n",
    "because it can copy from the first. Random tokens cannot be memorized,\n",
    "so any gain is genuine in-context copying.\n",
    "\n",
    "**Proving an absence requires the same care as proving a presence.**"
   ]
  },
  {
   "cell_type": "code",
   "id": "c036",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def copy_gain(half=32, trials=300, seed=0):\n",
    "    torch.manual_seed(seed)\n",
    "    first, second = 0.0, 0.0\n",
    "    for _ in range(trials):\n",
    "        base = torch.randint(0, 65, (half,))\n",
    "        seq = torch.cat([base, base])[None]\n",
    "        lp = F.cross_entropy(model(seq[:, :-1])[0][0], seq[0, 1:],\n",
    "                             reduction='none')\n",
    "        first  += lp[:half-1].mean().item()\n",
    "        second += lp[half:].mean().item()\n",
    "    return first/trials, second/trials\n",
    "\n",
    "f, s = copy_gain()\n",
    "print(f'loss on first  (novel)  copy : {f:.4f} nats')\n",
    "print(f'loss on second (repeat) copy : {s:.4f} nats')\n",
    "print(f'in-context copying gain      : {f - s:+.4f} nats')\n",
    "print(f'(for scale, ln V = {math.log(65):.4f}; these are >> that because\\n'\n",
    "      f' random uniform text is wildly off this model\\'s distribution)')"
   ]
  },
  {
   "cell_type": "code",
   "id": "c037",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Per-position loss across the repeated half: an induction head would make\n",
    "# this DROP sharply after the first couple of tokens. It stays flat.\n",
    "@torch.no_grad()\n",
    "def per_position(half=32, trials=400, seed=0):\n",
    "    torch.manual_seed(seed)\n",
    "    acc = torch.zeros(half)\n",
    "    for _ in range(trials):\n",
    "        base = torch.randint(0, 65, (half,))\n",
    "        seq = torch.cat([base, base])[None]\n",
    "        lp = F.cross_entropy(model(seq[:, :-1])[0][0], seq[0, 1:],\n",
    "                             reduction='none')\n",
    "        acc += lp[half-1:2*half-1]\n",
    "    return acc / trials\n",
    "\n",
    "pp = per_position()\n",
    "print('per-position loss in the repeated half:')\n",
    "print('  ', '  '.join(f'{v:.1f}' for v in pp[:8].tolist()), '...')\n",
    "print(f'  flat at ~{pp.mean():.1f} nats \u2014 no drop, so no copying')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c038",
   "metadata": {},
   "source": [
    "**The verdict, stated as carefully as a positive result.** The copying\n",
    "gain is $\\approx +0.02$ nats \u2014 indistinguishable from zero against the\n",
    "$\\sim 9.5$-nat scale \u2014 and the per-position loss is *flat* across the\n",
    "repeated half, where an induction head would produce a sharp drop after\n",
    "the first token or two. **This model has no working induction head.**\n",
    "\n",
    "That is the expected result and not a disappointment. Induction heads\n",
    "emerge with scale and depth (Olsson et al. found their onset is a\n",
    "*phase change* during training of larger models); a 4-layer,\n",
    "816,640-parameter character model on 1 MB of text is below that\n",
    "threshold. What it *does* have is the previous-token head of C.2 \u2014 the\n",
    "documented **prerequisite** for the two-head induction circuit. The\n",
    "model has assembled the first component and not the second, which is a\n",
    "more informative place to have stopped than either 'found it' or\n",
    "'found nothing.'\n",
    "\n",
    "**This is the standard of evidence Lecture 8 holds interpretability\n",
    "to**: a stable pattern, a causal intervention (C.2), and \u2014 for a\n",
    "claimed absence \u2014 a behavioural test *and* a mechanistic one that\n",
    "agree (C.3)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c039",
   "metadata": {},
   "source": [
    "## Tracks A and B \u2014 the mathematics, unexecuted\n",
    "\n",
    "### Track A: LoRA\n",
    "Freeze $W_0$, train $W = W_0 + BA$ with $B \\in \\mathbb R^{d\\times r}$,\n",
    "$A \\in \\mathbb R^{r\\times d}$, $B$ initialized to zero so $W = W_0$ at\n",
    "the start. Sweep $r\\in\\{1,2,4,8,16\\}$ and plot target-domain loss\n",
    "against trainable parameters. **The hypothesis under test is that the\n",
    "*update* is low rank** \u2014 not that $W_0$ is \u2014 so the value of $r$ where\n",
    "quality saturates is a measurement of the intrinsic rank of your\n",
    "adaptation.\n",
    "\n",
    "```python\n",
    "class LoRALinear(nn.Module):\n",
    "    def __init__(self, base, r):\n",
    "        super().__init__()\n",
    "        self.base = base\n",
    "        for p in self.base.parameters():\n",
    "            p.requires_grad = False\n",
    "        self.A = nn.Parameter(torch.randn(r, base.in_features) * 0.01)\n",
    "        self.B = nn.Parameter(torch.zeros(base.out_features, r))\n",
    "    def forward(self, x):\n",
    "        return self.base(x) + (x @ self.A.T) @ self.B.T\n",
    "```\n",
    "\n",
    "### Track B: KL-regularized preference optimization\n",
    "Define a programmatic reward $r(y)$ (e.g. $+1$ per line of the right\n",
    "syllable count, or a penalty for a banned word). The objective\n",
    "$$\\max_\\pi \\mathbb E_\\pi[r] - \\beta\\,D_{\\mathrm{KL}}(\\pi\\|\\pi_{\\mathrm{ref}})$$\n",
    "has the closed-form optimum $\\pi^\\ast \\propto \\pi_{\\mathrm{ref}}\n",
    "e^{r/\\beta}$ (Lecture 8, Thm 3.2). You cannot sample from it directly \u2014\n",
    "$Z$ is intractable \u2014 but **best-of-$n$ rejection sampling approximates\n",
    "it cheaply**: draw $n$ completions from $\\pi_{\\mathrm{ref}}$, keep the\n",
    "highest-reward one. Sweep $n$ (equivalently, sweep $\\beta$), and for\n",
    "each plot mean reward against the empirical\n",
    "$D_{\\mathrm{KL}}(\\pi\\|\\pi_{\\mathrm{ref}})$. You will get a\n",
    "reward-versus-KL trade-off curve, and \u2014 if your reward is at all naive\n",
    "\u2014 you will find the exploit that games it. Describing that exploit\n",
    "precisely is the most valuable paragraph you can write for the\n",
    "showcase.\n",
    "\n",
    "## Your bits/char table\n",
    "\n",
    "| step | model | bits/char |\n",
    "|---|---|---|\n",
    "| 1 | bigram counts | 3.54 |\n",
    "| 3 | Bengio MLP, $k=3$ | 2.86 |\n",
    "| 6 | GPT, 816,640 params | **2.19** |\n",
    "| \u2014 | Shannon's estimate for English | ~1 |\n",
    "\n",
    "Put that on your last slide. It is the story of the course in four\n",
    "rows: from a table of pair counts to a transformer, closing roughly\n",
    "half the gap to a human, with every line of code written by you."
   ]
  }
 ],
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  }
 }
}