Reward-Model Margin Loss

Background

The reward model (RM) in RLHF is trained on preference pairs. The standard objective is the Bradley-Terry-style log-sigmoid margin loss: the RM should score the chosen completion above the rejected one, with the loss vanishing as the margin grows.

Problem statement

Implement rm_loss(r_chosen, r_rejected):

$$\mathcal{L} \;=\; \frac{1}{N}\sum_i -\log \sigma\!\bigl(r^+_i - r^-_i\bigr) \;=\; \frac{1}{N}\sum_i \mathrm{softplus}\!\bigl(-(r^+_i - r^-_i)\bigr)$$

Use the numerically stable softplus identity to avoid overflow when the margin is very negative.

Input

• r_chosen — 1-D array of floats, length $N$. RM scores on the chosen completions.
• r_rejected — 1-D array of floats, length $N$. RM scores on the rejected completions.

Output

Returns a plain Python float >= 0.

Examples

Example 1 — chosen well above rejected → near-zero loss

Input:  r_chosen=[10, 10, ...], r_rejected=[-10, -10, ...]
Output: ~0

Example 2 — chosen below rejected → large loss

Input:  r_chosen=[-10, ...], r_rejected=[10, ...]
Output: ~20

Constraints

• Length mismatch raises ValueError.
• Use np.logaddexp(0.0, -(r_chosen - r_rejected)) for stability — direct $-\log \sigma(x)$ overflows when $x \ll 0$.
• Returns a plain Python float.

Notes

• Margin interpretation. The gradient is largest at margin 0 and decays as the chosen score moves above rejected — pushes the RM to develop a clear margin, then leaves it alone.
• DPO connection. DPO replaces the explicit RM with this same loss form, taking $r = \beta(\log\pi_\theta - \log\pi_{\text{ref}})$.