Reward Model Training Loss

Problem

Compute the Bradley-Terry reward model loss: the mean of softplus(-(r_preferred - r_rejected)). Use the numerically stable form max(y, 0) + log1p(exp(-|y|)) with y = -(r_preferred - r_rejected). An empty batch returns 0.0.

Hint

Training the reward model is simply logistic regression on preference pairs.

Test cases

case 1

in: [[2,0],[0,1]]

expected: 0.7200948492805976

case 2

in: [[1],[1]]

expected: 0.6931471805599453

case 3

in: [[],[]]

expected: 0

Starter codePython

import math
def reward_model_loss(r_preferred, r_rejected):
    # Your code here
    pass

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