DPO loss and implicit rewards · AI Engineering

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.

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.
dpo_loss(...) returns the mean DPO loss over the batch.
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

L = \frac{1}{N} \sum_{i=1}^{N} -\log \sigma(z_i), \qquad \sigma(z) = \frac{1}{1 + e^{-z}}.

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.

Name	Type	Shape	Meaning
`policy_logprob`	float array	`(N,)`	log probability of a response under the policy
`ref_logprob`	float array	`(N,)`	log probability of the same response under the reference
`policy_logprob_preferred`	float array	`(N,)`	policy log probability of each preferred response
`policy_logprob_rejected`	float array	`(N,)`	policy log probability of each rejected response
`ref_logprob_preferred`	float array	`(N,)`	reference log probability of each preferred response
`ref_logprob_rejected`	float array	`(N,)`	reference log probability of each rejected response
`beta`	float		the strength of the reference constraint, above 0

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

python

policy_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)

text

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

python

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

text

0.6802247996122515

python

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

text

grad_preferred
[-0.01500553 -0.0183278  -0.01583403]
grad_rejected
[0.01500553 0.0183278  0.01583403]