{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Module 13: Reasoning & Test-Time Compute"
      ],
      "id": "ef2fb5f1-139b-4fe9-9c37-5301a4a92327"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [],
      "id": "6db69ce2-a9d6-40fc-9084-4bccec9f9352"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "9543269c-5bc5-4ff9-9b6b-0a2cfb5f5f46"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Introduction\n",
        "\n",
        "Every module so far spends a **fixed** amount of compute per answer: run\n",
        "the model once, decode a token at a time (m08), stop. The answer you get\n",
        "is the answer the model happened to produce on its single pass. But hard\n",
        "problems — multi-step arithmetic, logic, code — are exactly the ones\n",
        "where a single pass is most likely to slip.\n",
        "\n",
        "**Test-time compute** is the frontier’s newest lever: spend *more*\n",
        "compute at **inference** to get a *better* answer, with the weights\n",
        "frozen. No retraining, no bigger model — just let the model think\n",
        "longer, or think several times and reconcile. This is the idea behind\n",
        "OpenAI’s o1 and DeepSeek-R1, and this module builds its most important,\n",
        "most learnable form from scratch:\n",
        "\n",
        "- **Chain-of-thought (CoT):** let the model write intermediate steps\n",
        "  before the final answer, turning one hard leap into several easy ones.\n",
        "- **Self-consistency:** sample *many* independent chains, read off each\n",
        "  final answer, and take the **majority vote**. One idea, a large,\n",
        "  reliable gain.\n",
        "- **Best-of-N:** sample N candidates and let a **verifier** pick the\n",
        "  best.\n",
        "\n",
        "Why it matters for LLMs:\n",
        "\n",
        "- Self-consistency alone lifted GSM8K accuracy by **+17.9 points** over\n",
        "  a single chain-of-thought pass (Wang et al., 2022) — no new\n",
        "  parameters.\n",
        "- The reasoning models topping today’s benchmarks are, at their core,\n",
        "  models that were taught (with RL) to *use* test-time compute well.\n",
        "  Understanding the parallel, no-training version first makes the RL\n",
        "  version legible.\n",
        "\n",
        "### What You’ll Learn\n",
        "\n",
        "After this module, you can:\n",
        "\n",
        "- Explain the difference between train-time and **test-time** compute,\n",
        "  and why spending more of the latter raises accuracy.\n",
        "- Build **self-consistency** from scratch: extract answers, tally a\n",
        "  plurality vote, and aggregate many sampled chains.\n",
        "- Prove *why* voting works with **Condorcet’s jury theorem**, and see\n",
        "  the exact condition it needs (the correct answer is the **mode**,\n",
        "  which is weaker than “right more than half the time”).\n",
        "- Implement **best-of-N** with a verifier and see when it beats voting.\n",
        "- Read the **compute–accuracy curve** — the empirical signature of\n",
        "  test-time scaling — and place chain-of-thought and RL-for-reasoning\n",
        "  (o1 / R1) on it.\n",
        "- Build **GRPO** from scratch — the critic-free RL recipe behind\n",
        "  DeepSeek-R1 — and train a toy policy to answer correctly from a\n",
        "  **verifiable reward** alone, watching the group-relative advantage do\n",
        "  the work of a value network.\n",
        "\n",
        "### Prerequisites\n",
        "\n",
        "This module requires familiarity with:\n",
        "\n",
        "- [Module 08: Generation](../m08_generation/lesson.qmd) — temperature\n",
        "  sampling and the `generate` loop that produces the diverse chains we\n",
        "  vote over.\n",
        "- [Module 12: Alignment](../m12_alignment/lesson.qmd) — the aligned\n",
        "  model that actually produces answers to reason about.\n",
        "\n",
        "## Intuition: Think Longer, Answer Better\n",
        "\n",
        "Ask a person a hard arithmetic question and make them answer instantly,\n",
        "and they’ll often be wrong. Give them scratch paper — *chain-of-thought*\n",
        "— and they do better. Let them solve it **three different ways** and go\n",
        "with the answer they reached most often — *self-consistency* — and they\n",
        "do better still. None of this changes *who* is solving the problem; it\n",
        "changes *how much thinking* they spend.\n",
        "\n",
        "Language models are the same. A single greedy decode is the instant\n",
        "answer. Two cheap upgrades, neither touching the weights:\n",
        "\n",
        "1.  **Chain-of-thought** — prompt the model to emit reasoning steps\n",
        "    *before* the answer. Each token now conditions on written-out\n",
        "    intermediate work instead of leaping straight to the result.\n",
        "2.  **Self-consistency** — because sampling (m08) makes each run take a\n",
        "    *different* reasoning path, run it many times and **vote**. A\n",
        "    correct answer is a fixed target many paths land on; wrong answers\n",
        "    scatter.\n",
        "\n",
        "Step through it: one prompt, many sampled chains, one vote."
      ],
      "id": "47500875-4106-4b4f-9b6d-36c1d82c82dc"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "343d42da-9ffd-4362-b2bc-91f134c8bbcc"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "0429efcb-2810-42e3-8bb4-e0167772a352"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "a2e49f5d-c6bd-4c10-aa83-bd11df45b358"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "baed2c15-643c-4804-91a0-2445a7ddfa82"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Key Insight**\n",
        ">\n",
        "> Self-consistency needs the model to be *unreliable but not\n",
        "> adversarial*: as long as the correct answer is the single **most\n",
        "> likely** one per sample, more samples sharpen the vote toward it. The\n",
        "> reasoning path is a means to an end — we **marginalize it out** and\n",
        "> keep only where the chains agree.\n",
        "\n",
        "## The Math: Why Voting Helps (Condorcet)\n",
        "\n",
        "Model each sampled chain as an independent voter that is correct with\n",
        "probability $p$. If it is wrong, take the worst case first: every wrong\n",
        "vote lands on the *same* wrong answer. Then the majority is correct\n",
        "exactly when more than half of $n$ voters are correct:\n",
        "\n",
        "$$P_\\text{correct}(n) = \\sum_{k > n/2} \\binom{n}{k}\\, p^{k} (1-p)^{n-k}.$$\n",
        "\n",
        "This is **Condorcet’s jury theorem**. Its behavior is a sharp phase\n",
        "transition at $p = \\tfrac12$:\n",
        "\n",
        "- If $p > \\tfrac12$, $P_\\text{correct}(n) \\to 1$ as $n \\to \\infty$ —\n",
        "  more votes, more accuracy.\n",
        "- If $p < \\tfrac12$, it goes to $0$ — voting *amplifies* a bad model’s\n",
        "  errors.\n",
        "- If $p = \\tfrac12$, it stays at $\\tfrac12$ forever.\n",
        "\n",
        "Real self-consistency does **better** than this bound, because wrong\n",
        "answers do not collude — they scatter across many values. Then the\n",
        "correct answer only has to beat each wrong answer *individually* (be the\n",
        "**mode**), a much weaker condition than $p > \\tfrac12$. Drive $p$ and\n",
        "watch the curve swing from “amplifies errors” to “converges to certain”:"
      ],
      "id": "ce37b811-a217-48cf-9353-ad3710318f12"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "55b6c584-3cac-4ce9-a9ac-86dabf8f6741"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        },
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        },
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        },
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "4471be34-ca53-4945-959c-eb327aa98ccb"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "b8063560-af8c-46d2-9846-4a9028b3eab7"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Code: Self-Consistency from Scratch\n",
        "\n",
        "`reasoning.py` builds the whole pipeline as small, model-agnostic\n",
        "functions — they wrap *any* generator, so we can test them without a\n",
        "trained model. First, pull the answer out of a chain. Following GSM8K,\n",
        "the model marks its final answer with `####`:"
      ],
      "id": "7ede3320-3a38-4237-ae30-b575cdb089f5"
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "42\n",
            "1024\n",
            "None"
          ]
        }
      ],
      "source": [
        "from reasoning import extract_answer\n",
        "\n",
        "print(extract_answer(\"6 eggs/day * 7 days = 42 eggs.\\n#### 42\"))\n",
        "print(extract_answer(\"...so there are 1,024 bytes. The answer is 1,024.\"))\n",
        "print(extract_answer(\"I'm not sure.\"))   # nothing to extract"
      ],
      "id": "6dac5084"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Then tally a **plurality** vote — the most common answer wins, `None`s\n",
        "ignored:"
      ],
      "id": "35dec8f7-1312-40ff-894d-b33d8c969367"
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "distribution: {'42': 4, '36': 1, '17': 1}\n",
            "vote winner:  42"
          ]
        }
      ],
      "source": [
        "from reasoning import majority_vote, vote_distribution\n",
        "\n",
        "answers = [\"42\", \"36\", \"42\", \"42\", \"17\", \"42\"]\n",
        "print(\"distribution:\", vote_distribution(answers))\n",
        "print(\"vote winner: \", majority_vote(answers))"
      ],
      "id": "c5d85e3f"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "`self_consistency` puts them together: sample `n` chains, extract each\n",
        "answer, return the vote and the tally. Here we drive it with\n",
        "`NoisyReasoner`, a controllable stand-in whose per-sample accuracy is\n",
        "exactly `p` — no model needed to see the effect (swap in a\n",
        "temperature-sampled `generate` call for real use):"
      ],
      "id": "bf72a7e9-27ef-4810-abec-69e62dd84fe6"
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "n= 1: voted '3' ✗   tally={'3': 1}\n",
            "n= 5: voted '42' ✓   tally={'42': 2, '3': 1, '2': 2}\n",
            "n=41: voted '42' ✓   tally={'1': 6, '42': 22, '2': 5, '3': 4, '0': 4}"
          ]
        }
      ],
      "source": [
        "from reasoning import self_consistency, NoisyReasoner\n",
        "\n",
        "gen = NoisyReasoner(correct=\"42\", p=0.55, num_distractors=4, seed=0)\n",
        "\n",
        "for n in (1, 5, 41):\n",
        "    answer, dist = self_consistency(gen.sample, n)\n",
        "    correct = \"✓\" if answer == \"42\" else \"✗\"\n",
        "    print(f\"n={n:>2}: voted {answer!r} {correct}   tally={dist}\")"
      ],
      "id": "db1cfa87"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "A single sample is a coin-flip near `p`; by 41 samples the correct\n",
        "answer is the clear plurality even though the generator is right only\n",
        "55% of the time. That is the phase transition, made concrete.\n",
        "\n",
        "## Watch Accuracy Climb with Compute\n",
        "\n",
        "The headline of test-time compute: **accuracy is a rising function of\n",
        "how many samples you draw**. `demonstrate_self_consistency` runs\n",
        "thousands of problems at each sample count and reports the fraction\n",
        "solved — the *empirical* compute–accuracy curve — beside the Condorcet\n",
        "bound."
      ],
      "id": "b807eae7-21e8-439a-a768-073cb2718a73"
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "============================================================\n",
            "SELF-CONSISTENCY: accuracy vs. samples (test-time compute)\n",
            "============================================================\n",
            "  per-sample accuracy p = 0.55, distractors = 4, trials = 3000\n",
            "\n",
            "     n   empirical   Condorcet\n",
            "     1       0.560       0.550\n",
            "     3       0.655       0.575\n",
            "     5       0.779       0.593\n",
            "    11       0.931       0.633\n",
            "    21       0.992       0.679\n",
            "    41       1.000       0.741\n",
            "\n",
            "  More samples -> higher accuracy (since p > 1/2); the empirical\n",
            "  curve beats the binary Condorcet bound because wrong answers scatter."
          ]
        }
      ],
      "source": [
        "from reasoning import demonstrate_self_consistency\n",
        "\n",
        "accuracy = demonstrate_self_consistency(p=0.55, num_distractors=4, trials=3000, seed=0)"
      ],
      "id": "4c324a4c"
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [],
      "source": [
        "from reasoning import condorcet_majority_prob\n",
        "\n",
        "# Bridge the empirical curve (and the theoretical bound) to the plot below.\n",
        "sc_points = [\n",
        "    {\"n\": n, \"empirical\": acc, \"condorcet\": condorcet_majority_prob(0.55, n)}\n",
        "    for n, acc in accuracy.items()\n",
        "]\n",
        "ojs_define(sc_points = sc_points)"
      ],
      "id": "f9518882"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "643b2f0d-d701-4eec-a8e2-ef4cf51e86b3"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Try This**\n",
        ">\n",
        "> The measured curve (orange) sits **above** the Condorcet bound (blue).\n",
        "> That gap is the scattering effect: with four distractors, the wrong\n",
        "> 45% splits four ways (~11% each), so the correct 55% is the runaway\n",
        "> plurality. Real reasoning benchmarks have *many* possible wrong\n",
        "> answers, so self-consistency is even more forgiving than the binary\n",
        "> theorem predicts.\n",
        "\n",
        "## Best-of-N and Verifiers\n",
        "\n",
        "Voting assumes the right answer is common. But sometimes the model finds\n",
        "the right answer *rarely* — it just needs help recognizing it.\n",
        "**Best-of-N** samples `N` candidates and keeps the one a **verifier** (a\n",
        "learned scorer / reward model, m12) rates highest, so a single good\n",
        "sample can win even if it is outvoted:"
      ],
      "id": "b16026b7-6928-42af-8ecc-c3ccd2070393"
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "verifier picked:\n",
            " 6 * 7 = 42, so the total is 42.\n",
            "#### 42"
          ]
        }
      ],
      "source": [
        "from reasoning import best_of_n\n",
        "\n",
        "# Toy verifier: prefers solutions that \"show their work\" (longer, with an '=').\n",
        "def verifier(chain: str) -> float:\n",
        "    return len(chain) + (5.0 if \"=\" in chain else 0.0)\n",
        "\n",
        "candidates = iter([\n",
        "    \"#### 42\",\n",
        "    \"6 * 7 = 42, so the total is 42.\\n#### 42\",\n",
        "    \"idk maybe 40\\n#### 40\",\n",
        "])\n",
        "best = best_of_n(lambda: next(candidates), score_fn=verifier, n=3)\n",
        "print(\"verifier picked:\\n\", best)"
      ],
      "id": "d1588ff9"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Voting and best-of-N are the two **parallel** ways to spend test-time\n",
        "compute: voting is a verifier-free majority; best-of-N trades the vote\n",
        "for a scorer that can spot a rare gem. Production systems blend them —\n",
        "e.g. weight each vote by its verifier score.\n",
        "\n",
        "## Chain-of-Thought & the Frontier\n",
        "\n",
        "Two loose ends connect this from-scratch core to the models making\n",
        "headlines.\n",
        "\n",
        "**Chain-of-thought is the enabler.** Self-consistency only helps if the\n",
        "sampled chains are *diverse yet mostly-correct* — which is exactly what\n",
        "CoT prompting produces. Asking for steps (“Let’s think step by step”)\n",
        "both raises per-sample accuracy $p$ and creates the path diversity\n",
        "voting feeds on.\n",
        "\n",
        "**RL for reasoning is the sequential counterpart.** Everything above is\n",
        "*parallel* test-time compute: draw independent samples, aggregate. The\n",
        "other axis is *sequential* — train the model (with RL against verifiable\n",
        "rewards) to produce one long, self-correcting chain that thinks longer\n",
        "on harder problems. That is the o1 / DeepSeek-R1 recipe; its **GRPO**\n",
        "objective is the natural from-scratch sequel to the DPO you built in m12\n",
        "— and you build it next, below.\n",
        "\n",
        "> **Key Insight**\n",
        ">\n",
        "> There are two knobs for test-time compute: **parallel** (sample many,\n",
        "> aggregate — this module) and **sequential** (one longer,\n",
        "> self-correcting chain — RL-trained). Both trade inference FLOPs for\n",
        "> accuracy, and both leave the pre-trained weights of the base model\n",
        "> exactly where scaling laws (m07) left them.\n",
        "\n",
        "## The Sequential Branch: Training to Reason with GRPO\n",
        "\n",
        "Self-consistency spends compute at *inference* on a frozen model. The\n",
        "reasoning models that top today’s benchmarks do something more: they are\n",
        "**trained** to spend that compute well — to emit one long,\n",
        "self-correcting chain and only then commit to an answer. The recipe that\n",
        "made this reproducible and open is **GRPO** (Group Relative Policy\n",
        "Optimization), introduced in DeepSeekMath (2024) and used to train\n",
        "DeepSeek-R1 (2025).\n",
        "\n",
        "GRPO is reinforcement learning, and it is the direct sequel to the DPO\n",
        "you built in m12. Both push a policy toward better outputs using only a\n",
        "*preference* or *reward* signal, no labeled target text. The twist that\n",
        "makes GRPO special is how it estimates whether an answer was good.\n",
        "\n",
        "**The problem RL has to solve.** To improve, the policy needs to know\n",
        "whether a sampled answer was *better than expected*. Classic PPO trains\n",
        "a second network — a **critic** / value model $V(s)$, as big as the\n",
        "policy — just to predict that expected reward, so the “advantage” is\n",
        "$r - V(s)$. That doubles the memory.\n",
        "\n",
        "**GRPO’s one idea: let the group be the baseline.** For a prompt $q$,\n",
        "sample a whole *group* of $G$ answers $o_1, \\dots, o_G$ from the current\n",
        "policy, score each with a reward $r_i$, and use the **group’s own mean\n",
        "reward** as the expectation. An answer that beat its group’s average\n",
        "gets pushed up; one below average gets pushed down. No critic, no value\n",
        "network — just the samples you already drew.\n",
        "\n",
        "> **Key Insight**\n",
        ">\n",
        "> The reward can be **verifiable**: for math or code you don’t need a\n",
        "> learned reward model at all — you can *check* whether the final answer\n",
        "> is correct. DeepSeek-R1 trained on exactly this (a correctness reward\n",
        "> plus a format bonus), and *R1-Zero* reached frontier reasoning with\n",
        "> **no supervised reasoning traces at all** — pure RL from a rule that\n",
        "> grades the answer.\n",
        "\n",
        "## The Math: Group-Relative Advantage\n",
        "\n",
        "Score the group, then turn raw rewards into **advantages** by\n",
        "standardizing within the group:\n",
        "\n",
        "$$A_i = \\frac{r_i - \\operatorname{mean}(r_1, \\dots, r_G)}{\\operatorname{std}(r_1, \\dots, r_G) + \\varepsilon}.$$\n",
        "\n",
        "Subtracting the mean centers the signal (this is the critic-free\n",
        "baseline); dividing by the std puts every prompt on the same scale. A\n",
        "group that is *all* correct or *all* wrong has zero std → zero advantage\n",
        "— correctly, since there is nothing *relative* to learn from it.\n",
        "\n",
        "Every token in output $o_i$ shares that one scalar $A_i$. The policy is\n",
        "then nudged with the same **clipped surrogate** PPO uses, with a KL\n",
        "leash to a frozen reference policy $\\pi_\\text{ref}$:\n",
        "\n",
        "$$\\mathcal{J}_\\text{GRPO} = \\frac{1}{G}\\sum_{i=1}^{G}\n",
        "  \\min\\!\\Big(\\rho_i A_i,\\ \\operatorname{clip}(\\rho_i,\\, 1-\\varepsilon,\\, 1+\\varepsilon)\\,A_i\\Big)\n",
        "  \\;-\\; \\beta\\, \\mathbb{D}_\\text{KL}\\!\\left(\\pi_\\theta \\,\\|\\, \\pi_\\text{ref}\\right),\n",
        "\\qquad\n",
        "\\rho_i = \\frac{\\pi_\\theta(o_i \\mid q)}{\\pi_{\\theta_\\text{old}}(o_i \\mid q)}.$$\n",
        "\n",
        "Two guards keep the update honest. The **clip** is a trust region: once\n",
        "a sample’s probability has moved more than $\\varepsilon$ in the helpful\n",
        "direction, the clipped branch wins the $\\min$ and its gradient goes to\n",
        "zero, so one lucky group can’t yank the policy off a cliff. The **KL\n",
        "penalty** keeps $\\pi_\\theta$ near the reference so it doesn’t forget how\n",
        "to write — or learn to reward-hack. GRPO uses Schulman’s unbiased,\n",
        "always-non-negative **k3** estimator,\n",
        "$\\mathbb{D}_\\text{KL} = e^{d} - d - 1$ with\n",
        "$d = \\log\\pi_\\text{ref} - \\log\\pi_\\theta$.\n",
        "\n",
        "### Step by Step"
      ],
      "id": "ec0ec63c-5fd0-49d0-8ee4-18ef0e6116a4"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "f96de3cb-45fa-4b4f-b7cf-304da3a0bb1e"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "d706c4c2-d089-4a0e-b8e1-6ca7622792b3"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "ce5e1a37-020f-4766-ba33-23f3d3795da4"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Code: GRPO from Scratch\n",
        "\n",
        "The whole algorithm is three small functions plus a loop. It lives in\n",
        "`grpo.py`; here we build the pieces inline. Start with the advantage —\n",
        "the one line that replaces PPO’s critic:"
      ],
      "id": "4743bb2d-9df8-43bd-be82-59e14c59bc4f"
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "rewards   : [1.0, 0.0, 0.0, 0.0]\n",
            "advantages: [1.732, -0.577, -0.577, -0.577]\n",
            "mean ~ 0  : 0.0  std ~ 1: 1.0\n",
            "all-correct advantages: [0.0, 0.0, 0.0, 0.0]"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "from grpo import group_relative_advantages, kl_divergence\n",
        "\n",
        "# A group of 4 answers to one prompt; only the first was correct (verifiable reward).\n",
        "rewards = torch.tensor([1.0, 0.0, 0.0, 0.0])\n",
        "adv = group_relative_advantages(rewards)\n",
        "print(\"rewards   :\", rewards.tolist())\n",
        "print(\"advantages:\", [round(a, 3) for a in adv.tolist()])\n",
        "print(\"mean ~ 0  :\", round(float(adv.mean()), 4), \" std ~ 1:\", round(float(adv.std(unbiased=False)), 4))\n",
        "\n",
        "# All-correct (or all-wrong) group -> no relative signal -> zero advantage.\n",
        "print(\"all-correct advantages:\", group_relative_advantages(torch.ones(4)).tolist())"
      ],
      "id": "a10b21e6"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The correct answer gets a positive advantage, the wrong ones negative,\n",
        "and an all-equal group yields zeros. Now the KL leash — the k3\n",
        "estimator, non-negative by construction:"
      ],
      "id": "d6ba1fb6-a6c8-41b0-af4a-16626c9ea1a7"
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "KL when unchanged: 0.000\n",
            "KL after drifting: 0.718  (always >= 0)"
          ]
        }
      ],
      "source": [
        "same = kl_divergence(torch.tensor(-1.0), torch.tensor(-1.0))   # policy == reference\n",
        "moved = kl_divergence(torch.tensor(-2.0), torch.tensor(-1.0))  # policy drifted away\n",
        "print(f\"KL when unchanged: {float(same):.3f}\")\n",
        "print(f\"KL after drifting: {float(moved):.3f}  (always >= 0)\")"
      ],
      "id": "b036b53b"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Put them together with the clipped surrogate and you have the GRPO loss\n",
        "— exactly the objective `loss.backward()` differentiates. To keep it\n",
        "hand-followable we make each “answer” a **single token**: a categorical\n",
        "policy over `K` candidate answers, where an output’s log-prob is one\n",
        "entry of `log_softmax`. (For a real chain you sum the per-token\n",
        "log-probs of the generated sequence; the loss above is unchanged.)"
      ],
      "id": "39c22c5b-c40b-4258-a330-30041156d7eb"
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "before: p(correct) = 0.167  (chance = 0.167)\n",
            "after : p(correct) = 0.993  (learned from a verifiable reward — no critic, no labels)"
          ]
        }
      ],
      "source": [
        "from grpo import ToyReasoningPolicy, grpo_step\n",
        "\n",
        "# A uniform policy over 6 possible answers; answer 3 is correct.\n",
        "policy = ToyReasoningPolicy(num_answers=6, seed=0)\n",
        "print(f\"before: p(correct) = {float(policy.probabilities()[3]):.3f}  (chance = {1/6:.3f})\")\n",
        "\n",
        "for _ in range(30):\n",
        "    grpo_step(policy, correct=3, group_size=32, lr=0.5)\n",
        "\n",
        "print(f\"after : p(correct) = {float(policy.probabilities()[3]):.3f}  \"\n",
        "      \"(learned from a verifiable reward — no critic, no labels)\")"
      ],
      "id": "29e0e87a"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Each `grpo_step` samples a group from the current policy, grades it,\n",
        "standardizes the rewards into advantages, and takes a few\n",
        "clipped-surrogate gradient steps toward the above-average answers — the\n",
        "full loop in `grpo.py`.\n",
        "\n",
        "## Watch a Policy Learn to Reason\n",
        "\n",
        "`demonstrate_grpo` runs that loop from a uniform start and records the\n",
        "curve. The policy climbs from chance ($1/K$) toward near-certainty on\n",
        "the correct answer, driven only by a reward that says *right* or *wrong*\n",
        "— a from-scratch miniature of the R1 recipe."
      ],
      "id": "79e3b730-344b-400d-8b8a-2988b4027520"
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {},
      "outputs": [],
      "source": [
        "from grpo import demonstrate_grpo\n",
        "\n",
        "grpo_hist = demonstrate_grpo(num_answers=6, correct=3, steps=40,\n",
        "                             group_size=32, lr=0.5, seed=0, verbose=False)\n",
        "\n",
        "# Bridge the learning curve to the interactive plot below.\n",
        "ojs_define(grpo_curve = grpo_hist)"
      ],
      "id": "9cfd69a7"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "a881266c-ae18-4a39-8755-0d2e104d4cdc"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Try This**\n",
        ">\n",
        "> 1.  **Kill the reward signal.** In the `demonstrate_grpo` call, set\n",
        ">     `correct` to an answer, then imagine every reward were identical —\n",
        ">     the advantages would all be zero and the curve would stay flat.\n",
        ">     This is why a group that is *all* right or *all* wrong teaches\n",
        ">     nothing.\n",
        "> 2.  **Shrink the group.** Re-run with `group_size=4` vs\n",
        ">     `group_size=64`. Smaller groups give noisier advantage estimates\n",
        ">     (the mean baseline is shakier), so the climb is bumpier — the same\n",
        ">     bias/variance trade a critic would smooth.\n",
        "> 3.  **Turn up the leash.** Raise `beta` toward `1.0`. The KL penalty\n",
        ">     fights the reward, slowing how fast the policy departs its uniform\n",
        ">     reference — a stronger leash trades learning speed for staying\n",
        ">     close to the starting model.\n",
        "\n",
        "## Common Pitfalls\n",
        "\n",
        "When implementing test-time compute, watch out for:\n",
        "\n",
        "1.  **Voting with greedy decoding.** Self-consistency *needs* stochastic\n",
        "    sampling (temperature \\> 0, top-p/top-k from m08). Greedy makes\n",
        "    every chain identical, so the “vote” is `n` copies of one answer —\n",
        "    no gain.\n",
        "2.  **Voting when `p < ½` and the answer space is small.** With a single\n",
        "    dominant wrong answer, more samples make you *more* confidently\n",
        "    wrong. Voting helps only when the correct answer is the mode.\n",
        "3.  **Comparing raw strings.** `\"42\"`, `\"42.\"`, and `\"1,024\"` vs\n",
        "    `\"1024\"` must be normalized before tallying, or correct chains split\n",
        "    their own vote. That is why `extract_answer` strips commas and pulls\n",
        "    a clean number.\n",
        "4.  **Counting compute as free.** N-sample self-consistency costs ~N×\n",
        "    the FLOPs of one pass. The compute–accuracy curve has diminishing\n",
        "    returns; past some N the next point of accuracy is not worth the\n",
        "    tokens.\n",
        "5.  **Trusting a weak verifier in best-of-N.** Best-of-N is only as good\n",
        "    as its scorer — an exploitable verifier makes the model\n",
        "    **reward-hack** (m12), picking samples that fool the scorer rather\n",
        "    than solve the problem.\n",
        "6.  **A degenerate GRPO group teaches nothing.** If every sampled answer\n",
        "    earns the same reward (all correct, or all wrong), the group’s std\n",
        "    is zero and *every* advantage is zero — no gradient. Reward design\n",
        "    and a policy that isn’t already saturated matter: you need a\n",
        "    *spread* of outcomes for the group baseline to have signal.\n",
        "7.  **Dropping GRPO’s KL leash.** Without the\n",
        "    $\\beta\\,\\mathbb{D}_\\text{KL}$ term the policy is free to chase the\n",
        "    reward straight into reward-hacking or degenerate text. The clipped\n",
        "    surrogate bounds *one* step; the KL keeps the model near a sane\n",
        "    reference across *many* steps.\n",
        "\n",
        "## Exercises\n",
        "\n",
        "### Exercise 1: The break-even sample count"
      ],
      "id": "2dda8575-375b-46d1-bac9-953197b838da"
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {},
      "outputs": [],
      "source": [
        "from reasoning import condorcet_majority_prob\n",
        "\n",
        "# For p = 0.55, how many samples n does the Condorcet bound need to first exceed\n",
        "# 0.90? (Loop odd n and check.) Then confirm real scattered-distractor accuracy\n",
        "# reaches it sooner.\n",
        "\n",
        "# Your implementation here:"
      ],
      "id": "9d601b42"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Exercise 2: Self-consistency vs. best-of-N"
      ],
      "id": "b2fbdb63-eee5-42bb-8b8a-d4d87aee4a27"
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {},
      "outputs": [],
      "source": [
        "from reasoning import self_consistency, best_of_n, NoisyReasoner\n",
        "\n",
        "# Build a generator that is correct only 30% of the time but whose correct\n",
        "# chains are always the longest. Show best_of_n(score_fn=len) beats\n",
        "# self_consistency at the same N. When does voting win instead?\n",
        "\n",
        "# Your implementation here:"
      ],
      "id": "9f16d2a7"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Exercise 3: Weighted voting"
      ],
      "id": "8fadd763-626a-47ff-9b64-cd7242363539"
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {},
      "outputs": [],
      "source": [
        "from collections import Counter\n",
        "\n",
        "# Combine voting and verification: instead of one vote per chain, weight each\n",
        "# chain's vote by a verifier score, then take the argmax. Implement\n",
        "# weighted_vote(answers, scores) and compare to plain majority_vote.\n",
        "\n",
        "# Your implementation here:"
      ],
      "id": "8364f62b"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Exercise 4: GRPO with a shaped reward"
      ],
      "id": "b1cae6ea-b053-4169-9ba2-40b1184ea591"
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {},
      "outputs": [],
      "source": [
        "from grpo import ToyReasoningPolicy, group_relative_advantages, grpo_step\n",
        "\n",
        "# The verifiable_reward above is 0/1. Give partial credit instead: reward the\n",
        "# correct answer 1.0, a \"close\" answer (index correct-1 or correct+1) 0.5, else 0.\n",
        "# Write shaped_reward(answers, correct, num_answers) -> tensor, then run GRPO with\n",
        "# it (adapt grpo_step, or roll your own loop). Does the policy still converge on\n",
        "# the exactly-correct answer, or does partial credit slow it down?\n",
        "\n",
        "# Your implementation here:"
      ],
      "id": "33af600b"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Summary\n",
        "\n",
        "Key takeaways:\n",
        "\n",
        "1.  **Test-time compute is a third axis.** Beyond bigger models and more\n",
        "    data (m07 scaling laws), you can spend more compute *at inference*\n",
        "    to raise accuracy — with the weights frozen.\n",
        "2.  **Self-consistency is sample-and-vote.** Draw many chain-of-thought\n",
        "    samples, extract each answer, take the plurality. It lifted GSM8K by\n",
        "    +17.9 points with no new parameters.\n",
        "3.  **Condorcet explains it.** If each sample is correct with\n",
        "    probability $p > \\tfrac12$, the majority of $n$ is correct with\n",
        "    probability $\\to 1$; below $\\tfrac12$ voting backfires. The real\n",
        "    requirement is weaker — the correct answer just has to be the\n",
        "    **mode**, since wrong answers scatter.\n",
        "4.  **Accuracy rises with samples.** The compute–accuracy curve climbs\n",
        "    (with diminishing returns), and empirically beats the binary\n",
        "    Condorcet bound.\n",
        "5.  **Best-of-N uses a verifier instead of a vote,** so it can surface a\n",
        "    *rare* correct answer — as good as its scorer, and vulnerable to\n",
        "    reward hacking.\n",
        "6.  **Parallel vs. sequential.** Sampling-and-aggregating is the\n",
        "    parallel branch; RL-trained long self-correcting chains (o1 / R1)\n",
        "    are the sequential branch. Both are test-time compute.\n",
        "7.  **GRPO trains reasoning without a critic.** Sample a *group* of\n",
        "    answers, use their mean reward as the baseline\n",
        "    ($A_i = (r_i - \\text{mean})/\\text{std}$), and nudge the policy with\n",
        "    a clipped surrogate plus a KL leash. With a **verifiable** reward\n",
        "    (is the answer correct?) it needs no reward model and no labeled\n",
        "    traces — the DeepSeek-R1 recipe, built from scratch in `grpo.py`.\n",
        "\n",
        "## What’s Next\n",
        "\n",
        "You now have **both** halves of the reasoning story: the parallel branch\n",
        "(sample-and-aggregate) and the sequential branch (GRPO — training a\n",
        "model to reason from verifiable rewards, no critic). From here the book\n",
        "turns to serving these compute-hungry models efficiently — [Module 14:\n",
        "Quantization](../m14_quantization/lesson.qmd) shrinks the weights, and\n",
        "speculative decoding speeds the decode — and to [Module 15:\n",
        "Retrieval-Augmented Generation](../m15_rag/lesson.qmd) and agents that\n",
        "put reasoning to work on live knowledge.\n",
        "\n",
        "### Going Deeper\n",
        "\n",
        "**Core Papers:**\n",
        "\n",
        "- [Self-Consistency Improves Chain of Thought Reasoning in Language\n",
        "  Models](https://arxiv.org/abs/2203.11171) — Wang et al. (2022), the\n",
        "  sample-and-vote method built here (+17.9 on GSM8K).\n",
        "- [Chain-of-Thought Prompting Elicits Reasoning in Large Language\n",
        "  Models](https://arxiv.org/abs/2201.11903) — Wei et al. (2022),\n",
        "  reasoning-before-answering.\n",
        "- [Training Verifiers to Solve Math Word\n",
        "  Problems](https://arxiv.org/abs/2110.14168) — Cobbe et al. (2021),\n",
        "  GSM8K, the `####` answer format, and best-of-N with a verifier.\n",
        "- [Scaling LLM Test-Time Compute Optimally can be More Effective than\n",
        "  Scaling Model Parameters](https://arxiv.org/abs/2408.03314) — Snell et\n",
        "  al. (2024), the compute-optimal proposer/verifier view.\n",
        "- [DeepSeekMath: Pushing the Limits of Mathematical\n",
        "  Reasoning](https://arxiv.org/abs/2402.03300) — Shao et al. (2024),\n",
        "  where **GRPO** is introduced: the group-relative advantage, the\n",
        "  clipped surrogate, and the critic-free objective built here.\n",
        "- [DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via\n",
        "  Reinforcement Learning](https://arxiv.org/abs/2501.12948) —\n",
        "  DeepSeek-AI (2025), GRPO with verifiable rewards; *R1-Zero* reasons\n",
        "  with no supervised traces at all.\n",
        "\n",
        "**Practical Resources:**\n",
        "\n",
        "- [Approximating KL Divergence](http://joschu.net/blog/kl-approx.html) —\n",
        "  Schulman (2020), the unbiased, non-negative **k3** estimator GRPO uses\n",
        "  for its KL leash.\n",
        "- [Condorcet’s Jury\n",
        "  Theorem](https://en.wikipedia.org/wiki/Condorcet%27s_jury_theorem) —\n",
        "  the 1785 result that voting converges to correct when each voter beats\n",
        "  a coin flip."
      ],
      "id": "afb4a286-b40e-4852-b75a-f51866db9b19"
    }
  ],
  "nbformat": 4,
  "nbformat_minor": 5,
  "metadata": {
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "path": "/opt/buildhome/.asdf/installs/python/3.11.14/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.14"
    }
  }
}