Reward Model Loss

Problem

Compute the Bradley-Terry reward model loss: -mean(log sigmoid(r_chosen - r_rejected)), computed stably with softplus. The loss is 0.0 when chosen rewards exceed rejected rewards by a lot.

The signature is reward_model_loss(chosen_rewards, rejected_rewards).

Hint

Pairwise preference loss is softplus of the negated reward margin.

Test cases

case 1

in: [[1],[0]]

expected: 0.3132616875

case 2

in: [[0],[0]]

expected: 0.6931471806

case 3

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

expected: 0.4100375958

case 4

in: [[1],[-1]]

expected: 0.126928011

Starter codePython

import math
def reward_model_loss(chosen_rewards, rejected_rewards):
    # Your code here
    pass

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