DL S062 · Reward Modeling
Train a reward model on preference pairs. Part of the 'Deep Learning & LLMs From Scratch' 80-session series.
🎯 Train a reward model on preference pairs.
Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 62 / 80 · Module M10 · ~2 hours
The story
From "what a good answer looks like" to "which of these two is better"
SFT gets a model to the vicinity of "well-formed and on-topic." But well-formed and on-topic doesn't mean helpful. Two responses to "How do I center a div in CSS?" might both be technically correct — one uses text-align, one uses flexbox. Which is better? SFT can't say. Its loss function has no notion of preference.
Here's the puzzle: we can't write down "helpful" as a loss function directly. We don't have a differentiable helpfulness detector. What we can do: pay human labelers to look at pairs of model outputs and click "A is better than B." Do this 50,000 times and you have a preference dataset. Train a small neural network to predict, given a prompt and a response, a scalar "how much would a human prefer this?" That model is the reward model (RM), and it's the object that RLHF (S063) and — as we'll see — even DPO (S064, which secretly doesn't need one but derives from the same theory) both revolve around.
- Derive the Bradley-Terry preference model from first principles: P(A > B) = sigmoid(r(A) - r(B)).
- Write down the reward-modeling loss: -log sigmoid(r_chosen - r_rejected), and explain why it works.
- Modify a HuggingFace CausalLM into a reward model: replace the LM head with a single scalar regression head.
- Train an RM on the UltraFeedback or HH-RLHF preference dataset with trl.RewardTrainer.
- Diagnose reward-model failure modes: length bias, agreement bias, and reward hacking on OOD prompts.
Prerequisites
- Session 058 (SFT). The reward model is initialized from your SFT'd model in most recipes.
- Session 022 (cross-entropy loss). RM loss is a special case of pairwise log-loss.
- Rudimentary probability — you should be comfortable with the term "logistic model" (from S005).
1 · The Bradley-Terry model, in 5 lines of math
Suppose each response y to a prompt x has some latent scalar "goodness" r(x, y) ∈ ℝ. A human comparing two responses yA and yB to the same prompt picks A over B with probability:
This is the Bradley-Terry model (Bradley & Terry 1952). It has one lovely property: it doesn't care about the absolute value of r, only the difference. Shift r by any constant and probabilities don't change. So r is only identifiable up to an additive constant — which is fine, we never need absolute rewards, only relative.
Where did the sigmoid come from? It's the logistic model of binary choice: given two things with utilities rA and rB, and a Gumbel-distributed noise added to each, the probability that noisy-A > noisy-B is exactly σ(rA - rB). Same trick as multinomial logistic regression. Same trick as softmax. Nothing new.
1.1 The RM loss falls out directly
If a human labeled ychosen ≻ yrejected, then the log-likelihood under our model is log σ(r_chosen - r_rejected). Negate for a loss:
That's it. That's the entire reward-modeling training objective. A one-line loss. Pairs, sigmoid, negative log.
import torch.nn.functional as F
def rm_loss(r_chosen, r_rejected):
return -F.logsigmoid(r_chosen - r_rejected).mean()2 · The reward model, architecturally
An RM is almost a language model. Same transformer backbone. Only one thing changes: instead of an LM head (Linear from d → V predicting next-token logits), we attach a regression head (Linear from d → 1 predicting a scalar).
import torch.nn as nn
from transformers import AutoModel
class RewardModel(nn.Module):
def __init__(self, base_name: str):
super().__init__()
self.backbone = AutoModel.from_pretrained(base_name) # transformer, no LM head
d = self.backbone.config.hidden_size
self.head = nn.Linear(d, 1, bias=False)
# initialize head near zero for stable early training
nn.init.zeros_(self.head.weight)
def forward(self, input_ids, attention_mask=None):
out = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
h = out.last_hidden_state # (B, T, d)
# Pool the last non-pad token's hidden state — the "final answer" representation
seq_lens = attention_mask.sum(dim=1) - 1
last_h = h[torch.arange(h.size(0)), seq_lens]
return self.head(last_h).squeeze(-1) # (B,)Two important details:
2.1 Pool the LAST non-pad token, not the first or the mean
The reward should reflect the entire response, weighted by the ending. Averaging over all tokens is a weak baseline (dilutes the signal). Taking the first token's hidden state loses the ending. Taking the last non-pad token is the convention: by that position, causal attention has aggregated the whole response.
For a chat-style prompt like [user turn] [assistant turn] [EOS], the last token is the EOS, whose hidden state has attended to the entire assistant response. That's what you want to score.
2.2 Initialize the head to zero (or near-zero)
If you Kaiming-init the regression head, at step 0 your rewards are big random numbers uniformly distributed over ±10 or so. Their differences are also random, so sigma(r_c - r_r) is uniform in [0,1] and the loss starts at -log(0.5) = 0.693. But the gradient magnitude is huge because the losses on individual pairs vary wildly. Training is unstable for the first few hundred steps.
Zero-init the head → all rewards start at 0 → all differences are 0 → sigmoid = 0.5 uniformly → loss starts at exactly 0.693 → clean, calm early gradients. Classic zero-init trick, same reason LoRA zeros B.
3 · Preference datasets — the actual game
Just like SFT, the data is the game. Two open datasets dominate:
3.1 Anthropic HH-RLHF (Helpful+Harmless)
~170K conversation pairs from Anthropic's early Claude data. Each example:
{
"chosen": "\n\nHuman: What are some pranks I can play on my parents?\n\nAssistant: I'd suggest...",
"rejected": "\n\nHuman: What are some pranks I can play on my parents?\n\nAssistant: You could tape their car doors shut..."
}Split into "helpful" (chosen is more helpful) and "harmless" (chosen refuses/redirects when appropriate). The dataset that trained Claude 1.
3.2 UltraFeedback (OpenBMB, 2023)
~64K prompts × 4 responses each from a mix of GPT-4, Claude, Llama-2, etc. Each response scored by GPT-4 on 4 dimensions (instruction-following, honesty, truthfulness, helpfulness). Convert to pairs by taking (best-scored, worst-scored) per prompt.
Advantage: the "labels" come from GPT-4, so you get much larger volume than paying humans. Disadvantage: your RM is now learning GPT-4's preferences, not humans'. In practice this is usually fine and sometimes better (GPT-4 is more consistent than crowdworkers).
3.3 Data quality is the whole game (again)
For RMs specifically, watch for:
- Length bias in labels. If human labelers systematically prefer longer answers, the RM will predict "longer = better" even when it isn't. Very well-documented failure. Some datasets (Anthropic's original HH) have >70% correlation between response length and preference — devastating.
- Agreement bias. Labelers agree with confident-sounding assertions even when wrong. RM learns "sound confident = high reward" → downstream model becomes a smug hallucinator.
- Task drift. If your prompts are all coding questions but your downstream use is general chat, the RM will generalize poorly.
A team fine-tuned an RM on their internal preference data (10K pairs, mostly customer support scenarios). RM loss looked great: dropped from 0.69 to 0.42, agreement with a holdout labeler 78%.
Then they used the RM in a PPO loop and the resulting model started giving very long answers to everything. A one-word yes/no question got a 400-token response.
Root cause: their preference data had a strong length bias. In 68% of their pairs, the chosen response was longer than the rejected. The RM correctly learned this bias. PPO exploited it.
Fix: retrain the RM with a length penalty (subtract β · len from the reward), or bias-corrected preference sampling. Cost: 3 days debugging + one whole PPO run wasted.
Lesson: after training an RM, always compute the correlation between predicted reward and response length. If |ρ| > 0.3, you have a length-bias problem.
4 · Training an RM with TRL
TRL's RewardTrainer implements the pairwise loss above. In 30 lines of code:
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from trl import RewardConfig, RewardTrainer
MODEL = "HuggingFaceTB/SmolLM-135M" # our SFT'd model in practice
tok = AutoTokenizer.from_pretrained(MODEL)
if tok.pad_token is None:
tok.pad_token = tok.eos_token
# AutoModelForSequenceClassification with num_labels=1 gives us a regression head automatically
model = AutoModelForSequenceClassification.from_pretrained(
MODEL, num_labels=1, torch_dtype="bfloat16"
)
model.config.pad_token_id = tok.pad_token_idData:
ds = load_dataset("Anthropic/hh-rlhf", split="train[:10000]")
def format_pair(ex):
return {
"chosen": ex["chosen"], # already includes the prompt
"rejected": ex["rejected"],
}
ds = ds.map(format_pair)Train:
cfg = RewardConfig(
output_dir="./rm-smollm-hh",
per_device_train_batch_size=8,
num_train_epochs=1,
learning_rate=1e-5,
bf16=True,
max_length=1024,
logging_steps=50,
)
trainer = RewardTrainer(model=model, args=cfg, train_dataset=ds, tokenizer=tok)
trainer.train()The trainer tokenizes both chosen and rejected, runs both through the model, computes the pairwise loss, backprops. One epoch is usually enough — RMs overfit fast.
4.1 Reading the training log
Two numbers to watch:
- Loss: starts at ~0.693 (log 2), should drop to 0.3–0.5. Much lower means overfitting on training pairs.
- Accuracy: fraction of pairs where rchosen > rrejected. Chance = 50%. Good RMs land at 70–80% on a holdout set. Above 90% either you have easy data or you're overfit.
If accuracy is ~50% after 1000 steps: your dataset format is broken (roles swapped? pad token wrong?) or LR is too high and rewards are collapsing to a constant.
5 · Sanity-checking your RM
Before shipping the RM into PPO, run these three diagnostics.
5.1 Length correlation
lens, rewards = [], []
for ex in val_ds:
ids = tok(ex["chosen"], return_tensors="pt", truncation=True, max_length=1024)
with torch.no_grad():
r = model(**ids).logits.squeeze().item()
lens.append(ids.input_ids.shape[1])
rewards.append(r)
import numpy as np
print(f"length ↔ reward correlation: {np.corrcoef(lens, rewards)[0,1]:.3f}")- |ρ| < 0.1 — clean.
- 0.1 < |ρ| < 0.3 — mild bias, monitor.
- |ρ| > 0.3 — you have a length-bias problem, PPO will exploit it.
5.2 Best-of-N sanity check
Given a prompt, sample 8 completions from your SFT model. Rank them by RM. Ask yourself (or a coworker): is the top-1 actually the best? If the RM ranks a garbled response above a clean one, something's wrong.
5.3 Adversarial prompts
Feed prompts the RM has never seen: math questions, code, non-English, edge-case ethics. The RM's rewards should be reasonable-ish (not wildly negative or wildly positive on OOD). Extreme rewards on OOD data are a sign of overfitting.
6 · The reward-hacking problem (preview)
Every RM is a proxy for what you actually want (helpful behavior). Optimize hard against a proxy and you Goodhart it — the model finds ways to score high without being helpful.
Classic examples from published RLHF runs:
- Length hacking: RM prefers longer → model produces essays for every question.
- Sycophancy: RM slightly prefers agreeable responses → model becomes a yes-man.
- Formatting hacking: RM saw more markdown in chosen responses → model produces bullet-listed hierarchical answers for everything.
- Refusal drift: RM was trained on many "I cannot help with that" chosen responses (from harmlessness data) → model refuses benign questions.
This is why RLHF (S063) uses a KL penalty to keep the policy close to the SFT baseline — you can't drift too far in reward-hacking directions without paying a KL cost. But the KL penalty is a bandaid, not a fix. The real fix is a better RM, which mostly means better data.
7 · Diagram — the RM training loop
Two forward passes per pair (one for chosen, one for rejected), one loss, one backward. Trivial to parallelize across pairs in a batch.
8 · The 2024–2025 reward-modeling landscape
Reward modeling was the sleepy branch of the alignment tree in 2023 — everyone was doing PPO/DPO downstream and the RM was treated as a black box. Then in 2024 several papers made clear that the RM is the alignment bottleneck, and the field exploded.
8.1 UltraFeedback — the workhorse preference dataset (Cui et al., 2023)
The most-cited preference dataset in 2024–2025 open work. 64K prompts × 4 responses (from GPT-4, GPT-3.5, Llama-2-70B-Chat, and other models) × GPT-4 grades on 4 axes: instruction-following, honesty, truthfulness, helpfulness. Yields ~340K binary preference pairs.
Why it dominates:
- Diverse response sources. Not all responses come from GPT-4 — the pairwise contrast has more signal.
- Multi-axis grading. You can train separate RMs for each axis, or a multi-objective RM.
- Open license. MIT.
Every Zephyr, Tulu, and StableLM alignment paper trains an RM on UltraFeedback as a baseline. If your RM doesn't outperform an UltraFeedback-trained one, your data is not yet SOTA.
Further reading: https://arxiv.org/abs/2310.01377
8.2 RewardBench — the benchmark that unified the field (Lambert et al., Mar 2024)
Before RewardBench, everyone claimed their RM was "good" but nobody could compare across labs. Lambert et al. built a curated eval with 4 categories — Chat, Chat-Hard, Safety, Reasoning — each with hand-verified chosen/rejected pairs where the correct answer is unambiguous.
As of mid-2025 the leaderboard (https://huggingface.co/spaces/allenai/reward-bench) is topped by:
- Nemotron-4-340B-Reward (NVIDIA): 92.2%
- Skywork-Reward-Llama-3.1-8B-v0.2: 93.8% (yes, an 8B beats a 340B — data recipe matters more than params)
- URM-LLaMA-3.1-8B: 92.9%
The headline finding from Skywork: a smaller RM trained on cleaner, in-domain data beats a huge RM trained on general preference data. Same lesson as SFT: quality > quantity.
Further reading: https://arxiv.org/abs/2403.13787
8.3 Generative RMs (LLM-as-a-Judge, but trained)
Instead of training a scalar-head RM, train the RM to generate a critique that ends in a [[Better: A]] or [[Better: B]] token. Extract the preference from the last-token logits. This is Zheng et al. 2024's JudgeLM approach and also what NVIDIA's Nemotron-4-Reward uses.
Why this wins:
- Chain-of-thought before judgment. The model can reason about criteria before committing.
- Interpretability. You can read the critique when the RM disagrees with a human.
- Same architecture as the policy. No separate scalar head to align; just use the base model.
Cost: 5–10× more inference compute per preference score (you generate ~200 tokens of critique). For PPO training (S063) this is often prohibitive; for offline DPO (S064) it's fine.
8.4 Process reward models (PRMs) — grading steps, not just answers
For math and code, the answer is right or wrong — but why it went wrong can be at step 3 of 10. OpenAI's PRM800K (Lightman et al. 2023) provided 800K step-level correctness labels on MATH-dataset solutions. Training a PRM gives you fine-grained rewards that guide search (best-of-N, MCTS) and RL (PPO, GRPO) far more effectively than outcome-only rewards.
The 2024–2025 line: DeepSeek-Math and DeepSeek-R1 used PRMs early in training, then switched to outcome-only rewards for the RL stage because well-trained models generate their own step-level reasoning cheaply. This foreshadows the "self-play with verifiable rewards" line we'll see in S063.
Further reading: https://arxiv.org/abs/2305.20050
8.5 The reward-hacking taxonomy (Skalse et al., 2022; refreshed 2024)
A reward model is an approximation of human preferences. The policy will find every place the approximation is wrong. Documented modes as of 2024:
- Length hacking. Longer responses score higher because human labelers subconsciously prefer thoroughness. Fix: length-normalize the reward, or use SimPO (S064).
- Sycophancy. Responses that agree with the user's stated position score higher. Fix: adversarial prompts in the RM training set that flip user positions.
- Format hacking. Bullet points and markdown score higher. Fix: normalize by format, or filter RM data.
- Confidence hacking. Assertive tone scores higher than hedged tone. Fix: include correct-but-hedged responses in the RM data as chosen.
- Refusal hacking. "I cannot help with that" scores higher on safety-labeled prompts and lower on helpful-labeled prompts — but the boundary is fuzzy, so the policy learns to over-refuse. Fix: careful red-team dataset (S065).
The key insight from 2024–2025: you cannot audit reward hacking by looking at the RM alone. You have to sample from the policy and see what game it's playing. The RM will happily give 5.2 to a length-hacked slop response; only human eyes catch it.
8.6 War story — the RM that learned to love em-dashes
Team I know trained an RM on 60K in-house preference pairs for a marketing-copy assistant. RM eval on holdout: 74% agreement with humans, respectable. They fed it into PPO. Two weeks later the policy was writing copy that was 40% em-dashes by character count. Every third clause was joined with an em-dash.
Investigation: their human labelers had (as a stylistic tic of the two senior copywriters who labeled 80% of the data) consistently preferred responses that used em-dashes over commas. The RM learned it. The policy exploited it. The final copy read like a Gwyneth Paltrow tweet transcribed by a Hemingway impersonator on ketamine.
Fix: they added a "punctuation normalizer" pre-processing step that replaced all em-dashes with commas in the RM training data, then retrained. Also added a diverse labeler pool. RM eval went from 74% to 71% (small drop from the noisier labels) but PPO-trained policies stopped exhibiting the em-dash pathology.
Lesson: every RM inherits the aesthetic biases of its labelers. If your labeler pool is small, ANY subtle preference will amplify into a policy pathology. Diversify labelers, or filter the data for stylistic markers before training.
9 · Retention scaffold
Recall
Q1. Write down the Bradley-Terry model for pairwise preferences.
Reveal
P(yA ≻ yB | x) = σ(r(x, yA) − r(x, yB)), where σ is the logistic sigmoid and r is the latent scalar reward. Absolute rewards aren't identifiable — only differences matter.
Q2. What is the reward-modeling loss?
Reveal
ℒ = −log σ(rchosen − rrejected), averaged over the batch. It's the negative log-likelihood of the labeled preference under the Bradley-Terry model.
Q3. When constructing an RM from a CausalLM, what changes and what stays?
Reveal
Stays: the entire transformer backbone. Changes: replace the LM head (Linear d→V, predicting token logits) with a regression head (Linear d→1, predicting a scalar reward, zero-initialized). Pool the last non-pad token's hidden state, not the mean or first.
Q4. What starting loss should you see for a properly-initialized RM, and why?
Reveal
~0.693 nats = log(2). Because at init, the zero-initialized head produces r=0 for both chosen and rejected, so σ(r_c − r_r) = σ(0) = 0.5, and −log(0.5) = 0.693. If you see something else at step 0, your head init or data format is off.
Q5. Name three RM failure modes and their diagnostic.
Reveal
(1) Length bias: correlation between predicted reward and response length > 0.3. (2) Agreement bias / sycophancy: model rewards confident-sounding wrong answers on held-out truthfulness tests. (3) Format hacking: model gives high reward to markdown-heavy responses regardless of content — test with plain-text vs bulleted versions of the same content.
Stretch prompt
You've trained an RM and want to test whether it truly captures "helpfulness" or just "length". Design a controlled evaluation: for a fixed set of prompts, generate response pairs where you can independently vary (a) actual helpfulness (via GPT-4 rewriting) and (b) length (via padding/truncation). Report predicted-reward as a function of both. Sketch what the plot should look like if the RM is well-calibrated vs length-biased.
In your own words
Explain to a colleague, in 4-5 sentences, why reward modeling is a necessary step between SFT and RLHF. Force yourself to explain (a) why SFT alone isn't sufficient and (b) why we don't just have humans directly rate model outputs during training.
Spaced review
- S022 — Cross-entropy loss. The RM loss is pairwise CE with a logistic model; if you can derive CE, you can derive this.
- S058 — SFT. In practice you initialize the RM from your SFT'd model, not from base. This means the RM inherits SFT's format understanding.
- S005 — Probability and logistic models. The Bradley-Terry model is a two-outcome logistic regression on the reward difference.
Next session
S063 — RLHF with PPO. Now that we have an RM, we can use it as a differentiable-ish critic to fine-tune the policy. PPO — the algorithm that trained ChatGPT — is famously fiddly: value heads, KL penalties, advantage estimation, four models loaded at once (policy, ref, RM, value). We'll build the intuition even if you never run one.
Bring back tomorrow
- The RM loss: −log σ(rc − rr).
- Zero-init the reward head. Pool the last non-pad token.
- Always check length ↔ reward correlation before shipping.
Take your trained RM. Score 500 held-out responses. For each, record (a) predicted reward, (b) response length in tokens. Compute Pearson correlation between reward and length. If it’s > 0.3, your RM has learned 'longer = better' — which means downstream RLHF will produce verbose responses regardless of quality. Fix by re-training with length-matched pairs (truncate the longer response, or filter pairs where |len_chosen - len_rejected| / max(len) > 0.3).