4. Reward Model Forward Pass with Bradley-Terry Preference Loss

Problem

RLHF starts by training a reward model on human preference pairs: for each pair, an annotator picked a "chosen" response over a "rejected" one, and the reward model should learn to score the chosen response higher. The standard objective for this is the Bradley-Terry pairwise-preference loss: minimize -log(sigmoid(r_chosen - r_rejected)), which pushes the score gap wider without needing an absolute reward scale, only a consistent relative ordering.

Implement both the reward-head forward pass and the loss computation together, given already-pooled embeddings for the chosen and rejected response in each pair.

Source: src/4_reward_model_bradley_terry.py

def reward_model_preference_loss(
    reward_head: torch.nn.Module,
    chosen_embeddings: torch.Tensor,
    rejected_embeddings: torch.Tensor,
) -> torch.Tensor: ...

>>> head = nn.Linear(4, 1)
>>> emb = torch.randn(3, 4)
>>> reward_model_preference_loss(head, emb, emb).item()  # identical pair -> log(2)
0.6931471824645996

Step-by-Step Approach

1. Run reward_head on chosen_embeddings to get raw scores of shape (batch, 1), then .squeeze(-1) to (batch,).
2. Do the same for rejected_embeddings.
3. Compute the score gap chosen_scores - rejected_scores.
4. Apply F.logsigmoid to the gap rather than torch.log(torch.sigmoid(...)) — the fused version stays accurate even when the gap is a large negative number.
5. Negate and average over the batch to get the scalar Bradley-Terry loss.
6. Sanity check: if chosen and rejected embeddings are identical, the score gap is exactly 0 regardless of the head's weights, so the loss must equal -log(sigmoid(0)) = log(2) ≈ 0.6931.

The key insight is the identical-pair sanity check: because the gap is exactly zero whenever chosen equals rejected, the expected loss (log 2) doesn't depend on the reward head's random initialization at all, which makes it a strong, implementation-agnostic test to reach for whenever you're debugging a Bradley-Terry loss.

Reference solution

import torch
import torch.nn as nn
import torch.nn.functional as F


def reward_model_preference_loss(
    reward_head: nn.Module,
    chosen_embeddings: torch.Tensor,
    rejected_embeddings: torch.Tensor,
) -> torch.Tensor:
    chosen_scores = reward_head(chosen_embeddings).squeeze(-1)  # (batch,)
    rejected_scores = reward_head(rejected_embeddings).squeeze(-1)  # (batch,)
    # F.logsigmoid(x) computes log(sigmoid(x)) via -softplus(-x) internally,
    # which stays numerically stable even when x is very negative -- plain
    # log(sigmoid(x)) would underflow to log(0) = -inf for large negative x
    return -F.logsigmoid(chosen_scores - rejected_scores).mean()

Key Functions & Tricks

• reward_head(embeddings).squeeze(-1) — a scalar-output head naturally produces shape (batch, 1); squeeze to (batch,) before subtracting.
• F.logsigmoid(x) — numerically stable log-sigmoid; internally equivalent to -F.softplus(-x), avoiding underflow that plain log(sigmoid(x)) suffers for very negative x.
• .mean() over the batch — reduces per-pair losses to the single scalar an optimizer needs.
• The identical-pair sanity check (log(2)) — a weight-independent regression test worth keeping in any reward-model test suite.

How to Recognize This Pattern