A minimal preference-tuning pass (DPO)

The pretrained model from the previous lessons imitates its training distribution; it has no notion of preferred output. Alignment closes that gap. This lesson runs the smallest faithful version of Direct Preference Optimization[1] on the tiny GPT: a reference copy, a handful of preference pairs, and one gradient step on the DPO loss. It is the offline counterpart to the human-feedback fine-tuning that InstructGPT introduced[2].

1. The setup DPO needs

DPO optimizes a policy $\pi_\theta$ against a frozen reference $\pi_{\text{ref}}$ using pairs $(x, y_w, y_l)$ where $y_w$ is preferred over $y_l$ for prompt $x$. It needs two ingredients beyond the base model:

• a frozen copy of the pretrained model as the reference, and
• preference pairs. At toy scale these can be synthetic: pick a property to prefer (say, lines that end in punctuation over lines that do not) and label pairs accordingly.

import copy
import torch
import torch.nn.functional as F

policy = model                       # the pretrained tiny GPT
ref = copy.deepcopy(model)           # frozen reference
for p in ref.parameters():
    p.requires_grad_(False)
ref.eval()

2. Sequence log-probabilities

DPO scores a completion by the policy's total log-probability of its tokens. Run the model, take log-softmax over the vocabulary, and gather the log-prob of each realized next token.

def seq_logprob(m, idx, targets):
    logits, _ = m(idx)                                  # (B,T,V)
    logp = F.log_softmax(logits, dim=-1)
    tok_logp = logp.gather(-1, targets.unsqueeze(-1)).squeeze(-1)  # (B,T)
    return tok_logp.sum(dim=-1)                          # (B,) sum over the sequence

3. The DPO loss

The objective raises the policy's log-prob of the winner relative to the loser, each measured against the reference, inside a logistic loss with temperature $\beta$:

The implicit reward $r_\theta$ is the log-ratio of policy to reference, which is why DPO needs no separately trained reward model. In code:

def dpo_loss(policy, ref, win, lose, beta=0.1):
    # win, lose: (idx, targets) tensors for preferred and dispreferred completions
    pi_w = seq_logprob(policy, *win)
    pi_l = seq_logprob(policy, *lose)
    with torch.no_grad():
        ref_w = seq_logprob(ref, *win)
        ref_l = seq_logprob(ref, *lose)
    logits = beta * ((pi_w - ref_w) - (pi_l - ref_l))
    return -F.logsigmoid(logits).mean()

opt = torch.optim.AdamW(policy.parameters(), lr=1e-5)
loss = dpo_loss(policy, ref, win_batch, lose_batch, beta=0.1)
opt.zero_grad(set_to_none=True)
loss.backward()
opt.step()

The reference terms are inside torch.no_grad() because only the policy is trained. The $\beta$ term controls how far the policy may move from the reference: small $\beta$ keeps it close, large $\beta$ lets preferences dominate, the same KL-control knob that the RLHF KL penalty provides in the PPO route.

4. What to expect

After a few steps on the synthetic pairs, sample again and compare to the reference: the policy shifts toward the preferred property while staying fluent, because the reference ratio penalizes drifting away from the pretrained distribution. This is the entire shape of preference tuning, minus the scale and the human labels.

Exercises

Exercise 1. At the start of training the policy equals the reference, so the implicit reward $r_\theta(x, y) = \log \frac{\pi_\theta(y \mid x)}{\pi_{\text{ref}}(y \mid x)}$ is zero for every completion. Compute the DPO loss in that initial state, and explain why this value is the natural baseline that any successful preference step must drive the loss below.

Exercise 2 (code). Implement the DPO loss in numpy from sequence log-probabilities (skip the network and feed the log-probs directly). Verify the $\log 2$ baseline from Exercise 1, and confirm that increasing the winner's policy log-prob while holding everything else fixed lowers the loss.

import numpy as np

def logsigmoid(z):
    # numerically stable log-sigmoid
    return -np.logaddexp(0.0, -z)

def dpo_loss(pi_w, pi_l, ref_w, ref_l, beta=0.1):
    r_w = pi_w - ref_w          # implicit reward for the winner
    r_l = pi_l - ref_l          # implicit reward for the loser
    z = beta * (r_w - r_l)
    return -logsigmoid(z).mean()

# At init the policy equals the reference: all log-ratios are zero.
ref_w = np.array([-3.0, -4.0])
ref_l = np.array([-5.0, -6.0])
pi_w0, pi_l0 = ref_w.copy(), ref_l.copy()
base = dpo_loss(pi_w0, pi_l0, ref_w, ref_l, beta=0.1)
assert np.isclose(base, np.log(2.0))

# Raise the winner's policy log-prob; the loss must drop below the baseline.
pi_w1 = ref_w + 2.0             # policy now prefers the winner
better = dpo_loss(pi_w1, pi_l0, ref_w, ref_l, beta=0.1)
assert better < base

print("baseline loss:", round(float(base), 4), "(log 2)")
print("after preferring winner:", round(float(better), 4))

Exercise 3 (code). The $\beta$ term is the KL-control knob: small $\beta$ keeps the policy near the reference, large $\beta$ lets the preference dominate. Show numerically that for a fixed reward gap $r_w - r_l > 0$, the loss decreases as $\beta$ grows, and explain what that means for how far the policy is pushed.

import numpy as np

def logsigmoid(z):
    return -np.logaddexp(0.0, -z)

# Fix a positive reward gap (winner already favored by the implicit reward).
gap = 1.5                       # r_w - r_l
betas = np.array([0.05, 0.1, 0.5, 1.0])
losses = np.array([-logsigmoid(b * gap) for b in betas])

# Larger beta -> smaller loss for a fixed positive gap.
assert np.all(np.diff(losses) < 0)

print("betas:  ", betas)
print("losses: ", np.round(losses, 4))

The model has now been pretrained, modernized, and aligned. The final lesson closes the loop by re-reading the original Transformer paper against the code written across this capstone.