{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e1b20e91",
   "metadata": {},
   "source": [
    "# HW 3: Understanding a GPT, line by line\n",
    "\n",
    "In this assignment you will take apart a complete, working GPT — the same kind of model behind ChatGPT, just very, very small.\n",
    "The code is Andrej Karpathy's [**microgpt**](https://gist.github.com/karpathy/8627fe009c40f57531cb18360106ce95) (see also his [blog post](http://karpathy.github.io/2026/02/12/microgpt/)): about 200 lines of plain Python, with no libraries, that trains a model on 32,000 baby names and then invents new ones.\n",
    "We have split his code into notebook cells and added questions between them. Lines we added or changed are marked `# [HW3]`.\n",
    "\n",
    "**This homework is not about writing code.** You will *read*, *predict*, *run*, *investigate*, and *modify* code. The goal is for you to understand how a GPT works well enough to explain it to someone else.\n",
    "\n",
    "### How you will be graded\n",
    "\n",
    "This notebook should be submitted, but it will not be graded. Instead, your grade will come from a **12–15 minute one-on-one check-off** with an instructor, where you will answer questions about **four topics chosen by the instructor** from this list:\n",
    "\n",
    "- Your **Part 6 experiments** (explain one: what you predicted, what happened, and why)\n",
    "- **Tokenization** (Part 2)\n",
    "- **Embeddings** (Part 4)\n",
    "- **Attention** (Part 4)\n",
    "- **MLP, residual connections, and normalization** (Part 4)\n",
    "- **Training** (Parts 3 and 5)\n",
    "- **Generation and sampling** (Part 5)\n",
    "\n",
    "Questions come in three flavors:\n",
    "- **Trace:** \"What is the shape of this list?\" \"What are the inputs and targets for the name `emma`?\"\n",
    "- **Explain why:** \"Why does the model need position embeddings?\"\n",
    "- **Predict:** \"What would happen if we removed this line?\"\n",
    "\n",
    "Each topic is scored 0–3 (12 points total). There are no retakes, so prepare well. The **✅ Check yourself** boxes throughout the notebook are good practice. You may bring this notebook to the check-off, but you should be able to explain things in your own words.\n",
    "\n",
    "### How to work\n",
    "\n",
    "- **Predict before you run.** Whenever a cell asks for a prediction, write it down *first*. Being wrong is fine and useful — that's where the learning happens. Being \"right\" because you ran the cell first teaches you nothing.\n",
    "- **Answer the written questions** in the \"Your answer\" cells. Short answers (1–3 sentences) are fine.\n",
    "- **Using AI:** You are encouraged to use the course tutor chatbot (or other AI tools) to *understand* the code. But remember that you will be on your own at the check-off. If an AI explains something to you, close the chat and try to explain it back in your own words.\n",
    "- **Pace yourself.** Parts 0–3 go quickly. **Part 4 (the model) is the heart of the assignment** and is where you should expect to spend most of your time: it is long, it is dense, and it is where the ideas from the videos become concrete. Don't try to do it in one sitting. Parts 5 and 6 are mostly running code and thinking about what it shows.\n",
    "\n",
    "### Contents\n",
    "\n",
    "| Part | Topic |\n",
    "|---|---|\n",
    "| 0 | Python you'll need (in class) |\n",
    "| 1 | The big picture |\n",
    "| 2 | Data and the tokenizer |\n",
    "| 3 | Autograd as a black box |\n",
    "| 4 | The model |\n",
    "| 5 | Training and generating |\n",
    "| 6 | Experiments |\n",
    "| 7 | Going further (optional) |"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70752c77",
   "metadata": {},
   "source": [
    "---\n",
    "# Part 0: Python you'll need\n",
    "\n",
    "*(We will work through this part together in class.)*\n",
    "\n",
    "microgpt uses only a handful of Python features, but it uses them *a lot*, often packed into a single line.\n",
    "Every example below is taken from (or modeled on) a real line of microgpt. For each cell: **read it, write your prediction in the comment, then run it.**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "baba0cfa",
   "metadata": {},
   "source": [
    "## 0.1 Lists and lists of lists\n",
    "\n",
    "microgpt stores every vector as a Python **list** of numbers, and every matrix as a **list of lists** (a list of rows)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "336b5e46",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = [0.5, -1.0, 2.0]\n",
    "W = [[1, 2, 3],\n",
    "     [4, 5, 6]]\n",
    "\n",
    "# Predict each value BEFORE running:\n",
    "# len(x)     -> ?\n",
    "# len(W)     -> ?\n",
    "# len(W[0])  -> ?\n",
    "# W[1]       -> ?\n",
    "# W[1][2]    -> ?\n",
    "print(len(x), len(W), len(W[0]), W[1], W[1][2])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "deb1f99c",
   "metadata": {},
   "source": [
    "**Slicing.** `a[start:stop]` gives the elements from index `start` up to *but not including* `stop`.\n",
    "In the attention code, a vector `q` of length 16 is cut into 4 \"heads\" of size 4 like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16b5a422",
   "metadata": {},
   "outputs": [],
   "source": [
    "q = list(range(16))   # [0, 1, 2, ..., 15]\n",
    "head_dim = 4\n",
    "h = 2\n",
    "hs = h * head_dim\n",
    "\n",
    "# Predict: what is q[hs:hs+head_dim]?  -> ?\n",
    "print(q[hs:hs+head_dim])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bfb33b7",
   "metadata": {},
   "source": [
    "## 0.2 List comprehensions\n",
    "\n",
    "A **list comprehension** builds a new list by looping over an old one, in a single line. These two cells do the same thing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4284edcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = [1.0, 2.0, 3.0]\n",
    "\n",
    "# The long way\n",
    "doubled = []\n",
    "for xi in x:\n",
    "    doubled.append(xi * 2)\n",
    "print(doubled)\n",
    "\n",
    "# The list-comprehension way\n",
    "doubled = [xi * 2 for xi in x]\n",
    "print(doubled)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63674446",
   "metadata": {},
   "source": [
    "microgpt also uses `_` as a variable name when it doesn't care about the loop variable, and it nests comprehensions to build matrices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88caefc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Predict the output of each line BEFORE running:\n",
    "print([0 for _ in range(3)])                     # -> ?\n",
    "print([[0 for _ in range(3)] for _ in range(2)]) # -> ?\n",
    "print([xi for xi in [3, -1, 4, -1, 5] if xi > 0])  # -> ?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67a97380",
   "metadata": {},
   "source": [
    "## 0.3 `zip`, `enumerate`, and `sum`\n",
    "\n",
    "- `zip(a, b)` walks through two lists **side by side**, giving pairs `(a[0], b[0]), (a[1], b[1]), ...`\n",
    "- `enumerate(a)` gives pairs `(index, element)`.\n",
    "- `sum(...)` adds everything up.\n",
    "\n",
    "microgpt adds a token embedding and a position embedding like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d210ef93",
   "metadata": {},
   "outputs": [],
   "source": [
    "tok_emb = [1, 2, 3]\n",
    "pos_emb = [10, 20, 30]\n",
    "\n",
    "# Predict -> ?\n",
    "x = [t + p for t, p in zip(tok_emb, pos_emb)]\n",
    "print(x)\n",
    "\n",
    "for i, value in enumerate(['a', 'b', 'c']):\n",
    "    print(i, value)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f05d820f",
   "metadata": {},
   "source": [
    "Here is the most important one-liner in microgpt. It computes a **dot product** of two vectors:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e68cecda",
   "metadata": {},
   "outputs": [],
   "source": [
    "w = [1, 2, 3]\n",
    "x = [4, 5, 6]\n",
    "\n",
    "# Predict (compute by hand) -> ?\n",
    "print(sum(wi * xi for wi, xi in zip(w, x)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdf49a56",
   "metadata": {},
   "source": [
    "## 0.4 Functions: matrix-vector multiplication\n",
    "\n",
    "Put the dot product inside a list comprehension and you get **matrix-vector multiplication**. This is microgpt's `linear` function, exactly as Karpathy wrote it.\n",
    "Each row `wo` of the matrix `w` is dotted with the input `x`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb9ea8fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def linear(x, w):\n",
    "    return [sum(wi * xi for wi, xi in zip(wo, x)) for wo in w]\n",
    "\n",
    "W = [[1, 0, 0],\n",
    "     [1, 1, 1]]\n",
    "x = [2, 3, 4]\n",
    "\n",
    "# Predict (by hand): linear(x, W) -> ?\n",
    "# If W has 2 rows and 3 columns, how long is the output? -> ?\n",
    "print(linear(x, W))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "427bfbed",
   "metadata": {},
   "source": [
    "**Remember this:** if `w` has `nout` rows and `nin` columns, then `linear(x, w)` takes a vector of length `nin` and returns a vector of length `nout`. You will use this fact many times."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54db8dbe",
   "metadata": {},
   "source": [
    "## 0.5 Dictionaries and f-strings\n",
    "\n",
    "microgpt keeps all of its weight matrices in a **dictionary** called `state_dict`, looked up by name. Names are built with **f-strings**, which substitute variables inside `{}`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac5202eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "state_dict = {'wte': [[0.1, 0.2], [0.3, 0.4]], 'layer0.attn_wq': [[1, 2], [3, 4]]}\n",
    "\n",
    "i = 0\n",
    "name = f'layer{i}.attn_wq'\n",
    "print(name)\n",
    "print(state_dict[name])\n",
    "\n",
    "# Predict -> ?\n",
    "print(state_dict['wte'][1])\n",
    "print(list(state_dict.keys()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee66be0a",
   "metadata": {},
   "source": [
    "## 0.6 Classes: what does `a + b` mean?\n",
    "\n",
    "When `a` and `b` are ordinary numbers, `a + b` is ordinary addition. But Python lets a **class** decide what `+` means for its objects, by defining a method named `__add__`.\n",
    "When you write `a + b`, Python secretly calls `a.__add__(b)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c07c5079",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Vec2:\n",
    "    def __init__(self, x, y):\n",
    "        self.x = x\n",
    "        self.y = y\n",
    "\n",
    "    def __add__(self, other):\n",
    "        print(\"  (__add__ was called!)\")\n",
    "        return Vec2(self.x + other.x, self.y + other.y)\n",
    "\n",
    "a = Vec2(1, 2)\n",
    "b = Vec2(10, 20)\n",
    "c = a + b\n",
    "\n",
    "# Predict: c.x and c.y -> ?\n",
    "print(c.x, c.y)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c43dda2",
   "metadata": {},
   "source": [
    "This is exactly how microgpt's `Value` class works (Part 3): every `+`, `*`, `exp`, etc. on a `Value` both computes the answer *and* secretly remembers how the answer was computed."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a6f493b",
   "metadata": {},
   "source": [
    "## 0.7 Random choices\n",
    "\n",
    "To generate text, microgpt picks the next character at random, with some characters more likely than others. `random.choices(options, weights=...)` does this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b44c8850",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "from collections import Counter\n",
    "\n",
    "options = ['a', 'b', 'c']\n",
    "weights = [0.7, 0.2, 0.1]\n",
    "picks = [random.choices(options, weights=weights)[0] for _ in range(1000)]\n",
    "\n",
    "# Predict: roughly how many of the 1000 picks will be 'a'? -> ?\n",
    "print(Counter(picks))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74fc25c2",
   "metadata": {},
   "source": [
    "> **✅ Check yourself** — these are the kinds of questions you'll be asked at the check-off. Can you answer them out loud, in your own words, without looking at your notes?\n",
    ">\n",
    "> - If `w` is a list of 64 lists, each of length 16, what is `len(linear(x, w))`? What must `len(x)` be?\n",
    "> - What does `[t + p for t, p in zip(a, b)]` compute?\n",
    "> - When `a` and `b` are objects of some class, what actually happens when Python runs `a * b`?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a9cd012",
   "metadata": {},
   "source": [
    "---\n",
    "# Part 1: The big picture\n",
    "\n",
    "Before looking at any details, let's watch the whole thing work.\n",
    "\n",
    "**Do this now:** From the menu, choose **Kernel → Restart Kernel and Run All Cells**. The whole notebook takes about 2 minutes to run (most of that is training in Part 5). Nothing in this notebook will stop with an error — the answer cells are just text.\n",
    "\n",
    "While it runs, scroll down to **Part 5** and watch the training loss go down. Then look at the names the model generates at the end of Part 5."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bdd2a79",
   "metadata": {},
   "source": [
    "## 1.1 A map of microgpt\n",
    "\n",
    "Every GPT, from microgpt to the largest models, is made of the same pieces. Here is a map of the code you'll be reading, in order."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd3b502b",
   "metadata": {},
   "source": [
    "| # | Section | Where | In one sentence |\n",
    "|---|---|---|---|\n",
    "| 1 | **Dataset** | Part 2 | A list of 32,000 names. |\n",
    "| 2 | **Tokenizer** | Part 2 | ? |\n",
    "| 3 | **Autograd** (`Value`) | Part 3 | ? |\n",
    "| 4 | **Parameters** (`state_dict`) | Part 4 | ? |\n",
    "| 5 | **Model** (`gpt`): embeddings → attention → MLP → logits | Part 4 | ? |\n",
    "| 6 | **Training loop**: loss → backward → update | Part 5 | ? |\n",
    "| 7 | **Inference**: sampling new names | Part 5 | ? |\n",
    "\n",
    "**Before reading further**, write your best guess for what each section does, based on the videos you've watched and what you saw when you ran the notebook. Don't worry about being right — you'll come back and revise this in Part 7."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27b78340",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "658d277c",
   "metadata": {},
   "source": [
    "---\n",
    "# Part 2: Data and the tokenizer\n",
    "\n",
    "We begin with Karpathy's code. (Karpathy's original downloads `names.txt` from the internet; we've included it in the homework folder instead.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d38ea0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os       # os.path.exists\n",
    "import math     # math.log, math.exp\n",
    "import random   # random.seed, random.choices, random.gauss, random.shuffle\n",
    "import time     # [HW3] for timing\n",
    "random.seed(42) # Let there be order among chaos\n",
    "\n",
    "# Let there be a Dataset `docs`: list[str] of documents (e.g. a list of names)\n",
    "docs = [line.strip() for line in open('names.txt') if line.strip()]  # [HW3] load the local copy\n",
    "random.shuffle(docs)\n",
    "print(f\"num docs: {len(docs)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70b04ce6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Look at the data\n",
    "print(docs[:10])\n",
    "print(\"longest name:\", max(docs, key=len), len(max(docs, key=len)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b5d2a3d",
   "metadata": {},
   "source": [
    "## 2.1 The tokenizer\n",
    "\n",
    "Neural networks work with numbers, not letters. A **tokenizer** converts text to a list of integers (\"tokens\") and back.\n",
    "Real GPTs use tokens that are chunks of words; microgpt uses one token per character.\n",
    "\n",
    "**Predict:** The names contain only lowercase letters `a`–`z`. After you read the next cell, predict what `vocab_size` will be *before* running it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c82e530",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let there be a Tokenizer to translate strings to sequences of integers (\"tokens\") and back\n",
    "uchars = sorted(set(''.join(docs))) # unique characters in the dataset become token ids 0..n-1\n",
    "BOS = len(uchars) # token id for a special Beginning of Sequence (BOS) token\n",
    "vocab_size = len(uchars) + 1 # total number of unique tokens, +1 is for BOS\n",
    "print(f\"vocab size: {vocab_size}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45612b28",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a90001d",
   "metadata": {},
   "source": [
    "## 2.2 Encoding and decoding\n",
    "\n",
    "microgpt doesn't define separate `encode`/`decode` functions — it does this inline, in the training loop and in inference. Here they are pulled out so you can experiment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eaa6b8d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode(name):\n",
    "    return [BOS] + [uchars.index(ch) for ch in name] + [BOS]   # exactly as in the training loop\n",
    "\n",
    "def decode(tokens):\n",
    "    return ''.join(uchars[t] for t in tokens if t != BOS)\n",
    "\n",
    "print(uchars)\n",
    "print(\"BOS =\", BOS)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bacf9e2",
   "metadata": {},
   "source": [
    "**Predict**, then run: What list of integers does `encode('emma')` return? How long is it?\n",
    "Then replace `'emma'` with **your own first name** (lowercase, letters only)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "786a7543",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokens = encode('emma')\n",
    "print(tokens, len(tokens))\n",
    "print(decode(tokens))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48725c4e",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bd88452",
   "metadata": {},
   "source": [
    "**Question:** The token `BOS` is put at *both* the beginning and the end of every name. What job does the BOS token do at the beginning? What job does it do at the end? *(Hint: think about what the model has to do when it generates a brand-new name.)*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d53a780",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "405d24a0",
   "metadata": {},
   "source": [
    "**Question:** What would happen if one of the names in the file contained a character that isn't a lowercase letter, like `'-'` in `mary-jane`?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c20a8f60",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "636bb809",
   "metadata": {},
   "source": [
    "> **✅ Check yourself** — these are the kinds of questions you'll be asked at the check-off. Can you answer them out loud, in your own words, without looking at your notes?\n",
    ">\n",
    "> - What is the vocabulary size, and where does each part of that number come from?\n",
    "> - Write out the token sequence for a short name. Why does BOS appear twice?\n",
    "> - How does the model 'know' to stop generating a name?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d90c81d",
   "metadata": {},
   "source": [
    "---\n",
    "# Part 3: Autograd as a black box\n",
    "\n",
    "In Lab 2 you computed gradients by hand, deriving the backpropagation equations for every layer. That doesn't scale: a real GPT has hundreds of different operations.\n",
    "Every modern framework (PyTorch, JAX, TensorFlow) instead uses **automatic differentiation (\"autograd\")**:\n",
    "\n",
    "1. **Forward pass:** As you compute with special number objects, each result remembers which inputs it came from and the *local* derivative of that one operation. Together these form a **computation graph**.\n",
    "2. **Backward pass:** Calling `.backward()` on the final result walks the graph backwards, applying the chain rule, and fills in `.grad` on every number in the graph.\n",
    "\n",
    "microgpt's `Value` class does this for single numbers. **In this homework, you do not need to understand how `Value` works inside.** Treat it as a black box with this contract:\n",
    "\n",
    "> `Value(x)` behaves like the number `x`. After computing a result `L` from Values and calling `L.backward()`, every Value `v` that `L` depends on has `v.grad` = $\\dfrac{\\partial L}{\\partial v}$.\n",
    "\n",
    "Run the cell below (you can skip reading it — or read it in Part 7 if you're curious)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "237146fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let there be Autograd to recursively apply the chain rule through a computation graph\n",
    "class Value:\n",
    "    __slots__ = ('data', 'grad', '_children', '_local_grads') # Python optimization for memory usage\n",
    "\n",
    "    def __init__(self, data, children=(), local_grads=()):\n",
    "        self.data = data                # scalar value of this node calculated during forward pass\n",
    "        self.grad = 0                   # derivative of the loss w.r.t. this node, calculated in backward pass\n",
    "        self._children = children       # children of this node in the computation graph\n",
    "        self._local_grads = local_grads # local derivative of this node w.r.t. its children\n",
    "\n",
    "    def __add__(self, other):\n",
    "        other = other if isinstance(other, Value) else Value(other)\n",
    "        return Value(self.data + other.data, (self, other), (1, 1))\n",
    "\n",
    "    def __mul__(self, other):\n",
    "        other = other if isinstance(other, Value) else Value(other)\n",
    "        return Value(self.data * other.data, (self, other), (other.data, self.data))\n",
    "\n",
    "    def __pow__(self, other): return Value(self.data**other, (self,), (other * self.data**(other-1),))\n",
    "    def log(self): return Value(math.log(self.data), (self,), (1/self.data,))\n",
    "    def exp(self): return Value(math.exp(self.data), (self,), (math.exp(self.data),))\n",
    "    def relu(self): return Value(max(0, self.data), (self,), (float(self.data > 0),))\n",
    "    def __neg__(self): return self * -1\n",
    "    def __radd__(self, other): return self + other\n",
    "    def __sub__(self, other): return self + (-other)\n",
    "    def __rsub__(self, other): return other + (-self)\n",
    "    def __rmul__(self, other): return self * other\n",
    "    def __truediv__(self, other): return self * other**-1\n",
    "    def __rtruediv__(self, other): return other * self**-1\n",
    "\n",
    "    def backward(self):\n",
    "        topo = []\n",
    "        visited = set()\n",
    "        def build_topo(v):\n",
    "            if v not in visited:\n",
    "                visited.add(v)\n",
    "                for child in v._children:\n",
    "                    build_topo(child)\n",
    "                topo.append(v)\n",
    "        build_topo(self)\n",
    "        self.grad = 1\n",
    "        for v in reversed(topo):\n",
    "            for child, local_grad in zip(v._children, v._local_grads):\n",
    "                child.grad += local_grad * v.grad\n",
    "\n",
    "    def __repr__(self): return f\"Value(data={self.data:.4f}, grad={self.grad:.4f})\"  # [HW3] nicer printing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0acf38e6",
   "metadata": {},
   "source": [
    "## 3.1 Using the black box\n",
    "\n",
    "Let $L = a \\cdot b + c$ with $a = 2$, $b = -3$, $c = 10$.\n",
    "\n",
    "**Predict by hand** (this is calculus you already know): What is $L$? What are $\\dfrac{\\partial L}{\\partial a}$, $\\dfrac{\\partial L}{\\partial b}$, and $\\dfrac{\\partial L}{\\partial c}$?\n",
    "Write your predictions in the answer cell, then run the code."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "544c731c",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b26d7dde",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = Value(2.0)\n",
    "b = Value(-3.0)\n",
    "c = Value(10.0)\n",
    "L = a * b + c\n",
    "\n",
    "print(\"before backward:\", a, b, c)\n",
    "L.backward()\n",
    "print(\"L =\", L.data)\n",
    "print(\"after backward: \", a, b, c)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4706526",
   "metadata": {},
   "source": [
    "Now a slightly harder one. Let $L = (a \\cdot b + c)^2$, using the same values.\n",
    "\n",
    "**Predict by hand** using the chain rule, then run."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f04054c9",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b489f9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = Value(2.0)\n",
    "b = Value(-3.0)\n",
    "c = Value(10.0)\n",
    "L = (a * b + c) ** 2\n",
    "L.backward()\n",
    "print(\"L =\", L.data)\n",
    "print(a, b, c)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7a3db0b",
   "metadata": {},
   "source": [
    "**Question:** A gradient like $\\partial L/\\partial a = -24$ tells us something useful. If we want to make $L$ *smaller*, should we increase or decrease $a$? What does the size of the number (24) tell us?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7841ae30",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c5b09fa",
   "metadata": {},
   "source": [
    "**That's all you need.** In microgpt, the loss is built from thousands of Value operations. One call to `loss.backward()` gives us the gradient for all 4,192 parameters at once, and the optimizer uses those gradients exactly the way you just described. **From here on, trust `backward()`.**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b92fe43",
   "metadata": {},
   "source": [
    "> **✅ Check yourself** — these are the kinds of questions you'll be asked at the check-off. Can you answer them out loud, in your own words, without looking at your notes?\n",
    ">\n",
    "> - In one or two sentences, what does autograd do for us? What is stored in `p.grad` after `loss.backward()`?\n",
    "> - If a parameter's gradient is positive, which direction should we move it to reduce the loss?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6149f2d0",
   "metadata": {},
   "source": [
    "---\n",
    "# Part 4: The model\n",
    "\n",
    "This is the heart of the assignment, and it is the part that will take you the longest. The videos you watched moved fast; this is where you can slow down and stare at each piece until it makes sense. We have tried to explain everything in detail. Read slowly, run every cell, and don't move on from a section until you could explain it to a classmate.\n",
    "\n",
    "## 4a. Parameters\n",
    "\n",
    "A model's **parameters** are the numbers that get learned. microgpt stores them as matrices (lists of lists of `Value`s) in `state_dict`, starting out as small random numbers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ed6a26d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the parameters, to store the knowledge of the model\n",
    "n_layer = 1     # depth of the transformer neural network (number of layers)\n",
    "n_embd = 16     # width of the network (embedding dimension)\n",
    "block_size = 16 # maximum context length of the attention window (note: the longest name is 15 characters)\n",
    "n_head = 4      # number of attention heads\n",
    "head_dim = n_embd // n_head # derived dimension of each head\n",
    "matrix = lambda nout, nin, std=0.08: [[Value(random.gauss(0, std)) for _ in range(nin)] for _ in range(nout)]\n",
    "state_dict = {'wte': matrix(vocab_size, n_embd), 'wpe': matrix(block_size, n_embd), 'lm_head': matrix(vocab_size, n_embd)}\n",
    "for i in range(n_layer):\n",
    "    state_dict[f'layer{i}.attn_wq'] = matrix(n_embd, n_embd)\n",
    "    state_dict[f'layer{i}.attn_wk'] = matrix(n_embd, n_embd)\n",
    "    state_dict[f'layer{i}.attn_wv'] = matrix(n_embd, n_embd)\n",
    "    state_dict[f'layer{i}.attn_wo'] = matrix(n_embd, n_embd)\n",
    "    state_dict[f'layer{i}.mlp_fc1'] = matrix(4 * n_embd, n_embd)\n",
    "    state_dict[f'layer{i}.mlp_fc2'] = matrix(n_embd, 4 * n_embd)\n",
    "params = [p for mat in state_dict.values() for row in mat for p in row] # flatten params into a single list[Value]\n",
    "print(f\"num params: {len(params)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb229bff",
   "metadata": {},
   "source": [
    "`matrix(nout, nin)` builds a matrix with `nout` rows and `nin` columns, filled with small random numbers. (The `lambda` is just a one-line way to define a function: `matrix(3, 5)` is the same as calling `def matrix(nout, nin): ...` with those arguments.)\n",
    "\n",
    "So, for example, `state_dict['wte'] = matrix(vocab_size, n_embd)` is a matrix with `vocab_size` = 27 rows and `n_embd` = 16 columns: one row of 16 numbers for each token.\n",
    "\n",
    "**Exercise:** Fill in the shape (rows × columns) and the number of parameters for each matrix **by hand**, using `vocab_size = 27`, `n_embd = 16`, `block_size = 16`. Add up the counts and check that you get the number printed above.\n",
    "\n",
    "| Matrix | Built by | Shape (rows × cols) | # params |\n",
    "|---|---|---|---|\n",
    "| `wte` | `matrix(vocab_size, n_embd)` | ? | ? |\n",
    "| `wpe` | `matrix(block_size, n_embd)` | ? | ? |\n",
    "| `layer0.attn_wq` | `matrix(n_embd, n_embd)` | ? | ? |\n",
    "| `layer0.attn_wk` | `matrix(n_embd, n_embd)` | ? | ? |\n",
    "| `layer0.attn_wv` | `matrix(n_embd, n_embd)` | ? | ? |\n",
    "| `layer0.attn_wo` | `matrix(n_embd, n_embd)` | ? | ? |\n",
    "| `layer0.mlp_fc1` | `matrix(4 * n_embd, n_embd)` | ? | ? |\n",
    "| `layer0.mlp_fc2` | `matrix(n_embd, 4 * n_embd)` | ? | ? |\n",
    "| `lm_head` | `matrix(vocab_size, n_embd)` | ? | ? |\n",
    "| **Total** | — | — | ? |"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5beb7184",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e00b988",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check your table\n",
    "import pandas as pd\n",
    "shapes = pd.DataFrame(\n",
    "    [(name, len(mat), len(mat[0]), len(mat) * len(mat[0])) for name, mat in state_dict.items()],\n",
    "    columns=['matrix', 'rows', 'cols', '# params'],\n",
    ")\n",
    "shapes.loc[len(shapes)] = ['total', '', '', shapes['# params'].sum()]\n",
    "shapes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9edd954a",
   "metadata": {},
   "source": [
    "## 4b. Helper functions\n",
    "\n",
    "The model uses three small helper functions. Run the cell, then read the explanations that follow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9468a010",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the model architecture: a function mapping tokens and parameters to logits over what comes next\n",
    "# Follow GPT-2, blessed among the GPTs, with minor differences: layernorm -> rmsnorm, no biases, GeLU -> ReLU\n",
    "def linear(x, w):\n",
    "    return [sum(wi * xi for wi, xi in zip(wo, x)) for wo in w]\n",
    "\n",
    "def softmax(logits):\n",
    "    max_val = max(val.data for val in logits)\n",
    "    exps = [(val - max_val).exp() for val in logits]\n",
    "    total = sum(exps)\n",
    "    return [e / total for e in exps]\n",
    "\n",
    "def rmsnorm(x):\n",
    "    ms = sum(xi * xi for xi in x) / len(x)\n",
    "    scale = (ms + 1e-5) ** -0.5\n",
    "    return [xi * scale for xi in x]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a54b17b",
   "metadata": {},
   "source": [
    "### The dot product, and what it means\n",
    "\n",
    "Almost everything in this model is built from one operation: the **dot product** (also called the **inner product**) of two vectors of the same length. You met it in Part 0:\n",
    "\n",
    "```python\n",
    "sum(wi * xi for wi, xi in zip(w, x))      # = w[0]*x[0] + w[1]*x[1] + ... + w[n-1]*x[n-1]\n",
    "```\n",
    "\n",
    "The formula is simple, but it's worth building an intuition for what the number *means*, because attention is built on it.\n",
    "\n",
    "**A dot product measures how much two vectors \"agree.\"** Think of two vectors as arrows:\n",
    "- If they point in the **same direction**, the products `wi * xi` are all positive, and the dot product is **large and positive**.\n",
    "- If they point in **unrelated (perpendicular) directions**, the positive and negative products cancel out, and the dot product is **near zero**.\n",
    "- If they point in **opposite directions**, the dot product is **large and negative**.\n",
    "\n",
    "(Precisely: $w \\cdot x = \\|w\\| \\, \\|x\\| \\cos\\theta$, where $\\theta$ is the angle between them. So for vectors of a fixed length, the dot product is largest when they point the same way.)\n",
    "\n",
    "So you can read `dot(w, x)` as **\"how similar is `x` to `w`?\"** Run the cell below and check that the numbers match the three cases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e28663cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dot(w, x):\n",
    "    return sum(wi * xi for wi, xi in zip(w, x))\n",
    "\n",
    "w = [1.0, 0.0]\n",
    "print(\"same direction:     dot(w, [ 0.9,  0.1]) =\", dot(w, [0.9, 0.1]))\n",
    "print(\"perpendicular:      dot(w, [ 0.0,  1.0]) =\", dot(w, [0.0, 1.0]))\n",
    "print(\"opposite direction: dot(w, [-0.9, -0.1]) =\", dot(w, [-0.9, -0.1]))\n",
    "print(\"similar, and longer: dot(w, [ 3.0,  0.5]) =\", dot(w, [3.0, 0.5]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff12ca7d",
   "metadata": {},
   "source": [
    "### `linear`: a stack of dot products\n",
    "\n",
    "`linear(x, w)` computes the dot product of `x` with **every row** of the matrix `w`, and returns the results as a list:\n",
    "\n",
    "```python\n",
    "def linear(x, w):\n",
    "    return [sum(wi * xi for wi, xi in zip(wo, x)) for wo in w]   # one dot product per row wo\n",
    "#           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  = dot(wo, x)\n",
    "```\n",
    "\n",
    "So `linear(x, w)` is the same as `[dot(wo, x) for wo in w]`. (In matrix notation, it's the matrix-vector product $Wx$.) Two useful ways to think about it:\n",
    "\n",
    "1. **Each row asks a question.** Row `wo` asks \"how much does `x` look like me?\" and the answer is `dot(wo, x)`. A matrix with 64 rows asks 64 questions and returns 64 answers. Which questions to ask is exactly what training learns.\n",
    "2. **Shapes.** If `w` has `nout` rows and `nin` columns, then `x` must have length `nin`, and the output has length `nout`. `linear` is the model's way of turning a vector of one length into a vector of another length."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7491b06",
   "metadata": {},
   "outputs": [],
   "source": [
    "W = [[1, 0, 0],      # row 0 asks: \"what is x[0]?\"\n",
    "     [1, 1, 1],      # row 1 asks: \"what is the sum of x?\"\n",
    "     [0, 0, -1]]     # row 2 asks: \"what is -x[2]?\"\n",
    "x = [2, 3, 4]\n",
    "\n",
    "# Predict (by hand): linear(x, W) -> ?\n",
    "print(linear(x, W))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de5b35e3",
   "metadata": {},
   "source": [
    "### `softmax`: scores → probabilities\n",
    "\n",
    "`softmax` turns a list of arbitrary scores (called **logits**) into a probability distribution: all outputs are positive and they add up to 1. Bigger logits get bigger probabilities.\n",
    "\n",
    "$$\\text{softmax}(z)_i = \\frac{e^{z_i}}{\\sum_j e^{z_j}}$$\n",
    "\n",
    "In the code: `exps` is the list of $e^{z_i}$, `total` is their sum, and the result divides each one by the total. (Subtracting `max_val` first doesn't change the answer — it cancels out in the division — it just prevents `exp` from overflowing on large inputs.)\n",
    "\n",
    "The model uses softmax in two places: to turn attention scores into attention weights, and to turn the final logits into next-token probabilities.\n",
    "\n",
    "**Predict:** For logits `[1, 2, 3]`: which output is largest? What do the outputs sum to? What happens to the probabilities if you add 100 to every logit? If you multiply every logit by 10?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3289e98",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e49ba85b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def show(values):\n",
    "    return [round(v.data, 5) for v in values]\n",
    "\n",
    "logits = [Value(1.0), Value(2.0), Value(3.0)]\n",
    "print(\"softmax([1,2,3])       =\", show(softmax(logits)))\n",
    "print(\"sum                    =\", sum(show(softmax(logits))))\n",
    "print(\"softmax([101,102,103]) =\", show(softmax([l + 100 for l in logits])))\n",
    "print(\"softmax([10,20,30])    =\", show(softmax([l * 10 for l in logits])))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d18e228",
   "metadata": {},
   "source": [
    "### `rmsnorm`: keeping numbers a sensible size\n",
    "\n",
    "`rmsnorm` rescales a vector so that its **root-mean-square** (RMS) is 1: it computes the mean of the squares (`ms`), then multiplies every element by $1/\\sqrt{\\text{ms}}$. (The tiny `1e-5` just avoids dividing by zero.)\n",
    "\n",
    "Why bother? As numbers pass through many layers of multiplications and additions, they can drift toward being enormous or tiny, which makes training unstable. Normalizing before each block keeps the vectors flowing through the network at a consistent size. RMSNorm doesn't change the *direction* of the vector — the relative sizes of the entries — only its overall scale.\n",
    "\n",
    "**Predict by hand:** What is `rmsnorm([3, 4])`? *(RMS = $\\sqrt{(3^2 + 4^2)/2}$.)* What is `rmsnorm([30, 40])`?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02657908",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fc3c7fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(show(rmsnorm([Value(3.0), Value(4.0)])))\n",
    "print(show(rmsnorm([Value(30.0), Value(40.0)])))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67e8668a",
   "metadata": {},
   "source": [
    "## 4c. The `gpt` function\n",
    "\n",
    "This is the model itself. Run the cell, skim the code once to get the shape of it, and then read on: we will walk through it a few lines at a time.\n",
    "\n",
    "Two important design choices to hold in your head:\n",
    "\n",
    "- **`gpt` processes one token at a time.** You give it the current token and its position, and it returns a **logit for every possible next token** (a list of length `vocab_size`). To process a whole name, you call `gpt` once per position, in order.\n",
    "- **The `keys` and `values` lists are the model's memory of earlier tokens.** Each call to `gpt` appends one entry to each, so the current token can \"see\" everything that came before it. This is called a **KV cache**.\n",
    "\n",
    "We added one optional argument, `attn_record`, so you can peek at the attention weights. It doesn't change what the model computes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e904300",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gpt(token_id, pos_id, keys, values, attn_record=None):  # [HW3] added attn_record\n",
    "    tok_emb = state_dict['wte'][token_id] # token embedding\n",
    "    pos_emb = state_dict['wpe'][pos_id] # position embedding\n",
    "    x = [t + p for t, p in zip(tok_emb, pos_emb)] # joint token and position embedding\n",
    "    x = rmsnorm(x) # note: not redundant due to backward pass via the residual connection\n",
    "\n",
    "    for li in range(n_layer):\n",
    "        # 1) Multi-head Attention block\n",
    "        x_residual = x\n",
    "        x = rmsnorm(x)\n",
    "        q = linear(x, state_dict[f'layer{li}.attn_wq'])\n",
    "        k = linear(x, state_dict[f'layer{li}.attn_wk'])\n",
    "        v = linear(x, state_dict[f'layer{li}.attn_wv'])\n",
    "        keys[li].append(k)\n",
    "        values[li].append(v)\n",
    "        x_attn = []\n",
    "        for h in range(n_head):\n",
    "            hs = h * head_dim\n",
    "            q_h = q[hs:hs+head_dim]\n",
    "            k_h = [ki[hs:hs+head_dim] for ki in keys[li]]\n",
    "            v_h = [vi[hs:hs+head_dim] for vi in values[li]]\n",
    "            attn_logits = [sum(q_h[j] * k_h[t][j] for j in range(head_dim)) / head_dim**0.5 for t in range(len(k_h))]\n",
    "            attn_weights = softmax(attn_logits)\n",
    "            if attn_record is not None:                                   # [HW3]\n",
    "                attn_record.append((li, h, [w.data for w in attn_weights]))  # [HW3]\n",
    "            head_out = [sum(attn_weights[t] * v_h[t][j] for t in range(len(v_h))) for j in range(head_dim)]\n",
    "            x_attn.extend(head_out)\n",
    "        x = linear(x_attn, state_dict[f'layer{li}.attn_wo'])\n",
    "        x = [a + b for a, b in zip(x, x_residual)]\n",
    "        # 2) MLP block\n",
    "        x_residual = x\n",
    "        x = rmsnorm(x)\n",
    "        x = linear(x, state_dict[f'layer{li}.mlp_fc1'])\n",
    "        x = [xi.relu() for xi in x]\n",
    "        x = linear(x, state_dict[f'layer{li}.mlp_fc2'])\n",
    "        x = [a + b for a, b in zip(x, x_residual)]\n",
    "\n",
    "    logits = linear(x, state_dict['lm_head'])\n",
    "    return logits"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b1c81a1",
   "metadata": {},
   "source": [
    "Here is the whole function as a picture. Every arrow carries a vector of length 16, except the last one (length 27).\n",
    "\n",
    "```\n",
    "token_id ──► wte row ──┐\n",
    "                       ├──► x = tok_emb + pos_emb ──► rmsnorm\n",
    "pos_id   ──► wpe row ──┘                                │\n",
    "                                                        ▼\n",
    "                   ┌────────────────────────────────────┤  x_residual\n",
    "                   │  rmsnorm ─► ATTENTION (looks at earlier tokens via keys/values)\n",
    "                   │                        │\n",
    "                   └──────────── + ◄────────┘        (residual connection)\n",
    "                                 │\n",
    "                   ┌─────────────┤  x_residual\n",
    "                   │  rmsnorm ─► MLP (fc1 → relu → fc2)\n",
    "                   │                        │\n",
    "                   └──────────── + ◄────────┘        (residual connection)\n",
    "                                 │\n",
    "                                 ▼\n",
    "                          lm_head ──► logits (length 27)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "465bdcd1",
   "metadata": {},
   "source": [
    "### Step 1: Embeddings\n",
    "\n",
    "The first three lines of `gpt` build the input vector `x`:\n",
    "\n",
    "```python\n",
    "tok_emb = state_dict['wte'][token_id]              # one row of wte: a vector of length 16\n",
    "pos_emb = state_dict['wpe'][pos_id]                # one row of wpe: a vector of length 16\n",
    "x = [t + p for t, p in zip(tok_emb, pos_emb)]      # add them element-wise: still length 16\n",
    "```\n",
    "\n",
    "**The token embedding.** `wte` has one row per token (27 rows). Row `token_id` is the vector the model uses to *represent* that character — its internal description of \"the letter m.\" These rows are parameters: they start out random, and training gradually moves them so that letters which behave similarly end up with similar vectors. This is the model's only source of information about *which* character it is looking at.\n",
    "\n",
    "**The position embedding.** `wpe` has one row per position (16 rows, because `block_size = 16`). Row `pos_id` is a learned vector that represents \"position number `pos_id`.\" Adding it to the token embedding gives the model a vector that encodes both *what* the current character is and *where* it is in the name.\n",
    "\n",
    "**Why does the model need a position embedding at all?** This is a subtle point that the videos go over quickly, so it's worth spelling out.\n",
    "\n",
    "In a Transformer, the current token gathers information from the tokens before it through attention (Step 3, below). But attention is a *weighted average* over those earlier tokens, and a weighted average is the same no matter what order you list the items in. So from the point of view of attention, **the earlier tokens are an unordered set** — like a bag of letters. On its own, attention can't tell `\"ann\"` from `\"nan\"`, and it can't tell whether the current letter is the 2nd or the 8th in the name.\n",
    "\n",
    "This might sound like a flaw, but it's a deliberate design choice, and it is the reason Transformers took over. Because attention treats the context as a set, **every position in a sequence can be computed at the same time, in parallel**, during training. Older models (recurrent networks) processed one token after another, each step waiting for the previous one to finish, which is far slower. (microgpt happens to loop over positions one at a time because it is written for clarity, not speed — but nothing about the math requires that.)\n",
    "\n",
    "The position embedding is how we get order back *without* giving up the parallelism: instead of changing how attention works, we stamp each token's input vector with a learned \"I am at position `pos_id`\" vector. The position information travels along inside `x`, and attention can learn to use it (for example, to notice \"I'm 8 letters in — time to end this name\").\n",
    "\n",
    "**Quick check** (no writing needed): if the letter `m` appears at positions 2 and 3 of `emma`, what is the same about the two input vectors `x`, and what is different?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3fd053b",
   "metadata": {},
   "source": [
    "### Step 2: Queries, keys, and values\n",
    "\n",
    "Now we enter the attention block. First, `x` is normalized, and then it is passed through three *different* `linear` layers:\n",
    "\n",
    "```python\n",
    "q = linear(x, state_dict[f'layer{li}.attn_wq'])    # query: length 16\n",
    "k = linear(x, state_dict[f'layer{li}.attn_wk'])    # key:   length 16\n",
    "v = linear(x, state_dict[f'layer{li}.attn_wv'])    # value: length 16\n",
    "keys[li].append(k)                                 # remember this token's key ...\n",
    "values[li].append(v)                               # ... and its value, for later tokens to use\n",
    "```\n",
    "\n",
    "The three matrices are `attn_wq` (the **query matrix**), `attn_wk` (the **key matrix**), and `attn_wv` (the **value matrix**). Each is 16 × 16, so `q`, `k`, and `v` are three different vectors of length 16, all computed from the same `x`. All three are learned during training. Here is what each one is *for*. A useful analogy is searching a library:\n",
    "\n",
    "- The **query** `q` is what the current token is *looking for* in the earlier tokens. (\"I'm an `m` — was there a vowel just before me?\")\n",
    "- The **key** `k` is how a token *advertises itself* to future tokens that will be searching. (\"I'm a vowel, near the start of the name.\")\n",
    "- The **value** `v` is the information a token *hands over* if someone decides to attend to it.\n",
    "\n",
    "The key and value of the current token are appended to the cache so that *later* tokens can search over them. Notice that `keys[li]` is a list of vectors, one for each token processed so far: after the 3rd call to `gpt`, `keys[0]` contains 3 vectors of length 16.\n",
    "\n",
    "Why three separate matrices instead of just comparing `x` vectors directly? Because \"what I'm looking for\" and \"what I'm offering\" are different things. A consonant might be looking for the previous vowel, but it advertises itself as a consonant. The three matrices let the model learn these roles separately."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b903e78a",
   "metadata": {},
   "source": [
    "### Step 3: Splitting into heads\n",
    "\n",
    "```python\n",
    "for h in range(n_head):                             # n_head = 4\n",
    "    hs = h * head_dim                               # head_dim = 4, so hs = 0, 4, 8, 12\n",
    "    q_h = q[hs:hs+head_dim]                         # the 4 entries of q that belong to head h\n",
    "    k_h = [ki[hs:hs+head_dim] for ki in keys[li]]   # the same 4 entries of EVERY cached key\n",
    "    v_h = [vi[hs:hs+head_dim] for vi in values[li]] # ... and of every cached value\n",
    "```\n",
    "\n",
    "Rather than doing one big attention computation with 16-dimensional vectors, the model chops `q`, `k`, and `v` into 4 pieces of length 4 and runs 4 small, independent attention computations, called **heads**. Head 0 uses entries 0–3 of every vector, head 1 uses entries 4–7, and so on (this is the slicing from Part 0).\n",
    "\n",
    "So inside the loop, for head `h`:\n",
    "- `q_h` is a vector of length 4 (the current token's query, head `h`'s slice),\n",
    "- `k_h` is a *list* of vectors of length 4, one per cached token,\n",
    "- `v_h` is a list of vectors of length 4, one per cached token.\n",
    "\n",
    "Each head gets to compute its own attention weights, so different heads can look at different earlier tokens for different reasons at the same time. We'll come back to why that's useful."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b39add80",
   "metadata": {},
   "source": [
    "### Step 4: Attention scores — \"how relevant is each earlier token?\"\n",
    "\n",
    "This is the line that does the searching:\n",
    "\n",
    "```python\n",
    "attn_logits = [sum(q_h[j] * k_h[t][j] for j in range(head_dim)) / head_dim**0.5 for t in range(len(k_h))]\n",
    "```\n",
    "\n",
    "It's dense, so let's unpack it into loops. It does exactly the same thing as:\n",
    "\n",
    "```python\n",
    "attn_logits = []\n",
    "for t in range(len(k_h)):                        # for each cached token t ...\n",
    "    score = 0\n",
    "    for j in range(head_dim):                    #   dot product of q_h with token t's key\n",
    "        score += q_h[j] * k_h[t][j]\n",
    "    attn_logits.append(score / head_dim**0.5)    #   scaled down a bit (divide by sqrt(4) = 2)\n",
    "```\n",
    "\n",
    "So `attn_logits[t]` is **the dot product of the current query with token `t`'s key**. Remember what a dot product means: it measures how much two vectors agree. So `attn_logits[t]` is a **relevance score**: it is large when what the current token is looking for (`q_h`) matches what token `t` is advertising (`k_h[t]`). There is one score per cached token, so `attn_logits` has the same length as `keys[li]`.\n",
    "\n",
    "(The division by $\\sqrt{\\text{head\\_dim}}$ just keeps the scores from growing too large as `head_dim` gets bigger, which would make the softmax in the next step too sharp.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a65ccbfc",
   "metadata": {},
   "source": [
    "### Step 5: Attention weights — turning scores into a distribution\n",
    "\n",
    "```python\n",
    "attn_weights = softmax(attn_logits)\n",
    "```\n",
    "\n",
    "The scores can be any numbers. Softmax turns them into **attention weights**: positive numbers that add up to 1, with higher scores getting more of the weight. You can read `attn_weights[t]` as **\"the fraction of its attention that the current token pays to token `t`.\"** If the current token is at position 2, there are 3 cached tokens (positions 0, 1, 2 — it can attend to itself), so `attn_weights` has 3 entries that sum to 1."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11ecbd50",
   "metadata": {},
   "source": [
    "### Step 6: The head's output — a weighted average of values\n",
    "\n",
    "```python\n",
    "head_out = [sum(attn_weights[t] * v_h[t][j] for t in range(len(v_h))) for j in range(head_dim)]\n",
    "```\n",
    "\n",
    "Unpacked:\n",
    "\n",
    "```python\n",
    "head_out = []\n",
    "for j in range(head_dim):                        # for each of the 4 entries of the output ...\n",
    "    total = 0\n",
    "    for t in range(len(v_h)):                    #   add up entry j of every cached value vector,\n",
    "        total += attn_weights[t] * v_h[t][j]     #   weighted by how much we attend to that token\n",
    "    head_out.append(total)\n",
    "```\n",
    "\n",
    "So `head_out` is a **weighted average of the value vectors of all cached tokens**, where the weights are the attention weights. If the current token pays 90% of its attention to token 1, then `head_out` is 90% token 1's value vector (plus small amounts of the others). This is how information from earlier tokens flows into the current position — and it is the *only* place in the entire model where tokens exchange information with each other.\n",
    "\n",
    "`head_out` has length `head_dim` = 4."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "107bf68e",
   "metadata": {},
   "source": [
    "### Step 7: Combining the heads\n",
    "\n",
    "```python\n",
    "    x_attn.extend(head_out)                                   # inside the loop: append this head's 4 numbers\n",
    "x = linear(x_attn, state_dict[f'layer{li}.attn_wo'])          # after the loop: mix the heads together\n",
    "x = [a + b for a, b in zip(x, x_residual)]                    # residual connection (explained below)\n",
    "```\n",
    "\n",
    "**Gluing the heads together.** Each head produces 4 numbers, and `extend` appends them one after another. After the loop, `x_attn` is a vector of length 16, laid out in slots:\n",
    "\n",
    "```\n",
    "x_attn = [ head 0's 4 numbers | head 1's 4 numbers | head 2's 4 numbers | head 3's 4 numbers ]\n",
    "           entries 0-3          entries 4-7          entries 8-11         entries 12-15\n",
    "```\n",
    "\n",
    "At this point the heads have not interacted at all. Each one computed its own attention weights and its own weighted average independently, and its result sits in its own slot.\n",
    "\n",
    "**The output projection.** The next line, `x = linear(x_attn, attn_wo)`, multiplies `x_attn` by the 16 × 16 matrix `attn_wo`. Recall from 4b that `linear` computes one dot product per row. Each row of `attn_wo` is 16 numbers long, so each entry of the result is a weighted combination of *all 16* entries of `x_attn`, which means all four heads' outputs at once. This is the first place the heads' results are combined.\n",
    "\n",
    "A cleaner way to see it: split the 16 columns of `attn_wo` into four blocks of 4 columns, one block per head, so $W_O = [\\,W_O^{(0)} \\mid W_O^{(1)} \\mid W_O^{(2)} \\mid W_O^{(3)}\\,]$, with each block 16 × 4. Multiplying a matrix by a stacked vector works block by block, so\n",
    "\n",
    "$$W_O \\, x_{\\text{attn}} = W_O^{(0)}\\,\\text{head\\_out}_0 + W_O^{(1)}\\,\\text{head\\_out}_1 + W_O^{(2)}\\,\\text{head\\_out}_2 + W_O^{(3)}\\,\\text{head\\_out}_3$$\n",
    "\n",
    "In words: **each head has its own 16 × 4 matrix that turns its 4 numbers into a 16-number vector, and the four results are added up.** This mirrors Step 3. There, the *rows* of `attn_wq` split into one block per head. Here, the *columns* of `attn_wo` split into one block per head.\n",
    "\n",
    "Why is this step needed, instead of using `x_attn` directly?\n",
    "- **Every head can affect every entry of `x`.** Without `attn_wo`, head 0 could only ever change entries 0–3 of `x`, and head 1 only entries 4–7. That restriction is just a side effect of how we sliced the vectors into heads. The projection lets each head write its result anywhere in the 16-dimensional vector.\n",
    "- **It converts the result back into the coordinates of `x`.** Each head's output is a weighted average of *value* vectors. Those were produced by the value matrix `attn_wv` back in Step 2, so they are expressed in whatever coordinates `attn_wv` learned to use, not in the coordinates of `x` itself. `attn_wo` is learned, so training can find how to turn each head's result into a sensible change to add to `x`.\n",
    "\n",
    "Finally, `x = [a + b for a, b in zip(x, x_residual)]` adds this length-16 result onto the vector that entered the attention block. So the attention block doesn't replace `x`. It computes an *update*, and the update is added on. (This residual connection is explained below, after Step 8.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79cf1012",
   "metadata": {},
   "source": [
    "### A worked example, by hand\n",
    "\n",
    "Let's do one head's attention with numbers small enough to check by hand. Suppose `head_dim = 2`, three tokens have been cached, and the current token's query is `q_h = [1, 0]`. The cached keys and values are:\n",
    "\n",
    "| token `t` | key `k_h[t]` | value `v_h[t]` |\n",
    "|---|---|---|\n",
    "| 0 | `[1, 0]` | `[10, 0]` |\n",
    "| 1 | `[0, 1]` | `[0, 10]` |\n",
    "| 2 | `[-1, 0]` | `[5, 5]` |\n",
    "\n",
    "**Predict by hand, before running:**\n",
    "1. The three dot products `q_h · k_h[t]` (ignore the $\\sqrt{\\text{head\\_dim}}$ scaling).\n",
    "2. Which token gets the *most* attention weight? Which gets the least?\n",
    "3. Roughly what will `head_out` look like — closer to `[10, 0]`, `[0, 10]`, or `[5, 5]`?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd4301d2",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bcdbd49",
   "metadata": {},
   "outputs": [],
   "source": [
    "q_h = [1.0, 0.0]\n",
    "k_h = [[1.0, 0.0], [0.0, 1.0], [-1.0, 0.0]]\n",
    "v_h = [[10.0, 0.0], [0.0, 10.0], [5.0, 5.0]]\n",
    "\n",
    "# the three steps of attention, with plain numbers (no Value objects)\n",
    "scores = [dot(q_h, k_h[t]) for t in range(len(k_h))]\n",
    "exps = [math.exp(s) for s in scores]\n",
    "weights = [e / sum(exps) for e in exps]\n",
    "head_out = [sum(weights[t] * v_h[t][j] for t in range(len(v_h))) for j in range(2)]\n",
    "\n",
    "print(\"scores  :\", scores)\n",
    "print(\"weights :\", [round(w, 3) for w in weights], \" (sum =\", round(sum(weights), 3), \")\")\n",
    "print(\"head_out:\", [round(h, 2) for h in head_out])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f230104a",
   "metadata": {},
   "source": [
    "### Tracing the shapes in the real model\n",
    "\n",
    "Now let's trace the real thing. The cell below feeds the first two tokens of `emma` through `gpt` normally (filling the KV cache), and then does the **third** call — position 2, the first `m` — *by hand*, step by step, printing the length of every variable along the way. This is the same code as inside `gpt`, just pulled out so we can look at it.\n",
    "\n",
    "**Before you run it**, use what you've read above to fill in your predictions in the table. Then run the cell and compare.\n",
    "\n",
    "| Variable | Predicted length | Actual |\n",
    "|---|---|---|\n",
    "| `x` | ? | ? |\n",
    "| `q`, `k`, `v` | ? | ? |\n",
    "| `keys[0]` (number of cached vectors) | ? | ? |\n",
    "| `q_h` | ? | ? |\n",
    "| `k_h` (number of vectors, and length of each) | ? | ? |\n",
    "| `attn_logits` | ? | ? |\n",
    "| `attn_weights` | ? | ? |\n",
    "| `head_out` | ? | ? |\n",
    "| `x_attn` (after all 4 heads) | ? | ? |\n",
    "| `logits` | ? | ? |"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44157a9b",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "255d8126",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokens = encode('emma')                     # [BOS, e, m, m, a, BOS]\n",
    "keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]\n",
    "\n",
    "# Positions 0 and 1: run gpt normally. This fills the cache with 2 keys and 2 values.\n",
    "for pos_id in range(2):\n",
    "    gpt(tokens[pos_id], pos_id, keys, values)\n",
    "print(f\"after positions 0 and 1: len(keys[0]) = {len(keys[0])}, len(values[0]) = {len(values[0])}\")\n",
    "\n",
    "# Position 2: the same computation as gpt(), step by step.\n",
    "pos_id, token_id = 2, tokens[2]\n",
    "print(f\"\\n--- position {pos_id}: token_id = {token_id} ('{uchars[token_id]}') ---\")\n",
    "tok_emb = state_dict['wte'][token_id]\n",
    "pos_emb = state_dict['wpe'][pos_id]\n",
    "x = [t + p for t, p in zip(tok_emb, pos_emb)]\n",
    "x = rmsnorm(x)\n",
    "print(f\"len(tok_emb) = {len(tok_emb)}, len(pos_emb) = {len(pos_emb)}, len(x) = {len(x)}\")\n",
    "\n",
    "x_residual = x\n",
    "x = rmsnorm(x)\n",
    "q = linear(x, state_dict['layer0.attn_wq'])\n",
    "k = linear(x, state_dict['layer0.attn_wk'])\n",
    "v = linear(x, state_dict['layer0.attn_wv'])\n",
    "keys[0].append(k)\n",
    "values[0].append(v)\n",
    "print(f\"len(q) = {len(q)}, len(k) = {len(k)}, len(v) = {len(v)}\")\n",
    "print(f\"len(keys[0]) = {len(keys[0])} cached keys, each of length {len(keys[0][0])}\")\n",
    "\n",
    "x_attn = []\n",
    "for h in range(n_head):\n",
    "    hs = h * head_dim\n",
    "    q_h = q[hs:hs+head_dim]\n",
    "    k_h = [ki[hs:hs+head_dim] for ki in keys[0]]\n",
    "    v_h = [vi[hs:hs+head_dim] for vi in values[0]]\n",
    "    attn_logits = [sum(q_h[j] * k_h[t][j] for j in range(head_dim)) / head_dim**0.5 for t in range(len(k_h))]\n",
    "    attn_weights = softmax(attn_logits)\n",
    "    head_out = [sum(attn_weights[t] * v_h[t][j] for t in range(len(v_h))) for j in range(head_dim)]\n",
    "    x_attn.extend(head_out)\n",
    "    if h == 0:\n",
    "        print(f\"\\nhead {h}: slice [{hs}:{hs+head_dim}]\")\n",
    "        print(f\"  len(q_h) = {len(q_h)}\")\n",
    "        print(f\"  len(k_h) = {len(k_h)} vectors, each of length {len(k_h[0])}\")\n",
    "        print(f\"  len(attn_logits) = {len(attn_logits)}   values: {show(attn_logits)}\")\n",
    "        print(f\"  len(attn_weights) = {len(attn_weights)}  values: {show(attn_weights)}  (sum = {sum(show(attn_weights)):.3f})\")\n",
    "        print(f\"  len(head_out) = {len(head_out)}\")\n",
    "print(f\"\\nafter all {n_head} heads: len(x_attn) = {len(x_attn)}\")\n",
    "\n",
    "x = linear(x_attn, state_dict['layer0.attn_wo'])\n",
    "x = [a + b for a, b in zip(x, x_residual)]\n",
    "print(f\"after attn_wo and the residual connection: len(x) = {len(x)}\")\n",
    "\n",
    "# MLP block\n",
    "x_residual = x\n",
    "x = rmsnorm(x)\n",
    "x = linear(x, state_dict['layer0.mlp_fc1'])\n",
    "print(f\"after mlp_fc1: len(x) = {len(x)}\")\n",
    "x = [xi.relu() for xi in x]\n",
    "x = linear(x, state_dict['layer0.mlp_fc2'])\n",
    "x = [a + b for a, b in zip(x, x_residual)]\n",
    "print(f\"after mlp_fc2 and the residual connection: len(x) = {len(x)}\")\n",
    "\n",
    "logits = linear(x, state_dict['lm_head'])\n",
    "print(f\"len(logits) = {len(logits)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "553d54cc",
   "metadata": {},
   "source": [
    "### Now, in your own words\n",
    "\n",
    "You've seen every line of attention. Answer these in your own words — these are the kinds of things you'll be asked at the check-off.\n",
    "\n",
    "1. What does `attn_logits[t]` measure, and how is it computed? What does `attn_weights[t]` mean?\n",
    "2. `head_out` is a weighted sum of what? Why is that useful — what does it let the current token do?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "653c673a",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f674c05",
   "metadata": {},
   "source": [
    "### Causality: no peeking at the future\n",
    "\n",
    "Look back at the printed output for position 2: `attn_weights` had exactly 3 entries, for positions 0, 1, and 2. The model was *not* able to look at position 3 — even though, during training, we know what the token at position 3 is.\n",
    "\n",
    "This matters because of what the model is being asked to do. At position 2, the model's job is to predict the token at position 3. If it could look at position 3, it could just copy the answer — it would get a perfect score during training and learn nothing useful. A model that can only look backward is called **causal** (or \"autoregressive\"), and it is what makes generation possible: when generating a new name, there *is* no future to look at.\n",
    "\n",
    "**Question:** In the very first call to `gpt` (position 0), every head's attention weights are exactly `[1.0]`. Why? And more generally: which specific lines of code guarantee that a token can never attend to a future token? *(Hint: what is in `keys[li]` at the moment the attention scores are computed, and how did it get there?)*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19a345e2",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15adbfbd",
   "metadata": {},
   "source": [
    "### Why several heads?\n",
    "\n",
    "We split the 16-dimensional `q`, `k`, `v` into 4 heads of size 4, and each head computes its *own* attention weights over the cached tokens. A single set of attention weights can only express one \"pattern of looking\": for example, mostly at the previous letter. But the current letter might benefit from several kinds of context at once: what the previous letter was, whether the name started with a vowel, how far in it is.\n",
    "\n",
    "**Question:** In your own words, why might it be useful to have several heads instead of one?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "495236c7",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7fb58b5",
   "metadata": {},
   "source": [
    "### Step 8: The MLP block\n",
    "\n",
    "```python\n",
    "x_residual = x\n",
    "x = rmsnorm(x)\n",
    "x = linear(x, state_dict[f'layer{li}.mlp_fc1'])    # 16 -> 64\n",
    "x = [xi.relu() for xi in x]                        # ReLU, element by element\n",
    "x = linear(x, state_dict[f'layer{li}.mlp_fc2'])    # 64 -> 16\n",
    "x = [a + b for a, b in zip(x, x_residual)]         # residual connection\n",
    "```\n",
    "\n",
    "After attention, each position goes through an **MLP** — a small two-layer neural network, the same kind you built in Lab 2. It expands the vector from 16 to 64 numbers (`mlp_fc1` has 64 rows), applies ReLU, and shrinks it back to 16 (`mlp_fc2` has 16 rows).\n",
    "\n",
    "The MLP and attention have complementary jobs. **Attention is where tokens communicate**: it is the only step that combines information across positions. **The MLP is where each token computes**: it takes the information that the current position has gathered and transforms it, on its own, without looking at any other position. Every position goes through the *same* MLP with the *same* weights. A helpful slogan: *attention = communication, MLP = computation.*\n",
    "\n",
    "Why is there a ReLU in the middle? You saw the reason in Lab 2: two `linear` layers back to back with nothing in between collapse into a single linear layer. The nonlinearity is what lets the MLP compute things a single matrix can't."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf36dc74",
   "metadata": {},
   "source": [
    "### Residual connections\n",
    "\n",
    "Twice in `gpt`, you see this pattern:\n",
    "\n",
    "```python\n",
    "x_residual = x          # save the input to the block\n",
    "x = ... block ...       # compute something from it\n",
    "x = [a + b for a, b in zip(x, x_residual)]     # ADD the block's output to the saved input\n",
    "```\n",
    "\n",
    "This is a **residual connection**: the block's output is *added* to its input instead of replacing it. Each block only needs to learn a *change* to make to `x`, rather than having to rebuild everything from scratch — and the original information in `x` is always carried forward, so a block can never accidentally destroy it.\n",
    "\n",
    "Residual connections also help on the backward pass. Recall from Part 3 that the local gradient of `a + b` with respect to either input is 1. So gradients flow straight through the `+` back to `x_residual`, unchanged, no matter what the block in between did. In deep networks with many blocks, this keeps the gradients from shrinking away to nothing by the time they reach the early layers — which is what makes it possible to train very deep Transformers at all.\n",
    "\n",
    "**Questions:**\n",
    "1. **Predict:** Suppose the attention block happened to output all zeros (so `x` is all zeros right before the first residual line). What would `x` be after the residual line?\n",
    "2. In your own words: why do residual connections make a deep network easier to train?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3056c86b",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39ca4e37",
   "metadata": {},
   "source": [
    "### Step 9: Logits\n",
    "\n",
    "```python\n",
    "logits = linear(x, state_dict['lm_head'])\n",
    "return logits\n",
    "```\n",
    "\n",
    "The last step turns the final 16-dimensional `x` into a score for every possible next token. `lm_head` has `vocab_size` = 27 rows, so this `linear` produces 27 numbers.\n",
    "\n",
    "**Question:** What does each entry of `logits` represent? What do we do with `logits` to get a probability for each possible next character?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1d313fd",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f8cb0bc",
   "metadata": {},
   "source": [
    "> **✅ Check yourself** — these are the kinds of questions you'll be asked at the check-off. Can you answer them out loud, in your own words, without looking at your notes?\n",
    ">\n",
    "> - Walk through the `gpt` function at a high level: token + position → ? → ? → ? → logits.\n",
    "> - What is a query, a key, and a value? How are the attention weights computed from them, and what do they mean?\n",
    "> - What does a dot product measure, and where does attention use one?\n",
    "> - How does microgpt guarantee that a token cannot attend to future tokens? Why does that matter?\n",
    "> - Why does the model need position embeddings, given how attention works?\n",
    "> - What is the difference between what attention does and what the MLP does?\n",
    "> - What would go wrong if we removed the residual connections?\n",
    "> - At position 3 of a name, what is the length of `attn_weights`? Of `x_attn`? Of `logits`?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1460c04",
   "metadata": {},
   "source": [
    "---\n",
    "# Part 5: Training and generating\n",
    "\n",
    "## 5.1 Inputs and targets\n",
    "\n",
    "A GPT is trained on one task: **predict the next token**. For the name `emma`, with tokens `[BOS, e, m, m, a, BOS]`, each position is a separate prediction.\n",
    "\n",
    "**Fill in by hand:** For each position, what is the input token and what is the target (the token the model should predict)?\n",
    "\n",
    "| `pos_id` | input token | target token |\n",
    "|---|---|---|\n",
    "| 0 | ? | ? |\n",
    "| 1 | ? | ? |\n",
    "| 2 | ? | ? |\n",
    "| 3 | ? | ? |\n",
    "| 4 | ? | ? |\n",
    "\n",
    "How many predictions does one name of $L$ letters give us?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a569799",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "233c976d",
   "metadata": {},
   "source": [
    "## 5.2 The loss before training\n",
    "\n",
    "For each prediction, the loss is $-\\log(p_\\text{target})$: the negative log of the probability the model gave to the correct next token. The loss for a name is the average over its positions. (This is the **cross-entropy loss** from Lab 2.)\n",
    "\n",
    "**Predict:** Before training, the weights are small and random, so the model's predictions are close to *uniform* — every one of the 27 tokens gets probability about 1/27. What loss do you expect? Put your prediction in the cell below and run it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b72a23e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "my_initial_loss_prediction = None   # <-- replace None with a number\n",
    "\n",
    "# Compute the loss of the (untrained) model on the first few names\n",
    "losses = []\n",
    "for doc in docs[:20]:\n",
    "    tokens = encode(doc)\n",
    "    n = min(block_size, len(tokens) - 1)\n",
    "    keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]\n",
    "    for pos_id in range(n):\n",
    "        probs = softmax(gpt(tokens[pos_id], pos_id, keys, values))\n",
    "        losses.append(-math.log(probs[tokens[pos_id + 1]].data))\n",
    "print(f\"your prediction:    {my_initial_loss_prediction}\")\n",
    "print(f\"actual average loss: {sum(losses) / len(losses):.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8fdeef1",
   "metadata": {},
   "source": [
    "**Question:** A model that guesses uniformly at random gets a loss of about 3.30. What loss would a *perfect* model get? Is a perfect model possible for this data? Why or why not?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afa3240d",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4fd78f3",
   "metadata": {},
   "source": [
    "## 5.3 The training loop\n",
    "\n",
    "Read the training loop below, then answer the questions. It takes about a minute to run.\n",
    "We made two small changes: we record the loss at each step (so we can plot it), and we print less often."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e360a0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let there be Adam, the blessed optimizer and its buffers\n",
    "learning_rate, beta1, beta2, eps_adam = 0.01, 0.85, 0.99, 1e-8\n",
    "m = [0.0] * len(params) # first moment buffer\n",
    "v = [0.0] * len(params) # second moment buffer\n",
    "\n",
    "# Repeat in sequence\n",
    "num_steps = 1000 # number of training steps\n",
    "loss_history = []   # [HW3]\n",
    "t0 = time.time()    # [HW3]\n",
    "for step in range(num_steps):\n",
    "\n",
    "    # Take single document, tokenize it, surround it with BOS special token on both sides\n",
    "    doc = docs[step % len(docs)]\n",
    "    tokens = [BOS] + [uchars.index(ch) for ch in doc] + [BOS]\n",
    "    n = min(block_size, len(tokens) - 1)\n",
    "\n",
    "    # Forward the token sequence through the model, building up the computation graph all the way to the loss\n",
    "    keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]\n",
    "    losses = []\n",
    "    for pos_id in range(n):\n",
    "        token_id, target_id = tokens[pos_id], tokens[pos_id + 1]\n",
    "        logits = gpt(token_id, pos_id, keys, values)\n",
    "        probs = softmax(logits)\n",
    "        loss_t = -probs[target_id].log()\n",
    "        losses.append(loss_t)\n",
    "    loss = (1 / n) * sum(losses) # final average loss over the document sequence. May yours be low.\n",
    "\n",
    "    # Backward the loss, calculating the gradients with respect to all model parameters\n",
    "    loss.backward()\n",
    "\n",
    "    # Adam optimizer update: update the model parameters based on the corresponding gradients\n",
    "    lr_t = learning_rate * (1 - step / num_steps) # linear learning rate decay\n",
    "    for i, p in enumerate(params):\n",
    "        m[i] = beta1 * m[i] + (1 - beta1) * p.grad\n",
    "        v[i] = beta2 * v[i] + (1 - beta2) * p.grad ** 2\n",
    "        m_hat = m[i] / (1 - beta1 ** (step + 1))\n",
    "        v_hat = v[i] / (1 - beta2 ** (step + 1))\n",
    "        p.data -= lr_t * m_hat / (v_hat ** 0.5 + eps_adam)\n",
    "        p.grad = 0\n",
    "\n",
    "    loss_history.append(loss.data)   # [HW3]\n",
    "    if (step + 1) % 100 == 0:        # [HW3] print every 100 steps\n",
    "        print(f\"step {step+1:4d} / {num_steps:4d} | loss {loss.data:.4f} | \"\n",
    "              f\"avg of last 100: {sum(loss_history[-100:]) / 100:.4f} | {time.time() - t0:.0f}s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51525138",
   "metadata": {},
   "source": [
    "**Adam** is an improved version of the gradient descent you used in Lab 1. The key line is:\n",
    "\n",
    "```python\n",
    "p.data -= lr_t * m_hat / (v_hat ** 0.5 + eps_adam)\n",
    "```\n",
    "\n",
    "where `m_hat` is (roughly) a running average of recent gradients, and `v_hat` is a running average of recent *squared* gradients. So, like plain gradient descent, each parameter takes a small step in the direction *opposite* its gradient. The division by `(v_hat ** 0.5 + eps_adam)` normalizes each step by the typical size of that parameter's gradients: in effect, every parameter gets its own learning rate, which adapts to how large its gradients have been, instead of one global learning rate for all 4,192 parameters. The tiny `eps_adam` avoids numerical problems when the denominator is close to zero.\n",
    "\n",
    "**Questions:**\n",
    "1. Describe **one training step** in plain words, in 3–4 sentences. Use the words *forward pass*, *loss*, *backward*, and *update*.\n",
    "2. At the end of the update loop, the code sets `p.grad = 0`. What would go wrong if we deleted that line? *(Hint: look at how `backward()` updates `child.grad` in the `Value` class — it uses `+=`.)*\n",
    "3. `lr_t` gets smaller as training goes on. Why might that be a good idea?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5eb260b2",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "483cdc85",
   "metadata": {},
   "source": [
    "## 5.4 The loss curve\n",
    "\n",
    "The loss for any *single* step is noisy (some names are just harder to predict than others), so we also plot a moving average."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68ca756a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def moving_average(xs, window=50):\n",
    "    return [sum(xs[max(0, i - window + 1):i + 1]) / (i + 1 - max(0, i - window + 1)) for i in range(len(xs))]\n",
    "\n",
    "plt.figure(figsize=(8, 4))\n",
    "plt.plot(loss_history, alpha=0.3, label='loss per step')\n",
    "plt.plot(moving_average(loss_history), label='moving average (50 steps)')\n",
    "plt.axhline(math.log(vocab_size), color='gray', linestyle='--', label='uniform guessing: log(27)')\n",
    "plt.xlabel('step'); plt.ylabel('loss'); plt.legend(); plt.show()\n",
    "print(f\"average loss over the last 100 steps: {sum(loss_history[-100:]) / 100:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c33716c8",
   "metadata": {},
   "source": [
    "**Question:** Look at the curve. Roughly where does the loss start and end? Is it still going down at the end? What would you try if you wanted a lower loss?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13af60ed",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6eeb533b",
   "metadata": {},
   "source": [
    "## 5.5 Generating new names\n",
    "\n",
    "Karpathy's inference code, wrapped in a function so you can call it with different settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e178b9b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inference: may the model babble back to us\n",
    "def generate(temperature=0.5, num_samples=20):   # [HW3] wrapped in a function\n",
    "    samples = []\n",
    "    for sample_idx in range(num_samples):\n",
    "        keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]\n",
    "        token_id = BOS\n",
    "        sample = []\n",
    "        for pos_id in range(block_size):\n",
    "            logits = gpt(token_id, pos_id, keys, values)\n",
    "            probs = softmax([l / temperature for l in logits])\n",
    "            token_id = random.choices(range(vocab_size), weights=[p.data for p in probs])[0]\n",
    "            if token_id == BOS:\n",
    "                break\n",
    "            sample.append(uchars[token_id])\n",
    "        samples.append(''.join(sample))\n",
    "    return samples\n",
    "\n",
    "random.seed(0)   # [HW3] so everyone sees the same samples\n",
    "names = generate(temperature=0.5)\n",
    "print(names)\n",
    "print(f\"{sum(name in set(docs) for name in names)} of {len(names)} generated names appear in the training data\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9717b247",
   "metadata": {},
   "source": [
    "**Questions:**\n",
    "1. Starting from `token_id = BOS`, describe how one new name is generated.\n",
    "2. Some generated names are real names from the training set, and others are new but still *look like* names. What has the model learned, and what hasn't it learned? Is it just copying the training data?\n",
    "3. Generation stops when the model samples `BOS`, *or* after `block_size` tokens. Why is the second condition needed?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b92f1fea",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33daea6f",
   "metadata": {},
   "source": [
    "## 5.6 Attention after training\n",
    "\n",
    "Now that the model is trained, let's see what the attention heads actually look at.\n",
    "\n",
    "Recall from Part 4 what the attention weights are: at each position, each head computes one weight per cached token, the weights are positive and **sum to 1**, and `attn_weights[t]` is the fraction of the current token's attention that goes to token `t` — equivalently, how much of token `t`'s *value* vector gets mixed into `head_out`. The function below feeds a name through the trained model, records those weights at every position, and draws them as a heatmap, one per head:\n",
    "\n",
    "- each **row** is a current token (the position being processed),\n",
    "- each **column** is a token it could attend to,\n",
    "- the **color** of the cell at (row, column) is the attention weight: bright = a big share of this row's attention, dark = almost none.\n",
    "\n",
    "So every row of every heatmap sums to 1, and a bright cell means \"when the model was at this token, this head was mostly looking at that token.\" Try a few different names by changing `name`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23e3f145",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_attention(name):\n",
    "    tokens = encode(name)                       # [BOS, letters..., BOS]\n",
    "    labels = ['BOS'] + list(name)               # tick labels: BOS followed by the letters\n",
    "    n = len(tokens) - 1                         # number of positions the model processes (last BOS is only a target)\n",
    "    keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]   # a fresh, empty KV cache\n",
    "\n",
    "    # grids[h] is an n x n table for head h: grids[h][row][col] = attention weight\n",
    "    # that the token at position `row` puts on the token at position `col`.\n",
    "    # Start with all zeros; cells for future tokens (col > row) will stay zero.\n",
    "    grids = [[[0.0] * n for _ in range(n)] for _ in range(n_head)]\n",
    "\n",
    "    for pos_id in range(n):\n",
    "        record = []                             # gpt() will append (layer, head, weights) tuples here\n",
    "        gpt(tokens[pos_id], pos_id, keys, values, attn_record=record)\n",
    "        for li, h, w in record:\n",
    "            if li == 0:                         # (only one layer in this model)\n",
    "                grids[h][pos_id][:len(w)] = w   # w has pos_id + 1 weights: one per cached token\n",
    "\n",
    "    # Draw one heatmap per head. Rows = current token, columns = token attended to.\n",
    "    fig, axes = plt.subplots(1, n_head, figsize=(3.2 * n_head, 3.2))\n",
    "    for h in range(n_head):\n",
    "        axes[h].imshow(grids[h], cmap='viridis', vmin=0, vmax=1)   # 0 = dark, 1 = bright\n",
    "        axes[h].set_xticks(range(n)); axes[h].set_xticklabels(labels[:n])\n",
    "        axes[h].set_yticks(range(n)); axes[h].set_yticklabels(labels[:n])\n",
    "        axes[h].set_title(f'head {h}')\n",
    "        axes[h].set_xlabel('attends to')\n",
    "    axes[0].set_ylabel('current token')\n",
    "    plt.tight_layout(); plt.show()\n",
    "\n",
    "name = 'isabella'\n",
    "plot_attention(name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d718f37",
   "metadata": {},
   "source": [
    "**Questions:**\n",
    "1. Why is everything above the diagonal exactly zero (dark)?\n",
    "2. Pick a bright cell in one of the heatmaps and say, in a sentence, what it means in terms of queries, keys, and values.\n",
    "3. Describe what at least two of the heads seem to be doing. (For example: does a head mostly look at the current token? The previous token? The BOS token? Vowels?) Try at least two different names to see if the pattern holds."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1116a75",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b72e897d",
   "metadata": {},
   "source": [
    "> **✅ Check yourself** — these are the kinds of questions you'll be asked at the check-off. Can you answer them out loud, in your own words, without looking at your notes?\n",
    ">\n",
    "> - What are the inputs and targets for a given name? How many predictions does one name give us?\n",
    "> - What loss does an untrained model get, and why that number?\n",
    "> - Describe one training step. Why do we reset `p.grad = 0`?\n",
    "> - How does the model generate a new name, step by step?\n",
    "> - Look at the attention heatmaps for a name. Point to a pattern in one head and explain what it means and why it might be useful for predicting the next letter."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c158ecc",
   "metadata": {},
   "source": [
    "---\n",
    "# Part 6: Experiments\n",
    "\n",
    "Now you'll change one thing at a time and see what happens. **At your check-off, you may be asked to explain one of these experiments.**\n",
    "\n",
    "For every experiment:\n",
    "1. **Predict** what will happen *before* running, and write it down.\n",
    "2. **Run** it and record what actually happened (numbers and example names).\n",
    "3. **Explain** why, connecting it to what you learned in Parts 2–5. If your prediction was wrong, explaining *why* it was wrong is the most valuable part.\n",
    "\n",
    "To make experiments easy, the cell below packages all the code you've read into a single function, `run_experiment`. **It is the same code** — skim it and convince yourself of that. The arguments let you change one setting at a time; everything else stays at Karpathy's defaults. The only new option is `use_pos_emb`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21502455",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_experiment(dataset='names.txt', n_layer=1, n_embd=16, n_head=4, block_size=16,\n",
    "                   learning_rate=0.01, num_steps=1000, use_pos_emb=True,\n",
    "                   temperature=0.5, num_samples=10, seed=42):\n",
    "    t0 = time.time()\n",
    "    random.seed(seed)\n",
    "\n",
    "    # Dataset and tokenizer\n",
    "    docs = [line.strip() for line in open(dataset) if line.strip()]\n",
    "    random.shuffle(docs)\n",
    "    uchars = sorted(set(''.join(docs)))\n",
    "    BOS = len(uchars)\n",
    "    vocab_size = len(uchars) + 1\n",
    "\n",
    "    # Parameters\n",
    "    head_dim = n_embd // n_head\n",
    "    matrix = lambda nout, nin, std=0.08: [[Value(random.gauss(0, std)) for _ in range(nin)] for _ in range(nout)]\n",
    "    state_dict = {'wte': matrix(vocab_size, n_embd), 'wpe': matrix(block_size, n_embd), 'lm_head': matrix(vocab_size, n_embd)}\n",
    "    for i in range(n_layer):\n",
    "        state_dict[f'layer{i}.attn_wq'] = matrix(n_embd, n_embd)\n",
    "        state_dict[f'layer{i}.attn_wk'] = matrix(n_embd, n_embd)\n",
    "        state_dict[f'layer{i}.attn_wv'] = matrix(n_embd, n_embd)\n",
    "        state_dict[f'layer{i}.attn_wo'] = matrix(n_embd, n_embd)\n",
    "        state_dict[f'layer{i}.mlp_fc1'] = matrix(4 * n_embd, n_embd)\n",
    "        state_dict[f'layer{i}.mlp_fc2'] = matrix(n_embd, 4 * n_embd)\n",
    "    params = [p for mat in state_dict.values() for row in mat for p in row]\n",
    "\n",
    "    # Model (same as gpt() above, except for the use_pos_emb option)\n",
    "    def gpt(token_id, pos_id, keys, values):\n",
    "        tok_emb = state_dict['wte'][token_id]\n",
    "        if use_pos_emb:\n",
    "            pos_emb = state_dict['wpe'][pos_id]\n",
    "            x = [t + p for t, p in zip(tok_emb, pos_emb)]\n",
    "        else:\n",
    "            x = tok_emb\n",
    "        x = rmsnorm(x)\n",
    "        for li in range(n_layer):\n",
    "            x_residual = x\n",
    "            x = rmsnorm(x)\n",
    "            q = linear(x, state_dict[f'layer{li}.attn_wq'])\n",
    "            k = linear(x, state_dict[f'layer{li}.attn_wk'])\n",
    "            v = linear(x, state_dict[f'layer{li}.attn_wv'])\n",
    "            keys[li].append(k)\n",
    "            values[li].append(v)\n",
    "            x_attn = []\n",
    "            for h in range(n_head):\n",
    "                hs = h * head_dim\n",
    "                q_h = q[hs:hs+head_dim]\n",
    "                k_h = [ki[hs:hs+head_dim] for ki in keys[li]]\n",
    "                v_h = [vi[hs:hs+head_dim] for vi in values[li]]\n",
    "                attn_logits = [sum(q_h[j] * k_h[t][j] for j in range(head_dim)) / head_dim**0.5 for t in range(len(k_h))]\n",
    "                attn_weights = softmax(attn_logits)\n",
    "                head_out = [sum(attn_weights[t] * v_h[t][j] for t in range(len(v_h))) for j in range(head_dim)]\n",
    "                x_attn.extend(head_out)\n",
    "            x = linear(x_attn, state_dict[f'layer{li}.attn_wo'])\n",
    "            x = [a + b for a, b in zip(x, x_residual)]\n",
    "            x_residual = x\n",
    "            x = rmsnorm(x)\n",
    "            x = linear(x, state_dict[f'layer{li}.mlp_fc1'])\n",
    "            x = [xi.relu() for xi in x]\n",
    "            x = linear(x, state_dict[f'layer{li}.mlp_fc2'])\n",
    "            x = [a + b for a, b in zip(x, x_residual)]\n",
    "        return linear(x, state_dict['lm_head'])\n",
    "\n",
    "    # Training\n",
    "    beta1, beta2, eps_adam = 0.85, 0.99, 1e-8\n",
    "    m = [0.0] * len(params)\n",
    "    v = [0.0] * len(params)\n",
    "    loss_history = []\n",
    "    for step in range(num_steps):\n",
    "        doc = docs[step % len(docs)]\n",
    "        tokens = [BOS] + [uchars.index(ch) for ch in doc] + [BOS]\n",
    "        n = min(block_size, len(tokens) - 1)\n",
    "        keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]\n",
    "        losses = []\n",
    "        for pos_id in range(n):\n",
    "            probs = softmax(gpt(tokens[pos_id], pos_id, keys, values))\n",
    "            losses.append(-probs[tokens[pos_id + 1]].log())\n",
    "        loss = (1 / n) * sum(losses)\n",
    "        loss.backward()\n",
    "        lr_t = learning_rate * (1 - step / num_steps)\n",
    "        for i, p in enumerate(params):\n",
    "            m[i] = beta1 * m[i] + (1 - beta1) * p.grad\n",
    "            v[i] = beta2 * v[i] + (1 - beta2) * p.grad ** 2\n",
    "            m_hat = m[i] / (1 - beta1 ** (step + 1))\n",
    "            v_hat = v[i] / (1 - beta2 ** (step + 1))\n",
    "            p.data -= lr_t * m_hat / (v_hat ** 0.5 + eps_adam)\n",
    "            p.grad = 0\n",
    "        loss_history.append(loss.data)\n",
    "\n",
    "    # Generation\n",
    "    random.seed(0)\n",
    "    samples = []\n",
    "    for _ in range(num_samples):\n",
    "        keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]\n",
    "        token_id, sample = BOS, []\n",
    "        for pos_id in range(block_size):\n",
    "            probs = softmax([l / temperature for l in gpt(token_id, pos_id, keys, values)])\n",
    "            token_id = random.choices(range(vocab_size), weights=[p.data for p in probs])[0]\n",
    "            if token_id == BOS:\n",
    "                break\n",
    "            sample.append(uchars[token_id])\n",
    "        samples.append(''.join(sample))\n",
    "\n",
    "    result = {\n",
    "        'num_params': len(params),\n",
    "        'final_loss': sum(loss_history[-200:]) / len(loss_history[-200:]),   # average of the last 200 steps\n",
    "        'seconds': time.time() - t0,\n",
    "        'samples': samples,\n",
    "        'loss_history': loss_history,\n",
    "    }\n",
    "    print(f\"params: {result['num_params']} | avg loss (last 200 steps): {result['final_loss']:.4f} | \"\n",
    "          f\"time: {result['seconds']:.0f}s\")\n",
    "    print(\"samples:\", samples)\n",
    "    return result\n",
    "\n",
    "\n",
    "def compare(results):\n",
    "    # Plot the moving-average loss curves of several experiments on one graph.\n",
    "    plt.figure(figsize=(8, 4))\n",
    "    for label, r in results.items():\n",
    "        plt.plot(moving_average(r['loss_history']), label=f\"{label} ({r['final_loss']:.3f})\")\n",
    "    plt.xlabel('step'); plt.ylabel('loss (moving average)'); plt.legend(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adc2d97c",
   "metadata": {},
   "source": [
    "First, run the **baseline** — Karpathy's default settings. Every other experiment will be compared to this. (About 1 minute.)\n",
    "\n",
    "The baseline trains exactly the same model as Part 5 (same random seed, same names in the same order), so you should see the same generated names. The loss number is a little different only because `run_experiment` averages the last 200 steps instead of the last 100. Always compare experiments using the numbers from `run_experiment`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7c04d94",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = {}\n",
    "results['baseline'] = run_experiment()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "146e5624",
   "metadata": {},
   "source": [
    "## Experiment 1: Temperature\n",
    "\n",
    "This one doesn't require retraining: we'll sample from the model you trained in Part 5, dividing the logits by different `temperature` values before the softmax.\n",
    "\n",
    "**Predict:** What will names look like with `temperature=0.1`? With `temperature=2.0`? Why? *(Hint: remember what happened in Part 4b when you multiplied all the logits by 10.)*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "406f071d",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb699ae2",
   "metadata": {},
   "outputs": [],
   "source": [
    "for temperature in [0.1, 0.5, 1.0, 2.0]:\n",
    "    random.seed(0)\n",
    "    print(f\"temperature {temperature}:\", generate(temperature=temperature, num_samples=12))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e197f16",
   "metadata": {},
   "source": [
    "**Result and explanation:** What happened? Why? Which temperature would you use if you wanted realistic names? Unusual, creative names?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d736518",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbadd24e",
   "metadata": {},
   "source": [
    "## Experiment 2: Remove the position embedding\n",
    "\n",
    "**Predict:** With `use_pos_emb=False`, the model has no information about positions. Will the loss be higher, lower, or about the same as the baseline? By a lot or a little? What will the generated names look like? Write your prediction *before* running (about 1 minute)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bce68b35",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c69c0f2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "results['no position embedding'] = run_experiment(use_pos_emb=False)\n",
    "compare({k: results[k] for k in ['baseline', 'no position embedding']})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d8eb4b0",
   "metadata": {},
   "source": [
    "**Result and explanation:** How big was the effect? Was it bigger or smaller than you expected? What can the model still figure out without position embeddings, and what can't it? *(Hint: without a position embedding, what information does the model have at each step? Think about the current token, and the set of tokens it can attend to.)*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c081e67",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "659dd019",
   "metadata": {},
   "source": [
    "## Experiment 3: Learning rate\n",
    "\n",
    "**Predict:** What will happen to the loss curve if the learning rate is **10× smaller** (`0.001`)? **10× bigger** (`0.1`)? Write predictions for both before running (about 2 minutes total)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c26878f5",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44ec19eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "results['lr = 0.001'] = run_experiment(learning_rate=0.001)\n",
    "results['lr = 0.1'] = run_experiment(learning_rate=0.1)\n",
    "compare({k: results[k] for k in ['baseline', 'lr = 0.001', 'lr = 0.1']})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "420aeba8",
   "metadata": {},
   "source": [
    "**Result and explanation:** What happened in each case, and why? Connect your answer to what a single update step does."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9535a388",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a131db90",
   "metadata": {},
   "source": [
    "## Experiment 4: Model size\n",
    "\n",
    "**Predict:** What will happen to the **loss**, the **number of parameters**, and the **training time** if we make the model wider (`n_embd=32`)? Deeper (`n_layer=2`)? *(Warning: these are slow. `n_embd=32` takes about 7 minutes and `n_layer=2` about 3 minutes on a fast laptop. Start them and take a break.)*\n",
    "\n",
    "Before running, **calculate by hand** how many parameters each model will have (use your table from Part 4a)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7f57036",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d04cfa21",
   "metadata": {},
   "outputs": [],
   "source": [
    "results['n_embd = 32'] = run_experiment(n_embd=32)\n",
    "results['n_layer = 2'] = run_experiment(n_layer=2)\n",
    "compare({k: results[k] for k in ['baseline', 'n_embd = 32', 'n_layer = 2']})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce84f477",
   "metadata": {},
   "source": [
    "**Result and explanation:** Did the bigger models help? Was it worth the extra time? Why do you think the improvement was the size it was? *(Hint: how many names has the model seen after 1000 steps, out of 32,000?)*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0324177",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "683c858a",
   "metadata": {},
   "source": [
    "## Experiment 5: A different dataset\n",
    "\n",
    "`words.txt` contains 10,000 English words (3–10 letters long) from Webster's dictionary. **Predict:** What will the vocabulary size be? Will the loss be higher or lower than for names? What will the generated \"words\" look like?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbf12f34",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "698de0cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "results['words'] = run_experiment(dataset='words.txt')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8c27e19",
   "metadata": {},
   "source": [
    "**Result and explanation:** Was the loss higher or lower than for names? Do the generated words look like English? Why do you think the loss was different? *(Nothing about the model changed — only the data.)*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88085f9f",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8174f6d1",
   "metadata": {},
   "source": [
    "## Summary of your experiments\n",
    "\n",
    "Run the cell below for a summary table of all the experiments you ran."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f78c6e85",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"{'experiment':25s} {'params':>8s} {'loss':>8s} {'time (s)':>9s}\")\n",
    "for label, r in results.items():\n",
    "    print(f\"{label:25s} {r['num_params']:8d} {r['final_loss']:8.4f} {r['seconds']:9.0f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72f451bc",
   "metadata": {},
   "source": [
    "> **✅ Check yourself** — these are the kinds of questions you'll be asked at the check-off. Can you answer them out loud, in your own words, without looking at your notes?\n",
    ">\n",
    "> - Pick one experiment. What did you predict, what happened, and why?\n",
    "> - Why does temperature change the output without changing the model?\n",
    "> - What happens when the learning rate is too small? Too big?\n",
    "> - Why didn't the bigger model help much?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a65ea06",
   "metadata": {},
   "source": [
    "---\n",
    "# Part 7: Going further (optional)\n",
    "\n",
    "These are for students who want more. They are **not** required, and the check-off will not ask about them unless you bring them up.\n",
    "\n",
    "1. **Look inside autograd.** Read the `Value` class in Part 3. For `__mul__`, why are the local gradients `(other.data, self.data)`? What does `backward()` do, and why does it need to visit the nodes in a particular order (a \"topological sort\")? Karpathy's video [*The spelled-out intro to neural networks and backpropagation: building micrograd*](https://www.youtube.com/watch?v=VMj-3S1tku0) explains this in detail.\n",
    "2. **Top-k sampling.** Modify `generate` so that it only samples from the `k` most likely tokens (set all the others' probabilities to zero). How do the names change for `k = 1`, `3`, `10`?\n",
    "3. **Your own dataset.** Make a text file with one item per line (city names, Pokémon, band names, ...) and train on it with `run_experiment(dataset=...)`. A few hundred items is enough to see something.\n",
    "4. **Train longer.** Run `run_experiment(num_steps=5000)` (about 5 minutes). How much does the loss improve?\n",
    "5. **Visualize the embeddings.** Each row of `state_dict['wte']` is a 16-number vector for one character. Use PCA (e.g., `numpy.linalg.svd`) to project them to 2-D and plot them with their letters. Do vowels cluster together?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ff4114e",
   "metadata": {},
   "source": [
    "---\n",
    "# Wrap-up: revisit the map\n",
    "\n",
    "Go back to your answer in **Part 1.1** and revise your one-sentence description of each section of microgpt. What did you get right? What changed?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69636980",
   "metadata": {},
   "source": [
    "> **✏️ Your answer:**\n",
    ">\n",
    "> *Double-click this cell, replace this line with your answer, then press Shift+Enter.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61e3b089",
   "metadata": {},
   "source": [
    "---\n",
    "# Check-off preparation\n",
    "\n",
    "At your check-off, you will answer questions about **four topics, chosen by the instructor** from:\n",
    "\n",
    "   - **Your Part 6 experiments** (explain one: what you predicted, what happened, and why)\n",
    "   - **Tokenization** (Part 2)\n",
    "   - **Embeddings** (Part 4)\n",
    "   - **Attention** (Part 4)\n",
    "   - **MLP, residual connections, and normalization** (Part 4)\n",
    "   - **Training** (Parts 3 and 5)\n",
    "   - **Generation and sampling** (Part 5)\n",
    "\n",
    "Each topic is scored from 0 to 3:\n",
    "\n",
    "| Score | What it looks like |\n",
    "|---|---|\n",
    "| **3** | Clear, correct explanation of *how and why* it works, without prompting. Can handle a follow-up question (\"what if...?\"). |\n",
    "| **2** | Correct explanation of the mechanism, but needs some prompting or has small gaps. |\n",
    "| **1** | Knows the vocabulary but can't explain the mechanism (\"attention lets the model pay attention to important things\"). |\n",
    "| **0** | No answer, or an incorrect explanation. |\n",
    "\n",
    "**How to prepare:** Go through every **✅ Check yourself** box and answer each question *out loud*, without notes. If you get stuck, go back to that part of the notebook, or ask the tutor chatbot to quiz you. Practicing with a classmate — taking turns asking each other questions — is one of the best ways to prepare."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
