Reward Modeling › Code Implementation
Michael BrenndoerferDecember 24, 202514 min read
Language AI HandbookAll chapters Part 37 of 67Alignment and RLHF
Reward Modeling › Code ImplementationOne section · Sections in this chapter7 / 10 ↓
Sections in this chapter10
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
Let's build a reward model from scratch using a small pretrained transformer. We'll implement the architecture, loss function, and training loop, then assess the model's ability to generalize to new preference pairs.
Setting Up the Environment
In[7]:
Code!uv pip install transformers torch matplotlib numpy
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
import torch.optim
from transformers import AutoModel, AutoTokenizer
import numpy as np
import matplotlib.pyplot as plt
from dataclasses import dataclass
from typing import List, Tuple, Optional
import warnings
warnings.filterwarnings('ignore')Reward Model Architecture
We'll build a reward model by adding a value head to a pretrained transformer. The value head is a simple linear layer that projects the final hidden state to a scalar.
In[8]:
Codeclass RewardModel(nn.Module):
"""
Reward model built on a pretrained transformer.
Takes (prompt, response) pairs and outputs scalar rewards.
"""
def __init__(self, model_name: str, dropout: float = 0.1):
super().__init__()
# Load pretrained transformer as the backbone
self.backbone = AutoModel.from_pretrained(model_name)
hidden_size = self.backbone.config.hidden_size
# Value head: projects final hidden state to scalar reward
self.value_head = nn.Sequential(
nn.Dropout(dropout), nn.Linear(hidden_size, 1)
)
def forward(
self,
input_ids: torch.Tensor,
attention_mask: torch.Tensor,
) -> torch.Tensor:
"""
Compute reward for input sequences.
Args:
input_ids: Token IDs [batch_size, seq_len]
attention_mask: Attention mask [batch_size, seq_len]
Returns:
Tensor: Scalar rewards [batch_size]
"""
# Get hidden states from backbone
outputs = self.backbone(
input_ids=input_ids, attention_mask=attention_mask
)
hidden_states = outputs.last_hidden_state # [batch, seq, hidden]
# Find position of last non-padding token for each sequence
# Sum attention mask to get sequence lengths
seq_lengths = attention_mask.sum(dim=1) - 1 # -1 for 0-indexing
batch_indices = torch.arange(
hidden_states.size(0), device=hidden_states.device
)
# Extract final token hidden state for each sequence
final_hidden = hidden_states[
batch_indices, seq_lengths
] # [batch, hidden]
# Project to scalar reward
rewards = self.value_head(final_hidden).squeeze(-1) # [batch]
return rewardsLet's verify the architecture outputs the expected shapes.
In[9]:
Code# Initialize model and tokenizer
model_name = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
reward_model = RewardModel(model_name)
# Create a sample input
sample_text = "What is machine learning? Machine learning is a field of AI."
inputs = tokenizer(sample_text, return_tensors="pt", padding=True)
# Forward pass
with torch.no_grad():
reward = reward_model(inputs["input_ids"], inputs["attention_mask"])Out[10]:
ConsoleInput shape: torch.Size([1, 15]) Reward shape: torch.Size([1]) Reward value: 0.2864
The output confirms that the model processes the tokenized sequence and produces a single scalar score (batch size 1, output dimension 1). This scalar represents the reward for the prompt-response pair, which will be used to rank different responses against each other.
Preference Dataset
Now let's create a dataset class that handles preference pairs. Each example contains a prompt with a chosen (preferred) and rejected response.
In[11]:
Code@dataclass
class PreferencePair:
prompt: str
chosen: str
rejected: str
class PreferenceDataset(Dataset):
"""Dataset of preference pairs for reward model training."""
def __init__(
self, pairs: List[PreferencePair], tokenizer, max_length: int = 256
):
self.pairs = pairs
self.tokenizer = tokenizer
self.max_length = max_length
def __len__(self):
return len(self.pairs)
def __getitem__(self, idx: int):
pair = self.pairs[idx]
# Tokenize prompt and chosen response
# Passing two strings automatically adds the separator token
chosen_enc = self.tokenizer(
pair.prompt,
pair.chosen,
max_length=self.max_length,
padding="max_length",
truncation=True,
return_tensors="pt",
)
# Tokenize prompt and rejected response
rejected_enc = self.tokenizer(
pair.prompt,
pair.rejected,
max_length=self.max_length,
padding="max_length",
truncation=True,
return_tensors="pt",
)
return {
"chosen_ids": chosen_enc["input_ids"].squeeze(0),
"chosen_mask": chosen_enc["attention_mask"].squeeze(0),
"rejected_ids": rejected_enc["input_ids"].squeeze(0),
"rejected_mask": rejected_enc["attention_mask"].squeeze(0),
}Let's create some synthetic preference data for demonstration.
In[12]:
Code# Synthetic preference pairs for demonstration
# 20 distinct pairs covering diverse topics so train and val sets are independent
all_preference_pairs = [
PreferencePair(
prompt="Explain gravity simply.",
chosen="Gravity is the force that pulls objects toward each other. The Earth pulls you down, which is why you stay on the ground.",
rejected="Gravity is described by Einstein's field equations relating the curvature of spacetime to energy-momentum.",
),
PreferencePair(
prompt="What is photosynthesis?",
chosen="Photosynthesis is how plants make food from sunlight, water, and carbon dioxide, creating oxygen as a byproduct.",
rejected="It's a plant thing.",
),
PreferencePair(
prompt="How do computers work?",
chosen="Computers process information using tiny electronic switches called transistors that can be on or off, representing 1s and 0s.",
rejected="Computers work by executing machine code instructions on a von Neumann architecture with fetch-decode-execute cycles.",
),
PreferencePair(
prompt="Why is the sky blue?",
chosen="Sunlight contains all colors. Blue light scatters more than other colors when it hits air molecules, so we see blue when we look up.",
rejected="The sky is blue.",
),
PreferencePair(
prompt="What causes rain?",
chosen="Water evaporates from oceans and lakes, rises as vapor, cools in clouds, and falls as rain when droplets get heavy enough.",
rejected="Rain is caused by the condensation of atmospheric water vapor into droplets when air masses cool below the dew point temperature.",
),
PreferencePair(
prompt="What is the immune system?",
chosen="Your immune system is your body's defense network. White blood cells recognize and destroy germs before they make you sick.",
rejected="The immune system is a complex network of cells and proteins.",
),
PreferencePair(
prompt="How does electricity work?",
chosen="Electricity is the flow of electrons through a wire, like water flowing through a pipe. A battery pushes those electrons to power devices.",
rejected="Electricity involves electromagnetic fields described by Maxwell's equations governing charge motion.",
),
PreferencePair(
prompt="What is DNA?",
chosen="DNA is a molecule shaped like a twisted ladder that holds the instructions for building every part of your body, stored in nearly every cell.",
rejected="dna stuff",
),
PreferencePair(
prompt="Why do we need sleep?",
chosen="Sleep lets your brain consolidate memories, repair cells, and clear out waste products that build up during the day.",
rejected="Sleep is a periodic state of rest involving complex neurological oscillations.",
),
PreferencePair(
prompt="What is climate change?",
chosen="Climate change refers to long-term shifts in global temperatures, largely driven by humans releasing greenhouse gases like CO2 into the atmosphere.",
rejected="It's when the weather changes.",
),
PreferencePair(
prompt="How do vaccines work?",
chosen="Vaccines teach your immune system to recognize a pathogen by exposing it to a harmless piece of the germ, so it's ready to fight the real thing.",
rejected="Vaccines involve antigen presentation and adaptive humoral immunity mediated by B-lymphocytes.",
),
PreferencePair(
prompt="What is inflation?",
chosen="Inflation is when prices rise over time, so the same amount of money buys less than it used to. Central banks try to keep it low and stable.",
rejected="inflation",
),
PreferencePair(
prompt="How does the internet work?",
chosen="Computers send data in small packets across a global network of cables and routers. Each packet finds its own path to the destination.",
rejected="The internet operates on TCP/IP protocols letting packet-switched communication across distributed nodes.",
),
PreferencePair(
prompt="What is evolution?",
chosen="Evolution is the process by which living things gradually change over generations. Individuals with helpful traits survive and pass them on.",
rejected="Evolution refers to allele frequency changes in populations across successive generations.",
),
PreferencePair(
prompt="Why do stars shine?",
chosen="Stars shine because they fuse hydrogen atoms into helium in their cores, releasing enormous amounts of energy as light and heat.",
rejected="stars are bright",
),
PreferencePair(
prompt="What is machine learning?",
chosen="Machine learning is when computers learn patterns from data to make predictions without being explicitly programmed for each task.",
rejected="ML utilizes gradient descent optimization on parameterized function approximators.",
),
PreferencePair(
prompt="How does a car engine work?",
chosen="A car engine burns fuel in cylinders, creating small explosions that push pistons down and turn the wheels through a series of gears.",
rejected="Internal combustion engines operate on thermodynamic cycles converting chemical energy to mechanical work via reciprocating pistons.",
),
PreferencePair(
prompt="What is democracy?",
chosen="Democracy is a system where citizens vote to choose their leaders and influence the laws that govern them.",
rejected="democracy",
),
PreferencePair(
prompt="How does sound travel?",
chosen="Sound travels as vibrations through air, like ripples on a pond. The vibrations reach your ears and your brain interprets them as sound.",
rejected="Sound propagates as longitudinal pressure waves through elastic media with frequency-dependent attenuation.",
),
PreferencePair(
prompt="What is gravity on the moon?",
chosen="The moon has about one-sixth of Earth's gravity because it is much less massive. That is why astronauts bounce when they walk there.",
rejected="Lunar surface gravity is approximately 1.62 m/s due to the moon's lower mass and radius.",
),
]
# Use 15 pairs for training and 5 held-out pairs for validation
train_preference_pairs = all_preference_pairs[:15]
val_preference_pairs = all_preference_pairs[15:]
# Duplicate training pairs to give enough batches for visible learning curves
extended_pairs = train_preference_pairs * 4 # 60 training examplesPreference Loss Function
The loss function implements the Bradley-Terry preference model. We compute rewards for both responses and maximize the probability of preferring the chosen response.
In[13]:
Codedef compute_preference_loss(
reward_model: nn.Module,
chosen_ids: torch.Tensor,
chosen_mask: torch.Tensor,
rejected_ids: torch.Tensor,
rejected_mask: torch.Tensor,
) -> Tuple[torch.Tensor, dict]:
"""
Compute the Bradley-Terry preference loss.
Args:
reward_model: The reward model
chosen_ids: Token IDs for chosen responses [batch, seq]
chosen_mask: Attention mask for chosen [batch, seq]
rejected_ids: Token IDs for rejected responses [batch, seq]
rejected_mask: Attention mask for rejected [batch, seq]
Returns:
loss: Scalar loss value
metrics: Dictionary with additional metrics
"""
# Compute rewards for both responses
chosen_rewards = reward_model(chosen_ids, chosen_mask)
rejected_rewards = reward_model(rejected_ids, rejected_mask)
# Preference loss: -log(sigmoid(r_chosen - r_rejected))
# Equivalent to: log(1 + exp(r_rejected - r_chosen))
reward_diff = chosen_rewards - rejected_rewards
loss = -F.logsigmoid(reward_diff).mean()
# Compute accuracy (how often chosen reward > rejected reward)
accuracy = (reward_diff > 0).float().mean()
# Average reward margin
margin = reward_diff.mean()
metrics = {
"loss": loss.item(),
"accuracy": accuracy.item(),
"margin": margin.item(),
"chosen_reward_mean": chosen_rewards.mean().item(),
"rejected_reward_mean": rejected_rewards.mean().item(),
}
return loss, metricsTraining Loop
Now we implement the complete training loop with logging and validation.
In[14]:
Codedef train_reward_model(
model: nn.Module,
train_loader: DataLoader,
val_loader: Optional[DataLoader],
num_epochs: int = 3,
learning_rate: float = 2e-5,
device: str = "cpu",
) -> dict:
"""
Train the reward model on preference data.
Args:
model: Reward model to train
train_loader: Training data loader
val_loader: Validation data loader (optional)
num_epochs: Number of training epochs
learning_rate: Learning rate for optimizer
device: Device to train on
Returns:
history: Dictionary with training metrics
"""
model = model.to(device)
optimizer = torch.optim.AdamW(
model.parameters(), lr=learning_rate, weight_decay=0.01
)
history = {
"train_loss": [],
"train_accuracy": [],
"val_loss": [],
"val_accuracy": [],
}
for epoch in range(num_epochs):
# Training phase
model.train()
train_metrics = {"loss": 0, "accuracy": 0, "count": 0}
for batch in train_loader:
chosen_ids = batch["chosen_ids"].to(device)
chosen_mask = batch["chosen_mask"].to(device)
rejected_ids = batch["rejected_ids"].to(device)
rejected_mask = batch["rejected_mask"].to(device)
optimizer.zero_grad()
loss, metrics = compute_preference_loss(
model, chosen_ids, chosen_mask, rejected_ids, rejected_mask
)
loss.backward()
# Gradient clipping for stability
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
train_metrics["loss"] += metrics["loss"] * len(chosen_ids)
train_metrics["accuracy"] += metrics["accuracy"] * len(chosen_ids)
train_metrics["count"] += len(chosen_ids)
# Compute epoch averages
train_loss = train_metrics["loss"] / train_metrics["count"]
train_acc = train_metrics["accuracy"] / train_metrics["count"]
history["train_loss"].append(train_loss)
history["train_accuracy"].append(train_acc)
# Validation phase
if val_loader is not None:
model.eval()
val_metrics = {"loss": 0, "accuracy": 0, "count": 0}
with torch.no_grad():
for batch in val_loader:
chosen_ids = batch["chosen_ids"].to(device)
chosen_mask = batch["chosen_mask"].to(device)
rejected_ids = batch["rejected_ids"].to(device)
rejected_mask = batch["rejected_mask"].to(device)
_, metrics = compute_preference_loss(
model,
chosen_ids,
chosen_mask,
rejected_ids,
rejected_mask,
)
val_metrics["loss"] += metrics["loss"] * len(chosen_ids)
val_metrics["accuracy"] += metrics["accuracy"] * len(
chosen_ids
)
val_metrics["count"] += len(chosen_ids)
val_loss = val_metrics["loss"] / val_metrics["count"]
val_acc = val_metrics["accuracy"] / val_metrics["count"]
history["val_loss"].append(val_loss)
history["val_accuracy"].append(val_acc)
return historyLet's train the model on our synthetic preference data.
In[15]:
Code# Create datasets -- val_preference_pairs are held-out pairs not seen during training
train_pairs = extended_pairs
val_pairs = val_preference_pairs
train_dataset = PreferenceDataset(train_pairs, tokenizer, max_length=128)
val_dataset = PreferenceDataset(val_pairs, tokenizer, max_length=128)
train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=8)
# Initialize fresh model for training
reward_model = RewardModel(model_name)
# Train
history = train_reward_model(
reward_model, train_loader, val_loader, num_epochs=5, learning_rate=2e-5
)Out[16]:
ConsoleTraining Results: ---------------------------------------- Epoch 1: Train Loss: 0.5974, Accuracy: 75.00% Val Loss: 0.3318, Accuracy: 100.00% Epoch 2: Train Loss: 0.1549, Accuracy: 100.00% Val Loss: 0.1605, Accuracy: 100.00% Epoch 3: Train Loss: 0.0074, Accuracy: 100.00% Val Loss: 0.0942, Accuracy: 100.00% Epoch 4: Train Loss: 0.0007, Accuracy: 100.00% Val Loss: 0.0688, Accuracy: 100.00% Epoch 5: Train Loss: 0.0002, Accuracy: 100.00% Val Loss: 0.0617, Accuracy: 100.00%
The model quickly learns to distinguish between preferred and rejected responses on this small dataset.
Visualizing Training Progress
Out[17]:
Visualization
The training curves demonstrate that the model effectively learns the preference ranking task, with validation accuracy tracking closely with training accuracy. This indicates the model is generalizing well to unseen preference pairs without significant overfitting.
Evaluating on New Examples
Let's assess the trained model on examples it hasn't seen during training.
In[18]:
Codedef score_response(
model, tokenizer, prompt: str, response: str, device: str = "cpu"
) -> float:
"""Score a single prompt-response pair."""
model.eval()
inputs = tokenizer(
prompt,
response,
return_tensors="pt",
padding=True,
truncation=True,
max_length=128,
)
with torch.no_grad():
reward = model(
inputs["input_ids"].to(device), inputs["attention_mask"].to(device)
)
return reward.item()
# Test on new examples
test_prompt = "What is machine learning?"
test_responses = [
(
"Good explanation",
"Machine learning is when computers learn patterns from data to make predictions without being explicitly programmed.",
),
(
"Too technical",
"ML utilizes gradient descent optimization on parameterized function approximators.",
),
("Too brief", "It's AI stuff."),
]
# Compute scores for all responses
scores = []
for label, response in test_responses:
score = score_response(reward_model, tokenizer, test_prompt, response)
scores.append((score, label, response))
# Sort by score (higher is better)
scores.sort(reverse=True)Out[19]:
ConsolePrompt: What is machine learning? Response Rankings: ------------------------------------------------------------ 1. [Good explanation] Score: -3.8208 "Machine learning is when computers learn patterns from data ..." 2. [Too technical] Score: -4.4391 "ML utilizes gradient descent optimization on parameterized f..." 3. [Too brief] Score: -4.6247 "It's AI stuff."
The model has learned to rank responses in a way that aligns with our training preferences: clear explanations over technical jargon or overly brief responses.
Analyzing Reward Distributions
A well-calibrated reward model should produce meaningful score separations between good and bad responses. Let's analyze the distribution of rewards on our validation set.
In[20]:
Code# Collect rewards from validation set
chosen_rewards = []
rejected_rewards = []
# Ensure model is on the correct device
device = next(reward_model.parameters()).device
reward_model.eval()
with torch.no_grad():
for batch in val_loader:
# Move inputs to the same device as the model
chosen_ids = batch["chosen_ids"].to(device)
chosen_mask = batch["chosen_mask"].to(device)
rejected_ids = batch["rejected_ids"].to(device)
rejected_mask = batch["rejected_mask"].to(device)
chosen_r = reward_model(chosen_ids, chosen_mask)
rejected_r = reward_model(rejected_ids, rejected_mask)
# Move to CPU for logging/plotting
chosen_rewards.extend(chosen_r.cpu().numpy().tolist())
rejected_rewards.extend(rejected_r.cpu().numpy().tolist())Out[21]:
Visualization
The clear separation between distributions indicates the model has learned meaningful preference distinctions.
Key Parameters
The key parameters for the Reward Model implementation are:
- model_name: The pretrained backbone (e.g.,
"distilbert-base-uncased"). Smaller models allow for faster iteration during experimentation. - dropout: Regularization applied to the value head (set to
0.1) to prevent overfitting on the small preference dataset. - learning_rate: A conservative rate (
2e-5) is used to fine-tune the backbone without destroying pretrained features. - batch_size: Set to
8for this demonstration, though larger batches are preferred for stability in full-scale training. - num_epochs: Training is limited to
5epochs to avoid overfitting on the small synthetic dataset.