Reward Model Loss (Bradley-Terry)

The problem

People find it easier to say which of two answers is better than to write the perfect answer. A reward model learns from such comparisons: it gives every answer a score, its reward, so that the answer people preferred (the chosen one) gets a higher score than the one they did not (the rejected one). You will write the loss that trains it, its gradient, and a way to measure how often it ranks the pair correctly.

The Bradley-Terry loss

The Bradley-Terry model says that the probability that the chosen answer beats the rejected one depends only on the difference of their rewards:

where $\sigma$ is the sigmoid, $r_c$ the reward of the chosen answer and $r_r$ the reward of the rejected one. The loss is the negative log of that probability, averaged over $N$ pairs:

Example. One pair.

• $r_c = 2$, $r_r = 0.5$: the difference is $d = 1.5$, $\sigma(1.5) = 0.82$, loss $-\log 0.82 = 0.20$. Correct ranking: small loss.
• $r_c = 0.5$, $r_r = 2$: $d = -1.5$, $\sigma(-1.5) = 0.18$, loss $1.70$. Wrong ranking: large loss.
• $r_c = 102$, $r_r = 100.5$: the same difference as the first case, the same loss $0.20$. Only the difference counts, never the absolute scale.

Why not compute $-\log \sigma(d)$ directly? For a very wrong ranking, such as $d = -1000$, $\sigma(d)$ rounds to 0 and the log is $-\infty$. The identity $-\log \sigma(d) = \log(1 + e^{-d})$ is computed safely by np.logaddexp(0, -d), which gives $1000$.

The gradient

With $d = r_c - r_r$, the derivative of $-\log \sigma(d)$ with respect to $d$ is $-\sigma(-d) = -(1 - \sigma(d))$: the more wrong the ranking, the closer to $-1$. So the gradient with respect to $r_c$ is that, and for $r_r$ it is the opposite, both divided by $N$ because of the mean.

Example. For the first pair above ($d = 1.5$): $-\sigma(-1.5) = -0.18$ for $r_c$ and $+0.18$ for $r_r$ (with $N = 1$). Training moves against the gradient, so it raises $r_c$ and lowers $r_r$, only a little: the ranking is already right. For the second pair ($d = -1.5$): $-0.82$ and $+0.82$, a strong push.

Accuracy

The accuracy is the fraction of pairs where the chosen answer gets the higher reward.

Example. Three pairs with r_chosen = [2, 1, 0] and r_rejected = [1, 2, -1]: the chosen answer wins in pairs 1 and 3, loses in pair 2: accuracy $2/3$.

What to implement

For arrays r_chosen and r_rejected of shape $(N,)$:

• preference_loss(r_chosen, r_rejected): the mean loss, stable for any rewards.
• preference_grads(r_chosen, r_rejected): the gradients with respect to r_chosen and to r_rejected, as a pair of arrays.
• preference_accuracy(r_chosen, r_rejected): the accuracy.

Requirements

• Use np.logaddexp(0, -d) for the loss, so that $d = -1000$ gives 1000 instead of inf.
• No loops.