{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preference tuning, in three tables\n",
    "\n",
    "Bradley-Terry turns a ranking into a loss, the KL penalty stops reward hacking, and DPO does it without a reward model.\n",
    "\n",
    "From chapter 20, [Adaptation and Alignment](https://nlp.jcrlabz.com/book/alignment/), of the course notes.\n",
    "\n",
    "Source: `book/code/worked_examples/alignment.py`\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What you need\n",
    "\n",
    "Nothing. This example uses only the Python standard library.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Nothing to install: this example uses only the Python standard library.\n",
    "import sys; print(sys.version.split()[0])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The script\n",
    "\n",
    "Everything the example defines, in one cell. Run it and the definitions\n",
    "are live in the notebook.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Preference tuning, worked as Chapter 20 works it.\n",
    "\n",
    "Three things, all reproducing the chapter's tables:\n",
    "\n",
    "    the reward model      Bradley-Terry turns a ranking into a loss\n",
    "    the KL penalty        why reward hacking needs a leash, in numbers\n",
    "    DPO                   the same preference, with no reward model\n",
    "\n",
    "    python3 alignment.py          # all three\n",
    "    python3 alignment.py --bt     # Bradley-Terry and the reward-model loss\n",
    "    python3 alignment.py --kl     # reward hacking and the penalty that stops it\n",
    "    python3 alignment.py --dpo    # direct preference optimisation\n",
    "\n",
    "Install:  nothing, the Python standard library is enough\n",
    "\"\"\"\n",
    "\n",
    "import argparse\n",
    "import math\n",
    "\n",
    "\n",
    "def sigmoid(z):\n",
    "    return 1.0 / (1.0 + math.exp(-z))\n",
    "\n",
    "\n",
    "# ------------------------------------------------------------ Bradley-Terry\n",
    "# Four preference pairs. r_w is the reward the model gives the response the\n",
    "# human preferred, r_l the one it gives the rejected response. The model is\n",
    "# right on the first two, wrong on the third, and undecided on the fourth.\n",
    "PAIRS = [\n",
    "    (\"clear answer\",      1.20, 0.40),\n",
    "    (\"answer with steps\", 2.10, 0.30),\n",
    "    (\"confident bluff\",   0.30, 1.10),\n",
    "    (\"two paraphrases\",   0.75, 0.75),\n",
    "]\n",
    "\n",
    "\n",
    "def show_bt():\n",
    "    print(\"\\nSTEP 1  the reward model learns from comparisons\\n\")\n",
    "    print(\"  Nobody can score helpfulness on an absolute scale, but people can\")\n",
    "    print(\"  say which of two answers they prefer. Bradley-Terry turns that\")\n",
    "    print(\"  comparison into a probability:\\n\")\n",
    "    print(\"      P(y_w > y_l) = sigma(r_w - r_l)\\n\")\n",
    "    print(\"  pair                 r_w     r_l   r_w-r_l   P(y_w>y_l)    loss\")\n",
    "    print(\"  \" + \"-\" * 64)\n",
    "    total = 0.0\n",
    "    for name, rw, rl in PAIRS:\n",
    "        d = rw - rl\n",
    "        p = sigmoid(d)\n",
    "        loss = -math.log(p)\n",
    "        total += loss\n",
    "        print(f\"  {name:18} {rw:5.2f}   {rl:5.2f}   {d:+7.2f}   {p:9.4f}  {loss:7.4f}\")\n",
    "    print(\"  \" + \"-\" * 64)\n",
    "    print(f\"  mean loss over the four pairs{' ' * 29}{total / len(PAIRS):7.4f}\\n\")\n",
    "    print(\"  Read the third row. The model prefers the bluff, so the pair costs\")\n",
    "    print(f\"  {-math.log(sigmoid(0.30 - 1.10)):.4f}, more than three times the first row. Training\")\n",
    "    print(\"  pushes r_w up and r_l down until that ordering flips.\\n\")\n",
    "    print(\"  Read the fourth. Equal rewards give exactly 0.5 and a loss of\")\n",
    "    print(f\"  {-math.log(0.5):.4f} = log 2, the most a single comparison can tell you\")\n",
    "    print(\"  when the model has no opinion.\")\n",
    "\n",
    "\n",
    "# ------------------------------------------------------------- the KL leash\n",
    "# Two candidate responses the policy could learn to produce. The first scores\n",
    "# well with the reward model and badly with people: that is reward hacking.\n",
    "CANDIDATES = [\n",
    "    (\"honest answer\",  4.00,  1.00),\n",
    "    (\"reward hack\",    5.00, 12.00),\n",
    "]\n",
    "BETAS = [0.00, 0.05, 0.10, 0.30]\n",
    "\n",
    "\n",
    "def show_kl():\n",
    "    print(\"\\nSTEP 2  the penalty is not an afterthought\\n\")\n",
    "    print(\"  The policy is optimised against the reward model, not against\")\n",
    "    print(\"  people. Left alone it finds text the reward model loves and humans\")\n",
    "    print(\"  do not. The objective keeps it on a leash:\\n\")\n",
    "    print(\"      maximise  E[r(x,y)]  -  beta * KL(pi_theta || pi_ref)\\n\")\n",
    "    print(\"  candidate          r      KL\" + \"\".join(f\"   beta={b:.2f}\" for b in BETAS))\n",
    "    print(\"  \" + \"-\" * (31 + 11 * len(BETAS)))\n",
    "    scores = {}\n",
    "    for name, r, kl in CANDIDATES:\n",
    "        row = [r - b * kl for b in BETAS]\n",
    "        scores[name] = row\n",
    "        cells = \"\".join(f\"   {v:9.2f}\" for v in row)\n",
    "        print(f\"  {name:15} {r:5.2f}  {kl:6.2f}{cells}\")\n",
    "    print(\"  \" + \"-\" * (31 + 11 * len(BETAS)))\n",
    "    print(\"\\n  winner per beta:  \", end=\"\")\n",
    "    flip = None\n",
    "    for i, b in enumerate(BETAS):\n",
    "        best = max(scores, key=lambda n: scores[n][i])\n",
    "        print(f\"beta={b:.2f} -> {best}   \", end=\"\")\n",
    "        if best == \"honest answer\" and flip is None:\n",
    "            flip = b\n",
    "    print()\n",
    "    print(f\"\\n  With no penalty the hack wins, because 5.00 beats 4.00 and the\")\n",
    "    print(f\"  reward model is the only judge. The ordering flips at beta = {flip:.2f},\")\n",
    "    print(\"  where the hack's distance from the starting model finally costs\")\n",
    "    print(\"  more than the extra reward it buys.\\n\")\n",
    "    print(\"  Too large a beta is its own failure. The policy then stays so close\")\n",
    "    print(\"  to the reference that preference training changes nothing.\")\n",
    "\n",
    "\n",
    "# ---------------------------------------------------------------------- DPO\n",
    "# Log-probabilities of the preferred and rejected response under the policy\n",
    "# being trained and under the frozen reference it started from.\n",
    "BETA_DPO = 0.1\n",
    "LOGP = [\n",
    "    (\"at the start\",   -2.30, -2.30, -2.20, -2.20),\n",
    "    (\"after one step\", -2.00, -2.30, -2.50, -2.20),\n",
    "    (\"over-corrected\", -1.40, -2.30, -3.40, -2.20),\n",
    "]\n",
    "\n",
    "\n",
    "def show_dpo():\n",
    "    print(\"\\nSTEP 3  the same preference, with no reward model\\n\")\n",
    "    print(\"  DPO drops the reward model and the reinforcement learning loop. It\")\n",
    "    print(\"  reads an implicit reward straight off the policy's own log-ratio:\\n\")\n",
    "    print(\"      r_hat(y) = beta * ( log pi_theta(y) - log pi_ref(y) )\")\n",
    "    print(\"      loss     = -log sigma( r_hat(y_w) - r_hat(y_l) )\\n\")\n",
    "    print(f\"  beta = {BETA_DPO}\\n\")\n",
    "    print(\"  stage             r_hat(y_w)  r_hat(y_l)   margin    sigma     loss\")\n",
    "    print(\"  \" + \"-\" * 66)\n",
    "    for name, tw, rw_, tl, rl_ in LOGP:\n",
    "        rhat_w = BETA_DPO * (tw - rw_)\n",
    "        rhat_l = BETA_DPO * (tl - rl_)\n",
    "        margin = rhat_w - rhat_l\n",
    "        s = sigmoid(margin)\n",
    "        loss = -math.log(s)\n",
    "        print(f\"  {name:16} {rhat_w:+10.4f}  {rhat_l:+10.4f}  {margin:+7.4f}  {s:7.4f}  {loss:7.4f}\")\n",
    "    print(\"  \" + \"-\" * 66)\n",
    "    print(\"\\n  At the start the policy is the reference, so both implicit rewards\")\n",
    "    print(f\"  are 0, the margin is 0 and the loss is log 2 = {-math.log(0.5):.4f}. That is the\")\n",
    "    print(\"  same number the undecided Bradley-Terry pair gave in Step 1, and for\")\n",
    "    print(\"  the same reason: no opinion yet.\\n\")\n",
    "    print(\"  One step raises the preferred response and lowers the rejected one,\")\n",
    "    print(\"  and the loss falls. The third row is the warning. Pushing the margin\")\n",
    "    print(\"  further keeps shrinking the loss, but the policy has now moved a long\")\n",
    "    print(\"  way from the reference, which is what beta is there to price.\")\n",
    "\n",
    "\n",
    "def main():\n",
    "    ap = argparse.ArgumentParser(description=__doc__)\n",
    "    ap.add_argument(\"--bt\", action=\"store_true\")\n",
    "    ap.add_argument(\"--kl\", action=\"store_true\")\n",
    "    ap.add_argument(\"--dpo\", action=\"store_true\")\n",
    "    a = ap.parse_args()\n",
    "    run_all = not (a.bt or a.kl or a.dpo)\n",
    "    if a.bt or run_all:\n",
    "        show_bt()\n",
    "    if a.kl or run_all:\n",
    "        show_kl()\n",
    "    if a.dpo or run_all:\n",
    "        show_dpo()\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running it\n",
    "\n",
    "The script is a command line tool, so each run below sets the arguments\n",
    "and calls `main()`. The heading shows the equivalent shell command.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "def run(*args):\n",
    "    sys.argv = [\"alignment.py\", *args]\n",
    "    main()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 alignment.py`\n",
    "\n",
    "All three stages of preference tuning.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 alignment.py --bt`\n",
    "\n",
    "Four preference pairs. The pair the model gets backwards costs 1.1711, more than three times the pair it gets right, and an undecided pair costs exactly log 2.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run(\"--bt\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 alignment.py --kl`\n",
    "\n",
    "A hack scoring 5.00 beats an honest answer scoring 4.00 until beta reaches 0.10, where its distance from the reference finally costs more than the reward it buys.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run(\"--kl\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 alignment.py --dpo`\n",
    "\n",
    "The implicit reward margin starts at 0 with a loss of log 2, the same number the undecided Bradley-Terry pair gives, and falls as the policy separates the two responses.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run(\"--dpo\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "Read the chapter this comes from: [Adaptation and Alignment](https://nlp.jcrlabz.com/book/alignment/).\n"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "alignment.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}