Notebook: plain-words narrative + 8 interactive widgets

tiny_llm_from_scratch.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "e2785f32",
+   "id": "980e5f0c",
    "metadata": {},
    "source": [
     "# Tiny Storyteller: building a language model from scratch\n",
@@ -22,6 +22,7 @@
     "## The road map\n",
     "| Part | What we build | Why |\n",
     "|---|---|---|\n",
+    "| 0 | The story first, plus a live flow diagram | PyTorch, models, attention and the KV cache in plain words |\n",
     "| 1 | Tensors and gradients | The two things every neural net is made of |\n",
     "| 2 | Data: TinyStories | Simple English stories, made for small models |\n",
     "| 3 | A bigram model | The simplest language model that can \"write\" |\n",
@@ -37,7 +38,277 @@
   },
   {
    "cell_type": "markdown",
-   "id": "5a109db5",
+   "id": "983741fb",
+   "metadata": {},
+   "source": [
+    "# Before any code: the story of what we are about to build\n",
+    "\n",
+    "*This section has no code. It is for anyone, technical or not. If you understand these\n",
+    "five ideas, every cell below will make sense.*\n",
+    "\n",
+    "## 1. What is PyTorch?\n",
+    "\n",
+    "Imagine a kitchen built for one job: cooking with enormous grids of numbers. PyTorch is that\n",
+    "kitchen. It gives you the pots and pans (ways to hold grids of numbers), the stove (a fast\n",
+    "way to do arithmetic on millions of numbers at once, using a graphics card), and one\n",
+    "magical appliance that nothing else in ordinary programming has: a device that, after you\n",
+    "cook a dish, tells you exactly *which ingredient to adjust, and by how much, to make it\n",
+    "taste better next time*. That appliance is called **autograd**, and it is the reason modern AI\n",
+    "is possible.\n",
+    "\n",
+    "PyTorch was built by Meta's AI lab and is the tool behind most of the models you have heard\n",
+    "of: ChatGPT's ancestors, Llama, Stable Diffusion, Whisper. Hugging Face, where we will\n",
+    "publish our model, is the library and community that grew around sharing PyTorch models.\n",
+    "\n",
+    "## 2. What is a \"model\", really?\n",
+    "\n",
+    "A model is a very large spreadsheet of numbers, called **weights**, plus a fixed recipe for\n",
+    "how to multiply and add them. Our tiny model has about six million weights. The famous big\n",
+    "ones have hundreds of billions. At the start, every one of those numbers is random noise:\n",
+    "the model knows nothing. **Training** is the process of nudging those numbers, a little at a\n",
+    "time, millions of times, until the recipe produces useful output.\n",
+    "\n",
+    "Real-world picture: a new employee at a help desk who starts by guessing every answer.\n",
+    "Each time a customer says \"no, that's wrong\", the employee adjusts slightly. After a\n",
+    "million customers, the guesses become good. The employee never memorized a rulebook. They\n",
+    "tuned their instincts. The weights are those instincts.\n",
+    "\n",
+    "## 3. What does a language model actually do?\n",
+    "\n",
+    "One thing only: **given some text, guess the next word.** That is the whole job.\n",
+    "\n",
+    "\"Once upon a time, a little\" → probably \"girl\", maybe \"boy\", maybe \"fox\".\n",
+    "\n",
+    "Writing a story is just doing that guess over and over: guess a word, add it to the text,\n",
+    "guess the next one, add it, and so on. Chatting with an AI assistant is the same trick. Your\n",
+    "question is the text so far, and its answer is a long run of next-word guesses.\n",
+    "\n",
+    "## 4. What happens when you send a question to an LLM?\n",
+    "\n",
+    "Say you type: *\"What is the capital of France?\"* Here is the journey, step by step.\n",
+    "\n",
+    "**Step 1, chopping: text becomes tokens.** The model cannot read letters. A **tokenizer**\n",
+    "chops your sentence into pieces called tokens (roughly word-sized: \"What\", \" is\", \" the\",\n",
+    "\" capital\", \" of\", \" France\", \"?\") and swaps each for a number from its dictionary. Your\n",
+    "question is now a short list of integers.\n",
+    "\n",
+    "**Step 2, meaning: each token becomes a point in space.** Each number is looked up in a\n",
+    "table and replaced by a long list of values (in our model, 256 of them). Think of it as\n",
+    "placing every word somewhere on a giant map, where words with similar meaning sit close\n",
+    "together: \"France\" near \"Spain\", \"capital\" near \"city\". The map was learned during training.\n",
+    "\n",
+    "**Step 3, looking around: attention.** This is the heart of it, and it deserves its own\n",
+    "section below. In short: each word looks back at all the words before it and asks \"which of\n",
+    "you matter to me right now?\" The word \"capital\" looks around, finds \"France\", and pulls\n",
+    "some of France's meaning into itself. After this step, \"capital\" no longer means \"capital\n",
+    "in general\". It means \"the capital of France\". This looking-around is repeated in several\n",
+    "layers, so meaning gets refined again and again.\n",
+    "\n",
+    "**Step 4, the guess.** The last word's refined meaning is compared against every word in the\n",
+    "dictionary, producing a score for each. The scores become probabilities. \"Paris\" gets, say,\n",
+    "92 percent. One word is picked.\n",
+    "\n",
+    "**Step 5, repeat.** \"Paris\" is appended to the text and the whole thing runs again to pick\n",
+    "the next word, until the model emits a special \"I'm done\" token. You see the words appear\n",
+    "one by one on screen because that is literally how they are produced.\n",
+    "\n",
+    "## 5. How attention \"looks things up\", in plain words\n",
+    "\n",
+    "Picture a meeting room. Every word in the sentence is a person sitting at the table. Each\n",
+    "person carries three cards:\n",
+    "\n",
+    "- a **query** card: \"here is what I am looking for\" (the word \"capital\" is looking for a\n",
+    "  country),\n",
+    "- a **key** card: \"here is what I can offer\" (the word \"France\" offers: I am a country),\n",
+    "- a **value** card: \"here is what I will hand over if you pick me\" (France's actual meaning).\n",
+    "\n",
+    "Each person holds up their query card and compares it with everyone else's key card. A good\n",
+    "match scores high, a poor match scores low. The scores are turned into percentages that add\n",
+    "up to 100, and the person collects a blend of everyone's value cards, weighted by those\n",
+    "percentages. \"capital\" might take 70 percent from \"France\", 20 percent from \"of\", and a\n",
+    "little from the rest. That blend becomes its new, more specific meaning.\n",
+    "\n",
+    "Two rules make this work for *writing*:\n",
+    "\n",
+    "- **No peeking ahead.** When the model is guessing word number 7, it may only look at words\n",
+    "  1 to 6. In code, this is a triangular mask: future words are blocked out.\n",
+    "- **Many heads.** The meeting runs several times in parallel with different cards, called\n",
+    "  \"heads\". One head might track grammar (which noun does this verb belong to), another might\n",
+    "  track names, another might track position. Our model uses 4 heads. Big models use dozens.\n",
+    "\n",
+    "## 6. The KV cache: why chat answers do not slow down as they get longer\n",
+    "\n",
+    "Remember Step 5: to produce each new word, the whole text is run through the model again.\n",
+    "Done naively, writing the 500th word means re-reading 499 words from scratch, and the 501st\n",
+    "means re-reading 500. The cost grows and grows.\n",
+    "\n",
+    "But notice: the key card and value card of a word that is already written **never change**.\n",
+    "\"France\" offers the same key and value whether the sentence is 7 words or 700 words long.\n",
+    "So the model simply keeps every finished word's key and value cards in a drawer, the\n",
+    "**KV cache** (K for key, V for value). When a new word arrives, only *its* three cards are\n",
+    "computed. It holds up its query, compares against all the stored keys, collects the stored\n",
+    "values, and done. Nothing is recomputed.\n",
+    "\n",
+    "That drawer is why a long chat stays fast, and it is also why long conversations use a lot\n",
+    "of GPU memory: every token ever seen keeps its cards in the drawer. When people talk about a\n",
+    "model's \"context window\" (how much text it can hold at once), they are partly talking about\n",
+    "how big that drawer is allowed to get.\n",
+    "\n",
+    "We demonstrate the cache with a few lines of code in Part 5. Our tiny model does not need it\n",
+    "for 200-word stories, but every production LLM has one.\n",
+    "\n",
+    "## The one-sentence version\n",
+    "\n",
+    "A language model is a huge pile of tuned numbers that turns text into tokens, places them on a\n",
+    "map of meaning, lets each word look back at the others to sharpen its meaning (attention),\n",
+    "and guesses the next word, one at a time, with a cache so it never redoes old work.\n",
+    "\n",
+    "Now let's build one."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5efd4788",
+   "metadata": {},
+   "source": [
+    "## Live flow diagram\n",
+    "\n",
+    "Run the cell below (it only draws, it trains nothing). Press **Play** and watch the model\n",
+    "produce one word at a time: the text is chopped into tokens, each token gets a place on the\n",
+    "meaning map, the newest token looks back at the others (the orange arcs are attention), the\n",
+    "dictionary is scored, a word is picked, and it joins the text. On the right, the KV cache\n",
+    "fills up one card per token, and only the newest token ever needs new cards."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b619ba27",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from IPython.display import HTML, display\n",
+    "\n",
+    "# Everything below is plain HTML + JavaScript, no libraries. It runs inside the notebook output.\n",
+    "display(HTML(r\"\"\"<div id=\"llmflow\" style=\"font-family:-apple-system,Segoe UI,Helvetica,Arial,sans-serif;color:#e8e8e8;background:#1e1e1e;border-radius:12px;padding:16px;max-width:980px\">\n",
+    "<style>\n",
+    "#llmflow .stage{display:inline-block;padding:8px 12px;border-radius:8px;border:2px solid #555;margin:0 4px;min-width:110px;text-align:center;font-size:13px;transition:all .3s;background:#2a2a2a}\n",
+    "#llmflow .stage.on{border-color:#ffb347;background:#3a2e1a;box-shadow:0 0 12px #ffb34788}\n",
+    "#llmflow .arrow{display:inline-block;color:#888;font-size:18px;vertical-align:middle}\n",
+    "#llmflow .chip{display:inline-block;padding:4px 8px;margin:3px;border-radius:6px;background:#2f4f6f;border:1px solid #5a8ab8;font-size:13px;transition:all .3s}\n",
+    "#llmflow .chip.new{background:#6a4a1a;border-color:#ffb347}\n",
+    "#llmflow .chip.q{outline:2px solid #ffb347}\n",
+    "#llmflow .slot{display:inline-block;width:22px;height:30px;margin:2px;border:1px dashed #666;border-radius:4px;vertical-align:middle;font-size:9px;text-align:center;line-height:30px;color:#aaa}\n",
+    "#llmflow .slot.full{background:#2f6f4f;border:1px solid #6fcf97;color:#fff}\n",
+    "#llmflow .bar{height:14px;background:#5a8ab8;border-radius:3px;margin:2px 0;transition:width .4s}\n",
+    "#llmflow button{background:#333;color:#eee;border:1px solid #777;border-radius:6px;padding:6px 12px;margin-right:6px;cursor:pointer;font-size:13px}\n",
+    "#llmflow button:hover{background:#444}\n",
+    "#llmflow .panel{background:#262626;border-radius:8px;padding:10px 12px;margin-top:10px;font-size:13px;line-height:1.45}\n",
+    "#llmflow .lbl{font-size:11px;color:#aaa;text-transform:uppercase;letter-spacing:.05em;margin:8px 0 4px}\n",
+    "</style>\n",
+    "\n",
+    "<div style=\"font-size:15px;font-weight:600;margin-bottom:8px\">How a language model writes: live flow</div>\n",
+    "<div style=\"font-size:12px;color:#aaa;margin-bottom:10px\">Press <b>Play</b> and watch one word get produced, then the next. Attention weights are illustrative.</div>\n",
+    "\n",
+    "<div id=\"stages\">\n",
+    "  <span class=\"stage\" data-s=\"0\">1. Text so far</span><span class=\"arrow\">&#8594;</span>\n",
+    "  <span class=\"stage\" data-s=\"1\">2. Tokenizer<br><small>chop into tokens</small></span><span class=\"arrow\">&#8594;</span>\n",
+    "  <span class=\"stage\" data-s=\"2\">3. Meaning map<br><small>token &#8594; vector</small></span><span class=\"arrow\">&#8594;</span>\n",
+    "  <span class=\"stage\" data-s=\"3\">4. Attention<br><small>look back</small></span><span class=\"arrow\">&#8594;</span>\n",
+    "  <span class=\"stage\" data-s=\"4\">5. Next-word scores</span><span class=\"arrow\">&#8617;</span>\n",
+    "  <span class=\"stage\" data-s=\"5\">6. Append word, repeat</span>\n",
+    "</div>\n",
+    "\n",
+    "<div class=\"lbl\">Text so far (each chip is one token)</div>\n",
+    "<div id=\"chips\" style=\"min-height:36px\"></div>\n",
+    "<svg id=\"arcs\" width=\"940\" height=\"70\" style=\"display:block\"></svg>\n",
+    "\n",
+    "<div style=\"display:flex;gap:20px;flex-wrap:wrap\">\n",
+    "  <div style=\"flex:1;min-width:280px\">\n",
+    "    <div class=\"lbl\">Next-word scores (top 5)</div>\n",
+    "    <div id=\"bars\"></div>\n",
+    "  </div>\n",
+    "  <div style=\"flex:1;min-width:280px\">\n",
+    "    <div class=\"lbl\">KV cache: stored key/value cards, one per token</div>\n",
+    "    <div id=\"cache\"></div>\n",
+    "    <div style=\"font-size:11px;color:#aaa;margin-top:4px\">Only the newest token computes its cards. Old ones are reused.</div>\n",
+    "  </div>\n",
+    "</div>\n",
+    "\n",
+    "<div class=\"panel\" id=\"explain\">Press Play.</div>\n",
+    "<div style=\"margin-top:10px\">\n",
+    "  <button onclick=\"llmPlay()\">&#9654; Play</button>\n",
+    "  <button onclick=\"llmStep()\">Step</button>\n",
+    "  <button onclick=\"llmReset()\">Reset</button>\n",
+    "  <label style=\"font-size:12px;color:#aaa;margin-left:8px\">Speed <input id=\"spd\" type=\"range\" min=\"1\" max=\"5\" value=\"3\" style=\"vertical-align:middle\"></label>\n",
+    "</div>\n",
+    "</div>\n",
+    "\n",
+    "<script>\n",
+    "(function(){\n",
+    "const prompt = [\"Once\",\" upon\",\" a\",\" time\",\",\",\" a\",\" little\",\" fox\"];\n",
+    "const gen = [\n",
+    "  {w:\" found\", top:[[\" found\",.41],[\" was\",.22],[\" saw\",.14],[\" lived\",.09],[\" had\",.06]], att:[.05,.03,.04,.06,.02,.08,.22,.50]},\n",
+    "  {w:\" a\",     top:[[\" a\",.55],[\" the\",.18],[\" some\",.08],[\" her\",.06],[\" his\",.05]], att:[.03,.02,.05,.03,.02,.06,.12,.27,.40]},\n",
+    "  {w:\" shiny\", top:[[\" shiny\",.31],[\" big\",.27],[\" red\",.16],[\" small\",.10],[\" little\",.06]], att:[.02,.02,.03,.02,.01,.04,.18,.33,.10,.25]},\n",
+    "  {w:\" ball\",  top:[[\" ball\",.46],[\" stone\",.15],[\" box\",.12],[\" hat\",.08],[\" key\",.05]], att:[.02,.01,.02,.02,.01,.03,.08,.28,.06,.12,.35]},\n",
+    "  {w:\".\",      top:[[\".\",.52],[\" in\",.14],[\" and\",.12],[\" on\",.08],[\" that\",.05]], att:[.02,.01,.02,.03,.02,.02,.05,.15,.05,.10,.18,.35]},\n",
+    "];\n",
+    "let tokens=[], gi=0, stage=-1, timer=null;\n",
+    "const $=id=>document.getElementById(id);\n",
+    "const txt={\n",
+    " 0:\"<b>1. Text so far.</b> Everything written so far, prompt plus generated words, is the input. Nothing else.\",\n",
+    " 1:\"<b>2. Tokenizer.</b> The text is chopped into tokens (roughly words, note the leading spaces) and each becomes a number from the model's dictionary.\",\n",
+    " 2:\"<b>3. Meaning map.</b> Each token number is swapped for a vector, a long list of values, placing it on a learned map where similar words sit close together.\",\n",
+    " 3:\"<b>4. Attention.</b> The newest token (orange outline) looks back at every earlier token. Arc thickness shows how much it attends to each. It blends their value cards into its own meaning. Future tokens do not exist yet, so there is nothing to peek at. Repeated in 6 layers with 4 heads each.\",\n",
+    " 4:\"<b>5. Next-word scores.</b> The refined meaning of the last token is compared with every word in the dictionary. Scores become probabilities; one word is sampled.\",\n",
+    " 5:\"<b>6. Append and repeat.</b> The chosen word is added to the text and its key/value cards go into the KV cache. The loop starts again from step 1, but only the new token needs computing.\",\n",
+    "};\n",
+    "function render(){\n",
+    "  document.querySelectorAll('#llmflow .stage').forEach(s=>s.classList.toggle('on',+s.dataset.s===stage));\n",
+    "  const c=$('chips'); c.innerHTML='';\n",
+    "  tokens.forEach((t,i)=>{const s=document.createElement('span');s.className='chip'+(i===tokens.length-1&&gi>0&&stage<5?' new':'')+(stage===3&&i===tokens.length-1?' q':'');s.textContent=t.replace(/^ /,'␣');c.appendChild(s);});\n",
+    "  const k=$('cache'); k.innerHTML='';\n",
+    "  const nCached = stage>=5 ? tokens.length : Math.max(0, tokens.length-(gi>0?1:0));\n",
+    "  tokens.forEach((t,i)=>{const s=document.createElement('span');s.className='slot'+(i<nCached?' full':'');s.textContent=i<nCached?'KV':'';k.appendChild(s);});\n",
+    "  const b=$('bars'); b.innerHTML='';\n",
+    "  if(stage>=4 && gi<gen.length){gen[gi].top.forEach(([w,p])=>{b.innerHTML+=`<div style=\"display:flex;align-items:center;gap:8px;font-size:12px\"><span style=\"width:60px;text-align:right\">${w.replace(/^ /,'␣')}</span><div class=\"bar\" style=\"width:${Math.round(p*260)}px\"></div><span style=\"color:#aaa\">${Math.round(p*100)}%</span></div>`;});}\n",
+    "  else b.innerHTML='<div style=\"color:#666;font-size:12px\">(appears at step 5)</div>';\n",
+    "  drawArcs();\n",
+    "  $('explain').innerHTML = stage<0 ? 'Press Play.' : txt[stage];\n",
+    "}\n",
+    "function drawArcs(){\n",
+    "  const svg=$('arcs'); svg.innerHTML='';\n",
+    "  if(stage!==3||gi>=gen.length) return;\n",
+    "  const chips=[...document.querySelectorAll('#chips .chip')]; if(chips.length<2) return;\n",
+    "  const base=$('chips').getBoundingClientRect(); const last=chips[chips.length-1].getBoundingClientRect();\n",
+    "  const x1=last.left-base.left+last.width/2; const w=gen[gi].att;\n",
+    "  chips.slice(0,-1).forEach((ch,i)=>{const r=ch.getBoundingClientRect();const x0=r.left-base.left+r.width/2;const a=w[i]||0.02;\n",
+    "    const p=document.createElementNS('http://www.w3.org/2000/svg','path');\n",
+    "    p.setAttribute('d',`M${x1},2 Q${(x0+x1)/2},${20+Math.abs(x1-x0)/6} ${x0},2`);\n",
+    "    p.setAttribute('fill','none');p.setAttribute('stroke','#ffb347');p.setAttribute('stroke-width',(0.5+a*10).toFixed(1));p.setAttribute('opacity',(0.25+a*1.5).toFixed(2));\n",
+    "    svg.appendChild(p);\n",
+    "    const t=document.createElementNS('http://www.w3.org/2000/svg','text');t.setAttribute('x',x0);t.setAttribute('y',60);t.setAttribute('fill','#ffb347');t.setAttribute('font-size','10');t.setAttribute('text-anchor','middle');t.textContent=Math.round(a*100)+'%';svg.appendChild(t);});\n",
+    "}\n",
+    "window.llmReset=function(){clearInterval(timer);timer=null;tokens=prompt.slice();gi=0;stage=-1;render();};\n",
+    "window.llmStep=function(){\n",
+    "  if(gi>=gen.length){$('explain').innerHTML='<b>Done.</b> The model wrote: <i>'+tokens.join('')+'</i>. Press Reset to watch again.';clearInterval(timer);timer=null;return;}\n",
+    "  stage++;\n",
+    "  if(stage===5){tokens.push(gen[gi].w);}\n",
+    "  if(stage>5){stage=0;gi++; if(gi>=gen.length){render();window.llmStep();return;}}\n",
+    "  render();\n",
+    "};\n",
+    "window.llmPlay=function(){if(timer){clearInterval(timer);timer=null;return;} const ms=[0,1600,1200,900,600,350][+$('spd').value]; timer=setInterval(window.llmStep,ms); window.llmStep();};\n",
+    "window.llmReset();\n",
+    "window.addEventListener('resize',drawArcs);\n",
+    "})();\n",
+    "</script>\n",
+    "\"\"\"))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b33f712a",
    "metadata": {},
    "source": [
     "## 0. Setup\n",
@@ -47,7 +318,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "00b4393c",
+   "id": "649f2bf7",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -57,7 +328,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e97e5f1c",
+   "id": "c60680a5",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -74,20 +345,141 @@
   },
   {
    "cell_type": "markdown",
-   "id": "5eecb952",
+   "id": "05aa805c",
    "metadata": {},
    "source": [
     "## 1. Tensors and gradients: the two building blocks\n",
     "\n",
-    "A **tensor** is just a grid of numbers with a shape. A vector is 1-D, a matrix is 2-D,\n",
-    "a batch of sentences is 3-D. Everything in a neural network (the input text, the weights,\n",
-    "the output probabilities) is a tensor."
+    "A **tensor** is just a grid of numbers with a shape. Think of three everyday objects:\n",
+    "\n",
+    "- a **shopping list** is 1-D: one line of items. Shape `(4,)`. We call it a vector.\n",
+    "- a **spreadsheet** is 2-D: rows and columns. Shape `(3, 4)`. We call it a matrix.\n",
+    "- a **stack of spreadsheets**, one per customer, is 3-D. Shape `(2, 3, 4)`.\n",
+    "\n",
+    "Everything in a neural network is one of these: the input text (one row per token), the\n",
+    "weights (big spreadsheets), the output probabilities (one row per position). Click through\n",
+    "the dimensions below and hover any cell to see its address."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "62b705cc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Interactive: 1-D, 2-D and 3-D tensors. Run this cell, then click and hover.\n",
+    "from IPython.display import HTML, display\n",
+    "display(HTML(r\"\"\"<div id=\"vt\" style=\"font-family:-apple-system,Segoe UI,Helvetica,Arial,sans-serif;color:#e8e8e8;background:#1e1e1e;border-radius:12px;padding:16px;max-width:900px\">\n",
+    "<style>\n",
+    "#vt button{background:#333;color:#eee;border:1px solid #777;border-radius:6px;padding:6px 12px;margin-right:6px;cursor:pointer}\n",
+    "#vt button.on{background:#6a4a1a;border-color:#ffb347}\n",
+    "#vt .cell{display:inline-block;width:38px;height:30px;line-height:30px;text-align:center;font-size:12px;border:1px solid #5a8ab8;background:#2f4f6f;margin:1px;border-radius:3px;cursor:crosshair}\n",
+    "#vt .cell:hover{background:#ffb347;color:#000}\n",
+    "#vt .grid{display:inline-block;padding:6px;background:#262626;border-radius:6px;margin:6px;vertical-align:top}\n",
+    "#vt .cap{font-size:13px;line-height:1.5;margin-top:10px;background:#262626;border-radius:8px;padding:10px}\n",
+    "</style>\n",
+    "<div style=\"font-size:15px;font-weight:600;margin-bottom:6px\">A tensor is a grid of numbers with a shape. Click a dimension, hover a cell.</div>\n",
+    "<div><button id=\"b1\" class=\"on\" onclick=\"vtShow(1)\">1-D: vector</button><button id=\"b2\" onclick=\"vtShow(2)\">2-D: matrix</button><button id=\"b3\" onclick=\"vtShow(3)\">3-D: batch</button></div>\n",
+    "<div id=\"vtArea\" style=\"margin-top:10px;min-height:120px\"></div>\n",
+    "<div id=\"vtHover\" style=\"font-size:12px;color:#ffb347;height:18px\"></div>\n",
+    "<div class=\"cap\" id=\"vtCap\"></div>\n",
+    "</div>\n",
+    "<script>\n",
+    "(function(){\n",
+    "const caps={\n",
+    "1:\"<b>Shape (4,)</b>: one row of numbers. Real-world: a shopping list, or one word's meaning as 4 numbers. In our model each token becomes a vector of <b>256</b> numbers.\",\n",
+    "2:\"<b>Shape (3, 4)</b>: rows and columns, like a spreadsheet. Real-world: one sentence, one row per token, one column per feature. Our model holds a 256-token story as a (256, 256) matrix.\",\n",
+    "3:\"<b>Shape (2, 3, 4)</b>: a stack of spreadsheets. Real-world: a batch of sentences processed at the same time, which is how GPUs stay busy. Training uses batches of 32 stories: shape (32, 256, 256).\"};\n",
+    "function grid(r,c,prefix){let h='<div class=\\\"grid\\\">';for(let i=0;i<r;i++){h+='<div>';for(let j=0;j<c;j++){const v=(Math.round((Math.sin(i*3+j*7+prefix.length)*5)*10)/10).toFixed(1);h+=`<span class=\\\"cell\\\" onmouseover=\\\"vtH('${prefix}${i},${j}', ${v})\\\" onmouseout=\\\"vtH('')\\\">${v}</span>`;}h+='</div>';}return h+'</div>';}\n",
+    "window.vtH=function(idx,v){document.getElementById('vtHover').textContent=idx?`index [${idx}]  value ${v}`:'';};\n",
+    "window.vtShow=function(d){[1,2,3].forEach(k=>document.getElementById('b'+k).classList.toggle('on',k===d));\n",
+    " const a=document.getElementById('vtArea');\n",
+    " if(d===1)a.innerHTML=grid(1,4,'');\n",
+    " if(d===2)a.innerHTML=grid(3,4,'');\n",
+    " if(d===3)a.innerHTML='<span style=\"font-size:12px;color:#aaa\">sentence 0</span>'+grid(3,4,'0,')+'<span style=\"font-size:12px;color:#aaa\">sentence 1</span>'+grid(3,4,'1,');\n",
+    " document.getElementById('vtCap').innerHTML=caps[d];};\n",
+    "window.vtShow(1);\n",
+    "})();\n",
+    "</script>\n",
+    "\"\"\"))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "690ef025",
+   "metadata": {},
+   "source": [
+    "### How matrix multiplication actually happens\n",
+    "\n",
+    "Almost all the work inside a neural network is one operation: multiplying two of those\n",
+    "spreadsheets together. People find it mysterious because textbooks show the formula first.\n",
+    "Here is the story instead. Two customers buy items. Two shops price those items differently.\n",
+    "What does each customer pay in each shop? Walk one row along one column, multiply the\n",
+    "pairs, add them up. That is one cell of the answer. Press **Next cell** and watch."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "3f77d7e0",
+   "id": "9be728bd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Interactive: matrix multiplication, one output cell at a time.\n",
+    "from IPython.display import HTML, display\n",
+    "display(HTML(r\"\"\"<div id=\"vm\" style=\"font-family:-apple-system,Segoe UI,Helvetica,Arial,sans-serif;color:#e8e8e8;background:#1e1e1e;border-radius:12px;padding:16px;max-width:960px\">\n",
+    "<style>\n",
+    "#vm button{background:#333;color:#eee;border:1px solid #777;border-radius:6px;padding:6px 12px;margin-right:6px;cursor:pointer}\n",
+    "#vm table{border-collapse:separate;border-spacing:3px;display:inline-table;vertical-align:middle;margin:0 6px}\n",
+    "#vm td,#vm th{width:46px;height:34px;text-align:center;font-size:13px;border-radius:5px}\n",
+    "#vm td{background:#2f4f6f;border:1px solid #5a8ab8;transition:all .25s}\n",
+    "#vm th{color:#aaa;font-weight:normal;font-size:11px;background:transparent}\n",
+    "#vm td.row{background:#6a4a1a;border-color:#ffb347}\n",
+    "#vm td.col{background:#2f6f4f;border-color:#6fcf97}\n",
+    "#vm td.out{background:#ffb347;color:#000;font-weight:600}\n",
+    "#vm td.empty{background:#262626;border:1px dashed #555;color:#666}\n",
+    "#vm .sym{font-size:22px;color:#aaa;vertical-align:middle;margin:0 4px}\n",
+    "#vm .cap{font-size:13px;line-height:1.55;margin-top:10px;background:#262626;border-radius:8px;padding:10px;min-height:60px}\n",
+    "#vm .work{font-family:monospace;font-size:14px;color:#ffb347;margin-top:6px}\n",
+    "</style>\n",
+    "<div style=\"font-size:15px;font-weight:600;margin-bottom:4px\">Matrix multiplication, one cell at a time</div>\n",
+    "<div style=\"font-size:12px;color:#aaa;margin-bottom:10px\">Story: two customers buy items. Two shops have different prices. What does each customer pay at each shop?</div>\n",
+    "<div id=\"vmArea\" style=\"overflow-x:auto;white-space:nowrap\"></div>\n",
+    "<div class=\"work\" id=\"vmWork\"></div>\n",
+    "<div class=\"cap\" id=\"vmCap\"></div>\n",
+    "<div style=\"margin-top:10px\"><button onclick=\"vmStep()\">Next cell</button><button onclick=\"vmAuto()\">Play all</button><button onclick=\"vmReset()\">Reset</button></div>\n",
+    "</div>\n",
+    "<script>\n",
+    "(function(){\n",
+    "const A=[[2,1,3],[1,4,0]];           // customers x items   (apples, bread, milk)\n",
+    "const B=[[1,2],[3,2],[2,4]];         // items x shops       (price at shop 1, shop 2)\n",
+    "const items=[\"apples\",\"bread\",\"milk\"], custs=[\"Ana\",\"Ben\"], shops=[\"Shop 1\",\"Shop 2\"];\n",
+    "const C=[[null,null],[null,null]]; let k=-1, timer=null;\n",
+    "function tbl(M,rh,ch,cls){let h='<table><tr><th></th>'+ch.map(c=>`<th>${c}</th>`).join('')+'</tr>';M.forEach((r,i)=>{h+=`<tr><th>${rh[i]}</th>`+r.map((v,j)=>`<td class=\"${cls(i,j,v)}\">${v===null?'?':v}</td>`).join('')+'</tr>';});return h+'</table>';}\n",
+    "function render(){const i=k<0?-1:Math.floor(k/2), j=k<0?-1:k%2;\n",
+    " document.getElementById('vmArea').innerHTML=\n",
+    "  tbl(A,custs,items,(r,c)=>r===i?'row':'')+'<span class=\"sym\">&times;</span>'+\n",
+    "  tbl(B,items,shops,(r,c)=>c===j?'col':'')+'<span class=\"sym\">=</span>'+\n",
+    "  tbl(C,custs,shops,(r,c,v)=>(r===i&&c===j)?'out':(v===null?'empty':''));\n",
+    " if(k<0){document.getElementById('vmWork').textContent='';document.getElementById('vmCap').innerHTML='<b>Left table:</b> how many of each item each customer buys (rows = customers, columns = items).<br><b>Right table:</b> the price of each item in each shop (rows = items, columns = shops).<br>Press <b>Next cell</b>. Each answer cell = one <span style=\"color:#ffb347\">row</span> from the left paired with one <span style=\"color:#6fcf97\">column</span> from the right: multiply matching pairs, add them up.';return;}\n",
+    " const terms=A[i].map((a,t)=>`${a}&times;${B[t][j]}`), prods=A[i].map((a,t)=>a*B[t][j]), tot=prods.reduce((x,y)=>x+y,0); C[i][j]=tot;\n",
+    " document.getElementById('vmArea').innerHTML=tbl(A,custs,items,(r,c)=>r===i?'row':'')+'<span class=\"sym\">&times;</span>'+tbl(B,items,shops,(r,c)=>c===j?'col':'')+'<span class=\"sym\">=</span>'+tbl(C,custs,shops,(r,c,v)=>(r===i&&c===j)?'out':(v===null?'empty':''));\n",
+    " document.getElementById('vmWork').innerHTML=`${custs[i]} at ${shops[j]}:  ${terms.join(' + ')}  =  ${prods.join(' + ')}  =  <b>${tot}</b>`;\n",
+    " document.getElementById('vmCap').innerHTML=`Take <span style=\"color:#ffb347\"><b>${custs[i]}'s row</b></span> (what she/he buys) and <span style=\"color:#6fcf97\"><b>${shops[j]}'s column</b></span> (its prices). Walk along both at the same time: ${A[i].map((a,t)=>`${a} ${items[t]} &times; ${B[t][j]} each`).join(', ')}. Add up: <b>${tot}</b>. That is one cell of the answer.<br><br>Why it matters: a neural-network layer is exactly this. Left table = your tokens (one row per token, one column per feature). Right table = the layer's weights. The answer = the tokens' new features. GPUs exist to do this billions of times per second.`;}\n",
+    "window.vmStep=function(){if(k>=3){return;}k++;render();};\n",
+    "window.vmAuto=function(){if(timer)return;timer=setInterval(()=>{if(k>=3){clearInterval(timer);timer=null;return;}vmStep();},1400);};\n",
+    "window.vmReset=function(){clearInterval(timer);timer=null;k=-1;C[0]=[null,null];C[1]=[null,null];render();};\n",
+    "render();\n",
+    "})();\n",
+    "</script>\n",
+    "\"\"\"))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "07dc3007",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -100,7 +492,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "2054c860",
+   "id": "dae67dc6",
    "metadata": {},
    "source": [
     "A **gradient** tells us how much the loss would change if we nudged a weight.\n",
@@ -108,13 +500,58 @@
     "`requires_grad=True`, do some math, call `.backward()`, and the gradient appears in `.grad`.\n",
     "\n",
     "Training a model is only this, repeated millions of times: compute a loss, get the gradients,\n",
-    "move every weight a tiny step against its gradient."
+    "move every weight a tiny step against its gradient.\n",
+    "\n",
+    "Picture a hiker in thick fog trying to reach the bottom of a valley. They cannot see the\n",
+    "valley, but they can feel which way the ground slopes under their feet. That slope is the\n",
+    "gradient. Step downhill, feel again, step again. Drag the slider or press **Roll**."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e0342de2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Interactive: gradient descent as a ball rolling downhill.\n",
+    "from IPython.display import HTML, display\n",
+    "display(HTML(r\"\"\"<div id=\"vg\" style=\"font-family:-apple-system,Segoe UI,Helvetica,Arial,sans-serif;color:#e8e8e8;background:#1e1e1e;border-radius:12px;padding:16px;max-width:900px\">\n",
+    "<style>#vg button{background:#333;color:#eee;border:1px solid #777;border-radius:6px;padding:6px 12px;margin-right:6px;cursor:pointer}#vg .cap{font-size:13px;line-height:1.5;margin-top:8px;background:#262626;border-radius:8px;padding:10px}</style>\n",
+    "<div style=\"font-size:15px;font-weight:600;margin-bottom:6px\">Gradient descent: a ball rolling downhill. Drag the weight, or let it roll.</div>\n",
+    "<svg id=\"vgSvg\" width=\"560\" height=\"260\" style=\"background:#262626;border-radius:8px\"></svg>\n",
+    "<div style=\"margin-top:8px\"><label style=\"font-size:13px\">weight w <input id=\"vgW\" type=\"range\" min=\"-4\" max=\"8\" step=\"0.05\" value=\"7\" style=\"width:260px;vertical-align:middle\" oninput=\"vgDraw(+this.value)\"></label>\n",
+    "<button onclick=\"vgStep()\">Take one step</button><button onclick=\"vgAuto()\">Roll (auto)</button></div>\n",
+    "<div class=\"cap\" id=\"vgCap\"></div>\n",
+    "</div>\n",
+    "<script>\n",
+    "(function(){\n",
+    "const svg=document.getElementById('vgSvg'),W=560,H=260;const lossOf=w=>0.5*(w-2)*(w-2)+1, gradOf=w=>(w-2);\n",
+    "const X=w=>40+(w+4)/12*(W-80), Y=l=>H-30-(l/19)*(H-60);\n",
+    "let timer=null;\n",
+    "window.vgDraw=function(w){document.getElementById('vgW').value=w;let d='';for(let x=-4;x<=8;x+=0.1){d+=(d?'L':'M')+X(x).toFixed(1)+','+Y(lossOf(x)).toFixed(1);}\n",
+    " const l=lossOf(w),g=gradOf(w),lr=0.3,wn=w-lr*g;\n",
+    " svg.innerHTML=`<path d=\"${d}\" fill=\"none\" stroke=\"#5a8ab8\" stroke-width=\"3\"/>\n",
+    " <text x=\"${W/2}\" y=\"${H-6}\" fill=\"#aaa\" font-size=\"12\" text-anchor=\"middle\">weight value w</text>\n",
+    " <text x=\"14\" y=\"${H/2}\" fill=\"#aaa\" font-size=\"12\" transform=\"rotate(-90 14,${H/2})\" text-anchor=\"middle\">loss (how wrong)</text>\n",
+    " <line x1=\"${X(w)}\" y1=\"${Y(l)}\" x2=\"${X(wn)}\" y2=\"${Y(l)}\" stroke=\"#ffb347\" stroke-width=\"3\" marker-end=\"url(#ar)\"/>\n",
+    " <defs><marker id=\"ar\" markerWidth=\"8\" markerHeight=\"8\" refX=\"6\" refY=\"4\" orient=\"auto\"><path d=\"M0,0 L8,4 L0,8 z\" fill=\"#ffb347\"/></marker></defs>\n",
+    " <circle cx=\"${X(w)}\" cy=\"${Y(l)}\" r=\"9\" fill=\"#ffb347\" stroke=\"#fff\" stroke-width=\"2\"/>\n",
+    " <text x=\"${X(2)}\" y=\"${Y(1)+22}\" fill=\"#6fcf97\" font-size=\"12\" text-anchor=\"middle\">best weight (lowest loss)</text>`;\n",
+    " document.getElementById('vgCap').innerHTML=`w = <b>${w.toFixed(2)}</b> &nbsp; loss = <b>${l.toFixed(2)}</b> &nbsp; gradient = <b>${g.toFixed(2)}</b><br>The gradient is the slope under the ball. Positive slope means \"loss goes up if w grows\", so we move w <i>down</i> by learning rate &times; gradient: ${w.toFixed(2)} &minus; 0.3 &times; ${g.toFixed(2)} = <b>${wn.toFixed(2)}</b>.<br>Real-world: a hiker in fog feeling which way the ground slopes and stepping downhill. Training does this for all 6 million weights at once, thousands of times.`;\n",
+    " return wn;};\n",
+    "window.vgStep=function(){const w=+document.getElementById('vgW').value;const wn=vgDraw(w);vgDraw(Math.abs(wn-2)<0.02?2:wn);};\n",
+    "window.vgAuto=function(){if(timer){clearInterval(timer);timer=null;return;}timer=setInterval(()=>{const w=+document.getElementById('vgW').value;if(Math.abs(w-2)<0.03){clearInterval(timer);timer=null;return;}vgStep();},350);};\n",
+    "vgDraw(7);\n",
+    "})();\n",
+    "</script>\n",
+    "\"\"\"))"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "0d5918ad",
+   "id": "ab57ff9e",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -132,7 +569,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "4c3d59ca",
+   "id": "4e717021",
    "metadata": {},
    "source": [
     "## 2. Data: TinyStories\n",
@@ -147,7 +584,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "6af6bd64",
+   "id": "67dfd24b",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -165,7 +602,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "ea6277bf",
+   "id": "3758528e",
    "metadata": {},
    "source": [
     "## 3. The simplest language model: a bigram\n",
@@ -176,13 +613,47 @@
     "\"after character A, how likely is character B\". That is a **bigram model**.\n",
     "It is just one `nn.Embedding` table of size *vocab x vocab*.\n",
     "\n",
-    "First we need a vocabulary. For now, one token = one character."
+    "First we need a vocabulary. For now, one token = one character.\n",
+    "\n",
+    "Before the code, play with the idea. The widget below counts, from a few sample sentences,\n",
+    "which letter tends to follow which. Click a letter. That table *is* the whole model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "69fb3e89",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Interactive: the bigram table, counted live from a few sentences.\n",
+    "from IPython.display import HTML, display\n",
+    "display(HTML(r\"\"\"<div id=\"vb\" style=\"font-family:-apple-system,Segoe UI,Helvetica,Arial,sans-serif;color:#e8e8e8;background:#1e1e1e;border-radius:12px;padding:16px;max-width:900px\">\n",
+    "<style>#vb .k{display:inline-block;width:26px;height:26px;line-height:26px;text-align:center;margin:2px;border-radius:5px;background:#2f4f6f;border:1px solid #5a8ab8;cursor:pointer;font-size:13px}#vb .k.on{background:#ffb347;color:#000}#vb .bar{height:14px;background:#6fcf97;border-radius:3px}#vb .cap{font-size:13px;line-height:1.5;margin-top:8px;background:#262626;border-radius:8px;padding:10px}</style>\n",
+    "<div style=\"font-size:15px;font-weight:600;margin-bottom:6px\">Bigram model: click a letter to see what it thinks comes next.</div>\n",
+    "<div id=\"vbKeys\"></div>\n",
+    "<div id=\"vbOut\" style=\"margin-top:10px;min-height:120px\"></div>\n",
+    "<div class=\"cap\">A bigram model is nothing but this table: for every character, how often each other character followed it in the training text. It is counted here live from a few sample sentences. It cannot remember anything further back than one letter, which is why its stories are gibberish. The GPT fixes that with attention.</div>\n",
+    "</div>\n",
+    "<script>\n",
+    "(function(){\n",
+    "const corpus=\"once upon a time there was a little fox. the fox liked to play in the park with her friend tom. one day they found a shiny red ball. they were very happy and played all day. the end. lily had a big red balloon. she gave it to her mom. her mom smiled and said thank you.\";\n",
+    "const keys=\"abcdefghijklmnopqrstuvwxyz. \".split('');const counts={};\n",
+    "for(let i=0;i<corpus.length-1;i++){const a=corpus[i],b=corpus[i+1];counts[a]=counts[a]||{};counts[a][b]=(counts[a][b]||0)+1;}\n",
+    "const kd=document.getElementById('vbKeys');keys.forEach(k=>{const s=document.createElement('span');s.className='k';s.textContent=k===' '?'␣':k;s.onclick=()=>vbShow(k);kd.appendChild(s);});\n",
+    "window.vbShow=function(k){[...kd.children].forEach((s,i)=>s.classList.toggle('on',keys[i]===k));const c=counts[k]||{};const tot=Object.values(c).reduce((a,b)=>a+b,0);\n",
+    " const rows=Object.entries(c).sort((a,b)=>b[1]-a[1]).slice(0,6);\n",
+    " document.getElementById('vbOut').innerHTML=tot?`<div style=\"font-size:13px;margin-bottom:6px\">After <b>\"${k===' '?'space':k}\"</b> (seen ${tot} times), the next character was:</div>`+rows.map(([ch,n])=>`<div style=\"display:flex;align-items:center;gap:8px;font-size:13px\"><span style=\"width:40px;text-align:right\">${ch===' '?'␣':ch}</span><div class=\"bar\" style=\"width:${Math.round(n/tot*300)}px\"></div><span style=\"color:#aaa\">${Math.round(n/tot*100)}%</span></div>`).join(''):`<div style=\"font-size:13px;color:#aaa\">\"${k}\" never appears in the sample text, so the model has no idea.</div>`;};\n",
+    "vbShow('t');\n",
+    "})();\n",
+    "</script>\n",
+    "\"\"\"))"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9a29bbc9",
+   "id": "b6b441fb",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -202,7 +673,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9990f156",
+   "id": "3301c521",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -251,7 +722,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "658f7e72",
+   "id": "073e9f1b",
    "metadata": {},
    "source": [
     "It learned *something* (the loss dropped, and it produces word-ish shapes with spaces),\n",
@@ -261,7 +732,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "763829c4",
+   "id": "07cea514",
    "metadata": {},
    "source": [
     "## 4. A real tokenizer: byte-pair encoding (BPE)\n",
@@ -271,13 +742,45 @@
     "repeatedly merges the most common pairs, so frequent words become single tokens and rare\n",
     "words split into pieces. Every big model (GPT, Llama, Qwen) uses some form of this.\n",
     "\n",
-    "We train our own tokenizer on the stories, with a small vocabulary of 4,096 tokens."
+    "We train our own tokenizer on the stories, with a small vocabulary of 4,096 tokens.\n",
+    "\n",
+    "Type your own sentence below and watch it get chopped into pieces and turned into numbers."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2d9f5023",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Interactive: type a sentence, see tokens and ids (illustrative splitter).\n",
+    "from IPython.display import HTML, display\n",
+    "display(HTML(r\"\"\"<div id=\"vk\" style=\"font-family:-apple-system,Segoe UI,Helvetica,Arial,sans-serif;color:#e8e8e8;background:#1e1e1e;border-radius:12px;padding:16px;max-width:900px\">\n",
+    "<style>#vk .chip{display:inline-block;padding:4px 8px;margin:3px;border-radius:6px;background:#2f4f6f;border:1px solid #5a8ab8;font-size:13px}#vk .id{display:inline-block;padding:4px 8px;margin:3px;border-radius:6px;background:#2f6f4f;border:1px solid #6fcf97;font-size:12px;font-family:monospace}#vk input{width:95%;padding:8px;border-radius:6px;border:1px solid #777;background:#262626;color:#eee;font-size:14px}#vk .cap{font-size:13px;line-height:1.5;margin-top:8px;background:#262626;border-radius:8px;padding:10px}</style>\n",
+    "<div style=\"font-size:15px;font-weight:600;margin-bottom:6px\">Tokenizer: type anything and watch it get chopped into tokens, then numbers.</div>\n",
+    "<input id=\"vkIn\" value=\"Once upon a time, a little fox found a shiny red ball.\" oninput=\"vkRun()\">\n",
+    "<div style=\"font-size:12px;color:#aaa;margin:8px 0 2px\">tokens (&#9251; marks a leading space)</div><div id=\"vkTok\"></div>\n",
+    "<div style=\"font-size:12px;color:#aaa;margin:8px 0 2px\">token ids (what the model actually receives)</div><div id=\"vkIds\"></div>\n",
+    "<div class=\"cap\">A tokenizer splits text into pieces it has seen often. Common words stay whole (\"little\"), rare or long words break into parts (\"un\" + \"believ\" + \"able\"), and punctuation is its own token. Each piece has a fixed number in the dictionary. <i>This demo mimics the idea in JavaScript; the real one is trained from data in Part 4.</i></div>\n",
+    "</div>\n",
+    "<script>\n",
+    "(function(){\n",
+    "function hash(s){let h=7;for(const ch of s)h=(h*31+ch.charCodeAt(0))%4096;return h;}\n",
+    "function split(t){const out=[];const re=/\\s?[A-Za-z]+|\\s?\\d+|\\s?[^\\sA-Za-z\\d]/g;let m;while((m=re.exec(t))){let p=m[0];if(p.replace(/^\\s/,'').length>7){const lead=p.startsWith(' ')?' ':'';const w=p.trim();out.push(lead+w.slice(0,Math.ceil(w.length/2)));out.push(w.slice(Math.ceil(w.length/2)));}else out.push(p);}return out;}\n",
+    "window.vkRun=function(){const t=document.getElementById('vkIn').value;const toks=split(t);\n",
+    " document.getElementById('vkTok').innerHTML=toks.map(p=>`<span class=\"chip\">${p.replace(/^ /,'␣')}</span>`).join('');\n",
+    " document.getElementById('vkIds').innerHTML=toks.map(p=>`<span class=\"id\">${hash(p)}</span>`).join('')+`<span style=\"font-size:12px;color:#aaa;margin-left:8px\">${t.length} characters &#8594; ${toks.length} tokens</span>`;};\n",
+    "vkRun();\n",
+    "})();\n",
+    "</script>\n",
+    "\"\"\"))"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8bcb2397",
+   "id": "e1d5392c",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -303,7 +806,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "23cb929f",
+   "id": "b38d367d",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -317,7 +820,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "d2e9b068",
+   "id": "3a3b651d",
    "metadata": {},
    "source": [
     "## 5. Attention, by hand\n",
@@ -338,7 +841,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "c68c6f5c",
+   "id": "459979b9",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -362,7 +865,51 @@
   },
   {
    "cell_type": "markdown",
-   "id": "0da003d3",
+   "id": "57ea4612",
+   "metadata": {},
+   "source": [
+    "### The KV cache, in twelve lines\n",
+    "\n",
+    "When the model writes, a new token arrives at every step. The keys and values of the tokens\n",
+    "already written **never change**, so we keep them in a store (the cache) and compute cards\n",
+    "only for the newcomer. Below: four tokens are already cached, a fifth arrives, and we check\n",
+    "that the cached shortcut gives exactly the same answer as recomputing all five from scratch."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "81bf3c28",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "Wq, Wk, Wv = nn.Linear(C, head, bias=False), nn.Linear(C, head, bias=False), nn.Linear(C, head, bias=False)\n",
+    "\n",
+    "x_old = torch.randn(1, 4, C)                  # 4 tokens already processed earlier\n",
+    "K_cache, V_cache = Wk(x_old), Wv(x_old)       # their key/value cards, stored once: \"the drawer\"\n",
+    "\n",
+    "x_new = torch.randn(1, 1, C)                  # the 5th token arrives\n",
+    "q_new, k_new, v_new = Wq(x_new), Wk(x_new), Wv(x_new)   # only ITS cards are computed\n",
+    "K_cache = torch.cat([K_cache, k_new], dim=1)  # the drawer grows by one\n",
+    "V_cache = torch.cat([V_cache, v_new], dim=1)\n",
+    "\n",
+    "att = F.softmax(q_new @ K_cache.transpose(-2, -1) / math.sqrt(head), dim=-1)  # 1 query vs 5 stored keys\n",
+    "out_cached = att @ V_cache\n",
+    "\n",
+    "# Check against the slow way: recompute attention for all 5 tokens from scratch, take the last row.\n",
+    "x_all = torch.cat([x_old, x_new], dim=1)\n",
+    "scores = Wq(x_all) @ Wk(x_all).transpose(-2, -1) / math.sqrt(head)\n",
+    "scores = scores.masked_fill(torch.tril(torch.ones(5, 5)) == 0, float(\"-inf\"))\n",
+    "out_full = (F.softmax(scores, dim=-1) @ Wv(x_all))[:, -1:]\n",
+    "\n",
+    "print(\"tokens stored in cache:\", K_cache.shape[1])\n",
+    "print(\"cached shortcut == full recompute:\", torch.allclose(out_cached, out_full, atol=1e-6))\n",
+    "print(\"work done this step: 1 token's cards instead of 5. For a 1,000-word chat, that is 1 instead of 1,000.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "190d030b",
    "metadata": {},
    "source": [
     "## 6. A tiny GPT\n",
@@ -377,13 +924,59 @@
     "   - layer norm and residual connections, which keep training stable.\n",
     "4. **Final layer norm + output head**: turn each vector into scores over the vocabulary.\n",
     "\n",
-    "We also write a `generate` method: predict the next token, append it, repeat."
+    "We also write a `generate` method: predict the next token, append it, repeat.\n",
+    "\n",
+    "Click any layer in the stack below to read what it does and why it is there."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b00876c0",
+   "id": "0c74a9de",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Interactive: the model's layers, bottom to top. Click a layer.\n",
+    "from IPython.display import HTML, display\n",
+    "display(HTML(r\"\"\"<div id=\"vs\" style=\"font-family:-apple-system,Segoe UI,Helvetica,Arial,sans-serif;color:#e8e8e8;background:#1e1e1e;border-radius:12px;padding:16px;max-width:900px\">\n",
+    "<style>#vs .box{display:block;width:300px;margin:4px auto;padding:8px;text-align:center;border-radius:8px;border:2px solid #5a8ab8;background:#2f4f6f;cursor:pointer;font-size:13px;transition:all .2s}#vs .box:hover,#vs .box.on{border-color:#ffb347;background:#3a2e1a}#vs .blk{border:2px dashed #6fcf97;border-radius:10px;padding:8px;width:340px;margin:6px auto;position:relative}#vs .tag{position:absolute;right:-70px;top:40%;color:#6fcf97;font-size:13px;font-weight:600}#vs .ar{text-align:center;color:#888;font-size:16px;line-height:14px}#vs .cap{font-size:13px;line-height:1.5;background:#262626;border-radius:8px;padding:10px;min-height:70px}</style>\n",
+    "<div style=\"font-size:15px;font-weight:600;margin-bottom:6px\">The tiny GPT, bottom to top. Click any layer.</div>\n",
+    "<div style=\"display:flex;gap:16px;flex-wrap:wrap\"><div style=\"flex:0 0 420px\">\n",
+    "<span class=\"box\" onclick=\"vsX(this,'out')\">next-word scores (4,096 numbers)</span><div class=\"ar\">&#8593;</div>\n",
+    "<span class=\"box\" onclick=\"vsX(this,'lnf')\">final LayerNorm + output head</span><div class=\"ar\">&#8593;</div>\n",
+    "<div class=\"blk\"><span class=\"tag\">&times; 6 blocks</span>\n",
+    "<span class=\"box\" onclick=\"vsX(this,'res2')\">+ add back (residual)</span><div class=\"ar\">&#8593;</div>\n",
+    "<span class=\"box\" onclick=\"vsX(this,'ff')\">feed-forward network</span><div class=\"ar\">&#8593;</div>\n",
+    "<span class=\"box\" onclick=\"vsX(this,'ln')\">LayerNorm</span><div class=\"ar\">&#8593;</div>\n",
+    "<span class=\"box\" onclick=\"vsX(this,'res1')\">+ add back (residual)</span><div class=\"ar\">&#8593;</div>\n",
+    "<span class=\"box\" onclick=\"vsX(this,'attn')\">multi-head attention (4 heads)</span><div class=\"ar\">&#8593;</div>\n",
+    "<span class=\"box\" onclick=\"vsX(this,'ln')\">LayerNorm</span></div>\n",
+    "<div class=\"ar\">&#8593;</div>\n",
+    "<span class=\"box\" onclick=\"vsX(this,'emb')\">token embedding + position embedding</span><div class=\"ar\">&#8593;</div>\n",
+    "<span class=\"box\" onclick=\"vsX(this,'tok')\">token ids (up to 256 of them)</span>\n",
+    "</div><div style=\"flex:1;min-width:260px\"><div class=\"cap\" id=\"vsCap\">Click a layer to read what it does and why it is there.</div></div></div>\n",
+    "</div>\n",
+    "<script>\n",
+    "(function(){const t={\n",
+    "tok:\"<b>Token ids.</b> The input: a list of up to 256 integers from the tokenizer. Real-world: the page numbers of words in a dictionary.\",\n",
+    "emb:\"<b>Embeddings.</b> Each id is swapped for a 256-number vector (its meaning), and a second vector for its position (1st, 2nd, 3rd...) is added, so the model knows word order. Real-world: each word gets a seat number and a name badge.\",\n",
+    "ln:\"<b>LayerNorm.</b> Rescales the numbers so they stay in a sensible range. It does no 'thinking'; it keeps training stable. Real-world: turning the volume to a steady level before each stage.\",\n",
+    "attn:\"<b>Multi-head attention.</b> Each token looks back at earlier tokens and blends in what matters. 4 heads run in parallel, each looking for something different (grammar, names, position...). This is the only place tokens talk to each other.\",\n",
+    "res1:\"<b>Residual (add back).</b> The attention output is <i>added</i> to the original, not replacing it. The token keeps its own meaning plus what it learned. It also gives gradients a direct path, which is what lets deep networks train at all.\",\n",
+    "ff:\"<b>Feed-forward network.</b> A small two-layer network applied to every token on its own, 256 &#8594; 1024 &#8594; 256. This is where facts and patterns are stored. Real-world: after listening to everyone (attention), each person goes away and thinks privately.\",\n",
+    "res2:\"<b>Residual again.</b> Same trick: add the feed-forward result back onto the stream.\",\n",
+    "lnf:\"<b>Final LayerNorm + output head.</b> One last rescale, then each token's 256-number vector is multiplied against the embedding table (shared weights, 'weight tying') to produce a score for all 4,096 dictionary entries.\",\n",
+    "out:\"<b>Scores.</b> 4,096 numbers per position. Softmax turns them into probabilities, and the next word is sampled from them. During training these scores are compared with the real next word to compute the loss.\"};\n",
+    "window.vsX=function(el,k){document.querySelectorAll('#vs .box').forEach(b=>b.classList.remove('on'));el.classList.add('on');document.getElementById('vsCap').innerHTML=t[k];};\n",
+    "})();\n",
+    "</script>\n",
+    "\"\"\"))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cc29c0f8",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -494,7 +1087,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "9c446a37",
+   "id": "515072d6",
    "metadata": {},
    "source": [
     "## 7. Training\n",
@@ -508,13 +1101,57 @@
     "- a **learning-rate schedule** (warm up, then cosine decay) so the start is gentle and the end is precise,\n",
     "- **mixed precision** on NVIDIA GPUs, which is roughly 2x faster.\n",
     "\n",
-    "About 1,500 steps takes a few minutes on a T4 or an M-series Mac. Watch the loss fall."
+    "About 1,500 steps takes a few minutes on a T4 or an M-series Mac. Watch the loss fall.\n",
+    "\n",
+    "Step through the loop below first to see the five stages and a loss curve drawing itself."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "69241e04",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Interactive: the training loop, stage by stage, with a loss curve.\n",
+    "from IPython.display import HTML, display\n",
+    "display(HTML(r\"\"\"<div id=\"vr\" style=\"font-family:-apple-system,Segoe UI,Helvetica,Arial,sans-serif;color:#e8e8e8;background:#1e1e1e;border-radius:12px;padding:16px;max-width:900px\">\n",
+    "<style>#vr button{background:#333;color:#eee;border:1px solid #777;border-radius:6px;padding:6px 12px;margin-right:6px;cursor:pointer}#vr .st{display:inline-block;padding:8px 10px;margin:3px;border-radius:8px;border:2px solid #555;background:#2a2a2a;font-size:13px;min-width:120px;text-align:center;transition:all .2s}#vr .st.on{border-color:#ffb347;background:#3a2e1a}#vr .cap{font-size:13px;line-height:1.5;margin-top:8px;background:#262626;border-radius:8px;padding:10px;min-height:48px}</style>\n",
+    "<div style=\"font-size:15px;font-weight:600;margin-bottom:6px\">The training loop. Step through it and watch the loss fall.</div>\n",
+    "<div id=\"vrStages\">\n",
+    "<span class=\"st\" data-i=\"0\">1. grab a batch<br><small>32 random story windows</small></span><span style=\"color:#888\">&#8594;</span>\n",
+    "<span class=\"st\" data-i=\"1\">2. forward<br><small>model guesses next tokens</small></span><span style=\"color:#888\">&#8594;</span>\n",
+    "<span class=\"st\" data-i=\"2\">3. loss<br><small>how wrong, on average</small></span><span style=\"color:#888\">&#8594;</span>\n",
+    "<span class=\"st\" data-i=\"3\">4. backward<br><small>autograd: gradient per weight</small></span><span style=\"color:#888\">&#8594;</span>\n",
+    "<span class=\"st\" data-i=\"4\">5. optimizer step<br><small>nudge all 6M weights</small></span><span style=\"color:#888\">&#8617;</span>\n",
+    "</div>\n",
+    "<div style=\"display:flex;gap:16px;flex-wrap:wrap;margin-top:10px\"><svg id=\"vrSvg\" width=\"460\" height=\"200\" style=\"background:#262626;border-radius:8px\"></svg>\n",
+    "<div style=\"flex:1;min-width:240px\"><div class=\"cap\" id=\"vrCap\">Press Step, or Run 20 steps.</div>\n",
+    "<div style=\"margin-top:8px\"><button onclick=\"vrStep()\">Step</button><button onclick=\"vrRun()\">Run 20 steps</button><button onclick=\"vrReset()\">Reset</button></div>\n",
+    "<div style=\"font-size:12px;color:#aaa;margin-top:6px\">step <b id=\"vrN\">0</b> &nbsp; loss <b id=\"vrL\">-</b></div></div></div>\n",
+    "</div>\n",
+    "<script>\n",
+    "(function(){\n",
+    "const txt=[\"<b>Grab a batch.</b> 32 random 256-token windows from the stories, plus the same windows shifted one token right: the correct answers.\",\"<b>Forward.</b> Run the batch through the model. Out come scores for every position's next token.\",\"<b>Loss.</b> Cross-entropy compares the scores with the true next tokens. One number: average surprise. Lower is better.\",\"<b>Backward.</b> Autograd walks the computation in reverse and gives every one of the 6 million weights its gradient: 'nudge me this way to lower the loss'.\",\"<b>Optimizer step.</b> AdamW moves every weight a tiny bit against its gradient. Then the loop repeats, 1,500 times.\"];\n",
+    "let stage=-1,n=0,pts=[];const svg=document.getElementById('vrSvg');let timer=null;\n",
+    "function loss(k){return 1.9+6.5*Math.exp(-k/6)+ (Math.sin(k*1.7)*0.12);}\n",
+    "function draw(){let d='';pts.forEach((l,i)=>{d+=(i?'L':'M')+(30+i*(410/60)).toFixed(1)+','+(180-(l-1.5)/7.5*160).toFixed(1);});\n",
+    " svg.innerHTML=`<text x=\"240\" y=\"196\" fill=\"#aaa\" font-size=\"11\" text-anchor=\"middle\">training steps</text><text x=\"12\" y=\"100\" fill=\"#aaa\" font-size=\"11\" transform=\"rotate(-90 12,100)\" text-anchor=\"middle\">loss</text><path d=\"${d}\" fill=\"none\" stroke=\"#6fcf97\" stroke-width=\"2.5\"/>`+(pts.length?`<circle cx=\"${30+(pts.length-1)*(410/60)}\" cy=\"${180-(pts[pts.length-1]-1.5)/7.5*160}\" r=\"4\" fill=\"#ffb347\"/>`:'');\n",
+    " document.querySelectorAll('#vr .st').forEach(s=>s.classList.toggle('on',+s.dataset.i===stage));\n",
+    " document.getElementById('vrCap').innerHTML=stage<0?'Press Step, or Run 20 steps.':txt[stage];document.getElementById('vrN').textContent=n;document.getElementById('vrL').textContent=pts.length?pts[pts.length-1].toFixed(2):'-';}\n",
+    "window.vrStep=function(){stage=(stage+1)%5;if(stage===2){n++;if(pts.length<60)pts.push(loss(n));}draw();};\n",
+    "window.vrRun=function(){if(timer)return;let c=0;timer=setInterval(()=>{vrStep();if(++c>=100){clearInterval(timer);timer=null;}},120);};\n",
+    "window.vrReset=function(){clearInterval(timer);timer=null;stage=-1;n=0;pts=[];draw();};\n",
+    "draw();\n",
+    "})();\n",
+    "</script>\n",
+    "\"\"\"))"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "222e1c51",
+   "id": "0f607250",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -581,7 +1218,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "4dee51b4",
+   "id": "e6723c30",
    "metadata": {},
    "source": [
     "**How to read the numbers.** The loss is the average \"surprise\" per token. *Perplexity* is\n",
@@ -591,7 +1228,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "284708da",
+   "id": "953a841e",
    "metadata": {},
    "source": [
     "## 8. Generate stories: the use case\n",
@@ -603,7 +1240,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ce19e9d9",
+   "id": "2cc94e4e",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -620,7 +1257,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "a3a6cfad",
+   "id": "75329d6b",
    "metadata": {},
    "source": [
     "## 9. Save and publish on the Hugging Face Hub\n",
@@ -636,7 +1273,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e5fc0838",
+   "id": "f26e61dc",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -677,7 +1314,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7e1ae430",
+   "id": "08db0245",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -698,7 +1335,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "34f8cabf",
+   "id": "b91cda98",
    "metadata": {},
    "source": [
     "## Loading it back later\n",
@@ -709,7 +1346,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "c5e86850",
+   "id": "c8a96007",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -730,7 +1367,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "c0f44cd1",
+   "id": "f2036bcb",
    "metadata": {},
    "source": [
     "## What we built, in one paragraph\n",