Direct Preference Optimization (DPO)

The problem

Learning from human preferences used to need two steps: train a reward model, then use reinforcement learning against it. DPO (Direct Preference Optimization) skips both and trains the language model directly on pairs of answers. For a prompt $x$ with a chosen answer $y_c$ and a rejected answer $y_r$, it needs a policy $\pi_\theta$ (the model being trained) and a frozen reference $\pi_{\text{ref}}$ (the model before this training, usually the SFT (supervised fine-tuning) model).

The loss

For an answer $y$, how much more likely has the policy made it, compared with the reference?

Here $\log \pi(y \mid x)$ is the sum of the log-probabilities of the answer's tokens. The loss wants $\Delta_c$ (chosen) to grow more than $\Delta_r$ (rejected):

where $\sigma$ is the sigmoid and $\beta$ (beta) is a constant that controls how strongly the policy may move away from the reference.

Example. At the start of training, $\pi_\theta = \pi_{\text{ref}}$, so $\Delta_c = \Delta_r = 0$ and the loss is $-\log \sigma(0) = \log 2 = 0.693$: a good first check.

Later, the sequence log-probabilities are: policy $-10$ for the chosen and $-12$ for the rejected; reference $-11$ for both. Then $\Delta_c = +1$ and $\Delta_r = -1$. With $\beta = 0.1$: $-\log \sigma(0.1 \cdot 2) = 0.598$, lower than $0.693$: the policy has moved in the right direction.

The length-normalized variant

A longer answer has a sum of more log-probabilities, so its $\Delta$ tends to be larger in size. OLMo 3 divides each $\Delta$ by the length of its answer before taking the difference:

Example. The same log-probabilities, with the chosen answer 10 tokens long and the rejected one 20, and $\beta = 5$. The normalized values are $\Delta_c / 10 = 0.1$ and $\Delta_r / 20 = -0.05$; the difference is $0.15$, and $-\log \sigma(5 \cdot 0.15) = 0.387$.

The log-probability of an answer

sequence_logprob computes $\log \pi(y \mid x)$ from the model's logits: for each position take the log-softmax (section 2), pick the entry of the actual token, and sum over the answer's positions only (where mask is true: the prompt is not part of it).

What to implement

• sequence_logprob(logits, tokens, mask): logits has shape $(T, V)$ and row $t$ is the model's distribution for tokens[t] (already aligned: no shift to do). Return the sum of the log-probabilities of the tokens where mask is true.
• dpo_loss(pi_c, pi_r, ref_c, ref_r, beta): arrays of shape $(N,)$ with the sequence log-probabilities of the chosen and rejected answers under the policy and the reference. Return the mean loss over the $N$ pairs.
• dpo_loss_length_normalized(pi_c, pi_r, ref_c, ref_r, len_c, len_r, beta): the variant above.

Requirements

• Stable for any values: $-\log\sigma(z)$ is np.logaddexp(0, -z).
• No loops over pairs or tokens.