{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Module 28: Knowledge Distillation"
   ],
   "id": "05f1089e-488f-4ac6-86b4-4074af93b651"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "A frontier model is expensive to serve. **Knowledge distillation** is\n",
    "the oldest and most widely-used way to shrink one: train a small\n",
    "**student** network to imitate a large **teacher** — not by copying its\n",
    "weights (that was m27), but by copying its *behavior*, the full\n",
    "probability distribution it puts over outputs.\n",
    "\n",
    "The trick is that a teacher’s soft probabilities say far more than a\n",
    "label does. Shown a photo of a car, a one-hot label says only `car`. The\n",
    "teacher’s distribution *also* says `truck` is far more plausible than\n",
    "`cat` — it has learned the similarity structure of the world. That extra\n",
    "signal, which Hinton called **dark knowledge**, is what a student learns\n",
    "from. A small model trained on soft targets can reach accuracy a small\n",
    "model trained on hard labels cannot.\n",
    "\n",
    "Why it matters for LLMs:\n",
    "\n",
    "- **Deployment.** Every small model you actually run — a phone-sized\n",
    "  assistant, a fast draft model (m16) — is likely a distilled one.\n",
    "- **Reasoning on a budget.** In 2025, DeepSeek showed that *distilling*\n",
    "  a giant reasoning model’s traces into a small one **beats** running\n",
    "  expensive RL (m13) on that small model directly.\n",
    "- **It is the same loss everywhere.** The soft-label idea reappears as\n",
    "  the tuned lens’s training objective (m21), as sequence-level KD for\n",
    "  translation, and as on-policy distillation for chat models.\n",
    "\n",
    "### What You’ll Learn\n",
    "\n",
    "After this module, you can:\n",
    "\n",
    "- Soften a distribution with **temperature** and read the **dark\n",
    "  knowledge** it exposes.\n",
    "- Write **Hinton’s distillation loss** from scratch and explain the\n",
    "  `$T^2$` factor.\n",
    "- Prove — numerically — that high-temperature distillation is **logit\n",
    "  matching**.\n",
    "- Distinguish **word-level** from **sequence-level** KD for a language\n",
    "  model, and see that sequence-level KD is just SFT on the teacher’s own\n",
    "  outputs.\n",
    "- Explain **on-policy distillation (GKD)** and why the *choice of\n",
    "  divergence* (forward vs reverse KL) trades mode-covering for\n",
    "  mode-seeking.\n",
    "- Retell the **R1 distillation** result and why it reshaped how small\n",
    "  reasoners are built.\n",
    "\n",
    "### Prerequisites\n",
    "\n",
    "This module requires familiarity with:\n",
    "\n",
    "- [Module 07: Training](../m07_training/lesson.qmd) — cross-entropy and\n",
    "  the training loop the student runs.\n",
    "- [Module 08: Generation](../m08_generation/lesson.qmd) — teachers\n",
    "  *generate* the sequences sequence-level KD trains on.\n",
    "- [Module 13: Reasoning](../m13_reasoning/lesson.qmd) — the RL baseline\n",
    "  that R1 distillation is measured against.\n",
    "\n",
    "## Intuition: Dark Knowledge\n",
    "\n",
    "Cross-entropy against a one-hot label only ever pushes up the\n",
    "probability of the correct class. It is indifferent to *how* the student\n",
    "distributes the leftover mass. But the teacher is not indifferent — it\n",
    "has opinions about the wrong answers, and those opinions encode real\n",
    "structure.\n",
    "\n",
    "Temperature is the dial that exposes them. Divide the logits by `$T$`\n",
    "before the softmax: at `$T=1$` you get the ordinary (peaky)\n",
    "distribution; as `$T$` grows the peak flattens and the plausible\n",
    "runner-ups rise into view. Drag the temperature below and watch the\n",
    "teacher’s faith in `truck` — not `cat` — emerge from behind the `car`\n",
    "spike."
   ],
   "id": "27e5a02d-194e-41f2-ab54-0f20a26177ef"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "learnllm-ojs-note": true
   },
   "source": [
    "> **Interactive figure**\n",
    ">\n",
    "> This visualization runs in the browser. Open the [HTML lesson](lesson.html) to explore it.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Try This**\n",
    ">\n",
    "> 1.  **Slide `$T$` from 1 to 12.** At `$T=1$` the `car` bar swallows\n",
    ">     everything. As `$T$` rises, `truck` climbs while `cat`, `plane`,\n",
    ">     `ship` stay low — the teacher is telling you a truck looks like a\n",
    ">     car and a cat does not.\n",
    "> 2.  **Watch the ratio** in the top-left. A one-hot label sets it to\n",
    ">     `0/0`. The soft target keeps `truck` several times larger than\n",
    ">     `cat` — that is the signal a student gets for free.\n",
    "\n",
    "## The Math: Temperature and the Distillation Loss\n",
    "\n",
    "Write the teacher’s logits `$v$` and the student’s `$z$`, both over a\n",
    "vocabulary of `$V$` classes. The **temperature-softened** distributions\n",
    "are\n",
    "\n",
    "$$p^{\\text{teacher}}_i = \\frac{\\exp(v_i / T)}{\\sum_j \\exp(v_j / T)}, \\qquad\n",
    "q^{\\text{student}}_i = \\frac{\\exp(z_i / T)}{\\sum_j \\exp(z_j / T)}.$$\n",
    "\n",
    "The **distillation loss** asks the student to reproduce the teacher’s\n",
    "soft distribution (via KL divergence), optionally blended with the\n",
    "ordinary hard-label cross-entropy `$H$`:\n",
    "\n",
    "$$\\mathcal{L} = \\underbrace{\\alpha \\, H(y,\\, q^{T=1})}_{\\text{hard label}}\n",
    "\\;+\\;\n",
    "\\underbrace{(1-\\alpha)\\, T^2 \\, \\mathrm{KL}\\!\\left(p^{\\text{teacher}} \\,\\|\\, q^{\\text{student}}\\right)}_{\\text{soft target}}.$$\n",
    "\n",
    "Two details earn their keep. The gradient of the soft term through the\n",
    "softmax scales like `$1/T^2$`, so the explicit `$T^2$` factor **rescales\n",
    "it back** to the same magnitude as the hard term — you can change `$T$`\n",
    "without retuning `$\\alpha$`. And `$\\mathrm{KL}$`, not plain\n",
    "cross-entropy, keeps the “match the whole distribution” reading exact.\n",
    "\n",
    "### Step by Step"
   ],
   "id": "45fe5fd1-0a7c-4e26-b11d-5208f2061d0b"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "learnllm-ojs-note": true
   },
   "source": [
    "> **Interactive figure**\n",
    ">\n",
    "> This visualization runs in the browser. Open the [HTML lesson](lesson.html) to explore it.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Code: Distillation From Scratch\n",
    "\n",
    "Everything above is a few lines of PyTorch. The reference implementation\n",
    "lives in `distillation.py`; here we build the core pieces inline."
   ],
   "id": "c8a09194-03a4-48b9-a91c-7a9f0b3d1a15"
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "T=1  (peaky): ['0.003', '0.004', '0.926', '0.056', '0.003', '0.007']\n",
      "T=4  (soft) : ['0.093', '0.102', '0.395', '0.196', '0.097', '0.116']\n",
      "truck/cat ratio at T=4: 2.1 x"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from distillation import softmax_T, soft_targets, distillation_loss, kl_divergence\n",
    "\n",
    "# A teacher confident it's a \"car\" (index 2), with a real hint toward \"truck\" (3).\n",
    "teacher_logits = torch.tensor([0.2, 0.6, 6.0, 3.2, 0.4, 1.1])\n",
    "\n",
    "hard = softmax_T(teacher_logits, temperature=1.0)\n",
    "soft = softmax_T(teacher_logits, temperature=4.0)\n",
    "print(\"T=1  (peaky):\", [f\"{v:.3f}\" for v in hard.tolist()])\n",
    "print(\"T=4  (soft) :\", [f\"{v:.3f}\" for v in soft.tolist()])\n",
    "print(\"truck/cat ratio at T=4:\", round((soft[3] / soft[0]).item(), 1), \"x\")"
   ],
   "id": "c16a6abb"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The loss itself is the hard/soft blend, with the `$T^2$` factor on the\n",
    "soft term:"
   ],
   "id": "5c8c608c-3f43-411e-b84c-e6c99ba34248"
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "distillation loss: 1.6313\n",
      "student receives gradients: True"
     ]
    }
   ],
   "source": [
    "# One batch of 4 examples over a 6-class vocabulary.\n",
    "torch.manual_seed(0)\n",
    "student_logits = torch.randn(4, 6, requires_grad=True)\n",
    "teacher_batch = torch.randn(4, 6)\n",
    "labels = torch.randint(0, 6, (4,))\n",
    "\n",
    "loss = distillation_loss(student_logits, teacher_batch, labels, temperature=4.0, alpha=0.5)\n",
    "loss.backward()\n",
    "print(f\"distillation loss: {loss.item():.4f}\")\n",
    "print(\"student receives gradients:\", student_logits.grad is not None)"
   ],
   "id": "eb3dba98"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### High Temperature *is* Logit Matching\n",
    "\n",
    "Hinton’s key result: at high `$T$` (with per-example zero-meaned logits)\n",
    "the gradient of the `$T^2$`-scaled soft loss w.r.t. a student logit\n",
    "approaches `$\\tfrac{1}{V}(z_i - v_i)$` — exactly the gradient of\n",
    "`$\\tfrac{1}{2}\\sum_i (z_i - v_i)^2$`. So distillation at high\n",
    "temperature is *regression onto the teacher’s logits*. We can check it\n",
    "numerically: the gap between the true autograd gradient and that\n",
    "approximation collapses as `$T$` grows."
   ],
   "id": "efaabbdd-c109-47b6-9d06-30669ae097ec"
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "T=   2.0   max|grad - logit-match|  =  3.73e-02\n",
      "T=  10.0   max|grad - logit-match|  =  8.26e-03\n",
      "T=  50.0   max|grad - logit-match|  =  1.67e-03\n",
      "T= 200.0   max|grad - logit-match|  =  4.18e-04"
     ]
    }
   ],
   "source": [
    "from distillation import verify_logit_matching\n",
    "\n",
    "torch.manual_seed(0)\n",
    "s = torch.randn(3, 8)\n",
    "t = torch.randn(3, 8)\n",
    "for T in (2.0, 10.0, 50.0, 200.0):\n",
    "    err = verify_logit_matching(s, t, temperature=T)\n",
    "    print(f\"T={T:6.1f}   max|grad - logit-match|  =  {err:.2e}\")"
   ],
   "id": "25263a68"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Key Insight**\n",
    ">\n",
    "> Temperature interpolates between two regimes. At `$T=1$` distillation\n",
    "> cares mostly about the top class — close to a hard label. At high\n",
    "> `$T$` it becomes pure logit regression, forcing the student to\n",
    "> reproduce the teacher’s *entire* score vector, dark knowledge and all.\n",
    "> The sweet spot (`$T \\in 2..10$` in most papers) keeps the runner-up\n",
    "> structure without drowning the correct answer.\n",
    "\n",
    "## Dark Knowledge, Measured\n",
    "\n",
    "Does the extra signal actually help a small student? It does exactly\n",
    "where you’d expect: when labels are **scarce**. Below, a wide teacher is\n",
    "trained on plenty of data; a narrow student then sees only 32 labeled\n",
    "points. The `hard` student gets just those one-hot labels; the `kd`\n",
    "student *also* sees the teacher’s soft targets on the same 32 inputs —\n",
    "and generalizes measurably better on held-out data."
   ],
   "id": "f6803836-42fa-4c55-99c0-e45a47d8bb6b"
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "teacher (trained on all data): 0.806\n",
      "student, hard labels only    : 0.725\n",
      "student, + soft targets (KD) : 0.819   (+0.094)"
     ]
    }
   ],
   "source": [
    "from distillation import train_student\n",
    "\n",
    "hard_run = train_student(\"hard\", seed=0)\n",
    "kd_run = train_student(\"kd\", seed=0)\n",
    "print(f\"teacher (trained on all data): {hard_run['teacher_acc']:.3f}\")\n",
    "print(f\"student, hard labels only    : {hard_run['final_acc']:.3f}\")\n",
    "print(f\"student, + soft targets (KD) : {kd_run['final_acc']:.3f}\"\n",
    "      f\"   (+{kd_run['final_acc'] - hard_run['final_acc']:.3f})\")"
   ],
   "id": "45efcc40"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "learnllm-ojs-note": true
   },
   "source": [
    "> **Interactive figure**\n",
    ">\n",
    "> This visualization runs in the browser. Open the [HTML lesson](lesson.html) to explore it.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Try This**\n",
    ">\n",
    "> 1.  **Read the gap.** The KD curve sits above the hard curve for\n",
    ">     essentially the whole run — same 32 labels, more knowledge per\n",
    ">     label.\n",
    "> 2.  In `distillation.py`, raise `n_labeled` toward the full set. The\n",
    ">     gap **shrinks**: dark knowledge matters most when labels are the\n",
    ">     bottleneck.\n",
    "\n",
    "## Distilling a Language Model: Word-level vs Sequence-level\n",
    "\n",
    "A language model outputs a next-token distribution at every position, so\n",
    "there are two honest ways to distill it.\n",
    "\n",
    "**Word-level KD** matches the teacher’s full distribution at each\n",
    "position — a dense KL over the whole vocabulary, every step. It is the\n",
    "direct generalization of what we did above:"
   ],
   "id": "08ba85e8-b9ac-47d3-9fe5-f5ea4ed1186a"
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "word-level KD (mean per-token KL): 0.5614"
     ]
    }
   ],
   "source": [
    "from distillation import word_level_kd_loss, sequence_level_kd_loss\n",
    "\n",
    "torch.manual_seed(0)\n",
    "# A short sequence: per-position logits over a 12-word vocabulary.\n",
    "teacher_seq = torch.randn(5, 12)\n",
    "student_seq = torch.randn(5, 12, requires_grad=True)\n",
    "\n",
    "wl = word_level_kd_loss(student_seq, teacher_seq, temperature=1.0)\n",
    "print(f\"word-level KD (mean per-token KL): {wl.item():.4f}\")"
   ],
   "id": "f6390234"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Sequence-level KD** (Kim & Rush, 2016) takes a different route: let\n",
    "the teacher *generate* an output sequence (e.g. by beam search), then\n",
    "train the student with ordinary next-token cross-entropy on those\n",
    "tokens. The student learns to reproduce the sequences the teacher\n",
    "actually produces — and crucially, **sequence-level KD is just SFT on\n",
    "the teacher’s outputs**:"
   ],
   "id": "5bb38450-db13-41e2-b77a-668ea1ef809b"
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "sequence-level KD : 2.8593\n",
      "plain SFT on teacher tokens: 2.8593\n",
      "identical: True"
     ]
    }
   ],
   "source": [
    "import torch.nn.functional as F\n",
    "\n",
    "# Pretend the teacher generated these tokens (here: its own argmax path).\n",
    "teacher_tokens = teacher_seq.argmax(dim=-1)\n",
    "\n",
    "seq_kd = sequence_level_kd_loss(student_seq, teacher_tokens)\n",
    "plain_sft = F.cross_entropy(student_seq, teacher_tokens)\n",
    "print(f\"sequence-level KD : {seq_kd.item():.4f}\")\n",
    "print(f\"plain SFT on teacher tokens: {plain_sft.item():.4f}\")\n",
    "print(\"identical:\", torch.allclose(seq_kd, plain_sft))"
   ],
   "id": "3d903ada"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Key Insight**\n",
    ">\n",
    "> Word-level KD gives the richest signal (the whole distribution, every\n",
    "> step) but requires the teacher’s logits at train time. Sequence-level\n",
    "> KD needs only the teacher’s *text* — cheaper, and it matches the\n",
    "> teacher’s *sequence* behavior rather than its per-token marginals. It\n",
    "> is the form behind most “trained on outputs of a bigger model”\n",
    "> recipes, R1 distillation included.\n",
    "\n",
    "## On-Policy Distillation (GKD)\n",
    "\n",
    "Both losses above train on a **fixed** set of sequences. But at\n",
    "inference the student decodes its *own* tokens, drifting into states the\n",
    "training set never covered — the **exposure bias** you met in m08.\n",
    "On-policy distillation (Agarwal et al.’s **GKD**) fixes this by training\n",
    "on sequences the *student itself* samples, scored by the teacher. The\n",
    "student is graded exactly where it actually goes.\n",
    "\n",
    "Sampling on-policy opens a second choice: *which divergence* to\n",
    "minimize. The forward and reverse KL pull in opposite directions."
   ],
   "id": "1c677c35-991b-4104-9dd9-ba3c3351ee98"
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "forward KL(p‖q) [mode-covering]: 0.1360\n",
      "reverse KL(q‖p) [mode-seeking] : 0.1816\n",
      "JS divergence   [symmetric]    : 0.0380"
     ]
    }
   ],
   "source": [
    "from distillation import forward_kl, reverse_kl, js_divergence\n",
    "\n",
    "p = torch.tensor([0.45, 0.10, 0.45])   # a bimodal teacher\n",
    "q = torch.tensor([0.34, 0.32, 0.34])   # a spread-out student\n",
    "print(f\"forward KL(p‖q) [mode-covering]: {forward_kl(p, q).item():.4f}\")\n",
    "print(f\"reverse KL(q‖p) [mode-seeking] : {reverse_kl(p, q).item():.4f}\")\n",
    "print(f\"JS divergence   [symmetric]    : {js_divergence(p, q).item():.4f}\")"
   ],
   "id": "668f26c7"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- **Forward KL** `$\\mathrm{KL}(p\\|q)$` is **mode-covering**: the student\n",
    "  is punished wherever the teacher has mass but the student does not, so\n",
    "  it spreads to cover *every* teacher mode (and smears probability\n",
    "  across the valleys between them).\n",
    "- **Reverse KL** `$\\mathrm{KL}(q\\|p)$` is **mode-seeking**: the student\n",
    "  is punished for mass where the teacher has none, so it commits to a\n",
    "  *single* mode. This yields decisive students — often what you want for\n",
    "  a chat model.\n",
    "\n",
    "Fit a single-peaked student to a two-peaked teacher under each\n",
    "divergence and the difference is stark:"
   ],
   "id": "db91f946-6e83-40a2-a57a-9d4b47c1fc12"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "learnllm-ojs-note": true
   },
   "source": [
    "> **Interactive figure**\n",
    ">\n",
    "> This visualization runs in the browser. Open the [HTML lesson](lesson.html) to explore it.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Try This**\n",
    ">\n",
    "> 1.  **Toggle the divergence.** Reverse KL parks a narrow student\n",
    ">     squarely on one peak; forward KL spreads a wide student across\n",
    ">     both, piling mass into the empty valley between them.\n",
    "> 2.  This is why the “right” objective depends on the goal: cover the\n",
    ">     teacher faithfully (forward / word-level KL) or imitate its\n",
    ">     decisiveness (reverse / on-policy).\n",
    "\n",
    "## Reasoning Distillation: The R1 Result\n",
    "\n",
    "In January 2025, DeepSeek trained **R1**, a strong reasoning model,\n",
    "largely with reinforcement learning (the GRPO lineage of m13). Then they\n",
    "did something that reset expectations for small models: they took R1’s\n",
    "~800K curated reasoning traces and ran **plain supervised fine-tuning**\n",
    "(sequence-level KD) on small, open base models — Qwen2.5 and Llama —\n",
    "with *no RL stage at all*.\n",
    "\n",
    "The distilled students beat the RL baseline they were compared against:\n",
    "\n",
    "| Model                    | How it was built | AIME 2024 (pass@1) |\n",
    "|--------------------------|------------------|--------------------|\n",
    "| QwQ-32B-Preview          | large-scale RL   | 50.0%              |\n",
    "| **R1-Distill-Qwen-1.5B** | SFT on R1 traces | 28.9%              |\n",
    "| **R1-Distill-Qwen-7B**   | SFT on R1 traces | 55.5%              |\n",
    "| **R1-Distill-Qwen-32B**  | SFT on R1 traces | **72.6%**          |\n",
    "\n",
    "The paper’s own conclusion is blunt: *“distilling more powerful models\n",
    "into smaller ones yields excellent results, whereas smaller models\n",
    "relying on large-scale RL … require enormous computational power and may\n",
    "not even achieve the performance of distillation.”* A 1.5B distilled\n",
    "model doing meaningful AIME math, and a 32B distilled model beating a\n",
    "32B RL model, is dark knowledge at the frontier — the teacher’s\n",
    "reasoning *traces* are a far richer signal than a scalar reward.\n",
    "\n",
    "> **Distillation inherits the teacher’s ceiling — and its flaws**\n",
    ">\n",
    "> A student trained purely on a teacher can rarely *exceed* it, and it\n",
    "> copies the teacher’s biases and mistakes wholesale. Distillation moves\n",
    "> capability *down* in size, it does not create new capability. To push\n",
    "> past the teacher you still need signal the teacher lacks — RL against\n",
    "> a verifier (m13), fresh data, or on-policy exploration.\n",
    "\n",
    "## Common Pitfalls\n",
    "\n",
    "| Pitfall | Why it bites | Fix |\n",
    "|-----------------------|------------------------------------|--------------|\n",
    "| Dropping the `$T^2$` factor | The soft gradient shrinks like `$1/T^2$`; without the factor the soft term vanishes at high `$T$` and `$\\alpha$` needs retuning per temperature. | Keep `$T^2$` on the soft term (as in `distillation_loss`). |\n",
    "| Mismatched temperatures | Softening the teacher but not the student (or vice-versa) makes the KL compare different distributions. | Use the **same** `$T$` for both; divide by `$T$` **before** the softmax. |\n",
    "| `$T=1$` “distillation” | At `$T=1$` the soft target is nearly one-hot — you’ve thrown away the dark knowledge you came for. | Use `$T \\in 2..10$`; verify the runner-ups are actually visible. |\n",
    "| Reading student logits as probabilities of correctness | The student learns the teacher’s *distribution*, calibration flaws included. | Evaluate the student’s own task metric, not its agreement with the teacher. |\n",
    "| Expecting the student to beat the teacher | Pure distillation is bounded by the teacher. | Add verifier-RL / fresh data if you need to surpass it. |\n",
    "\n",
    "## Exercises\n",
    "\n",
    "### Exercise 1: The temperature sweet spot\n",
    "\n",
    "Sweep the distillation temperature and watch how much dark knowledge\n",
    "survives. Report the runner-up (`truck`) probability at each `$T$`."
   ],
   "id": "af82003f-0a63-4d24-a14a-8bb23a1c6068"
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "T=1.0: car=0.926  truck=0.056\n",
      "T=2.0: car=0.660  truck=0.163\n",
      "T=4.0: car=0.395  truck=0.196\n",
      "T=8.0: car=0.268  truck=0.189"
     ]
    }
   ],
   "source": [
    "from distillation import softmax_T\n",
    "\n",
    "logits = torch.tensor([0.2, 0.6, 6.0, 3.2, 0.4, 1.1])  # cat dog car truck plane ship\n",
    "for T in (1.0, 2.0, 4.0, 8.0):\n",
    "    p = softmax_T(logits, T)\n",
    "    # Your turn: print T, the top class prob p[2], and the runner-up p[3].\n",
    "    print(f\"T={T}: car={p[2]:.3f}  truck={p[3]:.3f}\")"
   ],
   "id": "1145f643"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 2: Reverse-KL is mode-seeking\n",
    "\n",
    "Using `forward_kl` and `reverse_kl`, confirm that for a bimodal teacher\n",
    "a *broad* student scores better under forward KL while a *narrow,\n",
    "on-mode* student scores better under reverse KL."
   ],
   "id": "eed7e40e-9979-482b-9adb-8e79cb6d68db"
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "forward: broad 0.591 narrow 1.098\n",
      "reverse: broad 5.434 narrow 0.953"
     ]
    }
   ],
   "source": [
    "from distillation import forward_kl, reverse_kl\n",
    "\n",
    "teacher = torch.tensor([0.45, 0.05, 0.0, 0.05, 0.45])   # two modes at the ends\n",
    "broad = torch.tensor([0.20, 0.20, 0.20, 0.20, 0.20])    # covers everything\n",
    "narrow = torch.tensor([0.02, 0.06, 0.02, 0.10, 0.80])   # commits to one mode\n",
    "# Your turn: compare forward_kl(teacher, ·) and reverse_kl(teacher, ·) for both.\n",
    "print(\"forward: broad\", round(forward_kl(teacher, broad).item(), 3),\n",
    "      \"narrow\", round(forward_kl(teacher, narrow).item(), 3))\n",
    "print(\"reverse: broad\", round(reverse_kl(teacher, broad).item(), 3),\n",
    "      \"narrow\", round(reverse_kl(teacher, narrow).item(), 3))"
   ],
   "id": "bf455f2d"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 3: Distill your own tiny student\n",
    "\n",
    "Call `train_student` with a harder task (raise `noise`, lower\n",
    "`n_labeled`) and confirm the KD student’s advantage over the hard-label\n",
    "student grows as labels get scarcer."
   ],
   "id": "6d3598b4-c0c1-4c20-be8a-14796f539a60"
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "n_labeled=16: hard=0.569  kd=0.669  gap=+0.100\n",
      "n_labeled=32: hard=0.725  kd=0.819  gap=+0.094"
     ]
    }
   ],
   "source": [
    "from distillation import train_student\n",
    "\n",
    "# Your turn: try n_labeled in {16, 24, 48} and compare final_acc for \"hard\" vs \"kd\".\n",
    "for nl in (16, 32):\n",
    "    h = train_student(\"hard\", n_labeled=nl, seed=0)[\"final_acc\"]\n",
    "    k = train_student(\"kd\", n_labeled=nl, seed=0)[\"final_acc\"]\n",
    "    print(f\"n_labeled={nl}: hard={h:.3f}  kd={k:.3f}  gap={k - h:+.3f}\")"
   ],
   "id": "508eca44"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "Key takeaways:\n",
    "\n",
    "1.  **Distillation copies behavior, not weights.** A student learns from\n",
    "    the teacher’s full probability distribution — the “dark knowledge” a\n",
    "    one-hot label discards.\n",
    "2.  **Temperature is the dial.** `softmax(z/T)` softens the\n",
    "    distribution; higher `$T$` exposes the runner-up classes that carry\n",
    "    the similarity structure.\n",
    "3.  **The loss is a hard/soft blend with a `$T^2$` factor.**\n",
    "    `$\\mathcal{L} = \\alpha H + (1-\\alpha) T^2 \\mathrm{KL}(p\\|q)$`; the\n",
    "    `$T^2$` keeps the two gradients comparable across temperatures.\n",
    "4.  **High-`$T$` distillation is logit matching** — provable and\n",
    "    numerically verifiable: the student regresses onto the teacher’s raw\n",
    "    logits.\n",
    "5.  **Language models distill two ways.** Word-level KD matches the\n",
    "    per-token distribution; sequence-level KD is plain SFT on the\n",
    "    teacher’s generated tokens.\n",
    "6.  **On-policy distillation (GKD)** trains on the student’s own samples\n",
    "    to kill exposure bias, and the *choice of divergence* trades\n",
    "    mode-covering (forward KL) for mode-seeking (reverse KL).\n",
    "7.  **R1 reset the small-model playbook.** SFT on a strong reasoner’s\n",
    "    traces beat large-scale RL on the same small model — distillation\n",
    "    moves frontier capability down in size.\n",
    "\n",
    "## What’s Next\n",
    "\n",
    "You have now built the language model from tensors all the way to the\n",
    "current frontier — architecture, training, alignment, reasoning,\n",
    "efficiency, and the compression and knowledge-transfer techniques\n",
    "(quantization, merging, and now distillation) that put a frontier model\n",
    "on real hardware. Circle back to [Module 13:\n",
    "Reasoning](../m13_reasoning/lesson.qmd) to see the RL side of the\n",
    "distillation-vs-RL trade, or [Module 16: Speculative\n",
    "Decoding](../m16_speculative_decoding/lesson.qmd) where a small (often\n",
    "distilled) draft model makes a large one faster.\n",
    "\n",
    "### Going Deeper\n",
    "\n",
    "**Core Papers:**\n",
    "\n",
    "- [Distilling the Knowledge in a Neural Network (Hinton, Vinyals, Dean,\n",
    "  2015)](https://arxiv.org/abs/1503.02531) — soft targets, temperature,\n",
    "  and the high-`$T$` logit-matching result.\n",
    "- [Sequence-Level Knowledge Distillation (Kim & Rush,\n",
    "  2016)](https://arxiv.org/abs/1606.07947) — distilling a language model\n",
    "  by SFT on its generated sequences.\n",
    "- [On-Policy Distillation of Language Models: Learning from\n",
    "  Self-Generated Mistakes (Agarwal et al.,\n",
    "  2023)](https://arxiv.org/abs/2306.13649) — GKD: on-policy student\n",
    "  samples and generalized divergences.\n",
    "- [DeepSeek-R1 (DeepSeek-AI, 2025)](https://arxiv.org/abs/2501.12948) —\n",
    "  distilling reasoning traces into small models beats large-scale RL on\n",
    "  those models.\n",
    "\n",
    "**Practical Resources:**\n",
    "\n",
    "- [DistilBERT (Sanh et al., 2019)](https://arxiv.org/abs/1910.01108) — a\n",
    "  widely-used distilled encoder; 40% smaller, 97% of BERT’s performance."
   ],
   "id": "d3c91923-c108-4b5e-9b03-a20a53307ee5"
  }
 ],
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "path": "/usr/local/share/jupyter/kernels/python3"
  },
  "language_info": {
   "name": "python",
   "codemirror_mode": {
    "name": "ipython",
    "version": "3"
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.15"
  }
 }
}