Preference Learning and Direct Preference Optimization

Optimize relative preference between responses while measuring what the labels actually encode.

Demonstrations specify one acceptable response per prompt; preference pairs rank two. A preference example is (x, y_w, y_l): a prompt, a preferred response, and a rejected response. This chapter implements the numerical pieces behind a reward model and Direct Preference Optimization (DPO), then constructs a preference dataset whose labels deliberately contain a shortcut. The second evaluator stays untouched by optimization, so an improved training score can be compared with the behavior we actually intended.

The implementation uses small arrays rather than a language-model-sized network. That keeps every probability visible while preserving the same sequence-level objectives used by larger policies. Design rule: preference optimization is only as meaningful as the evaluator that remains outside the training loop.

Preference pairs

A reward model is trained to rank the preferred response above the rejected one:

\mathcal{L}_{\mathrm{RM}}(\phi) = -\log\sigma\left(r_\phi(x,y_w)-r_\phi(x,y_l)\right).

DPO removes the separately trained reward model. With a trainable policy \pi_\theta and frozen reference \pi_{\mathrm{ref}}, it optimizes the difference in log-ratios,

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

The reference is not a second label source. It is a trust-region-like anchor that makes the update relative to a known policy. The temperature \beta controls how sharply a preference margin is pursued.

The sequence log-probability is the shared primitive: it sums token log-probabilities over the response mask, exactly as the SFT mask in Chapter 08 selected response positions.

Sequence log-probabilities

A sequence probability is a product of token probabilities, so its log-probability is a sum. Work in log space to avoid underflow and use the response mask to exclude prompt and padding tokens. The function below returns one scalar per sequence; changing a masked token cannot change its value.

import numpy as np

rng = np.random.default_rng(9)


def log_softmax(logits, axis=-1):
    shifted = logits - logits.max(axis=axis, keepdims=True)
    return shifted - np.log(np.exp(shifted).sum(axis=axis, keepdims=True))


def sequence_logprob(logits, targets, mask):
    """Return summed log p(targets) over the unmasked positions."""
    log_probs = log_softmax(logits)
    selected = np.take_along_axis(log_probs, targets[..., None], axis=-1)[..., 0]
    return (selected * mask).sum(axis=-1)


# Two sequences share the same response tokens; only their padded positions differ.
logits = np.array([
    [[3.0, 0.0, 0.0], [0.0, 3.0, 0.0], [0.0, 0.0, 3.0], [9.0, 0.0, 0.0]],
    [[3.0, 0.0, 0.0], [0.0, 3.0, 0.0], [0.0, 0.0, 3.0], [0.0, 9.0, 0.0]],
])
targets = np.array([[0, 1, 2, 0], [0, 1, 2, 1]])
response_mask = np.array([[1, 1, 1, 0], [1, 1, 1, 0]])

logprobs = sequence_logprob(logits, targets, response_mask)
print("masked sequence log-probabilities:", logprobs.round(4).tolist())
np.testing.assert_allclose(logprobs[0], logprobs[1])

masked sequence log-probabilities: [-0.2848, -0.2848]

The fourth position has different logits and targets in the two rows, but it is padding and contributes zero. The equality check is an invariant for batching: preference scores must not change when examples are padded to a common length. Summing rather than averaging is intentional here because DPO compares the likelihood of complete responses; if response lengths are allowed to differ, report a separate length analysis rather than silently changing the objective.

Pairwise reward models

A reward model maps a prompt-response sequence to a scalar. The tiny bag-of-tokens model below is not intended to understand language; it makes the pairwise update inspectable. Its embedding pools the response tokens and a linear head produces one reward. The training target is only the ordering, not an absolute score.

VOCAB = 8
PAIRS_W = np.array([
    [1, 1, 3],
    [1, 3, 4],
    [1, 1, 4],
    [3, 1, 4],
])
PAIRS_L = np.array([
    [2, 2, 3],
    [2, 3, 4],
    [2, 2, 4],
    [3, 2, 4],
])
PAIR_MASK = np.ones_like(PAIRS_W)


def reward_scores(embedding, head, tokens, mask):
    pooled = (embedding[tokens] * mask[..., None]).sum(axis=1)
    pooled /= np.maximum(mask.sum(axis=1, keepdims=True), 1)
    return pooled @ head


def reward_loss_and_gradients(embedding, head, preferred, rejected, mask):
    n = len(preferred)
    preferred_hidden = (embedding[preferred] * mask[..., None]).sum(axis=1)
    rejected_hidden = (embedding[rejected] * mask[..., None]).sum(axis=1)
    lengths = np.maximum(mask.sum(axis=1, keepdims=True), 1)
    preferred_hidden /= lengths
    rejected_hidden /= lengths
    margins = (preferred_hidden - rejected_hidden) @ head
    sigmoid_neg = 1 / (1 + np.exp(margins))
    loss = np.logaddexp(0.0, -margins).mean()

    grad_head = ((-sigmoid_neg[:, None] / n) * (preferred_hidden - rejected_hidden)).sum(axis=0)
    grad_embedding = np.zeros_like(embedding)
    grad_hidden_preferred = (-sigmoid_neg[:, None] / n) * head
    grad_hidden_rejected = -grad_hidden_preferred
    for row in range(n):
        for token in preferred[row]:
            grad_embedding[token] += grad_hidden_preferred[row] / lengths[row, 0]
        for token in rejected[row]:
            grad_embedding[token] += grad_hidden_rejected[row] / lengths[row, 0]
    return loss, grad_embedding, grad_head, margins


embedding = rng.normal(scale=0.1, size=(VOCAB, 8))
head = rng.normal(scale=0.1, size=8)
for _ in range(180):
    loss, grad_embedding, grad_head, _ = reward_loss_and_gradients(
        embedding, head, PAIRS_W, PAIRS_L, PAIR_MASK
    )
    embedding -= 0.08 * grad_embedding
    head -= 0.08 * grad_head

loss, _, _, margins = reward_loss_and_gradients(
    embedding, head, PAIRS_W, PAIRS_L, PAIR_MASK
)
print("pairwise loss:", round(float(loss), 4))
print("preferred margins:", np.round(margins, 3).tolist())
print("pair accuracy:", round(float((margins > 0).mean()), 3))
assert np.all(margins > 0)

pairwise loss: 0.0556
preferred margins: [4.501, 2.251, 4.501, 2.251]
pair accuracy: 1.0

The reward model learns an ordering: all four preferred margins are positive on this tiny training set. That is a fit check, not evidence that the model learned helpfulness. The next experiment changes the data-generating process so the same optimization machinery can be shown to learn the wrong property.

Direct Preference Optimization (DPO)

For a full language model, the policy and reference log-probabilities come from two forward passes over the chosen and rejected responses. To isolate the DPO update, represent four complete responses by four trainable policy logits. The candidate log-probabilities below are the exact same quantities that the sequence log-probability function would return for a real decoder; only the model that produces them has been compressed to a lookup table.

def dpo_loss(policy_logits, reference_logits, chosen, rejected, beta=0.2):
    policy_logprobs = log_softmax(policy_logits)
    reference_logprobs = log_softmax(reference_logits)
    rows = np.arange(len(chosen))
    delta = (
        policy_logprobs[rows, chosen] - reference_logprobs[rows, chosen]
        - policy_logprobs[rows, rejected] + reference_logprobs[rows, rejected]
    )
    return np.logaddexp(0.0, -beta * delta).mean()


def dpo_gradient(policy_logits, reference_logits, chosen, rejected, beta=0.2):
    policy_logprobs = log_softmax(policy_logits)
    reference_logprobs = log_softmax(reference_logits)
    rows = np.arange(len(chosen))
    delta = (
        policy_logprobs[rows, chosen] - reference_logprobs[rows, chosen]
        - policy_logprobs[rows, rejected] + reference_logprobs[rows, rejected]
    )
    # d loss / d delta, averaged over the pair batch.
    coefficient = -beta / (1 + np.exp(beta * delta)) / len(rows)
    gradient = np.zeros_like(policy_logits)
    for i, row in enumerate(rows):
        d_logprob = np.zeros(policy_logits.shape[1])
        d_logprob[chosen[i]] += coefficient[i]
        d_logprob[rejected[i]] -= coefficient[i]
        probabilities = np.exp(policy_logprobs[row])
        gradient[row] += d_logprob - probabilities * d_logprob.sum()
    return gradient


# Candidate IDs stand for complete responses for two prompts.
chosen = np.array([0, 2])
rejected = np.array([1, 3])
reference_logits = np.array([
    [1.4, 1.1, 0.2, 0.0],
    [0.1, 0.0, 1.2, 1.0],
])
policy_logits = reference_logits.copy()
before_logprobs = log_softmax(policy_logits)
before = before_logprobs[np.arange(2), chosen] - before_logprobs[np.arange(2), rejected]

for _ in range(80):
    policy_logits -= 0.35 * dpo_gradient(
        policy_logits, reference_logits, chosen, rejected
    )

after_logprobs = log_softmax(policy_logits)
after = after_logprobs[np.arange(2), chosen] - after_logprobs[np.arange(2), rejected]
print("DPO loss:", round(float(dpo_loss(
    policy_logits, reference_logits, chosen, rejected
)), 4))
print("chosen-minus-rejected margins before:", np.round(before, 3).tolist())
print("chosen-minus-rejected margins after: ", np.round(after, 3).tolist())
assert np.all(after > before)

DPO loss: 0.4779
chosen-minus-rejected margins before: [0.3, 0.2]
chosen-minus-rejected margins after:  [2.749, 2.649]

DPO increases the policy’s relative likelihood of the preferred completions while the reference stays fixed. The loss does not say that the preferred response is factually correct; it says only that it should outrank the rejected response according to the labels. In practice, evaluate the policy at each checkpoint with held-out preference pairs and with a ground-truth evaluator that never supplies gradients.

Length bias

A preference dataset can make an accidental feature look like quality. Here every training winner is longer, so a reward model that sees only response length can achieve perfect training accuracy. The held-out pairs reverse the correlation: the preferred answer is concise. The evaluator is known in advance, which lets us measure the shortcut directly instead of debating whether a learned judge is trustworthy.

def pairwise_loss_from_lengths(weight, bias, preferred, rejected):
    reward_preferred = weight * preferred + bias
    reward_rejected = weight * rejected + bias
    return np.logaddexp(0.0, -(reward_preferred - reward_rejected)).mean()


train_preferred = np.array([6.0, 5.0, 7.0, 8.0])
train_rejected = np.array([3.0, 2.0, 4.0, 5.0])
heldout_preferred = np.array([2.0, 3.0, 1.0])
heldout_rejected = np.array([5.0, 6.0, 4.0])
length_weight = 0.0
length_bias = 0.0

for _ in range(120):
    differences = train_preferred - train_rejected
    margins = length_weight * differences
    derivative = -1 / (1 + np.exp(margins))
    gradient_weight = (derivative * differences).mean()
    length_weight -= 0.25 * gradient_weight

train_margin = length_weight * (train_preferred - train_rejected)
heldout_margin = length_weight * (heldout_preferred - heldout_rejected)
print("learned length coefficient:", round(float(length_weight), 3))
print("training preference accuracy:", round(float((train_margin > 0).mean()), 3))
print("held-out preference accuracy:", round(float((heldout_margin > 0).mean()), 3))
assert np.all(train_margin > 0)
assert np.all(heldout_margin < 0)

learned length coefficient: 1.868
training preference accuracy: 1.0
held-out preference accuracy: 0.0

The model reaches perfect training accuracy by assigning a positive value to length, then fails every held-out pair because the intended preference is reversed. This is the smallest possible overoptimization example: the optimizer is behaving correctly with respect to the labels, while the labels fail to express the intended property. A real study should vary length independently of quality, report response-length distributions, and keep a ground-truth evaluator outside the reward or DPO loop.

Reliability checks

At every preference-training checkpoint, report:

• held-out pairwise accuracy, separated by length-balanced and length-confounded subsets;
• the ground-truth evaluator score, never used for gradients;
• response length and token-level entropy;
• the distance from the reference policy, such as a KL estimate;
• SFT and general-suite scores from Chapters 07 and 08.

Noise, annotator disagreement, and reference choice should be swept explicitly. A falling DPO loss is compatible with a widening gap between reward-model score and the untouched evaluator. That gap is the measurement of proxy optimization, not an inconvenient outlier.

Shared DPO contract

The array implementation keeps the sequence objective visible; the reusable path takes chosen and rejected sequence log-probabilities relative to the same frozen reference. Because it consumes sequence scores rather than raw token positions, it can share the response-mask discipline from Chapter 08 and remain invariant to padding that is outside the response.

import torch

from proof_lm.posttraining import PreferencePair, dpo_loss

pair = PreferencePair(
    prompt="Explain entailment.",
    chosen="A premise guarantees the goal.",
    rejected="A premise is unrelated to the goal.",
)
policy_chosen = torch.tensor([-2.0, -1.8])
policy_rejected = torch.tensor([-2.5, -2.2])
reference_chosen = torch.tensor([-2.1, -1.9])
reference_rejected = torch.tensor([-2.4, -2.0])
per_example_dpo = dpo_loss(
    policy_chosen,
    policy_rejected,
    reference_chosen,
    reference_rejected,
    beta=0.2,
)
print("preference prompt:", pair.prompt)
print("DPO losses:", per_example_dpo.detach().round(decimals=4).tolist())
print("mean DPO loss:", round(float(per_example_dpo.mean().detach()), 4))
assert torch.isfinite(per_example_dpo).all()

preference prompt: Explain entailment.
DPO losses: [0.67330002784729, 0.6636000275611877]
mean DPO loss: 0.6685

Summary

• Pairwise reward modeling learns an ordering through -\log\sigma(r_w-r_l); it does not create an absolute notion of quality.
• DPO updates relative policy likelihoods against a frozen reference, with \beta controlling the strength of the preference margin.
• Sequence log-probabilities must use the same response mask discipline as SFT and must be invariant to padding.
• A length-only evaluator can fit every training preference and fail every held-out preference. Always report a ground-truth evaluator, length-balanced subsets, and the general regression suite.

Chapter 10 turns the proxy problem into a policy-gradient experiment with an executable checker and a deliberate loophole.

Exercises

Use the exercises to test the chapter’s invariants and connect the derivations to the reusable implementation. Solutions are hidden in the notebook source and are available through the course tooling when needed.

[P9.1] Length confound

Suppose every preferred response in a dataset is exactly twice the length of its rejected counterpart. Describe what a reward model trained on this data learns to score, and specify one held-out pair that detects the confound.

[P9.2]

Why does DPO compare policy and reference log-ratios for both the chosen and rejected responses? State what the beta parameter changes.