Direct Preference Optimisation

RLHF maximises expected reward under a KL penalty toward a reference policy. That problem is solved analytically by a Boltzmann tilt of the reference, which means any reward function can be written in terms of the optimal policy it induces.

Inverting eq. 1 gives $r(x,y) = \beta \log \frac{\pi(y \mid x)}{\pi_{\text{ref}}(y \mid x)} + \beta \log Z(x)$. Under Bradley–Terry only reward differences within a prompt matter, so the intractable partition function cancels and the preference likelihood becomes a logistic loss on the implicit reward margin.

$\beta$ is not a learning rate. It sets how far the policy is allowed to move from the reference, and it is the only handle on the failure mode where both chosen and rejected log-probabilities fall together while the margin technically improves.

import torch.nn.functional as F
from torch import Tensor


def dpo_loss(
    pi_chosen: Tensor, pi_rejected: Tensor,      # policy logprobs, summed
    ref_chosen: Tensor, ref_rejected: Tensor,    # reference logprobs, frozen
    beta: float = 0.1,
) -> Tensor:
    chosen_logratio = pi_chosen - ref_chosen
    rejected_logratio = pi_rejected - ref_rejected
    margin = beta * (chosen_logratio - rejected_logratio)
    return -F.logsigmoid(margin).mean()

The log-probabilities must be summed over completion tokens only, with the prompt masked out — averaging instead silently reweights the objective by completion length and biases the policy toward short answers. Track chosen and rejected log-ratios separately during training: a healthy run raises the margin, an unhealthy one lowers both terms and widens the gap by falling more slowly on the chosen side.