RLHF

Training a reward model on real preferences

The practice API has a preference dataset built for this lesson and the next. The chat model, tiny-gpt-chat, wrote four answers to each of 360 product prompts. Then a brand reviewer, POST /v1/brand-check, judged each answer against four guidelines:

• name the product that was asked about;
• don't repeat a word in a sentence;
• use a material the catalog uses for that kind of product;
• mention the repair promise.

Each pair puts the best of the four answers against the worst. The reward model below never sees those guidelines. It only sees which answer in each pair won.

Once trained, the reward model is used for best-of-n sampling. Generate several answers (here, $n = 4$), score each one, and keep the top scorer. It's like asking four people for a draft and keeping the one your editor likes most.

Here's what the code does, step by step:

1. It loads the chat model and the 293 preference pairs. Pairs about six product names are held back for testing. The rest are for training.
2. The features function turns each prompt and answer into a long list of numbers. It uses the chat model's hidden states, the numbers the model builds inside itself as it reads. It averages them over the answer and also keeps the ones at the answer's last token.
3. The reward model is one nn.Linear layer, a weighted sum of those numbers. This layer is called the scoring head. It trains for 201 steps with the Bradley-Terry loss.
4. It checks how often the reward model prefers the chosen answer in the held-back pairs.
5. For each held-back product, it writes four answers. The real brand reviewer counts each answer's flaws. Then it compares three ways of picking one answer: take the first, take the reward model's favorite, or take the reviewer's favorite.

It's slow, so the Run button counts the seconds.

import requests
import torch
import torch.nn as nn
import torch.nn.functional as F
from tinygpt import load_pretrained

lm, tok = load_pretrained("tiny-gpt-chat")
end = tok.token_id("<end>")
pairs = requests.get("https://api.gradient.test/v1/preferences").json()["data"]
held_out = {"Kestrel", "Lark", "River", "Sierra", "Summit", "Willow"}      # product names kept for testing
train = [p for p in pairs if p["prompt"].split()[2] not in held_out]
test = [p for p in pairs if p["prompt"].split()[2] in held_out]
print(f"{len(pairs)} preference pairs: {len(train)} to train on, {len(test)} held out")

@torch.no_grad()
def features(prompts, answers):
    # The chat model's final hidden states, averaged over the answer and taken at its last token.
    rows = [(tok.encode(f"<user> {p} <assistant>"), tok.encode(a) + [end]) for p, a in zip(prompts, answers)]
    rows = [((p + a)[:64], len(p)) for p, a in rows]
    T = max(len(ids) for ids, _ in rows)
    h = lm.hidden(torch.tensor([ids + [end] * (T - len(ids)) for ids, _ in rows]))
    in_answer = torch.tensor([[float(start <= t < len(ids)) for t in range(T)] for ids, start in rows]).unsqueeze(-1)
    last = h[torch.arange(len(rows)), torch.tensor([len(ids) - 1 for ids, _ in rows])]
    return torch.cat([(h * in_answer).sum(dim=1) / in_answer.sum(dim=1), last], dim=-1)

def side(ps, which):
    return features([p["prompt"] for p in ps], [p[which] for p in ps])

torch.manual_seed(0)
reward = nn.Linear(2 * lm.head.weight.shape[1], 1)          # the reward model's scoring head
chosen, rejected = side(train, "chosen"), side(train, "rejected")
opt = torch.optim.Adam(reward.parameters(), lr=1e-2)
for step in range(201):
    loss = -F.logsigmoid(reward(chosen) - reward(rejected)).mean()       # Bradley-Terry
    opt.zero_grad()
    loss.backward()
    opt.step()
with torch.no_grad():
    accuracy = (reward(side(test, "chosen")) > reward(side(test, "rejected"))).float().mean().item()
print(f"reward model prefers the chosen answer in {accuracy:.0%} of held-out pairs")

# best-of-4 on the held-out products, judged by the real brand reviewer
kinds = sorted({" ".join(p["prompt"].split()[3:-1]) for p in pairs})
torch.manual_seed(1)
picked = {"first answer": [], "reward model's best of 4": [], "reviewer's best of 4": []}
for name in sorted(held_out):
    for kind in kinds:
        ids = torch.tensor([tok.encode(f"<user> Describe the {name} {kind} . <assistant>")] * 4)
        out = lm.generate(ids, 60, temperature=0.8, stop_id=end)
        answers = [tok.decode(row).split("<end>")[0].strip() for row in out[:, ids.shape[1]:].tolist()]
        review = requests.post("https://api.gradient.test/v1/brand-check",
                               json={"items": [{"product": f"{name} {kind}", "text": a} for a in answers]}).json()["results"]
        flaws = [len(r["flaws"]) for r in review]
        with torch.no_grad():
            scores = reward(features([f"Describe the {name} {kind} ."] * 4, answers)).squeeze(-1)
        picked["first answer"].append(flaws[0])
        picked["reward model's best of 4"].append(flaws[int(scores.argmax())])
        picked["reviewer's best of 4"].append(min(flaws))
for how, f in picked.items():
    print(f"{how:<26} {sum(f) / len(f):.2f} guideline violations per answer")

The reward model agrees with the reviewer on 69% of the pairs it never trained on. That's in the range real reward models reach against held-out human judgments. It's a sobering number. The "reward" that RLHF pushes up is a noisy imitation of what people want.

Best-of-4 with the reward model cuts the flaws from 1.78 to 1.39 per answer. Picking with the reviewer itself would reach 0.99. The gap between 1.39 and 0.99 is the cost of the reward model's mistakes.

This reward model is a cheap version. It only trains a scoring head on top of the chat model's frozen hidden states. Production reward models are usually a full copy of the SFT model with a scoring head, and all of it is fine-tuned. (torch here is the course's in-browser PyTorch subset.)

Step 3: reinforcement learning, on a leash

Best-of-n improves answers, but every request now costs $n$ answers instead of one. RLHF goes further. It trains the policy to write high-scoring answers in the first place.

There's a catch. The reward model is only about 70% right here, and a training process hunts for anything that raises the score. Left free, it will find the reward model's mistakes and use them. So RLHF keeps the policy on a leash, like a dog on a walk. The dog can explore, but only so far from you. Here, "you" is the SFT model.

The goal can be written as one line of math. First, the symbols, one at a time:

• $x$ is the prompt.
• $y$ is an answer the model writes.
• $\pi$ (the Greek letter pi) is the policy. Here it has nothing to do with 3.14.
• $\pi_{\text{ref}}$ is the reference model: the SFT model, kept frozen.
• $r(x, y)$ is the reward model's score for answer $y$ to prompt $x$.
• $\mathbb{E}_{y \sim \pi}$ means "the average over answers the policy writes."
• $\mathrm{KL}$ is the KL divergence from Phase 2. It measures how far apart two sets of chances are, and it's 0 when they match.
• $\beta$ (beta) is a number you choose. It sets how tight the leash is.

In plain words: find the policy with the best average score after you take off a fine for drifting from the SFT model. The fine is the drift times $\beta$. A big $\beta$ means a short leash. A small $\beta$ lets the policy chase the score.

The best policy has a simple shape

This goal has an exact answer. Call the best policy $\pi^*$ (pi star). For one fixed prompt, the score depends only on the answer, so it's written $r(y)$.

Here $e$ is the number from the exponentials and logarithms lesson, about 2.718. The symbol $\propto$ means "is proportional to." You compute the right side for every answer. Then you divide each result by their total, so the chances add up to 1.

In plain words: start from how often the SFT model writes each answer. Multiply by a boost that grows with the answer's score. Then rescale so the chances add up to 1. Dividing the score by $\beta$ sets how strong the boost is.

Real models can't use this formula directly, because there are far too many possible answers to list. On a small made-up example, though, it shows exactly what $\beta$ does.

import math

p_ref = [0.8, 0.2]      # how often the SFT model writes "plain" and "punchy"
reward = [0.0, 1.0]     # the reward model's score for each
for beta in [10, 1, 0.5, 0.25]:
    weights = [p * math.exp(r / beta) for p, r in zip(p_ref, reward)]
    total = sum(weights)
    print(beta, [round(w / total, 3) for w in weights])

Watching reward hacking happen

Now try five possible answers. One of them games the reward model: it repeats the repair promise three times. The reward model loves that, and gives it the top score of 1.6. The brand team hates it. So the code keeps two lists of scores, the reward model's and the brand team's true opinion. For each $\beta$, it builds the best policy with the formula above and prints:

• the average reward model score;
• the average true quality;
• the KL, how far the policy has drifted from the SFT model;
• how often the policy writes the gaming answer.

import numpy as np

answers = ["on brand, clean", "on brand, plain", "repeats a word", "wrong material", "repair promise three times"]
p_ref = np.array([0.30, 0.35, 0.15, 0.15, 0.05])     # how likely the SFT model is to write each
reward = np.array([1.0, 0.7, -0.5, -0.8, 1.6])       # the reward model's score
true = np.array([1.0, 0.7, -0.5, -0.8, -1.0])        # what the brand team actually thinks

print(f"{'beta':>5}  {'reward score':>12}  {'true quality':>12}  {'KL':>5}   share of the gaming answer")
for beta in [10, 2, 1, 0.5, 0.25, 0.1]:
    p = p_ref * np.exp(reward / beta)                 # the optimal policy for this beta
    p /= p.sum()
    kl = float((p * np.log(p / p_ref)).sum())
    print(f"{beta:>5}  {p @ reward:>12.2f}  {p @ true:>12.2f}  {kl:>5.2f}   {p[4]:.0%}")

Read the table from top to bottom, as $\beta$ shrinks and the leash gets longer. The reward model's score climbs the whole way, from 0.48 to 1.59. True quality rises at first and peaks at 0.50 when $\beta = 1$. Then it collapses to −0.97. By then the policy writes the gaming answer 98% of the time.

This is called reward hacking, or over-optimization. It's Goodhart's law in action: "when a measure becomes a target, it ceases to be a good measure." A school that's judged only on test scores can start teaching to the test. The scores go up, and the learning doesn't.

Researchers at OpenAI measured the same shape in real models. They used a much larger "gold" reward model as a stand-in for people. As training pushed the policy further from the reference, the trained reward model's score kept rising. The gold score peaked and then fell.

How the RL step actually runs

With a real model, you can't list every possible answer, so you can't use the formula. Instead, RLHF uses a training method from reinforcement learning. The classic choice is PPO, short for proximal policy optimization. Each round goes like this:

1. The policy writes a batch of answers.
2. The reward model scores each one.
3. The KL fine is subtracted token by token, so each token that drifts from the SFT model costs a little.
4. The policy is nudged toward answers that scored better than expected. "Expected" comes from a fourth model, the value model. It predicts the score an answer is likely to get, from the part written so far.
5. Each update is clipped, meaning capped, so no single round changes the policy too much.

That means four models in memory at once: the policy, the frozen reference, the reward model and the value model. It also means lots of answer writing and many settings to tune. PPO works, but it's famously fussy.

That's why the next lesson's method, DPO, became popular. It skips the reward model and the answer writing. It's also why reasoning models are trained with simpler variants, such as GRPO. Those use rewards you can check, like whether a math answer is right, instead of a learned reward model.

For marketers: you can use best-of-n with a scorer today, without training anything. Write several ad variants. Score them with a rubric-based LLM judge or a classifier, and keep the best. Spot-check the winners, though. A scorer can be gamed by copy that hits its keywords without being good.