Capstone Lesson 40: Direct Preference Optimization from Scratch

You have an SFT model. It follows instructions, but its outputs are uneven: some completions are clear, others are wordy or wrong. You also have a small dataset of preference pairs, where for the same prompt a human marked one completion as chosen and the other as rejected.

The classical RLHF answer is a two-stage pipeline. Train a reward model on the preferences, then optimise the policy against that reward with PPO. It works, but it is expensive. PPO keeps two models in memory, you need KL control to hold the policy near the reference, and a brittle reward model invites reward hacking.

DPO replaces both stages with a single supervised loss. The reward model never exists explicitly. The policy trains directly on the preference pairs, with an explicit KL penalty toward the SFT reference. Under the Bradley-Terry preference model it reaches the same optimal solution with far less code.

Direct Preference Optimization (DPO) is a supervised loss that trains a policy straight from preference pairs, with no separate reward model and no PPO loop. It works by rewriting the RLHF objective so the reward can be expressed through the policy itself. The derivation starts from the Bradley-Terry model.

Given a prompt x and two completions y_w (chosen) and y_l (rejected), the probability the human prefers y_w is

P(y_w > y_l | x) = sigmoid( r(x, y_w) - r(x, y_l) )

where r is some latent reward function. RLHF first fits r from preferences, then trains a policy pi to maximise r with a KL anchor:

max_pi   E_{x, y~pi} [ r(x, y) ] - beta * KL(pi || pi_ref)

The DPO derivation observes that the optimal policy pi* under this objective has a closed form in terms of r:

pi*(y | x) = (1/Z(x)) * pi_ref(y | x) * exp( r(x, y) / beta )

Rearrange for r:

r(x, y) = beta * ( log pi*(y | x) - log pi_ref(y | x) ) + beta * log Z(x)

The log Z(x) term depends on x but not on y, so it is the same for y_w and y_l and cancels in the preference difference:

r(x, y_w) - r(x, y_l) = beta * ( log pi_theta(y_w|x) - log pi_ref(y_w|x)
                                - log pi_theta(y_l|x) + log pi_ref(y_l|x) )

Substitute into the Bradley-Terry sigmoid and take the negative log likelihood over preference pairs:

L_DPO(theta) = - E_{(x, y_w, y_l)} [
  log sigmoid( beta * ( log pi_theta(y_w|x) - log pi_ref(y_w|x)
                       - log pi_theta(y_l|x) + log pi_ref(y_l|x) ) )
]

That is the loss: a sigmoid over a single scalar per example, computed from four log-probabilities. There is no separate reward model and no PPO. The loss carries no KL term, because the KL constraint is baked into the closed-form derivation.

Check the sign of the gradient before any training run. Take the gradient with respect to log pi_theta(y_w | x):

d L_DPO / d log pi_theta(y_w | x) = - beta * (1 - sigmoid(z))

where z is the argument to the sigmoid. This is negative for all z, so increasing the policy's log-probability of the chosen completion decreases the loss. By symmetry, the gradient with respect to log pi_theta(y_l | x) is positive: raising the rejected log-probability raises the loss. Training pushes the chosen up and the rejected down. The reference is frozen and does not move.

The fixture holds twelve preference triples, each a (prompt, chosen, rejected). The chosen completion is short and precise. The rejected one is wordy, off-topic, or wrong. The pairs cover task families that also appear in lesson 39 (capital, arithmetic, list, define, Python snippet), so a policy that started from an SFT base has a reasonable starting point.

The fixture is small on purpose. DPO works on tens of thousands of pairs in production. Here the goal is to run the loss math and the loop end to end on a tiny dataset and watch the chosen-versus-rejected log-prob gap grow.

The reference model is the SFT model frozen in place, and a DPO implementation has to handle it carefully. Three properties must hold:

• The reference parameters never receive gradients.
• The reference log-probabilities never change between epochs.
• The policy starts from the same weights as the reference. (The optimal theta is the reference plus a learned update, so initialising the policy as a copy of the reference is the well-defined start.)

The implementation enforces these by:

• Computing the reference log-probabilities once, before training starts, under torch.no_grad(), and reusing those stored values every epoch.
• Setting requires_grad=False on every reference parameter (build_models does this, and run_demo re-applies it after warm-up).
• Constructing the policy via policy.load_state_dict(reference.state_dict()) after the reference is built.

One wrinkle: run_demo runs a short warm-up pretrain so the reference has non-trivial log-probabilities. For that step it temporarily sets requires_grad=True and reference.train(), then copies the warmed-up weights into the policy and re-freezes the reference with requires_grad=False and reference.eval() before DPO training begins.

The model is the same TinyGPT architecture used in lesson 39 (decoder-only, causal, byte tokeniser), minus lesson 39's padding mask, since DPO here runs one sequence at a time. Reference and policy share one architecture. The policy's weights drift from the reference under training while the reference stays fixed.

The implementation is one main.py plus tests.

1. InstructionTokenizer: byte tokeniser with INST and RESP specials. Same special tokens and vocabulary as lesson 39.
2. TinyGPT: decoder-only transformer. Same architecture as lesson 39 (minus the padding mask), so this build is self-contained even if you skipped 39.
3. make_preferences: returns twelve (prompt, chosen, rejected) triples.
4. sequence_log_prob: given the model, a prompt prefix, and a completion, returns the sum of next-token log-probabilities over the completion (no prompt-position contribution).
5. dpo_loss: takes the four log-probabilities and beta, returns the per-example loss tensor and the margin (the implicit reward difference before the beta scaling) for logging.
6. train_dpo: snapshots the reference log-probs once before training (under no_grad), then for each epoch loops over the triples, computes chosen and rejected log-probs under the policy, applies the loss, and steps Adam once per triple.
7. evaluate_margins: returns the mean chosen-rejected log-probability margin under the policy at any point.
8. run_demo: builds reference and policy from a small warm-up pretrain, copies weights, trains for 30 epochs (12 triples per epoch, one Adam step each, 360 optimizer steps in all), prints the per-epoch loss and margin, and exits zero on success.

Under the Bradley-Terry preference model, DPO is mathematically equivalent to RLHF, up to the parameterisation of the reward. The implicit reward r(x, y) = beta * (log pi(y|x) - log pi_ref(y|x)) is identifiable from preferences up to a function of x, and that function cancels in the difference. The closed-form policy lets you skip the explicit reward model. The KL constraint is enforced structurally: any deviation of pi from pi_ref makes the log-ratio larger, and the sigmoid saturates, which damps the gradient when the policy moves too far. The reference is your safety net.

• Study and extend the existing length_normalised_log_prob (log-probability sum divided by completion length), for example by training on it. Length bias is a known DPO failure mode: the model preferentially chooses shorter completions because their log-probabilities are larger in absolute terms.
• Study the existing ipo_loss, which replaces the sigmoid + log with a squared distance between the margin and a target of 1/(2*beta): (margin - 1/(2*beta))^2. Swap it into train_dpo and compare convergence on the fixture.
• Add a label-smoothing parameter that interpolates between the hard chosen-rejected label and a uniform 0.5.
• Replace the reference with a smaller cheaper model (knowledge distillation flavour).

main.py also ships margin_table and print_margin_table for per-triple debugging; extend them to report length-normalised margins.

The implementation gives you the loss, the reference invariance, and the training loop. The math is the point, and the code makes it concrete.