DPO loss and implicit rewards

Problem

Direct preference optimization (DPO) trains a language model on pairs of responses. A person preferred one response of each pair and rejected the other. DPO needs no separate reward model. The log probability ratio between the policy and a frozen reference model acts as the reward.

Task

Write three functions. Each input holds one sequence log probability for each pair in a batch of N pairs.

1. implicit_reward(policy_logprob, ref_logprob, beta) returns the implicit reward r = \beta \, (\log \pi(y \mid x) - \log \pi_{ref}(y \mid x)) for each pair.
2. dpo_loss(...) returns the mean DPO loss over the batch.
3. dpo_grad(...) returns the gradient of that mean loss with respect to the two policy log probabilities.

Let r_w be the implicit reward of the preferred response and r_l the implicit reward of the rejected response. The margin of pair i is z_i = r_{w,i} - r_{l,i}. The loss is

The gradient of L with respect to the preferred policy log probability of pair i is -\beta \, \sigma(-z_i) / N. The gradient with respect to the rejected policy log probability is +\beta \, \sigma(-z_i) / N. The reference log probabilities are constants.

Input and output

All four log probability arguments have the same order of names in dpo_loss and dpo_grad: policy_logprob_preferred, policy_logprob_rejected, ref_logprob_preferred, ref_logprob_rejected.

• implicit_reward returns a float array of shape (N,).
• dpo_loss returns a Python float.
• dpo_grad returns a tuple (grad_preferred, grad_rejected) of two float arrays of shape (N,).

Constraints

• 1 <= N <= 1000 and 0 < beta <= 10.
• A log probability can be as small as -5000. The margin can be as large as several thousand in either direction. The loss and the gradient must stay finite and exact.
• Do not add a small constant inside a log. Use a stable form of -\log \sigma(z).
• Do not change the input arrays.
• Use numpy only.

Examples

pythonCopypolicy_preferred = np.array([-12.0, -20.0, -8.0])
policy_rejected = np.array([-15.0, -18.0, -9.0])
ref_preferred = np.array([-13.0, -19.0, -8.5])
ref_rejected = np.array([-14.0, -19.0, -8.5])
implicit_reward(policy_preferred, ref_preferred, 0.1)
implicit_reward(policy_rejected, ref_rejected, 0.1)

textCopy[ 0.1  -0.1   0.05]
[-0.1   0.1  -0.05]

pythonCopydpo_loss(policy_preferred, policy_rejected, ref_preferred, ref_rejected, 0.1)

textCopy0.6802247996122515

pythonCopygrad_preferred, grad_rejected = dpo_grad(policy_preferred, policy_rejected, ref_preferred, ref_rejected, 0.1)

textCopygrad_preferred
[-0.01500553 -0.0183278  -0.01583403]
grad_rejected
[0.01500553 0.0183278  0.01583403]

Source paper Direct Preference Optimization: Your Language Model is Secretly a Reward Model, Rafailov et al., 2023.