DPO Loss Value

Problem

Compute the mean DPO loss over preference pairs.

With margin = beta * ((policy_chosen - ref_chosen) - (policy_rejected - ref_rejected)), each pair contributes softplus(-margin) = log(1 + exp(-margin)). Use the numerically stable form max(y, 0) + log1p(exp(-|y|)) with y = -margin; an empty batch returns 0.0.

Hint

DPO turns the reward-model softmax loss into a loss directly on policy log-probabilities.

Test cases

case 1

in: [[0],[0],[0],[0],0.1]

expected: 0.6931471805599453

case 2

in: [[0],[-1],[0],[0],0.1]

expected: 0.6443966600735709

case 3

in: [[0.5],[0],[0],[0],1]

expected: 0.4740769841801067

case 4

in: [[],[],[],[],0.5]

expected: 0

Starter codePython

import math
def dpo_loss(policy_chosen, policy_rejected, ref_chosen, ref_rejected, beta):
    # Your code here
    pass

4 test cases · run your code against them in the editor