# Lab 13: Direct preference optimization on TinyLM

You will write the DPO loss (and two variants) from four log-probabilities, check every piece against numbers from the chapter, then preference-tune a tiny language model and watch something surprising happen: with one common setting, DPO makes the model *less* likely to produce the answers it was told to prefer.

Chapter: <https://learn.xhelio.ai/rl/learn/13-dpo/>

## Files

| File | What it is |
|---|---|
| `dpo.py` | **Your file.** Four tasks marked `TODO`. `train_dpo` at the bottom is complete. |
| `prefs.py` | Given: builds the preference pairs, and computes per-token log-probs of *given* responses (`response_logprobs`). Read it once. |
| `tinylm.py` | Given: the tiny transformer and the addition task "37+48=" (same file as the other Part III labs). |
| `test_dpo.py` | Tests. |
| `run.py` | Trains four variants on the same pairs and saves `dpo_tinylm.png`. |
| `solutions/dpo.py` | Reference solution. Try not to peek. |

## The data: a hidden preference

TinyLM answers addition problems: the prompt is `37+48=`, the response is the digits of the answer followed by `<eos>`. After supervised fine-tuning (SFT) it gets about 76% right with greedy decoding and about 51% when sampling at temperature 1.

The "annotator" prefers a correct answer to a wrong one. `make_preference_pairs` samples 4 answers per prompt from the SFT model. Whenever at least one is right and one is wrong, it makes a pair `(prompt, chosen=a right answer, rejected=a wrong answer)`. From 3,000 prompts you get about 2,200 pairs. Prompts where all four samples agree give no pair, just as a human has nothing to compare when both answers are equally good.

Why correctness rather than the style preference of the Chapter 12 lab? Because then "is the model getting better?" has an exact answer, answer accuracy, that is independent of the loss. You can see directly whether satisfying the preference pairs also improved the behavior the pairs were meant to teach.

This dataset is built **once** from the SFT model and never refreshed: it is offline, like most DPO datasets in practice. (Building pairs from checkable answers is a real recipe too, e.g. the offline DPO stage in the Qwen2.5 report.)

## Tasks

Run the tests for one task at a time with `-k`, e.g. `python -m pytest -q -k task1`.

1. **`sequence_logprob`** (Task 1). Sum the per-token log-probs over real response tokens (mask), or average them. Warm-up, but every later number depends on it.
2. **`implicit_reward` and `dpo_loss`** (Task 2). $\hat r_\theta = \beta\log\frac{\pi_\theta}{\pi_\text{ref}}$ and $-\log\sigma(\hat r_w - \hat r_l)$. The test uses the chapter's worked example ($\beta = 0.1$, log-probs $-12, -15, -13, -14$, loss $0.598$) and checks the gradient is $\mp\beta\,\sigma(\hat r_l - \hat r_w)$.
3. **`ipo_loss` and `simpo_loss`** (Task 3). IPO regresses the log-ratio margin to $1/(2\beta)$; SimPO drops the reference and uses length-normalized log-probs with a margin.
4. **`dpo_step`** (Task 4). Policy log-probs with gradients, reference log-probs without, the loss, an optional NLL term on the chosen response, one optimizer step, and a dict of diagnostics. The last test runs 40 steps on 32 pairs and expects the implicit reward to rank them correctly.

```bash
python -m pytest -q                   # all tests; they fail until you fill in the TODOs
python -m pytest -q -k task2          # just Task 2
RL_SOLUTIONS=1 python -m pytest -q    # run the tests against the reference solution
```

The tests take about 10 seconds.

## Then run the experiment

```bash
python run.py            # about 1.5 minutes; saves dpo_tinylm.png
```

The first run trains the SFT model (about 15 s) and caches it in `sft_tinylm.pt`. Then it trains four copies of the SFT model on the **same** pairs for 300 steps (Adam, learning rate 3e-4, batch 64):

| Variant | What changes |
|---|---|
| DPO β = 0.1 | plain DPO with a common default β |
| DPO β = 1 | a stronger leash to the reference |
| DPO β = 0.1 + NLL | plain DPO plus `nll_coef=1.0`: an extra log-likelihood term on the chosen answer |
| IPO β = 0.1 | IPO instead of DPO |

**What you should see** (reference solution; your numbers may differ a little):

| Variant | greedy accuracy | sampled accuracy | held-out log π(chosen) | held-out log π(rejected) |
|---|---|---|---|---|
| start (SFT) | 0.76 | 0.51 | −0.67 | −2.33 |
| DPO β = 0.1 | **0.38** | 0.36 | **−2.49** | −9.79 |
| DPO β = 1 | 0.89 | 0.74 | −0.27 | −4.41 |
| DPO β = 0.1 + NLL | 0.91 | 0.82 | −0.18 | −5.94 |
| IPO β = 0.1 | 0.84 | 0.70 | −0.40 | −4.95 |

- Every variant learns to separate chosen from rejected: the training reward accuracy climbs above 90% for all four.
- **DPO with β = 0.1 gets worse at the task.** The log-probability of the *chosen* (correct) answers falls, from −0.67 to about −2.5. The loss only cares about the *gap* between chosen and rejected, and the cheapest way to widen the gap moved probability off both, onto other wrong answers that appear in no pair. This is **likelihood displacement** (chapter section "What DPO gives up").
- Why does β matter so much here? The DPO loss stops pushing a pair once its implicit-reward margin $\beta(h_w - h_l)$ is a few units. With β = 1 that happens when the log-ratio gap $h_w - h_l$ is a few nats; with β = 0.1 the gap has to reach tens of nats, so the optimizer keeps pushing, for many epochs over the same small dataset, long after the ranking is right.
- Each of the three fixes works: a larger β, an NLL term that holds the chosen answer up (Llama 3 used one with coefficient 0.2; Pang et al. call it crucial), or IPO's finite target.

Don't over-read this toy: a 100k-parameter model, almost nine epochs over 2,200 pairs, and a learning rate far larger than the ~1e-6 used for real LLMs. The *mechanism* (the loss is blind to the absolute likelihood of the chosen response) is the same one reported for real models.

## Things to try

- **Where did the probability go?** After training DPO β = 0.1, sample 10 answers for a few held-out prompts with `tinylm.sample` and `decode_responses`. Are they near-misses of the correct answer? (In the reference run, `41+19=` went from mostly `60` and `50` under the SFT model to `71` ten times out of ten, a wrong answer the SFT model never produced in its ten samples.)
- **Noisy annotators.** Rebuild the pairs with `make_preference_pairs(..., label_noise=0.2)`. Which variant suffers most? (Think about what DPO does with a pair whose label is wrong, versus IPO.)
- **SimPO.** Add `("SimPO", dict(loss="simpo", beta=2.0, simpo_margin=0.5))` to `VARIANTS`. With no reference model, what stops it from drifting?
- **Fewer epochs.** Set `steps=60` (about 1.7 epochs). Is DPO β = 0.1 still harmful?
- **One round of online DPO.** After training, call `make_preference_pairs` on the *trained* policy (new seed) and train again for 150 steps. Fresh pairs come from the policy's own current mistakes. Does accuracy improve beyond one offline round?