{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Module 12: Alignment"
      ],
      "id": "e01fc9ac-ae6c-4138-a237-ece7edf7014d"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [],
      "id": "562b022c-e812-4ce7-9b96-96139321cc6b"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "9cc1ad54-ead0-4c3e-a914-6bb1b59d1e9e"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Introduction\n",
        "\n",
        "You have trained a GPT (m07) and made it generate (m08). It predicts\n",
        "*likely* text — but likely is not the same as *helpful*. Ask a raw\n",
        "pretrained model a question and it is just as happy to continue with\n",
        "more questions, because that is what the internet looks like.\n",
        "**Alignment** is the post-training that turns a next-token predictor\n",
        "into an assistant that follows instructions and honors human\n",
        "preferences.\n",
        "\n",
        "This module builds the three canonical stages from scratch — and shows\n",
        "how the third one quietly deletes the machinery of the second:\n",
        "\n",
        "- **SFT** (supervised fine-tuning): teach the format by imitation, on\n",
        "  demonstrations.\n",
        "- **Reward modeling**: learn *what people prefer* from comparisons, not\n",
        "  demonstrations.\n",
        "- **DPO** (Direct Preference Optimization): fold the reward model and\n",
        "  the whole reinforcement-learning loop into a single classification\n",
        "  loss — *“your language model is secretly a reward model.”*\n",
        "\n",
        "Why it matters for LLMs:\n",
        "\n",
        "- Every assistant you have used (ChatGPT, Claude, Gemini) is a\n",
        "  pretrained model put through exactly this pipeline. Without it, the\n",
        "  base model is a brilliant autocomplete, not a helper.\n",
        "- DPO (and its relatives) is now the default way to align open models,\n",
        "  precisely because it removes the fragile RL step this module explains.\n",
        "\n",
        "### What You’ll Learn\n",
        "\n",
        "After this module, you can:\n",
        "\n",
        "- Explain the gap between **likelihood** and **helpfulness**, and the\n",
        "  three-stage pipeline that closes it.\n",
        "- Implement the **SFT loss** from scratch — masked next-token\n",
        "  cross-entropy that trains the completion, not the prompt.\n",
        "- Build a **reward model** by swapping an LM’s head, and train it with\n",
        "  the **Bradley-Terry** preference loss.\n",
        "- Derive the **KL-constrained RLHF objective**, its closed-form optimum,\n",
        "  and the reparameterization that turns it into **DPO**.\n",
        "- Implement and drive the **DPO loss**, and see the implicit reward push\n",
        "  the preferred answer up and the rejected one down.\n",
        "- Build **PPO** from scratch — the clipped surrogate, the value baseline\n",
        "  (critic), and the KL-shaped reward — and watch **reward\n",
        "  over-optimization** collapse the true reward when the KL leash is too\n",
        "  loose.\n",
        "\n",
        "### Prerequisites\n",
        "\n",
        "This module requires familiarity with:\n",
        "\n",
        "- [Module 07: Training](../m07_training/lesson.qmd) — the\n",
        "  language-modeling loss and gradient descent; SFT is that loss with a\n",
        "  mask.\n",
        "- [Module 08: Generation](../m08_generation/lesson.qmd) — sampling from\n",
        "  a policy; alignment changes *which* continuations the policy prefers.\n",
        "\n",
        "## Intuition: The Alignment Pipeline\n",
        "\n",
        "Alignment is a relay of three hand-offs, each fixing the previous\n",
        "stage’s limit. Pretraining gives a model that *knows language*. SFT\n",
        "shows it *the shape of a good answer*. But you cannot demonstrate the\n",
        "best answer to every prompt — so instead you let people **compare**\n",
        "answers, learn a notion of “better,” and push the model toward it. Step\n",
        "through the pipeline and watch what each stage consumes and what it\n",
        "changes:"
      ],
      "id": "0e6a2489-9352-4f7e-84b8-f7b22dc46ece"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "3293ffdb-b1cc-42fc-b50d-0fc1d6d58c43"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "5c82a757-f6a0-4cce-8070-da883461bc87"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "3ba4e051-4a36-4231-a204-06f66eabcb9f"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Key Insight**\n",
        ">\n",
        "> Only **SFT, RLHF, and DPO update the policy** — the model you ship.\n",
        "> The reward model is scaffolding: it exists to score answers during\n",
        "> RLHF. DPO’s whole trick is to notice you can express that score *using\n",
        "> the policy itself*, and throw the scaffolding away.\n",
        "\n",
        "## Stage 1 — Supervised Fine-Tuning\n",
        "\n",
        "SFT is the m07 language-modeling loss with one change: given a\n",
        "`(prompt, completion)` pair, you train the model to predict **only the\n",
        "completion tokens**. The prompt is context, not something to imitate, so\n",
        "its tokens are masked out of the loss.\n",
        "\n",
        "$$\\mathcal{L}_{\\text{SFT}} = - \\frac{1}{|C|} \\sum_{t \\in C} \\log \\pi_\\theta(y_t \\mid y_{<t})$$\n",
        "\n",
        "where $C$ is the set of completion-token positions. Everything else —\n",
        "the softmax, the cross-entropy, the next-token shift — is exactly m07.\n",
        "\n",
        "`alignment.py` implements this as `sft_loss`, built on a reusable\n",
        "`sequence_logprob` that computes $\\log \\pi_\\theta(y \\mid x)$ with a\n",
        "completion mask:"
      ],
      "id": "80791745-c9c6-4f7d-a5b8-ffd7a52806db"
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Loss over all tokens:        3.3155\n",
            "Loss over completion only:   3.5788\n",
            "They differ because the prompt tokens no longer count as targets."
          ]
        }
      ],
      "source": [
        "import torch\n",
        "from alignment import sft_loss, sequence_logprob\n",
        "\n",
        "torch.manual_seed(0)\n",
        "batch, seq, vocab = 2, 6, 20\n",
        "logits = torch.randn(batch, seq, vocab)\n",
        "labels = torch.randint(0, vocab, (batch, seq))\n",
        "\n",
        "# The first 3 tokens are the prompt (mask 0); the last 3 are the completion (1).\n",
        "completion_mask = torch.zeros(batch, seq)\n",
        "completion_mask[:, 3:] = 1\n",
        "\n",
        "loss_all = sft_loss(logits, labels, torch.ones(batch, seq))\n",
        "loss_completion = sft_loss(logits, labels, completion_mask)\n",
        "\n",
        "print(f\"Loss over all tokens:        {loss_all.item():.4f}\")\n",
        "print(f\"Loss over completion only:   {loss_completion.item():.4f}\")\n",
        "print(\"They differ because the prompt tokens no longer count as targets.\")"
      ],
      "id": "c7effcfe"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The mask is the whole point: without it, the model spends capacity\n",
        "learning to regurgitate prompts instead of answering them.\n",
        "`sequence_logprob` — the $\\log \\pi_\\theta(y \\mid x)$ it wraps — returns\n",
        "in a moment as the atom DPO is built from.\n",
        "\n",
        "## Stage 2 — Preferences and Reward Modeling\n",
        "\n",
        "SFT can only teach what a human can *write down*. For most prompts there\n",
        "is no single ideal answer, but people can reliably say **which of two\n",
        "answers is better**. That comparison is cheaper to collect and far more\n",
        "informative at the margin.\n",
        "\n",
        "The **Bradley-Terry** model turns a scalar reward into the probability\n",
        "that one answer beats another:\n",
        "\n",
        "$$p(y_w \\succ y_l \\mid x) = \\sigma\\!\\big(r(x, y_w) - r(x, y_l)\\big)$$\n",
        "\n",
        "so a **reward model** $r_\\phi(x, y)$ is trained by maximizing the\n",
        "probability it assigns to the human’s choice — equivalently, minimizing\n",
        "\n",
        "$$\\mathcal{L}_R = -\\,\\mathbb{E}_{(x, y_w, y_l)}\\Big[\\log \\sigma\\big(r_\\phi(x, y_w) - r_\\phi(x, y_l)\\big)\\Big].$$\n",
        "\n",
        "Only the *difference* of rewards ever appears, so the absolute scale is\n",
        "free — a fact DPO will exploit. A reward model is not a new\n",
        "architecture: it is a language-model backbone with the vocabulary head\n",
        "swapped for a single scalar head reading the last token. `RewardModel`\n",
        "does exactly that:"
      ],
      "id": "fa49bd8f-26c1-4570-898d-6f375b29acb6"
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Reward shape (one scalar per sequence): (4,)\n",
            "Bradley-Terry loss: 0.7040\n",
            "Ranking accuracy:   0.50"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "from alignment import RewardModel, bradley_terry_loss, preference_accuracy\n",
        "import importlib.util, sys\n",
        "from pathlib import Path\n",
        "\n",
        "# Load the m06 GPT to use as the reward model's backbone.\n",
        "_p = Path(\"../m06_transformer/transformer.py\").resolve()\n",
        "_s = importlib.util.spec_from_file_location(\"transformer\", _p)\n",
        "_t = importlib.util.module_from_spec(_s); sys.modules[\"transformer\"] = _t\n",
        "_s.loader.exec_module(_t)\n",
        "GPTModel = _t.GPTModel\n",
        "\n",
        "torch.manual_seed(0)\n",
        "backbone = GPTModel(vocab_size=50, embed_dim=32, num_heads=2, num_layers=2, max_seq_len=32)\n",
        "reward_model = RewardModel(backbone)\n",
        "\n",
        "chosen = torch.randint(0, 50, (4, 8))    # 4 preferred answers\n",
        "rejected = torch.randint(0, 50, (4, 8))  # 4 dispreferred answers\n",
        "\n",
        "r_chosen = reward_model(chosen)\n",
        "r_rejected = reward_model(rejected)\n",
        "print(f\"Reward shape (one scalar per sequence): {tuple(r_chosen.shape)}\")\n",
        "print(f\"Bradley-Terry loss: {bradley_terry_loss(r_chosen, r_rejected).item():.4f}\")\n",
        "print(f\"Ranking accuracy:   {preference_accuracy(r_chosen, r_rejected).item():.2f}\")"
      ],
      "id": "51b4b824"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The reward model is untrained here, so its accuracy is a coin flip —\n",
        "training it with `bradley_terry_loss` on real preference pairs is what\n",
        "teaches it to rank.\n",
        "\n",
        "## The Math: RLHF and Its Closed Form\n",
        "\n",
        "With a reward model in hand, the obvious move is to fine-tune the policy\n",
        "to score well. But a reward model is a *proxy*; chase it too hard and\n",
        "the policy drifts into gibberish that games the reward (**reward\n",
        "hacking**). So RLHF maximizes reward **minus a KL leash** to the\n",
        "reference (SFT) model:\n",
        "\n",
        "$$\\max_{\\pi_\\theta}\\; \\mathbb{E}_{x,\\, y \\sim \\pi_\\theta}\\big[r(x, y)\\big]\n",
        "\\;-\\; \\beta\\, \\mathrm{KL}\\!\\big[\\pi_\\theta(y \\mid x)\\,\\|\\,\\pi_{\\text{ref}}(y \\mid x)\\big].$$\n",
        "\n",
        "This objective has an exact, closed-form optimum — you can write down\n",
        "the best possible policy without any training:\n",
        "\n",
        "$$\\pi^*(y \\mid x) = \\frac{1}{Z(x)}\\, \\pi_{\\text{ref}}(y \\mid x)\\, \\exp\\!\\Big(\\tfrac{1}{\\beta}\\, r(x, y)\\Big).$$\n",
        "\n",
        "Read it as: **start from the reference and re-weight by reward**. The\n",
        "temperature is $\\beta$. Drive it and watch the optimal policy slide\n",
        "between “obey the reward” and “stay put”:"
      ],
      "id": "106f8392-2c25-4178-9e2c-cc621c6329c8"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "37693e72-8c83-4c1a-b324-ac1ea52f0f7a"
    },
    {
      "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": "a1e57989-30d4-4df8-920b-90e85ec535d5"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "a8dafb69-6cfe-482b-8710-08ce9f5cabdc"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "b6ff9eb7-f17e-4dab-874f-ead7a30f094d"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Key Insight**\n",
        ">\n",
        "> The optimum is *known* — so why is RLHF hard? Because $Z(x)$ sums over\n",
        "> every possible completion and is intractable, so you cannot just\n",
        "> sample from $\\pi^*$. Classic RLHF approximates it with **PPO**: roll\n",
        "> out completions, score them with the reward model, and nudge the\n",
        "> policy up the reward gradient. It works, but it needs a live reward\n",
        "> model, on-policy sampling, and careful tuning. The next section\n",
        "> removes all of it.\n",
        "\n",
        "## Stage 3 — DPO: Your Language Model Is the Reward Model\n",
        "\n",
        "Here is the move. Take the closed-form optimum and solve it *for the\n",
        "reward* instead of the policy — a line of algebra gives\n",
        "\n",
        "$$r(x, y) = \\beta \\log \\frac{\\pi^*(y \\mid x)}{\\pi_{\\text{ref}}(y \\mid x)} + \\beta \\log Z(x).$$\n",
        "\n",
        "The reward is just a **log-ratio of the policy to the reference**, plus\n",
        "a term that depends only on $x$. Now substitute this into the\n",
        "Bradley-Terry loss. Because only reward *differences* matter, the\n",
        "intractable $\\beta \\log Z(x)$ **cancels**, and the reward model vanishes\n",
        "with it. What remains is a loss on the policy alone:\n",
        "\n",
        "$$\\mathcal{L}_{\\text{DPO}} = -\\,\\mathbb{E}_{(x, y_w, y_l)}\\left[\\log \\sigma\\!\\left(\n",
        "\\beta \\log \\frac{\\pi_\\theta(y_w \\mid x)}{\\pi_{\\text{ref}}(y_w \\mid x)}\n",
        "- \\beta \\log \\frac{\\pi_\\theta(y_l \\mid x)}{\\pi_{\\text{ref}}(y_l \\mid x)}\n",
        "\\right)\\right].$$\n",
        "\n",
        "No reward model. No RL. No sampling. It is the *same supervised\n",
        "classification loss* as reward modeling — but the “reward” is the\n",
        "policy’s own **implicit reward**\n",
        "$\\hat r(x, y) = \\beta \\log \\frac{\\pi_\\theta(y|x)}{\\pi_{\\text{ref}}(y|x)}$.\n",
        "Both log-probs come straight from `sequence_logprob`; `dpo_loss`\n",
        "assembles them:"
      ],
      "id": "34bfa859-72a4-47d0-b8d8-cc0f2abeb07a"
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "DPO loss:               0.6583\n",
            "Implicit reward chosen:   [0.10000000149011612, 0.0]\n",
            "Implicit reward rejected: [-0.10000000149011612, 0.05000000074505806]\n",
            "Reward margin (want > 0): [0.20000000298023224, -0.05000000074505806]"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "from alignment import dpo_loss, implicit_reward\n",
        "\n",
        "# Per-sequence log πθ(y|x) and log π_ref(y|x) for a batch of preference pairs.\n",
        "# (In a real run these come from batch_sequence_logprob on the policy and a\n",
        "# frozen copy of the SFT model.)\n",
        "policy_chosen   = torch.tensor([-2.0, -3.0])\n",
        "policy_rejected = torch.tensor([-4.0, -2.5])\n",
        "ref_chosen      = torch.tensor([-3.0, -3.0])\n",
        "ref_rejected    = torch.tensor([-3.0, -3.0])\n",
        "\n",
        "loss, chosen_reward, rejected_reward = dpo_loss(\n",
        "    policy_chosen, policy_rejected, ref_chosen, ref_rejected, beta=0.1\n",
        ")\n",
        "print(f\"DPO loss:               {loss.item():.4f}\")\n",
        "print(f\"Implicit reward chosen:   {chosen_reward.tolist()}\")\n",
        "print(f\"Implicit reward rejected: {rejected_reward.tolist()}\")\n",
        "print(f\"Reward margin (want > 0): {(chosen_reward - rejected_reward).tolist()}\")"
      ],
      "id": "4b6cc02d"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The DPO gradient increases $\\log \\pi_\\theta(y_w)$ and decreases\n",
        "$\\log \\pi_\\theta(y_l)$, weighted by\n",
        "$\\sigma\\big(\\hat r(y_l) - \\hat r(y_w)\\big)$ — so pairs the model\n",
        "currently ranks **backwards** get the biggest push. Watch it run on a\n",
        "toy policy: `demonstrate_dpo` optimizes a 5-way softmax where response 0\n",
        "is preferred over response 1, against a frozen uniform reference."
      ],
      "id": "f684360e-b8ff-499b-a63a-1ce925440424"
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [],
      "source": [
        "from alignment import demonstrate_dpo\n",
        "\n",
        "history = demonstrate_dpo(beta=0.1, steps=200, lr=1.0, seed=0)\n",
        "steps = list(range(len(history[\"loss\"])))\n",
        "\n",
        "dpo_curve = [\n",
        "    {\"step\": s, \"chosen_reward\": history[\"chosen_reward\"][s],\n",
        "     \"rejected_reward\": history[\"rejected_reward\"][s],\n",
        "     \"chosen_prob\": history[\"chosen_prob\"][s]}\n",
        "    for s in steps\n",
        "]\n",
        "ojs_define(dpo_curve = dpo_curve)"
      ],
      "id": "5f3ae50c"
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "loss:        0.693 → 0.240\n",
            "chosen reward:   +0.000 → +0.161\n",
            "rejected reward: +0.000 → -1.144\n",
            "chosen P:    0.210 → 0.996"
          ]
        }
      ],
      "source": [
        "print(f\"loss:        {history['loss'][0]:.3f} → {history['loss'][-1]:.3f}\")\n",
        "print(f\"chosen reward:   {history['chosen_reward'][0]:+.3f} → {history['chosen_reward'][-1]:+.3f}\")\n",
        "print(f\"rejected reward: {history['rejected_reward'][0]:+.3f} → {history['rejected_reward'][-1]:+.3f}\")\n",
        "print(f\"chosen P:    {history['chosen_prob'][0]:.3f} → {history['chosen_prob'][-1]:.3f}\")"
      ],
      "id": "400de71e"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "0ac2e1d3-06e0-490c-bbc8-dde9700e0fc8"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The two curves split apart: DPO lifts the chosen response’s implicit\n",
        "reward and drives the rejected one down — the “push the winner up, the\n",
        "loser down” behavior, straight from a supervised loss.\n",
        "\n",
        "## Interactive Exploration\n",
        "\n",
        "The DPO loss depends on just two numbers per pair: how much the policy\n",
        "prefers the chosen answer over the reference,\n",
        "$\\Delta_w = \\log\\frac{\\pi_\\theta(y_w)}{\\pi_{\\text{ref}}(y_w)}$, and the\n",
        "same for the rejected answer, $\\Delta_l$. Drive both (and $\\beta$) and\n",
        "watch the implicit rewards, the loss, and — crucially — the **gradient\n",
        "weight** $\\sigma\\big(\\beta(\\Delta_l - \\Delta_w)\\big)$, which is largest\n",
        "exactly when the pair is ordered wrong."
      ],
      "id": "ab722dae-ea3a-40a6-a415-ee73382325ad"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "1f8b4bde-7e91-429e-9a3a-cbdbc4934faf"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "1575b4dc-7470-4b8f-98fd-41cec978b583"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "358a4184-819c-4603-90d0-1dcfc0700df7"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        },
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "a9068c0f-2f25-44e1-b0e1-be8e795331b9"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "162bb50c-4682-4f4f-bbbf-e1e0a734bd60"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "9f02d1d3-ec87-4b78-8ff2-6f7054209a7b"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Try This**\n",
        ">\n",
        "> 1.  Set $\\Delta_w$ **below** $\\Delta_l$ (drag chosen down, rejected\n",
        ">     up). The loss climbs and the gradient weight heads toward 1 — DPO\n",
        ">     spends its effort on pairs it currently gets wrong.\n",
        "> 2.  Push $\\Delta_w$ far above $\\Delta_l$. Loss collapses toward 0 and\n",
        ">     the gradient weight toward 0: an already-correct pair contributes\n",
        ">     almost nothing.\n",
        "> 3.  Shrink $\\beta$. The implicit rewards compress toward 0 — $\\beta$\n",
        ">     sets how strongly the log-ratios translate into reward, i.e. how\n",
        ">     hard the policy is allowed to move from the reference.\n",
        "\n",
        "## Beyond DPO: IPO, KTO, and ORPO\n",
        "\n",
        "DPO turns a pair $(y_w \\succ y_l)$ into one classification loss on the\n",
        "**margin**\n",
        "\n",
        "$$h = \\big(\\log\\pi_\\theta(y_w) - \\log\\pi_{\\text{ref}}(y_w)\\big)\n",
        "  - \\big(\\log\\pi_\\theta(y_l) - \\log\\pi_{\\text{ref}}(y_l)\\big),\n",
        "\\qquad L_\\text{DPO} = -\\log\\sigma(\\beta\\, h).$$\n",
        "\n",
        "That log-sigmoid never stops rewarding a bigger $h$: with clean, near-\n",
        "deterministic preferences it will drive the log-ratio toward infinity\n",
        "and overfit. The three successors each change **one** thing about DPO —\n",
        "and each is a one-function edit on the *same* margin you already built\n",
        "(`preference_margin` in `dpo_variants.py`).\n",
        "\n",
        "> **Key Insight**\n",
        ">\n",
        "> DPO opened a design space with three knobs: the **loss shape** (IPO\n",
        "> changes it), the **data** (KTO changes it), and the **machinery**\n",
        "> (ORPO changes it). All three still do the one thing DPO does — raise\n",
        "> the chosen answer, lower the rejected one.\n",
        "\n",
        "### IPO: A Bounded Target\n",
        "\n",
        "**Identity Preference Optimization** (Azar et al., 2023) keeps DPO’s\n",
        "margin but swaps the loss shape: instead of an ever-decreasing\n",
        "log-sigmoid, minimize the *squared distance* to a **finite** target\n",
        "margin $\\frac{1}{2\\tau}$:\n",
        "\n",
        "$$L_\\text{IPO} = \\Big(h - \\tfrac{1}{2\\tau}\\Big)^2.$$\n",
        "\n",
        "DPO’s gradient never reaches zero, so it keeps pushing; IPO’s gradient\n",
        "crosses zero exactly at $h = \\frac{1}{2\\tau}$ and **stops**. Watch the\n",
        "two loss shapes and their gradients side by side — DPO slides forever,\n",
        "IPO settles into its target:"
      ],
      "id": "b0bb6cff-4b24-47fa-86d7-d3ea46257c9b"
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "IPO target margin: 5.0\n",
            "IPO loss when margin == target: 0.0"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "from dpo_variants import ipo_loss, demonstrate_dpo_variants\n",
        "\n",
        "# Same margin, two losses. tau=0.1 -> IPO target margin 1/(2*0.1) = 5.0\n",
        "curves = demonstrate_dpo_variants(tau=0.1, beta=1.0)\n",
        "print(\"IPO target margin:\", curves[\"ipo_target\"][0])\n",
        "loss, _, _ = ipo_loss(torch.tensor([0.0]), torch.tensor([-5.0]),\n",
        "                      torch.tensor([0.0]), torch.tensor([0.0]), tau=0.1)\n",
        "print(\"IPO loss when margin == target:\", round(loss.item(), 4))"
      ],
      "id": "c148326d"
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Bridge the loss-vs-margin sweep to the plot below.\n",
        "variant_curve = [\n",
        "    {\"margin\": m, \"dpo\": d, \"ipo\": i, \"dpo_grad\": dg, \"ipo_grad\": ig}\n",
        "    for m, d, i, dg, ig in zip(curves[\"margin\"], curves[\"dpo\"], curves[\"ipo\"],\n",
        "                               curves[\"dpo_grad\"], curves[\"ipo_grad\"])\n",
        "]\n",
        "ipo_target = curves[\"ipo_target\"][0]\n",
        "ojs_define(variant_curve = variant_curve, ipo_target = ipo_target)"
      ],
      "id": "3d678d29"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "2e029fa0-8e07-49f0-8317-8660da61be60"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Try This**\n",
        ">\n",
        "> The DPO curve (orange) keeps sliding down as the margin grows — there\n",
        "> is no margin it considers “enough,” which is how DPO overfits crisp\n",
        "> preferences. IPO (blue) is a parabola with its floor at the green\n",
        "> target line: push past it and IPO pulls the margin *back*. Smaller\n",
        "> $\\tau$ moves the target further right (a wider chosen–rejected gap);\n",
        "> it is the direct dial on separation strength.\n",
        "\n",
        "### KTO: Learning from Unpaired Labels\n",
        "\n",
        "DPO and IPO need **pairs** — the same prompt answered two ways, ranked.\n",
        "But most feedback in the wild is a single thumbs-up or thumbs-down on\n",
        "*one* answer. **Kahneman-Tversky Optimization** (Ethayarajh et al.,\n",
        "2024) drops the pairing: each example carries one bit, *desirable* or\n",
        "*undesirable*, and a prospect-theory value pushes desirable rewards\n",
        "**above** a reference point $z_0$ and undesirable ones **below** it:\n",
        "\n",
        "$$v = \\begin{cases}\n",
        "      \\lambda_D\\,\\sigma\\big(\\beta(r - z_0)\\big) & \\text{desirable}\\\\\n",
        "      \\lambda_U\\,\\sigma\\big(\\beta(z_0 - r)\\big) & \\text{undesirable}\n",
        "    \\end{cases},\n",
        "\\qquad r = \\log\\frac{\\pi_\\theta(y|x)}{\\pi_{\\text{ref}}(y|x)},\n",
        "\\qquad L_\\text{KTO} = \\lambda_y - v.$$\n",
        "\n",
        "$z_0$ is the batch’s average KL (clamped at 0, **no gradient**) — the\n",
        "“neutral” reward each example is measured against."
      ],
      "id": "4394ec03-5d99-45ff-84ea-2ed4deaa9350"
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "implicit rewards β·(πθ−ref): [0.05000000074505806, -0.019999999552965164]\n",
            "KTO loss: 0.4913"
          ]
        }
      ],
      "source": [
        "from dpo_variants import kto_loss\n",
        "\n",
        "# One desirable and one undesirable example — no pairing needed.\n",
        "policy_logps = torch.tensor([-0.5, -0.8])   # log πθ(y|x)\n",
        "ref_logps    = torch.tensor([-1.0, -0.6])   # log π_ref(y|x)\n",
        "desirable    = torch.tensor([True, False])  # thumbs up / thumbs down\n",
        "\n",
        "loss, rewards = kto_loss(policy_logps, ref_logps, desirable, beta=0.1)\n",
        "print(\"implicit rewards β·(πθ−ref):\", rewards.round(decimals=3).tolist())\n",
        "print(\"KTO loss:\", round(loss.item(), 4))"
      ],
      "id": "994f4abb"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Key Insight**\n",
        ">\n",
        "> KTO trades DPO’s *pairwise* signal for a *pointwise* one, so it runs\n",
        "> on the abundant, cheap feedback (a like, a flag) that never came as a\n",
        "> clean A-vs-B comparison. The reference point $z_0$ is what lets an\n",
        "> unpaired label still say “better/worse than typical.”\n",
        "\n",
        "### ORPO: Dropping the Reference Model\n",
        "\n",
        "DPO, IPO, and KTO all need a **frozen reference** $\\pi_{\\text{ref}}$ — a\n",
        "second forward pass every step. **Odds Ratio Preference Optimization**\n",
        "(Hong et al., 2024) removes it: fold preference straight into SFT as the\n",
        "ordinary NLL on the chosen answer, plus a small **odds-ratio** penalty\n",
        "that separates chosen from rejected — using the policy *alone*:\n",
        "\n",
        "$$L_\\text{ORPO} = \\underbrace{-\\log P_\\theta(y_w)}_{\\text{SFT}}\n",
        "  + \\lambda\\underbrace{\\Big(\\!-\\log\\sigma\\big(\\log\\tfrac{\\text{odds}_\\theta(y_w)}{\\text{odds}_\\theta(y_l)}\\big)\\Big)}_{\\text{odds ratio}},\n",
        "\\qquad \\text{odds}(y) = \\frac{P(y)}{1 - P(y)}.$$\n",
        "\n",
        "Here $P_\\theta(y)$ is the **length-normalized** (average per-token)\n",
        "probability, so the inputs are *mean* log-probs. One model, one stage,\n",
        "no reference:"
      ],
      "id": "239ac0c3-bc90-4494-926a-37d56fdcfc89"
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "SFT part: 0.4000   odds-ratio part: 0.0929\n",
            "ORPO loss = SFT + λ·OR = 0.4929"
          ]
        }
      ],
      "source": [
        "from dpo_variants import orpo_loss\n",
        "\n",
        "# Length-normalized (average) log-probs — NOTE: no reference model anywhere.\n",
        "chosen_logps   = torch.tensor([-0.4])   # log P̄θ(y_w)\n",
        "rejected_logps = torch.tensor([-1.8])   # log P̄θ(y_l)\n",
        "\n",
        "loss, sft, orl = orpo_loss(chosen_logps, rejected_logps, lambda_or=1.0)\n",
        "print(f\"SFT part: {sft.item():.4f}   odds-ratio part: {orl.item():.4f}\")\n",
        "print(f\"ORPO loss = SFT + λ·OR = {loss.item():.4f}\")"
      ],
      "id": "9ba1a1d6"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Which One?**\n",
        ">\n",
        "> All four reduce to “raise the chosen, lower the rejected,” and differ\n",
        "> only in what they give up: **DPO** — the reward model; **IPO** — DPO’s\n",
        "> unbounded margin (a finite target instead); **KTO** — the need for\n",
        "> pairs (unpaired labels); **ORPO** — the reference model (fused into\n",
        "> SFT). Reach for KTO when your data is unpaired thumbs-up/down, IPO\n",
        "> when DPO overfits, ORPO when you want a single training stage.\n",
        "\n",
        "## PPO: The RL Detour DPO Removes\n",
        "\n",
        "We wrote the RLHF objective and its closed-form optimum\n",
        "$\\pi^* \\propto \\pi_{\\text{ref}}\\exp(r/\\beta)$, then took the **DPO**\n",
        "shortcut straight past the reinforcement learning. But that RL is what\n",
        "every pre-DPO aligned model (InstructGPT, the first ChatGPT) actually\n",
        "ran, and seeing it makes clear *what* DPO saves you from. The algorithm\n",
        "is **PPO** (Proximal Policy Optimization), and its loop has four steps:"
      ],
      "id": "3d639f97-d382-415d-b176-8339904017f9"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "ee501d71-4249-4977-861d-c0f74978e725"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "e87b1ddd-b132-4667-82cc-c3471177eb2c"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "781e59e4-69a0-4762-8b91-de0486640d8b"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The one structural difference from GRPO (m13) is **step 3**: PPO learns\n",
        "a **value network** $V$ (the *critic*) to estimate how good a response\n",
        "was expected to be, and the advantage is the surprise\n",
        "$A = \\text{return} - V$. GRPO throws the critic away and uses a group’s\n",
        "mean reward as the baseline — cheaper, but PPO’s critic is the classic\n",
        "choice. The surrogate (step 4) is the *same* clipped objective GRPO\n",
        "uses:\n",
        "\n",
        "$$L = -\\,\\mathbb{E}\\!\\left[\\min\\big(\\rho A,\\ \\text{clip}(\\rho, 1-\\epsilon, 1+\\epsilon)\\,A\\big)\\right]\n",
        "    + c_v\\,(V - \\text{return})^2, \\qquad \\rho = \\frac{\\pi_\\theta(a)}{\\pi_{\\theta_{\\text{old}}}(a)}.$$\n",
        "\n",
        "## Code: PPO on a Toy Policy\n",
        "\n",
        "`ppo.py` builds each piece. Advantages come from **GAE** (generalized\n",
        "advantage estimation) — a bias/variance knob over the value baseline —\n",
        "here on a one-step episode, so $A = \\text{return} - V$:"
      ],
      "id": "78702b21-7090-48e7-95c9-ffdb6ccb4a4c"
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "advantage: [0.75]\n",
            "clipped surrogate: -1.2"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "from ppo import compute_gae, clipped_surrogate, ActorCritic\n",
        "\n",
        "# GAE on one terminal step: reward 1, the critic expected 0.25 -> advantage 0.75.\n",
        "print(\"advantage:\", compute_gae([1.0], [0.25]))\n",
        "\n",
        "# The clipped surrogate is a loss to minimize; a huge ratio is capped at (1+eps)*A.\n",
        "lp_old = torch.tensor([-2.0]); lp_new = torch.tensor([0.0])  # ratio e^2 >> 1.2\n",
        "print(\"clipped surrogate:\", round(float(clipped_surrogate(lp_new, lp_old, torch.tensor([1.0]))), 3))"
      ],
      "id": "e1d168db"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The policy is `ActorCritic` — m13’s `ToyReasoningPolicy` with the value\n",
        "network added back: a categorical over `K` answers *plus* a scalar\n",
        "critic. `ppo_step` runs one full update (rollout → KL-shaped reward →\n",
        "advantage → clipped surrogate + value loss). Train it against a reward\n",
        "model and watch the reward climb:"
      ],
      "id": "1bb77af9-8f83-4f3b-9031-a5a0b154e723"
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "============================================================\n",
            "PPO on a toy policy: proxy reward vs. the true reward\n",
            "============================================================\n",
            "  K=6, good=1, hackable=4, beta=0.0\n",
            "\n",
            "   step    proxy     true       kl\n",
            "      0    0.300    0.167    0.000\n",
            "      7    0.453    0.234    0.049\n",
            "     14    0.703    0.312    0.385\n",
            "     21    0.853    0.266    0.762\n",
            "     28    0.916    0.109    1.135\n",
            "     35    0.916    0.109    1.120\n",
            "     42    0.900    0.031    1.179\n",
            "     49    0.978    0.031    1.546\n",
            "     56    1.000    0.000    1.732\n",
            "\n",
            "  proxy 0.300 -> 0.950, true 0.167 -> 0.016"
          ]
        }
      ],
      "source": [
        "from ppo import demonstrate_ppo\n",
        "\n",
        "# Train with a weak KL leash (beta=0). The PROXY reward model is what we optimize;\n",
        "# the TRUE reward is what we actually want (tracked, never optimized).\n",
        "history = demonstrate_ppo(steps=60, beta=0.0, seed=0)"
      ],
      "id": "d23c44ae"
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {},
      "outputs": [],
      "source": [],
      "id": "f0ebe09c"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "a54b6954-1b6f-40ce-87f4-38a73a5cb210"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Look at what happened: the **proxy** reward (blue) climbs to 1.0, but\n",
        "the **true** reward (orange) rises for a while and then *collapses back\n",
        "toward zero*. The policy learned to maximize the reward model — by\n",
        "exploiting an answer the reward model over-rates but that isn’t actually\n",
        "what we wanted. That is **reward over-optimization**.\n",
        "\n",
        "## Reward Over-Optimization\n",
        "\n",
        "The reward model is a *proxy* for human preference, and every proxy has\n",
        "flaws. Push the policy hard enough against it and the policy stops\n",
        "improving on the real goal and starts exploiting the proxy’s mistakes —\n",
        "**Goodhart’s law**: *when a measure becomes a target, it ceases to be a\n",
        "good measure*. The lever that controls this is the **KL leash** to the\n",
        "reference (the `beta` above): the further the policy is allowed to\n",
        "drift, the more it over-optimizes. Sweep `beta` and plot the **true\n",
        "reward against the KL** the policy reached (Gao et al., 2022):"
      ],
      "id": "41d1f1dc-0f9a-4e71-92f7-e5dafaf697a1"
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "============================================================\n",
            "Reward over-optimization: true reward vs. KL leash (beta)\n",
            "============================================================\n",
            "    beta    proxy     true       kl\n",
            "    0.00    1.000    0.000    1.756\n",
            "    0.10    0.953    0.078    1.338\n",
            "    0.20    0.906    0.234    1.003\n",
            "    0.40    0.797    0.391    0.620\n",
            "    0.80    0.537    0.281    0.140\n",
            "    1.60    0.409    0.219    0.030\n",
            "\n",
            "  Weak leash (small beta): proxy up, KL up, TRUE reward collapses."
          ]
        }
      ],
      "source": [
        "from ppo import demonstrate_overoptimization\n",
        "\n",
        "rows = demonstrate_overoptimization(steps=80, seed=0)"
      ],
      "id": "e622f01c"
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {},
      "outputs": [],
      "source": [],
      "id": "f6c8fb48"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "a1fc55ab-b1c2-4702-92d5-77808d884ed4"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The **proxy** reward (blue, dashed) rises monotonically as the policy\n",
        "drifts — more optimization always looks better *by the proxy*. But the\n",
        "**true** reward (orange) is a **hump**: it rises with a little\n",
        "optimization, peaks at a moderate KL, then falls as the policy\n",
        "over-optimizes the proxy’s flaws. The best model is not the one with the\n",
        "highest reward-model score — it’s the one at the top of the hump.\n",
        "\n",
        "> **Key Insight**\n",
        ">\n",
        "> This hump is *why DPO is attractive*. RLHF-PPO has to walk the\n",
        "> reward-vs-KL curve carefully — pick β wrong and you ship a\n",
        "> reward-hacked model — on top of running a full on-policy RL loop with\n",
        "> a reward model and a critic. DPO trades the whole loop for one\n",
        "> supervised loss on fixed preference pairs. It doesn’t make Goodhart go\n",
        "> away (your *preference data* is still a proxy), but it removes the\n",
        "> moving target the RL detour has to chase.\n",
        "\n",
        "> **Try This**\n",
        ">\n",
        "> 1.  **Find the peak.** In the sweep above, read off the β at the top\n",
        ">     of the true-reward hump. Re-run `demonstrate_overoptimization`\n",
        ">     with `betas` clustered around it to locate the compute-optimal KL\n",
        ">     more precisely.\n",
        "> 2.  **Break the leash.** Set `beta=0.0` in `demonstrate_ppo` and then\n",
        ">     a large `beta=0.8`. Watch the true reward collapse in the first\n",
        ">     and survive in the second — the same trade the DPO pitfall (β too\n",
        ">     small) warned about, now made visible.\n",
        "\n",
        "## Common Pitfalls\n",
        "\n",
        "1.  **Forgetting to mask the prompt in SFT.** Training on prompt tokens\n",
        "    wastes capacity teaching the model to echo inputs. Only completion\n",
        "    tokens belong in the loss (`completion_mask`).\n",
        "2.  **Letting the reference model drift.** $\\pi_{\\text{ref}}$ must be a\n",
        "    **frozen** copy of the SFT model. If you accidentally update it (or\n",
        "    forget `torch.no_grad` / `.detach()` on its log-probs), the KL\n",
        "    anchor moves and the implicit reward becomes meaningless.\n",
        "3.  **β too small or too large.** Small $\\beta$ lets the policy drift\n",
        "    far from the reference — faster preference-fitting but risk of\n",
        "    degenerate, reward-hacked text. Large $\\beta$ barely moves the\n",
        "    policy. Typical DPO runs use $\\beta \\in [0.1, 0.5]$.\n",
        "4.  **Comparing raw log-probs across different lengths.**\n",
        "    `sequence_logprob` sums over completion tokens, so longer\n",
        "    completions have more-negative log-probs. Always compare chosen\n",
        "    vs. rejected for the *same* prompt; DPO’s log-ratio against the\n",
        "    reference largely cancels the length effect, but a badly unbalanced\n",
        "    pair can still bias it.\n",
        "5.  **Reward hacking is not solved, only reframed.** DPO removes the\n",
        "    *separate* reward model, but the preference *data* is still a proxy\n",
        "    for what you want. Garbage preferences in, misaligned policy out.\n",
        "6.  **Feeding ORPO summed log-probs.** ORPO’s odds need\n",
        "    $P_\\theta(y) \\in (0,1)$, so it requires **length-normalized** (mean\n",
        "    per-token) log-probs — pass `sequence_logprob(..., average=True)`.\n",
        "    Summed log-probs make $P > 1$ possible and `log(1 - P)` undefined.\n",
        "7.  **Backpropagating through KTO’s $z_0$.** The reference point must be\n",
        "    `.detach()`ed (it only controls loss saturation). If gradients flow\n",
        "    through the batch-mean KL, the objective couples every example in\n",
        "    the batch and training destabilizes.\n",
        "8.  **Reading the reward model’s score as truth (PPO).** The reward\n",
        "    model is a proxy; its score rising is *not* proof the model got\n",
        "    better. Track a held-out true metric and watch for the\n",
        "    over-optimization hump — the best checkpoint is at the peak, not the\n",
        "    end.\n",
        "9.  **Optimizing without a KL leash (PPO).** With $\\beta \\to 0$ the\n",
        "    policy is free to reward-hack the proxy and collapse the true reward\n",
        "    (and often the text quality). The KL-to-reference penalty — folded\n",
        "    into the reward here — is what keeps RLHF-PPO from running off the\n",
        "    rails.\n",
        "\n",
        "## Exercises\n",
        "\n",
        "### Exercise 1: The SFT mask matters"
      ],
      "id": "ecf8faf5-c3f7-4e33-9675-ddb7a22d3415"
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {},
      "outputs": [],
      "source": [
        "import torch\n",
        "from alignment import sft_loss\n",
        "\n",
        "# Build a (1, 8) example. Compute the SFT loss twice: once over all tokens, once\n",
        "# over only the last 4 (the \"completion\"). Confirm they differ, and explain which\n",
        "# one trains the assistant behavior.\n",
        "\n",
        "# Your implementation here:\n",
        "# logits = torch.randn(1, 8, 30)\n",
        "# labels = torch.randint(0, 30, (1, 8))\n",
        "# ..."
      ],
      "id": "a443aee6"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Exercise 2: Bradley-Terry by hand"
      ],
      "id": "ffc8b984-fd5b-4388-88c6-d15428dbf833"
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {},
      "outputs": [],
      "source": [
        "import torch\n",
        "from alignment import bradley_terry_loss\n",
        "\n",
        "# For chosen reward 2.0 and rejected reward 0.0, verify by hand that the loss is\n",
        "# -log(sigmoid(2.0)). Then show the loss equals log(2) when the two rewards match.\n",
        "\n",
        "# Your implementation here:"
      ],
      "id": "24f0599d"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Exercise 3: DPO endpoints"
      ],
      "id": "925fc20a-8ad5-4d87-a46d-c1d700ab3e74"
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {},
      "outputs": [],
      "source": [
        "import torch\n",
        "from alignment import dpo_loss\n",
        "\n",
        "# Show that when the policy equals the reference (all four log-probs equal), the\n",
        "# DPO loss is exactly log(2) and both implicit rewards are 0 — DPO starts from\n",
        "# \"no preference\" and has to earn the margin.\n",
        "\n",
        "# Your implementation here:"
      ],
      "id": "2e8c7e72"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Exercise 4: IPO stops, DPO doesn’t"
      ],
      "id": "2577d6d6-a0de-4936-9a81-917178abb185"
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {},
      "outputs": [],
      "source": [
        "import torch\n",
        "from dpo_variants import ipo_loss, preference_margin\n",
        "from alignment import dpo_loss\n",
        "\n",
        "# For a pair whose margin is already huge (chosen far above rejected), compare the\n",
        "# gradient magnitude DPO vs IPO produces on the chosen log-prob. Confirm IPO's is\n",
        "# ~0 once the margin passes 1/(2τ) while DPO's stays positive — IPO knows when to\n",
        "# stop. (Set requires_grad on the chosen log-prob and inspect .grad.)\n",
        "\n",
        "# Your implementation here:"
      ],
      "id": "7f25c456"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Exercise 5: The over-optimization peak"
      ],
      "id": "218adabb-bac8-4992-b437-55a0fe603c34"
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {},
      "outputs": [],
      "source": [
        "from ppo import demonstrate_overoptimization\n",
        "\n",
        "# demonstrate_overoptimization sweeps the KL weight beta and reports final (proxy, true,\n",
        "# kl). Find the beta whose TRUE reward is highest, then sweep a finer grid of betas\n",
        "# around it to pin down the peak. Argue why the peak — not the highest proxy reward —\n",
        "# is the checkpoint you would actually ship.\n",
        "\n",
        "# Your implementation here:"
      ],
      "id": "742f2e7d"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Summary\n",
        "\n",
        "Key takeaways:\n",
        "\n",
        "1.  **Alignment closes the likelihood-vs-helpfulness gap** — a\n",
        "    pretrained model maximizes plausibility; SFT, reward modeling, and\n",
        "    DPO/RLHF make it *helpful*.\n",
        "2.  **SFT is masked next-token loss** — the m07 loss with the prompt\n",
        "    masked out, so the model learns to produce the completion, not echo\n",
        "    the prompt.\n",
        "3.  **Reward models learn from comparisons** — the Bradley-Terry loss\n",
        "    $-\\log\\sigma(r_w - r_l)$ trains a scalar-headed backbone to rank\n",
        "    answers; only reward *differences* matter.\n",
        "4.  **RLHF has a closed-form optimum** —\n",
        "    $\\pi^* \\propto \\pi_{\\text{ref}}\\exp(r/\\beta)$ — but its partition\n",
        "    function $Z(x)$ is intractable, so PPO approximates it with\n",
        "    on-policy RL.\n",
        "5.  **DPO deletes the reward model and the RL loop** — solving the\n",
        "    optimum for the reward gives\n",
        "    $\\hat r = \\beta\\log(\\pi_\\theta/\\pi_{\\text{ref}})$; substituted into\n",
        "    Bradley-Terry, the intractable term cancels and alignment becomes\n",
        "    one supervised loss on preference pairs.\n",
        "6.  **The DPO gradient is self-prioritizing** — it lifts the chosen\n",
        "    answer and lowers the rejected one, weighted by how wrong the\n",
        "    current ranking is.\n",
        "7.  **DPO opened a family, each dropping one thing** — **IPO** swaps the\n",
        "    unbounded log-sigmoid for a squared loss toward a finite target\n",
        "    margin $\\frac{1}{2\\tau}$; **KTO** drops the pairs to learn from\n",
        "    unpaired desirable/undesirable labels; **ORPO** drops the reference\n",
        "    model, fusing an odds-ratio penalty into SFT. All still raise the\n",
        "    chosen answer and lower the rejected one.\n",
        "8.  **PPO is the RL detour DPO removes** — rollout → KL-shaped reward →\n",
        "    advantage from a **critic** → clipped surrogate. It optimizes the\n",
        "    same objective directly, at the cost of an on-policy loop, a reward\n",
        "    model, and a value network (the critic m13’s GRPO drops). And it\n",
        "    must be walked carefully: **reward over-optimization** means the\n",
        "    *true* reward humps then falls as the policy over-fits the proxy —\n",
        "    the KL leash sets where you land on that hump.\n",
        "\n",
        "## What’s Next\n",
        "\n",
        "You now have the alignment toolkit: imitate with SFT, learn preferences\n",
        "with a reward model, and optimize them directly with DPO. From here the\n",
        "frontier builds on it — [Module 13: Reasoning & Test-Time\n",
        "Compute](../m13_reasoning/lesson.qmd) spends more compute at inference\n",
        "(chain-of-thought, self-consistency, best-of-N) to get better answers,\n",
        "and its sequential counterpart trains **reasoning models** with RL\n",
        "against verifiable rewards (o1/R1-style). The preference methods beyond\n",
        "DPO you just built — **IPO, KTO, ORPO** — reshape the same log-ratio\n",
        "loss, each dropping one of DPO’s ingredients (its unbounded margin, its\n",
        "pairs, or its reference model).\n",
        "\n",
        "### Going Deeper\n",
        "\n",
        "**Core Papers:**\n",
        "\n",
        "- [Direct Preference Optimization: Your Language Model is Secretly a\n",
        "  Reward Model](https://arxiv.org/abs/2305.18290) — Rafailov et\n",
        "  al. (2023), the DPO loss and its derivation.\n",
        "- [Proximal Policy Optimization\n",
        "  Algorithms](https://arxiv.org/abs/1707.06347) — Schulman et\n",
        "  al. (2017), **PPO** — the clipped surrogate built here.\n",
        "- [High-Dimensional Continuous Control Using Generalized Advantage\n",
        "  Estimation](https://arxiv.org/abs/1506.02438) — Schulman et\n",
        "  al. (2015), **GAE** — the advantage estimator.\n",
        "- [Scaling Laws for Reward Model\n",
        "  Overoptimization](https://arxiv.org/abs/2210.10760) — Gao, Schulman &\n",
        "  Hilton (2022), the true-reward-vs-KL **over-optimization** hump.\n",
        "- [Training language models to follow instructions with human\n",
        "  feedback](https://arxiv.org/abs/2203.02155) — Ouyang et al. (2022),\n",
        "  InstructGPT: the SFT → reward model → PPO pipeline.\n",
        "- [Deep reinforcement learning from human\n",
        "  preferences](https://arxiv.org/abs/1706.03741) — Christiano et\n",
        "  al. (2017), the Bradley-Terry reward model from comparisons.\n",
        "- [Learning to summarize from human\n",
        "  feedback](https://arxiv.org/abs/2009.01325) — Stiennon et al. (2020),\n",
        "  the KL-regularized RLHF objective in practice.\n",
        "- [A General Theoretical Paradigm to Understand Learning from Human\n",
        "  Preferences](https://arxiv.org/abs/2310.12036) — Azar et al. (2023),\n",
        "  **IPO** — the squared-loss bounded-margin fix for DPO’s overfitting.\n",
        "- [KTO: Model Alignment as Prospect Theoretic\n",
        "  Optimization](https://arxiv.org/abs/2402.01306) — Ethayarajh et\n",
        "  al. (2024), **KTO** — preference from unpaired desirable/undesirable\n",
        "  labels.\n",
        "- [ORPO: Monolithic Preference Optimization without Reference\n",
        "  Model](https://arxiv.org/abs/2403.07691) — Hong et al. (2024),\n",
        "  **ORPO** — the reference-free odds-ratio penalty fused into SFT.\n",
        "\n",
        "**Practical Resources:**\n",
        "\n",
        "- [Constitutional AI: Harmlessness from AI\n",
        "  Feedback](https://arxiv.org/abs/2212.08073) — Bai et al. (2022), RLAIF\n",
        "  — preferences from a model instead of humans.\n",
        "- [The Alignment\n",
        "  Handbook](https://github.com/huggingface/alignment-handbook) —\n",
        "  reference SFT + DPO training recipes."
      ],
      "id": "dd8710a2-5a46-44e0-be82-fd883b1959a2"
    }
  ],
  "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"
    }
  }
}