Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

RLHF Pipeline › Stage 2: Reward Model Training

Michael BrenndoerferDecember 29, 202513 min read

Language AI HandbookAll chapters

Part 37 of 67Alignment and RLHF

Chapter 280RLHF Pipeline › Stage 2: Reward Model Training

Reading section 4 of 11 · Stage 2: Reward Model Training · 13 minSection 4 of 11 · 13 min

Section 4 of 11 · 13 min read

RLHF Pipeline › Stage 2: Reward Model TrainingOne section · Sections in this chapter4 / 11 ↓

Sections in this chapter11

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

The reward model learns to predict which responses humans prefer. Building on the Bradley-Terry model from earlier chapters, it converts pairwise comparisons into a scalar reward signal that guides PPO optimization. This stage is arguably the most consequential in the entire pipeline, because any systematic error in the reward model will be amplified by PPO optimization. The policy will find the ways to maximize the reward model score, and if the reward model is wrong about what constitutes a good response in certain regions of the input space, PPO will drive the policy directly into those regions.

Think of the reward model as a compressed representation of annotator judgment. Thousands of pairwise comparisons, each representing a human's opinion about which of two responses is better, are distilled into a single neural network that can assign a quality score to any response in fractions of a second. This compression is both the strength and the weakness of the approach. The strength is that the reward model can generalize beyond the specific comparisons it was trained on. This provides quality judgments for entirely novel responses. The weakness is that the generalization may be imperfect, and the compression inevitably loses some of the nuance that human annotators would apply in novel situations.

The point is that the reward model is not trained to produce absolute quality scores, only to produce scores that correctly rank pairs of responses. Two reward models trained on the same data could assign completely different absolute values and yet behave identically in practice, because PPO optimization only cares about relative rankings, not absolute magnitudes. This property changes the interpretation for how you should interpret reward values and how you calibrate the KL penalty coefficient in stage 3.

Reward Model ArchitectureLink Copied

As discussed in the Reward Modeling chapter, the reward model typically shares architecture with the language model but replaces the language modeling head with a scalar output head. This architectural choice is deliberate: by starting from a language model that understands text in context, the reward model inherits much of its ability to interpret fine-grained meaning. The only modification is the final layer, which now outputs a single number representing quality rather than a distribution over vocabulary tokens.

In[11]:

Code

import torch
import torch.nn as nn


class RewardModel(nn.Module):
    """Reward model that outputs a scalar score for text quality."""

    def __init__(self, base_model, hidden_size):
        super().__init__()
        self.base_model = base_model

        self.reward_head = nn.Sequential(
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, 1),
        )

    def forward(self, input_ids, attention_mask):
        outputs = self.base_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=True,
        )

        last_hidden_state = outputs.hidden_states[-1]

        sequence_lengths = attention_mask.sum(dim=1) - 1
        batch_indices = torch.arange(input_ids.size(0), device=input_ids.device)
        final_hidden = last_hidden_state[batch_indices, sequence_lengths]

        reward = self.reward_head(final_hidden).squeeze(-1)
        return reward

The reward model extracts the hidden state at the final token position and passes it through a small neural network to produce a scalar reward. This design uses an important property of causal transformers: the final token's hidden state has attended to all previous tokens in the sequence, meaning it encapsulates information about the entire prompt-response pair. From this aggregated representation, the model can make an overall quality judgment that considers both what was asked and how well the response addresses it.

An important practical detail is that the reward model is initialized from the SFT model rather than the raw pretrained model. This initialization matters because the reward model needs to understand the specific response style established during SFT. If it were initialized from the pretrained model, it might assign high scores to responses that look like pretraining data (encyclopedia entries, forum posts) rather than assistant responses. Starting from the SFT model ensures the reward model operates in the same distribution as the responses it will be asked to score during PPO training.

The two-layer reward head sits on top of the transformer backbone. The first linear layer allows the network to learn a projection of the final hidden state that is maximally informative for quality prediction. The ReLU nonlinearity allows the network to model asymmetric quality differences. The second linear layer collapses this representation to a single scalar. Training this relatively small head while keeping the backbone frozen is one option, but joint fine-tuning of the entire model typically produces better reward models because it allows the transformer layers to reorganize their representations to make quality prediction easier.

Preference Data FormatLink Copied

Training data consists of prompts with two or more responses ranked by human annotators. Each training example captures a comparison: given the same prompt, which response did humans consider better? This pairwise structure is basic to the Bradley-Terry model, which we'll examine shortly.

In[12]:

Code

preference_data = [
    {
        "prompt": "How do I improve my writing skills?",
        "chosen": """Here are effective strategies to improve your writing:

1. **Read widely and actively** - Study how skilled authors structure sentences and arguments
2. **Write daily** - Even 15 minutes of practice builds consistency
3. **Seek feedback** - Join a writing group or find a trusted reader
4. **Edit ruthlessly** - Good writing is rewriting; cut unnecessary words
5. **Study grammar** - Understanding rules helps you break them purposefully

Start with one habit and build from there.""",
        "rejected": """Writing is hard but you can get better. Just practice a lot and read books. 
Maybe take a class or something. Good luck!""",
    },
    {
        "prompt": "Explain quantum entanglement simply",
        "chosen": """Quantum entanglement is like having two magic coins. When you flip them while they're together, then separate them, even across the universe, they always land the same way.

In physics terms: two particles become correlated so that measuring one instantly affects what you'll measure on the other, regardless of distance. Einstein called this "spooky action at a distance."

This isn't about sending information faster than light, but about correlations that can't be explained by classical physics.""",
        "rejected": """Quantum entanglement is a quantum mechanical phenomenon where particles become interconnected through their quantum states such that the quantum state of each particle cannot be described independently.""",
    },
]

Out[13]:

Console

Number of preference pairs: 2

Example prompt: How do I improve my writing skills?
Chosen response length: 481 chars
Rejected response length: 120 chars

The training examples typically show that chosen responses are longer, more detailed, and better structured than rejected ones. The reward model learns to associate these features, along with factual accuracy and tone, with higher scalar scores.

The quality of the preference data is the single biggest lever in determining how well the reward model generalizes. Constructing good preference data requires careful thought about how to generate diverse candidate responses and how to ensure annotators are evaluating the right properties. If all comparisons pit an obviously good response against an obviously bad one, the reward model learns only from extreme quality differences and may fail to distinguish between two good responses, which is precisely the judgment needed during PPO when the policy has already learned to produce reasonable outputs. Some comparison data should therefore deliberately contrast responses that are both reasonable but differ in subtle quality dimensions: accuracy, depth, appropriate confidence, and conciseness.

Annotator consistency is another important concern. Human annotators often disagree about which of two responses is better, especially when both responses are in the "good" quality range. The Bradley-Terry model assumes a total ordering of quality that may not exist. Two annotators might consistently prefer different response styles (one preferring brevity, another preferring depth), and training on their mixed judgments produces a reward model that has internalized a blend of their preferences. This blend may not correspond to any user's preference. Managing annotator consistency through training, calibration sessions, and inter-annotator agreement monitoring is an important part of high-quality preference data collection.

Reward Model Training LossLink Copied

The training objective for the reward model emerges from a probabilistic framework for modeling human preferences. We begin with a natural question: given two responses to the same prompt, how likely is it that a human prefers one over the other? The Bradley-Terry model provides an elegant answer by relating this preference probability to the difference in quality scores assigned to each response.

The core insight is that we can model preferences as arising from latent quality scores. If response $y_w$ has a higher quality score than response $y_l$ , then $y_w$ should be preferred more often. The Bradley-Terry model formalizes this intuition by expressing the preference probability as a function of the score difference:

P(y_w \succ y_l | x) = \sigma(r_\theta(x, y_w) - r_\theta(x, y_l))

where:

$P(y_w \succ y_l | x)$ : the probability that response $y_w$ is preferred over $y_l$ given prompt $x$
$\sigma$ : the logistic sigmoid function, $\sigma(z) = \frac{1}{1+e^{-z}}$
$r_\theta$ : the reward model with parameters $\theta$
$x$ : the input prompt
$y_w$ : the preferred ("winning") response
$y_l$ : the rejected ("losing") response

The sigmoid function $\sigma$ plays a useful role in this formulation. It converts the unbounded difference in reward scores into a probability between 0 and 1. This provides a smooth and differentiable mapping. When the reward for the winning response $r_\theta(x, y_w)$ is significantly larger than for the losing response, the difference becomes a large positive number, and the sigmoid approaches 1, showing near-certainty that $y_w$ would be preferred. Conversely, if the scores are equal, the sigmoid returns 0.5. This reflects maximum uncertainty. This elegant mathematical structure captures our intuition that larger quality differences should correspond to more decisive preferences.

Out[14]:

Visualization

Sigmoid curve showing how reward difference maps to preference probability in the Bradley-Terry model. — Bradley-Terry preference probability. The sigmoid function maps the reward difference between two responses to a probability of preference. A large positive difference implies near-certainty that the chosen response is better, while zero difference indicates equal preference likelihood.

To train the model, we minimize the negative log-likelihood of observing the preferences in our dataset. For a single preference pair, the loss is:

\mathcal{L}_{RM} = -\log \sigma(r_\theta(x, y_w) - r_\theta(x, y_l))

where:

$\mathcal{L}_{RM}$ : the scalar loss value to be minimized
$\sigma$ : the logistic sigmoid function, $\sigma(z) = \frac{1}{1+e^{-z}}$
$r_\theta$ : the reward model with parameters $\theta$ that assigns a scalar score to a prompt-response pair
$x$ : the input prompt or instruction
$y_w$ : the "winning" or preferred response
$y_l$ : the "losing" or rejected response
$r_\theta(x, y_w) - r_\theta(x, y_l)$ : the difference in reward scores (which we want to be positive)

Understanding why this loss works requires examining what happens during optimization. When the model correctly assigns a higher score to the preferred response (making the difference positive and large), the sigmoid outputs a value close to 1, and the negative log becomes small. When the model incorrectly ranks the responses (making the difference negative), the sigmoid outputs a value close to 0, and the negative log becomes very large, creating a strong gradient signal to correct this error. This objective therefore directly maximizes the likelihood that the model assigns a higher score to the preferred response $y_w$ than the rejected response $y_l$ , which is precisely what we want from a reward model.

In[15]:

Code

import torch


def compute_reward_model_loss(reward_model, batch, tokenizer, device):
    """Compute Bradley-Terry loss for preference learning."""

    prompts = batch["prompts"]
    chosen_responses = batch["chosen"]
    rejected_responses = batch["rejected"]

    chosen_texts = [p + c for p, c in zip(prompts, chosen_responses)]
    rejected_texts = [p + r for p, r in zip(prompts, rejected_responses)]

    chosen_tokens = tokenizer(
        chosen_texts, padding=True, return_tensors="pt"
    ).to(device)
    rejected_tokens = tokenizer(
        rejected_texts, padding=True, return_tensors="pt"
    ).to(device)

    chosen_rewards = reward_model(
        chosen_tokens["input_ids"], chosen_tokens["attention_mask"]
    )
    rejected_rewards = reward_model(
        rejected_tokens["input_ids"], rejected_tokens["attention_mask"]
    )

    loss = -torch.log(torch.sigmoid(chosen_rewards - rejected_rewards)).mean()

    accuracy = (chosen_rewards > rejected_rewards).float().mean()

    return loss, accuracy

Reward Model Training Metrics:

Loss: Bradley-Terry negative log-likelihood
Accuracy: Fraction where $r(\text{chosen}) > r(\text{rejected})$
Target accuracy: 70-80% (higher may indicate overfitting)

The 70-80% accuracy target deserves explanation, because it might seem counterintuitive to stop well below 100%. Higher accuracy on the training set typically means the model has overfit to the specific comparison pairs it was trained on and has lost the ability to generalize to novel responses. A reward model that achieves 95% accuracy on training data but only 65% on held-out data has learned to recognize specific text patterns rather than quality. In the downstream PPO stage, this overfitted reward model will assign high scores to responses that match patterns from the training comparison data rather than to better responses. Stopping at 70-80% accuracy is an empirical heuristic that tends to produce reward models with better generalization, though the right threshold varies with dataset size and diversity.

Reward Model CalibrationLink Copied

A necessary but often overlooked aspect is reward model calibration. The absolute reward values don't matter for ranking, since the Bradley-Terry model only uses differences between scores. However, the scale of these values significantly affects PPO training stability. If rewards are too large in magnitude, gradient updates can become unstable; if they vary too widely, optimization becomes difficult.

In[16]:

Code

import torch


def calibrate_reward_model(reward_model, calibration_data, tokenizer, device):
    """Calibrate reward model to have zero mean and unit variance."""
    reward_model.eval()

    all_rewards = []

    with torch.no_grad():
        for batch in calibration_data:
            texts = batch["texts"]
            tokens = tokenizer(texts, padding=True, return_tensors="pt").to(
                device
            )
            rewards = reward_model(
                tokens["input_ids"], tokens["attention_mask"]
            )
            all_rewards.append(rewards.cpu())

    all_rewards = torch.cat(all_rewards)
    mean_reward = all_rewards.mean().item()
    std_reward = all_rewards.std().item()

    return mean_reward, std_reward


def normalized_reward(raw_reward, mean, std):
    """Normalize reward to zero mean and unit variance."""
    return (raw_reward - mean) / (std + 1e-8)

Calibrating the reward model to have approximately zero mean and unit variance at the start of PPO training provides a stable foundation for the optimization. The zero mean ensures that the policy does not start with a systematic positive or negative reward bias, which could cause it to initially over-generate or under-generate regardless of quality. The unit variance ensures that the KL penalty coefficient $\beta$ has a consistent meaning across different training runs: a $\beta$ of 0.05 represents the same trade-off between reward and KL regularization regardless of whether the raw reward model happens to output values in the range $[-1, 1]$ or $[-10, 10]$ .

The calibration should be computed on a representative sample of prompt-response pairs that reflect the distribution the policy will encounter during PPO training. Using the SFT model to generate calibration responses is the most common approach, since the SFT model provides the starting point for the PPO policy. This ensures that the normalization statistics are computed at the operating point of the system, not on some other distribution.