Module 12: Alignment

Introduction

You have trained a GPT (m07) and made it generate (m08). It predicts likely text — but likely is not the same as helpful. Ask a raw pretrained model a question and it is just as happy to continue with more questions, because that is what the internet looks like. Alignment is the post-training that turns a next-token predictor into an assistant that follows instructions and honors human preferences.

This module builds the three canonical stages from scratch — and shows how the third one quietly deletes the machinery of the second:

  • SFT (supervised fine-tuning): teach the format by imitation, on demonstrations.
  • Reward modeling: learn what people prefer from comparisons, not demonstrations.
  • DPO (Direct Preference Optimization): fold the reward model and the whole reinforcement-learning loop into a single classification loss — “your language model is secretly a reward model.”

Why it matters for LLMs:

  • Every assistant you have used (ChatGPT, Claude, Gemini) is a pretrained model put through exactly this pipeline. Without it, the base model is a brilliant autocomplete, not a helper.
  • DPO (and its relatives) is now the default way to align open models, precisely because it removes the fragile RL step this module explains.

What You’ll Learn

After this module, you can:

  • Explain the gap between likelihood and helpfulness, and the three-stage pipeline that closes it.
  • Implement the SFT loss from scratch — masked next-token cross-entropy that trains the completion, not the prompt.
  • Build a reward model by swapping an LM’s head, and train it with the Bradley-Terry preference loss.
  • Derive the KL-constrained RLHF objective, its closed-form optimum, and the reparameterization that turns it into DPO.
  • Implement and drive the DPO loss, and see the implicit reward push the preferred answer up and the rejected one down.
  • Build PPO from scratch — the clipped surrogate, the value baseline (critic), and the KL-shaped reward — and watch reward over-optimization collapse the true reward when the KL leash is too loose.

Prerequisites

This module requires familiarity with:

  • Module 07: Training — the language-modeling loss and gradient descent; SFT is that loss with a mask.
  • Module 08: Generation — sampling from a policy; alignment changes which continuations the policy prefers.

Intuition: The Alignment Pipeline

Alignment is a relay of three hand-offs, each fixing the previous stage’s limit. Pretraining gives a model that knows language. SFT shows it the shape of a good answer. But you cannot demonstrate the best answer to every prompt — so instead you let people compare answers, learn a notion of “better,” and push the model toward it. Step through the pipeline and watch what each stage consumes and what it changes:

NoteKey Insight

Only SFT, RLHF, and DPO update the policy — the model you ship. The reward model is scaffolding: it exists to score answers during RLHF. DPO’s whole trick is to notice you can express that score using the policy itself, and throw the scaffolding away.

Stage 1 — Supervised Fine-Tuning

SFT is the m07 language-modeling loss with one change: given a (prompt, completion) pair, you train the model to predict only the completion tokens. The prompt is context, not something to imitate, so its tokens are masked out of the loss.

\[ \mathcal{L}_{\text{SFT}} = - \frac{1}{|C|} \sum_{t \in C} \log \pi_\theta(y_t \mid y_{<t}) \]

where \(C\) is the set of completion-token positions. Everything else — the softmax, the cross-entropy, the next-token shift — is exactly m07.

alignment.py implements this as sft_loss, built on a reusable sequence_logprob that computes \(\log \pi_\theta(y \mid x)\) with a completion mask:

import torch
from alignment import sft_loss, sequence_logprob

torch.manual_seed(0)
batch, seq, vocab = 2, 6, 20
logits = torch.randn(batch, seq, vocab)
labels = torch.randint(0, vocab, (batch, seq))

# The first 3 tokens are the prompt (mask 0); the last 3 are the completion (1).
completion_mask = torch.zeros(batch, seq)
completion_mask[:, 3:] = 1

loss_all = sft_loss(logits, labels, torch.ones(batch, seq))
loss_completion = sft_loss(logits, labels, completion_mask)

print(f"Loss over all tokens:        {loss_all.item():.4f}")
print(f"Loss over completion only:   {loss_completion.item():.4f}")
print("They differ because the prompt tokens no longer count as targets.")
Loss over all tokens:        3.3155
Loss over completion only:   3.5788
They differ because the prompt tokens no longer count as targets.

The mask is the whole point: without it, the model spends capacity learning to regurgitate prompts instead of answering them. sequence_logprob — the \(\log \pi_\theta(y \mid x)\) it wraps — returns in a moment as the atom DPO is built from.

Stage 2 — Preferences and Reward Modeling

SFT can only teach what a human can write down. For most prompts there is no single ideal answer, but people can reliably say which of two answers is better. That comparison is cheaper to collect and far more informative at the margin.

The Bradley-Terry model turns a scalar reward into the probability that one answer beats another:

\[ p(y_w \succ y_l \mid x) = \sigma\!\big(r(x, y_w) - r(x, y_l)\big) \]

so a reward model \(r_\phi(x, y)\) is trained by maximizing the probability it assigns to the human’s choice — equivalently, minimizing

\[ \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]. \]

Only the difference of rewards ever appears, so the absolute scale is free — a fact DPO will exploit. A reward model is not a new architecture: it is a language-model backbone with the vocabulary head swapped for a single scalar head reading the last token. RewardModel does exactly that:

import torch
from alignment import RewardModel, bradley_terry_loss, preference_accuracy
import importlib.util, sys
from pathlib import Path

# Load the m06 GPT to use as the reward model's backbone.
_p = Path("../m06_transformer/transformer.py").resolve()
_s = importlib.util.spec_from_file_location("transformer", _p)
_t = importlib.util.module_from_spec(_s); sys.modules["transformer"] = _t
_s.loader.exec_module(_t)
GPTModel = _t.GPTModel

torch.manual_seed(0)
backbone = GPTModel(vocab_size=50, embed_dim=32, num_heads=2, num_layers=2, max_seq_len=32)
reward_model = RewardModel(backbone)

chosen = torch.randint(0, 50, (4, 8))    # 4 preferred answers
rejected = torch.randint(0, 50, (4, 8))  # 4 dispreferred answers

r_chosen = reward_model(chosen)
r_rejected = reward_model(rejected)
print(f"Reward shape (one scalar per sequence): {tuple(r_chosen.shape)}")
print(f"Bradley-Terry loss: {bradley_terry_loss(r_chosen, r_rejected).item():.4f}")
print(f"Ranking accuracy:   {preference_accuracy(r_chosen, r_rejected).item():.2f}")
Reward shape (one scalar per sequence): (4,)
Bradley-Terry loss: 0.7040
Ranking accuracy:   0.50

The reward model is untrained here, so its accuracy is a coin flip — training it with bradley_terry_loss on real preference pairs is what teaches it to rank.

The Math: RLHF and Its Closed Form

With a reward model in hand, the obvious move is to fine-tune the policy to score well. But a reward model is a proxy; chase it too hard and the policy drifts into gibberish that games the reward (reward hacking). So RLHF maximizes reward minus a KL leash to the reference (SFT) model:

\[ \max_{\pi_\theta}\; \mathbb{E}_{x,\, y \sim \pi_\theta}\big[r(x, y)\big] \;-\; \beta\, \mathrm{KL}\!\big[\pi_\theta(y \mid x)\,\|\,\pi_{\text{ref}}(y \mid x)\big]. \]

This objective has an exact, closed-form optimum — you can write down the best possible policy without any training:

\[ \pi^*(y \mid x) = \frac{1}{Z(x)}\, \pi_{\text{ref}}(y \mid x)\, \exp\!\Big(\tfrac{1}{\beta}\, r(x, y)\Big). \]

Read it as: start from the reference and re-weight by reward. The temperature is \(\beta\). Drive it and watch the optimal policy slide between “obey the reward” and “stay put”:

NoteKey Insight

The optimum is known — so why is RLHF hard? Because \(Z(x)\) sums over every possible completion and is intractable, so you cannot just sample from \(\pi^*\). Classic RLHF approximates it with PPO: roll out completions, score them with the reward model, and nudge the policy up the reward gradient. It works, but it needs a live reward model, on-policy sampling, and careful tuning. The next section removes all of it.

Stage 3 — DPO: Your Language Model Is the Reward Model

Here is the move. Take the closed-form optimum and solve it for the reward instead of the policy — a line of algebra gives

\[ r(x, y) = \beta \log \frac{\pi^*(y \mid x)}{\pi_{\text{ref}}(y \mid x)} + \beta \log Z(x). \]

The reward is just a log-ratio of the policy to the reference, plus a term that depends only on \(x\). Now substitute this into the Bradley-Terry loss. Because only reward differences matter, the intractable \(\beta \log Z(x)\) cancels, and the reward model vanishes with it. What remains is a loss on the policy alone:

\[ \mathcal{L}_{\text{DPO}} = -\,\mathbb{E}_{(x, y_w, y_l)}\left[\log \sigma\!\left( \beta \log \frac{\pi_\theta(y_w \mid x)}{\pi_{\text{ref}}(y_w \mid x)} - \beta \log \frac{\pi_\theta(y_l \mid x)}{\pi_{\text{ref}}(y_l \mid x)} \right)\right]. \]

No reward model. No RL. No sampling. It is the same supervised classification loss as reward modeling — but the “reward” is the policy’s own implicit reward \(\hat r(x, y) = \beta \log \frac{\pi_\theta(y|x)}{\pi_{\text{ref}}(y|x)}\). Both log-probs come straight from sequence_logprob; dpo_loss assembles them:

import torch
from alignment import dpo_loss, implicit_reward

# Per-sequence log πθ(y|x) and log π_ref(y|x) for a batch of preference pairs.
# (In a real run these come from batch_sequence_logprob on the policy and a
# frozen copy of the SFT model.)
policy_chosen   = torch.tensor([-2.0, -3.0])
policy_rejected = torch.tensor([-4.0, -2.5])
ref_chosen      = torch.tensor([-3.0, -3.0])
ref_rejected    = torch.tensor([-3.0, -3.0])

loss, chosen_reward, rejected_reward = dpo_loss(
    policy_chosen, policy_rejected, ref_chosen, ref_rejected, beta=0.1
)
print(f"DPO loss:               {loss.item():.4f}")
print(f"Implicit reward chosen:   {chosen_reward.tolist()}")
print(f"Implicit reward rejected: {rejected_reward.tolist()}")
print(f"Reward margin (want > 0): {(chosen_reward - rejected_reward).tolist()}")
DPO loss:               0.6583
Implicit reward chosen:   [0.10000000149011612, 0.0]
Implicit reward rejected: [-0.10000000149011612, 0.05000000074505806]
Reward margin (want > 0): [0.20000000298023224, -0.05000000074505806]

The DPO gradient increases \(\log \pi_\theta(y_w)\) and decreases \(\log \pi_\theta(y_l)\), weighted by \(\sigma\big(\hat r(y_l) - \hat r(y_w)\big)\) — so pairs the model currently ranks backwards get the biggest push. Watch it run on a toy policy: demonstrate_dpo optimizes a 5-way softmax where response 0 is preferred over response 1, against a frozen uniform reference.

from alignment import demonstrate_dpo

history = demonstrate_dpo(beta=0.1, steps=200, lr=1.0, seed=0)
steps = list(range(len(history["loss"])))

dpo_curve = [
    {"step": s, "chosen_reward": history["chosen_reward"][s],
     "rejected_reward": history["rejected_reward"][s],
     "chosen_prob": history["chosen_prob"][s]}
    for s in steps
]
ojs_define(dpo_curve = dpo_curve)
print(f"loss:        {history['loss'][0]:.3f} → {history['loss'][-1]:.3f}")
print(f"chosen reward:   {history['chosen_reward'][0]:+.3f} → {history['chosen_reward'][-1]:+.3f}")
print(f"rejected reward: {history['rejected_reward'][0]:+.3f} → {history['rejected_reward'][-1]:+.3f}")
print(f"chosen P:    {history['chosen_prob'][0]:.3f} → {history['chosen_prob'][-1]:.3f}")
loss:        0.693 → 0.240
chosen reward:   +0.000 → +0.161
rejected reward: +0.000 → -1.144
chosen P:    0.210 → 0.996

The two curves split apart: DPO lifts the chosen response’s implicit reward and drives the rejected one down — the “push the winner up, the loser down” behavior, straight from a supervised loss.

Interactive Exploration

The DPO loss depends on just two numbers per pair: how much the policy prefers the chosen answer over the reference, \(\Delta_w = \log\frac{\pi_\theta(y_w)}{\pi_{\text{ref}}(y_w)}\), and the same for the rejected answer, \(\Delta_l\). Drive both (and \(\beta\)) and watch the implicit rewards, the loss, and — crucially — the gradient weight \(\sigma\big(\beta(\Delta_l - \Delta_w)\big)\), which is largest exactly when the pair is ordered wrong.

TipTry This
  1. Set \(\Delta_w\) below \(\Delta_l\) (drag chosen down, rejected up). The loss climbs and the gradient weight heads toward 1 — DPO spends its effort on pairs it currently gets wrong.
  2. Push \(\Delta_w\) far above \(\Delta_l\). Loss collapses toward 0 and the gradient weight toward 0: an already-correct pair contributes almost nothing.
  3. Shrink \(\beta\). The implicit rewards compress toward 0 — \(\beta\) sets how strongly the log-ratios translate into reward, i.e. how hard the policy is allowed to move from the reference.

Beyond DPO: IPO, KTO, and ORPO

DPO turns a pair \((y_w \succ y_l)\) into one classification loss on the margin

\[ h = \big(\log\pi_\theta(y_w) - \log\pi_{\text{ref}}(y_w)\big) - \big(\log\pi_\theta(y_l) - \log\pi_{\text{ref}}(y_l)\big), \qquad L_\text{DPO} = -\log\sigma(\beta\, h). \]

That log-sigmoid never stops rewarding a bigger \(h\): with clean, near- deterministic preferences it will drive the log-ratio toward infinity and overfit. The three successors each change one thing about DPO — and each is a one-function edit on the same margin you already built (preference_margin in dpo_variants.py).

NoteKey Insight

DPO opened a design space with three knobs: the loss shape (IPO changes it), the data (KTO changes it), and the machinery (ORPO changes it). All three still do the one thing DPO does — raise the chosen answer, lower the rejected one.

IPO: A Bounded Target

Identity Preference Optimization (Azar et al., 2023) keeps DPO’s margin but swaps the loss shape: instead of an ever-decreasing log-sigmoid, minimize the squared distance to a finite target margin \(\frac{1}{2\tau}\):

\[ L_\text{IPO} = \Big(h - \tfrac{1}{2\tau}\Big)^2. \]

DPO’s gradient never reaches zero, so it keeps pushing; IPO’s gradient crosses zero exactly at \(h = \frac{1}{2\tau}\) and stops. Watch the two loss shapes and their gradients side by side — DPO slides forever, IPO settles into its target:

import torch
from dpo_variants import ipo_loss, demonstrate_dpo_variants

# Same margin, two losses. tau=0.1 -> IPO target margin 1/(2*0.1) = 5.0
curves = demonstrate_dpo_variants(tau=0.1, beta=1.0)
print("IPO target margin:", curves["ipo_target"][0])
loss, _, _ = ipo_loss(torch.tensor([0.0]), torch.tensor([-5.0]),
                      torch.tensor([0.0]), torch.tensor([0.0]), tau=0.1)
print("IPO loss when margin == target:", round(loss.item(), 4))
IPO target margin: 5.0
IPO loss when margin == target: 0.0
# Bridge the loss-vs-margin sweep to the plot below.
variant_curve = [
    {"margin": m, "dpo": d, "ipo": i, "dpo_grad": dg, "ipo_grad": ig}
    for m, d, i, dg, ig in zip(curves["margin"], curves["dpo"], curves["ipo"],
                               curves["dpo_grad"], curves["ipo_grad"])
]
ipo_target = curves["ipo_target"][0]
ojs_define(variant_curve = variant_curve, ipo_target = ipo_target)
TipTry This

The DPO curve (orange) keeps sliding down as the margin grows — there is no margin it considers “enough,” which is how DPO overfits crisp preferences. IPO (blue) is a parabola with its floor at the green target line: push past it and IPO pulls the margin back. Smaller \(\tau\) moves the target further right (a wider chosen–rejected gap); it is the direct dial on separation strength.

KTO: Learning from Unpaired Labels

DPO and IPO need pairs — the same prompt answered two ways, ranked. But most feedback in the wild is a single thumbs-up or thumbs-down on one answer. Kahneman-Tversky Optimization (Ethayarajh et al., 2024) drops the pairing: each example carries one bit, desirable or undesirable, and a prospect-theory value pushes desirable rewards above a reference point \(z_0\) and undesirable ones below it:

\[ v = \begin{cases} \lambda_D\,\sigma\big(\beta(r - z_0)\big) & \text{desirable}\\ \lambda_U\,\sigma\big(\beta(z_0 - r)\big) & \text{undesirable} \end{cases}, \qquad r = \log\frac{\pi_\theta(y|x)}{\pi_{\text{ref}}(y|x)}, \qquad L_\text{KTO} = \lambda_y - v. \]

\(z_0\) is the batch’s average KL (clamped at 0, no gradient) — the “neutral” reward each example is measured against.

from dpo_variants import kto_loss

# One desirable and one undesirable example — no pairing needed.
policy_logps = torch.tensor([-0.5, -0.8])   # log πθ(y|x)
ref_logps    = torch.tensor([-1.0, -0.6])   # log π_ref(y|x)
desirable    = torch.tensor([True, False])  # thumbs up / thumbs down

loss, rewards = kto_loss(policy_logps, ref_logps, desirable, beta=0.1)
print("implicit rewards β·(πθ−ref):", rewards.round(decimals=3).tolist())
print("KTO loss:", round(loss.item(), 4))
implicit rewards β·(πθ−ref): [0.05000000074505806, -0.019999999552965164]
KTO loss: 0.4913
NoteKey Insight

KTO trades DPO’s pairwise signal for a pointwise one, so it runs on the abundant, cheap feedback (a like, a flag) that never came as a clean A-vs-B comparison. The reference point \(z_0\) is what lets an unpaired label still say “better/worse than typical.”

ORPO: Dropping the Reference Model

DPO, IPO, and KTO all need a frozen reference \(\pi_{\text{ref}}\) — a second forward pass every step. Odds Ratio Preference Optimization (Hong et al., 2024) removes it: fold preference straight into SFT as the ordinary NLL on the chosen answer, plus a small odds-ratio penalty that separates chosen from rejected — using the policy alone:

\[ L_\text{ORPO} = \underbrace{-\log P_\theta(y_w)}_{\text{SFT}} + \lambda\underbrace{\Big(\!-\log\sigma\big(\log\tfrac{\text{odds}_\theta(y_w)}{\text{odds}_\theta(y_l)}\big)\Big)}_{\text{odds ratio}}, \qquad \text{odds}(y) = \frac{P(y)}{1 - P(y)}. \]

Here \(P_\theta(y)\) is the length-normalized (average per-token) probability, so the inputs are mean log-probs. One model, one stage, no reference:

from dpo_variants import orpo_loss

# Length-normalized (average) log-probs — NOTE: no reference model anywhere.
chosen_logps   = torch.tensor([-0.4])   # log P̄θ(y_w)
rejected_logps = torch.tensor([-1.8])   # log P̄θ(y_l)

loss, sft, orl = orpo_loss(chosen_logps, rejected_logps, lambda_or=1.0)
print(f"SFT part: {sft.item():.4f}   odds-ratio part: {orl.item():.4f}")
print(f"ORPO loss = SFT + λ·OR = {loss.item():.4f}")
SFT part: 0.4000   odds-ratio part: 0.0929
ORPO loss = SFT + λ·OR = 0.4929
NoteWhich One?

All four reduce to “raise the chosen, lower the rejected,” and differ only in what they give up: DPO — the reward model; IPO — DPO’s unbounded margin (a finite target instead); KTO — the need for pairs (unpaired labels); ORPO — the reference model (fused into SFT). Reach for KTO when your data is unpaired thumbs-up/down, IPO when DPO overfits, ORPO when you want a single training stage.

PPO: The RL Detour DPO Removes

We wrote the RLHF objective and its closed-form optimum \(\pi^* \propto \pi_{\text{ref}}\exp(r/\beta)\), then took the DPO shortcut straight past the reinforcement learning. But that RL is what every pre-DPO aligned model (InstructGPT, the first ChatGPT) actually ran, and seeing it makes clear what DPO saves you from. The algorithm is PPO (Proximal Policy Optimization), and its loop has four steps:

The one structural difference from GRPO (m13) is step 3: PPO learns a value network \(V\) (the critic) to estimate how good a response was expected to be, and the advantage is the surprise \(A = \text{return} - V\). GRPO throws the critic away and uses a group’s mean reward as the baseline — cheaper, but PPO’s critic is the classic choice. The surrogate (step 4) is the same clipped objective GRPO uses:

\[ L = -\,\mathbb{E}\!\left[\min\big(\rho A,\ \text{clip}(\rho, 1-\epsilon, 1+\epsilon)\,A\big)\right] + c_v\,(V - \text{return})^2, \qquad \rho = \frac{\pi_\theta(a)}{\pi_{\theta_{\text{old}}}(a)}. \]

Code: PPO on a Toy Policy

ppo.py builds each piece. Advantages come from GAE (generalized advantage estimation) — a bias/variance knob over the value baseline — here on a one-step episode, so \(A = \text{return} - V\):

import torch
from ppo import compute_gae, clipped_surrogate, ActorCritic

# GAE on one terminal step: reward 1, the critic expected 0.25 -> advantage 0.75.
print("advantage:", compute_gae([1.0], [0.25]))

# The clipped surrogate is a loss to minimize; a huge ratio is capped at (1+eps)*A.
lp_old = torch.tensor([-2.0]); lp_new = torch.tensor([0.0])  # ratio e^2 >> 1.2
print("clipped surrogate:", round(float(clipped_surrogate(lp_new, lp_old, torch.tensor([1.0]))), 3))
advantage: [0.75]
clipped surrogate: -1.2

The policy is ActorCritic — m13’s ToyReasoningPolicy with the value network added back: a categorical over K answers plus a scalar critic. ppo_step runs one full update (rollout → KL-shaped reward → advantage → clipped surrogate + value loss). Train it against a reward model and watch the reward climb:

from ppo import demonstrate_ppo

# Train with a weak KL leash (beta=0). The PROXY reward model is what we optimize;
# the TRUE reward is what we actually want (tracked, never optimized).
history = demonstrate_ppo(steps=60, beta=0.0, seed=0)
============================================================
PPO on a toy policy: proxy reward vs. the true reward
============================================================
  K=6, good=1, hackable=4, beta=0.0

   step    proxy     true       kl
      0    0.300    0.167    0.000
      7    0.453    0.234    0.049
     14    0.703    0.312    0.385
     21    0.853    0.266    0.762
     28    0.916    0.109    1.135
     35    0.916    0.109    1.120
     42    0.900    0.031    1.179
     49    0.978    0.031    1.546
     56    1.000    0.000    1.732

  proxy 0.300 -> 0.950, true 0.167 -> 0.016

Look at what happened: the proxy reward (blue) climbs to 1.0, but the true reward (orange) rises for a while and then collapses back toward zero. The policy learned to maximize the reward model — by exploiting an answer the reward model over-rates but that isn’t actually what we wanted. That is reward over-optimization.

Reward Over-Optimization

The reward model is a proxy for human preference, and every proxy has flaws. Push the policy hard enough against it and the policy stops improving on the real goal and starts exploiting the proxy’s mistakes — Goodhart’s law: when a measure becomes a target, it ceases to be a good measure. The lever that controls this is the KL leash to the reference (the beta above): the further the policy is allowed to drift, the more it over-optimizes. Sweep beta and plot the true reward against the KL the policy reached (Gao et al., 2022):

from ppo import demonstrate_overoptimization

rows = demonstrate_overoptimization(steps=80, seed=0)
============================================================
Reward over-optimization: true reward vs. KL leash (beta)
============================================================
    beta    proxy     true       kl
    0.00    1.000    0.000    1.756
    0.10    0.953    0.078    1.338
    0.20    0.906    0.234    1.003
    0.40    0.797    0.391    0.620
    0.80    0.537    0.281    0.140
    1.60    0.409    0.219    0.030

  Weak leash (small beta): proxy up, KL up, TRUE reward collapses.

The proxy reward (blue, dashed) rises monotonically as the policy drifts — more optimization always looks better by the proxy. But the true reward (orange) is a hump: it rises with a little optimization, peaks at a moderate KL, then falls as the policy over-optimizes the proxy’s flaws. The best model is not the one with the highest reward-model score — it’s the one at the top of the hump.

NoteKey Insight

This hump is why DPO is attractive. RLHF-PPO has to walk the reward-vs-KL curve carefully — pick β wrong and you ship a reward-hacked model — on top of running a full on-policy RL loop with a reward model and a critic. DPO trades the whole loop for one supervised loss on fixed preference pairs. It doesn’t make Goodhart go away (your preference data is still a proxy), but it removes the moving target the RL detour has to chase.

TipTry This
  1. Find the peak. In the sweep above, read off the β at the top of the true-reward hump. Re-run demonstrate_overoptimization with betas clustered around it to locate the compute-optimal KL more precisely.
  2. Break the leash. Set beta=0.0 in demonstrate_ppo and then a large beta=0.8. Watch the true reward collapse in the first and survive in the second — the same trade the DPO pitfall (β too small) warned about, now made visible.

Common Pitfalls

  1. Forgetting to mask the prompt in SFT. Training on prompt tokens wastes capacity teaching the model to echo inputs. Only completion tokens belong in the loss (completion_mask).
  2. Letting the reference model drift. \(\pi_{\text{ref}}\) must be a frozen copy of the SFT model. If you accidentally update it (or forget torch.no_grad / .detach() on its log-probs), the KL anchor moves and the implicit reward becomes meaningless.
  3. β too small or too large. Small \(\beta\) lets the policy drift far from the reference — faster preference-fitting but risk of degenerate, reward-hacked text. Large \(\beta\) barely moves the policy. Typical DPO runs use \(\beta \in [0.1, 0.5]\).
  4. Comparing raw log-probs across different lengths. sequence_logprob sums over completion tokens, so longer completions have more-negative log-probs. Always compare chosen vs. rejected for the same prompt; DPO’s log-ratio against the reference largely cancels the length effect, but a badly unbalanced pair can still bias it.
  5. Reward hacking is not solved, only reframed. DPO removes the separate reward model, but the preference data is still a proxy for what you want. Garbage preferences in, misaligned policy out.
  6. Feeding ORPO summed log-probs. ORPO’s odds need \(P_\theta(y) \in (0,1)\), so it requires length-normalized (mean per-token) log-probs — pass sequence_logprob(..., average=True). Summed log-probs make \(P > 1\) possible and log(1 - P) undefined.
  7. Backpropagating through KTO’s \(z_0\). The reference point must be .detach()ed (it only controls loss saturation). If gradients flow through the batch-mean KL, the objective couples every example in the batch and training destabilizes.
  8. Reading the reward model’s score as truth (PPO). The reward model is a proxy; its score rising is not proof the model got better. Track a held-out true metric and watch for the over-optimization hump — the best checkpoint is at the peak, not the end.
  9. Optimizing without a KL leash (PPO). With \(\beta \to 0\) the policy is free to reward-hack the proxy and collapse the true reward (and often the text quality). The KL-to-reference penalty — folded into the reward here — is what keeps RLHF-PPO from running off the rails.

Exercises

Exercise 1: The SFT mask matters

import torch
from alignment import sft_loss

# Build a (1, 8) example. Compute the SFT loss twice: once over all tokens, once
# over only the last 4 (the "completion"). Confirm they differ, and explain which
# one trains the assistant behavior.

# Your implementation here:
# logits = torch.randn(1, 8, 30)
# labels = torch.randint(0, 30, (1, 8))
# ...

Exercise 2: Bradley-Terry by hand

import torch
from alignment import bradley_terry_loss

# For chosen reward 2.0 and rejected reward 0.0, verify by hand that the loss is
# -log(sigmoid(2.0)). Then show the loss equals log(2) when the two rewards match.

# Your implementation here:

Exercise 3: DPO endpoints

import torch
from alignment import dpo_loss

# Show that when the policy equals the reference (all four log-probs equal), the
# DPO loss is exactly log(2) and both implicit rewards are 0 — DPO starts from
# "no preference" and has to earn the margin.

# Your implementation here:

Exercise 4: IPO stops, DPO doesn’t

import torch
from dpo_variants import ipo_loss, preference_margin
from alignment import dpo_loss

# For a pair whose margin is already huge (chosen far above rejected), compare the
# gradient magnitude DPO vs IPO produces on the chosen log-prob. Confirm IPO's is
# ~0 once the margin passes 1/(2τ) while DPO's stays positive — IPO knows when to
# stop. (Set requires_grad on the chosen log-prob and inspect .grad.)

# Your implementation here:

Exercise 5: The over-optimization peak

from ppo import demonstrate_overoptimization

# demonstrate_overoptimization sweeps the KL weight beta and reports final (proxy, true,
# kl). Find the beta whose TRUE reward is highest, then sweep a finer grid of betas
# around it to pin down the peak. Argue why the peak — not the highest proxy reward —
# is the checkpoint you would actually ship.

# Your implementation here:

Summary

Key takeaways:

  1. Alignment closes the likelihood-vs-helpfulness gap — a pretrained model maximizes plausibility; SFT, reward modeling, and DPO/RLHF make it helpful.
  2. SFT is masked next-token loss — the m07 loss with the prompt masked out, so the model learns to produce the completion, not echo the prompt.
  3. Reward models learn from comparisons — the Bradley-Terry loss \(-\log\sigma(r_w - r_l)\) trains a scalar-headed backbone to rank answers; only reward differences matter.
  4. RLHF has a closed-form optimum — \(\pi^* \propto \pi_{\text{ref}}\exp(r/\beta)\) — but its partition function \(Z(x)\) is intractable, so PPO approximates it with on-policy RL.
  5. DPO deletes the reward model and the RL loop — solving the optimum for the reward gives \(\hat r = \beta\log(\pi_\theta/\pi_{\text{ref}})\); substituted into Bradley-Terry, the intractable term cancels and alignment becomes one supervised loss on preference pairs.
  6. The DPO gradient is self-prioritizing — it lifts the chosen answer and lowers the rejected one, weighted by how wrong the current ranking is.
  7. DPO opened a family, each dropping one thing — IPO swaps the unbounded log-sigmoid for a squared loss toward a finite target margin \(\frac{1}{2\tau}\); KTO drops the pairs to learn from unpaired desirable/undesirable labels; ORPO drops the reference model, fusing an odds-ratio penalty into SFT. All still raise the chosen answer and lower the rejected one.
  8. PPO is the RL detour DPO removes — rollout → KL-shaped reward → advantage from a critic → clipped surrogate. It optimizes the same objective directly, at the cost of an on-policy loop, a reward model, and a value network (the critic m13’s GRPO drops). And it must be walked carefully: reward over-optimization means the true reward humps then falls as the policy over-fits the proxy — the KL leash sets where you land on that hump.

What’s Next

You now have the alignment toolkit: imitate with SFT, learn preferences with a reward model, and optimize them directly with DPO. From here the frontier builds on it — Module 13: Reasoning & Test-Time Compute spends more compute at inference (chain-of-thought, self-consistency, best-of-N) to get better answers, and its sequential counterpart trains reasoning models with RL against verifiable rewards (o1/R1-style). The preference methods beyond DPO you just built — IPO, KTO, ORPO — reshape the same log-ratio loss, each dropping one of DPO’s ingredients (its unbounded margin, its pairs, or its reference model).

Going Deeper

Core Papers:

Practical Resources: