Module 7. Reward models and preference data
Read the explanations and code in the existing notebooks, or download the files. Code does not run on this site.
Download course files (.zip) 382.1 KBThe archive preserves the source folder structure, both languages, and supporting files.
Module 7 practice: reward models, calibration, and proxy overoptimization
Main module notebook
Download .ipynb module-07-reward-models.en.ipynb · 17.2 KB
Show notebook text
Module 7 practice: reward models, calibration, and proxy overoptimization
This notebook has two deliberately separate routes.
-
Offline synthetic route — runs on CPU with NumPy and PyTorch, trains a tiny pairwise scorer, checks prompt-level splitting, evaluates ranking and calibration, and reproduces the toy overoptimization curve.
-
Real-model route — an optional accelerator-backed plan using TRL, LoRA,
Qwen/Qwen3-0.6B, andtrl-lib/ultrafeedback_binarized.
The canonical notebook is stored unexecuted. Set RUN_REAL_MODEL=True only in a Colab runtime after reviewing the resource and provenance notes. The real branch was not executed for the course release.
Resource and provenance contract
- Smoke mode: 512–1,000 preference pairs,
max_length=512, one short epoch or a limited number of steps. - Full learning mode: 10k–50k pairs,
max_length=1024, and explicit evaluation/calibration slices. Feasibility depends on actual lengths, precision, hardware, and library revisions; measure memory rather than assuming a particular free runtime. - The UltraFeedback Binarized pairs are synthetic AI feedback: candidate completions were scored by GPT-4; they are not direct human preference labels.
- Save the tokenizer/chat template, model revision, dataset revision, truncation counts, and the scalar
scorehead together with the LoRA adapters. - A successful smoke run validates the data path. It is not evidence that the trained RM is robust under policy optimization.
from __future__ import annotations
from dataclasses import dataclass
import math
import random
import numpy as np
import torch
SEED = 7
RUN_REAL_MODEL = False
SMOKE_MODE = True
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
print({"torch": torch.__version__, "device": "cuda" if torch.cuda.is_available() else "cpu"})1. Synthetic prompt-level preference data
The tiny model below learns one scalar from interpretable response features. Each prompt has several candidates; all pairs from one prompt stay in one split. This is intentionally smaller than a language RM so that the statistical contract is visible.
@dataclass(frozen=True)
class PairBatch:
chosen: torch.Tensor
rejected: torch.Tensor
prompt_ids: list[str]
first: torch.Tensor
second: torch.Tensor
first_preferred: torch.Tensor
def make_synthetic_pairs(n_prompts: int = 800, pairs_per_prompt: int = 2, seed: int = SEED):
rng = np.random.default_rng(seed)
true_w = np.array([1.5, 0.8, -0.9, 0.4], dtype=np.float32)
rows = []
for prompt_index in range(n_prompts):
prompt_id = f"p{prompt_index:05d}"
prompt_shift = rng.normal(0.0, 0.5)
for _ in range(pairs_per_prompt):
a = rng.normal(size=4).astype(np.float32)
b = rng.normal(size=4).astype(np.float32)
utility_a = float(a @ true_w + prompt_shift + rng.normal(0.0, 0.35))
utility_b = float(b @ true_w + prompt_shift + rng.normal(0.0, 0.35))
chosen, rejected = (a, b) if utility_a >= utility_b else (b, a)
rows.append((prompt_id, chosen, rejected, a, b, float(utility_a >= utility_b)))
return rows
def split_by_prompt(rows, eval_fraction: float = 0.2, seed: int = SEED):
prompt_ids = sorted({row[0] for row in rows})
rng = np.random.default_rng(seed)
rng.shuffle(prompt_ids)
cut = int(round(len(prompt_ids) * (1.0 - eval_fraction)))
train_ids = set(prompt_ids[:cut])
eval_ids = set(prompt_ids[cut:])
assert train_ids.isdisjoint(eval_ids)
def pack(selected_ids):
chosen = np.stack([row[1] for row in rows if row[0] in selected_ids])
rejected = np.stack([row[2] for row in rows if row[0] in selected_ids])
ids = [row[0] for row in rows if row[0] in selected_ids]
first = np.stack([row[3] for row in rows if row[0] in selected_ids])
second = np.stack([row[4] for row in rows if row[0] in selected_ids])
labels = [row[5] for row in rows if row[0] in selected_ids]
return PairBatch(torch.tensor(chosen), torch.tensor(rejected), ids,
torch.tensor(first), torch.tensor(second), torch.tensor(labels))
return pack(train_ids), pack(eval_ids)
rows = make_synthetic_pairs()
train_batch, eval_batch = split_by_prompt(rows)
print({"train_pairs": len(train_batch.prompt_ids), "eval_pairs": len(eval_batch.prompt_ids)})2. Pairwise training
For scores and , the stable Bradley–Terry loss is
The model below has no language representation; it isolates the pairwise objective and score-centering convention.
class TinyRewardModel(torch.nn.Module):
def __init__(self, n_features: int = 4):
super().__init__()
self.score = torch.nn.Linear(n_features, 1)
def forward(self, features: torch.Tensor) -> torch.Tensor:
return self.score(features).squeeze(-1)
def pair_metrics(model: TinyRewardModel, batch: PairBatch) -> dict[str, float]:
model.eval()
with torch.no_grad():
delta = model(batch.chosen) - model(batch.rejected)
probability = torch.sigmoid(delta)
target = torch.ones_like(probability)
loss = torch.nn.functional.binary_cross_entropy(probability, target)
brier = torch.mean((probability - target) ** 2)
return {
"accuracy": float((delta > 0).float().mean()),
"log_loss": float(loss),
"brier": float(brier),
"mean_margin": float(delta.mean()),
}
model = TinyRewardModel()
optimizer = torch.optim.AdamW(model.parameters(), lr=0.03, weight_decay=1e-3)
steps = 120 if SMOKE_MODE else 600
batch_size = 256
rng = np.random.default_rng(SEED)
for step in range(steps):
index = torch.tensor(rng.integers(0, len(train_batch.prompt_ids), size=batch_size))
chosen_score = model(train_batch.chosen[index])
rejected_score = model(train_batch.rejected[index])
loss = torch.nn.functional.softplus(-(chosen_score - rejected_score)).mean()
# Optional gauge convention: discourage a drifting common score origin.
center = 0.5 * (chosen_score.mean() + rejected_score.mean())
total_loss = loss + 1e-2 * center.square()
optimizer.zero_grad(set_to_none=True)
total_loss.backward()
optimizer.step()
print("train", pair_metrics(model, train_batch))
print("eval ", pair_metrics(model, eval_batch))3. Reliability bins
Calibration is evaluated on predicted pairwise probabilities, not on raw scores. Keep the original A/B order chosen before observing the preference and use the corresponding binary label. Sorting every pair into chosen/rejected would condition on the outcome and make every target one; those bins do not estimate calibration of P(A preferred to B). BT log loss and Brier score are unchanged by a consistent swap of probability and label. Equal-width ECE is only one summary and depends on the binning rule.
def reliability_bins(probability: np.ndarray, label: np.ndarray, n_bins: int = 10):
probability = np.asarray(probability, dtype=np.float64).reshape(-1)
label = np.asarray(label, dtype=np.float64).reshape(-1)
if probability.shape != label.shape or probability.size == 0:
raise ValueError("aligned non-empty arrays required")
indices = np.minimum((probability * n_bins).astype(int), n_bins - 1)
bins = []
for bin_id in range(n_bins):
mask = indices == bin_id
if np.any(mask):
bins.append({
"bin": bin_id,
"count": int(mask.sum()),
"confidence": float(probability[mask].mean()),
"frequency": float(label[mask].mean()),
})
return bins
with torch.no_grad():
eval_delta = model(eval_batch.first) - model(eval_batch.second)
eval_probability = torch.sigmoid(eval_delta).numpy()
labels = eval_batch.first_preferred.numpy()
bins = reliability_bins(eval_probability, labels, n_bins=5)
ece = sum(item["count"] / len(labels) * abs(item["confidence"] - item["frequency"]) for item in bins)
print({"ece_5_bins": ece, "bins": bins})4. Proxy overoptimization
The next cell reproduces the module's teaching mechanism. It is not a language-model experiment and does not reproduce a published scaling law.
def toy_overoptimization(pool_sizes=(1, 2, 4, 8, 16, 32, 64, 128, 256, 512), trials=4000, seed=7):
rng = np.random.default_rng(seed)
proxy_means = []
target_means = []
for pool_size in pool_sizes:
useful = rng.normal(size=(trials, pool_size))
nuisance = rng.normal(size=(trials, pool_size))
proxy = useful + nuisance
target = useful - np.maximum(nuisance - 1.0, 0.0) ** 2
selected = np.argmax(proxy, axis=1)
rows = np.arange(trials)
proxy_means.append(proxy[rows, selected].mean())
target_means.append(target[rows, selected].mean())
return np.asarray(proxy_means), np.asarray(target_means)
pool_sizes = np.array([1, 2, 4, 8, 16, 32, 64, 128, 256, 512])
proxy, target = toy_overoptimization(pool_sizes)
# Finite Monte Carlo estimates need not be monotone for every seed/budget.
assert np.isfinite(proxy).all() and np.isfinite(target).all()
print({"best_target_pool": int(pool_sizes[np.argmax(target)]), "final_proxy": proxy[-1], "final_target": target[-1]})5. Optional real-model route with an accelerator
The following command is an unpinned installation example. Resolve and test compatible releases in a fresh accelerator-enabled runtime, then restart the runtime if the installer requests it:
pip install -U transformers datasets accelerate peft trlThe code follows the current TRL RewardTrainer interface at the factual snapshot. Library APIs can change. Pin the working environment in a released run. The score head is included in modules_to_save because it is newly initialized for sequence classification.
if RUN_REAL_MODEL:
import torch
from datasets import load_dataset
from peft import LoraConfig, TaskType
from trl import RewardConfig, RewardTrainer
MODEL_ID = "Qwen/Qwen3-0.6B"
DATASET_ID = "trl-lib/ultrafeedback_binarized"
TRAIN_PAIRS = 512 if SMOKE_MODE else 10_000
EVAL_PAIRS = 256 if SMOKE_MODE else 2_000
MAX_LENGTH = 512 if SMOKE_MODE else 1024
dataset = load_dataset(DATASET_ID)
train_dataset = dataset["train"].select(range(min(TRAIN_PAIRS, len(dataset["train"]))))
eval_split_name = "test" if "test" in dataset else next(name for name in dataset if name != "train")
eval_dataset = dataset[eval_split_name].select(range(min(EVAL_PAIRS, len(dataset[eval_split_name]))))
peft_config = LoraConfig(
task_type=TaskType.SEQ_CLS,
r=16,
lora_alpha=32,
lora_dropout=0.05,
target_modules=["q_proj", "v_proj"],
modules_to_save=["score"],
)
args = RewardConfig(
output_dir="module-07-rm-smoke" if SMOKE_MODE else "module-07-rm-full",
max_length=MAX_LENGTH,
learning_rate=1e-4,
per_device_train_batch_size=2,
per_device_eval_batch_size=2,
gradient_accumulation_steps=8,
max_steps=40 if SMOKE_MODE else -1,
num_train_epochs=1.0,
eval_strategy="steps",
eval_steps=20 if SMOKE_MODE else 200,
save_strategy="no" if SMOKE_MODE else "steps",
logging_steps=5,
fp16=True,
gradient_checkpointing=True,
center_rewards_coefficient=1e-2,
report_to="none",
seed=SEED,
)
trainer = RewardTrainer(
model=MODEL_ID,
args=args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
peft_config=peft_config,
)
trainer.train()
metrics = trainer.evaluate()
print(metrics)
else:
print("Real-model branch skipped. Set RUN_REAL_MODEL=True in a reviewed Colab runtime.")Publication checklist
Before publishing measured results from the real branch:
- pin model, dataset, tokenizer, TRL, Transformers, PEFT, Accelerate, and PyTorch revisions;
- report the number of pairs retained, filtered, and truncated;
- split by prompt and add generator/domain slices;
- report pair accuracy, log loss, Brier score, and a calibration plot;
- save LoRA adapters and the scalar score head;
- state the T4 type, wall-clock time, peak allocated memory, and random seeds;
- do not present UltraFeedback Binarized as direct human preference data;
- evaluate fresh responses from a policy that was not used to construct the training pairs.