5.5 The RLHF Pipeline & Reward Modeling¶

A base language model, freshly pretrained on a trillion tokens of internet text, is a magnificent next-token predictor and a frustrating assistant. Ask it “How do I make a loaf of sourdough?” and it might continue with three more questions in the same style, because forum posts that begin with a question are often followed by more questions. It has learned the distribution of text, not the intent to be helpful. Supervised Fine-Tuning & Instruction Tuning goes a long way toward fixing this: show the model tens of thousands of (instruction, good-response) pairs and it learns the format of being an assistant. But SFT has a ceiling. It can only imitate demonstrations a human bothered to write, it gives equal weight to every token of every demonstration regardless of quality, and it has no notion that one valid answer might be better than another equally valid answer.

The problem is that “good response” is not a thing you can easily write down. You cannot author a demonstration for every possible prompt, and for open-ended prompts (“write me a poem about the ocean”) there is no single correct target at all. But humans find it easy to compare: shown two ocean poems, a person can reliably say which they prefer, even when they could not have written either. Reinforcement Learning from Human Feedback (RLHF) is the machinery that converts that cheap, plentiful comparison signal into a training objective. We collect human preferences over model outputs, distill them into a learned reward model (RM) that scores any response with a scalar, and then optimize the language model — now the policy — to maximize that score while staying close to its SFT starting point.

This chapter is the cornerstone of Part V. We follow the InstructGPT recipe (Ouyang et al., 2022) end to end: how preference data is collected, the Bradley–Terry statistical model that turns pairwise comparisons into a continuous reward, how to train and evaluate the reward model (and how it gets hacked), and the four-model — actor, critic, reward, reference — apparatus that makes the optimization step work. The actual policy-optimization algorithm (PPO) is the subject of Policy Gradients & PPO for Language Models; the reward-free shortcut that skips the RM entirely is Direct Preference Optimization & Its Variants. Here we build everything up to and including the reward signal, plus the scaffolding around the optimizer.

The three-stage recipe and why it works¶

InstructGPT (and ChatGPT after it) is built on a three-stage pipeline. Each stage produces an artifact the next stage consumes.

The key conceptual move is decoupling the source of supervision from the form of supervision. Humans supply preferences — a comparison signal that is cheap, has low inter-annotator variance, and requires no expertise to author. The reward model turns those discrete comparisons into a dense, differentiable scalar defined on the entire space of possible responses, including responses no human has ever seen. The policy then gets to generate its own training data: it proposes responses, the RM grades them, and the policy is nudged toward the high-scoring region. This is why RLHF can exceed SFT — the policy explores outputs better than any demonstration in the dataset, and the RM, generalizing from preferences, can recognize and reward them.

Preference data collection¶

Everything downstream is only as good as the preference data. This is where most of the human cost and most of the subtle failure modes live.

What a labeler actually does¶

The labeling unit is a comparison. A prompt \(x\) is drawn (from real API traffic, in InstructGPT’s case, plus labeler-written prompts). The SFT model generates \(K\) candidate responses \(y_1, \dots, y_K\) (InstructGPT used \(K\) between 4 and 9). A human is shown the prompt and all \(K\) responses and produces a ranking — a total or partial order over the candidates. From a single ranking of \(K\) items you extract \(\binom{K}{2}\) pairwise preferences. Ranking \(K\) items at once is far more sample-efficient than \(\binom{K}{2}\) independent pairwise tasks, and it forces internal consistency: a labeler cannot say A>B, B>C, and C>A within one ranking.

Each extracted pair becomes a training row \((x, y_w, y_l)\) where \(y_w\) (“win”) is the preferred response and \(y_l\) (“lose”) is the dispreferred one. A modern open preference dataset row looks like:

{
  "prompt": "Explain why the sky is blue to a 10-year-old.",
  "chosen": "Sunlight is actually made of all the colors mixed together...",
  "rejected": "The sky is blue due to Rayleigh scattering, wherein the cross-section...",
  "metadata": {"annotator_id": 71, "confidence": "strong", "k_rank_position": [1, 3]}
}

You do not have to collect this yourself to get started. The open ecosystem publishes preference sets in exactly this (prompt, chosen, rejected) shape on the HuggingFace Hub, loadable with one line of datasets: Anthropic/hh-rlhf (the original helpfulness/harmlessness comparisons), HuggingFaceH4/ultrafeedback_binarized (GPT-4-scored responses from many models, binarized into pairs — the workhorse for open RM and DPO training), nvidia/HelpSteer2 (human ratings on several attributes, from which pairs are derived), and Skywork/Skywork-Reward-Preference-80K-v0.2 (a curated mixture behind several strong open RMs). These are also exactly the datasets Direct Preference Optimization & Its Variants consumes, since DPO eats the same rows.

The dimensions labelers are asked to judge¶

“Which is better?” is underspecified, so production labeling uses an explicit rubric. InstructGPT’s instructions asked labelers to weigh:

• Helpfulness — does it follow the instruction and actually solve the user’s problem (including inferring unstated intent)?
• Honesty / truthfulness — is it factually correct; does it avoid confidently stating falsehoods?
• Harmlessness — does it avoid toxic, dangerous, or biased content?

These can conflict: the most helpful answer to “how do I pick a lock” is harmful; the most harmless is unhelpful. InstructGPT told labelers to prioritize helpfulness during comparison labeling but harmlessness during a separate safety labeling pass — a deliberate decomposition that later work (Anthropic’s helpfulness/harmlessness split, then Constitutional AI, RLAIF & Self-Improvement) made explicit by training separate reward models per axis.

How much data¶

InstructGPT used on the order of tens of thousands of comparisons. Reward-model data is more leveraged than SFT data because each comparison is a relative judgment that constrains the whole reward surface, not a single absolute target. But preference data has a shelf life: once you optimize a policy against an RM, the new policy generates responses in a distribution the RM never saw during training. This is the distribution-shift problem and it motivates iterated RLHF — collect fresh preferences on the new policy’s outputs, retrain the RM, optimize again. LLaMA-2 ran roughly five such iterations; this loop is the seed of the Data Flywheels & Continuous Improvement story.

The Bradley–Terry model: from comparisons to a scalar¶

We have pairs \((x, y_w, y_l)\). We want a function \(r_\phi(x, y)\) — a single scalar score — such that higher means “more preferred.” How do we get a continuous score from discrete comparisons? The bridge is a 70-year-old idea from the statistics of paired comparisons: the Bradley–Terry model (Bradley & Terry, 1952), the same math behind Elo chess ratings.

The model¶

Assume each response \(y\) to a prompt \(x\) has a latent real-valued “strength” \(r(x,y)\). The Bradley–Terry assumption is that the probability a human prefers \(y_w\) over \(y_l\) is a logistic function of the difference in strengths:

where \(\sigma(z) = 1/(1 + e^{-z})\) is the logistic sigmoid. This is exactly a softmax over two items, or equivalently logistic regression on the score gap. Three properties make it the right choice:

1. Only differences matter. Adding a constant \(c\) to every score leaves all preference probabilities unchanged (\(\sigma((r_w + c) - (r_l + c)) = \sigma(r_w - r_l)\)). The reward is identified only up to an additive constant — a fact with real consequences (it is why you must not compare raw reward magnitudes across two separately-trained RMs, and why we mean-center rewards before PPO).
2. It is calibrated and monotone. Equal scores give a coin flip (\(\sigma(0)=0.5\)); a score gap of \(+2\) means \(\sigma(2)\approx 0.88\), i.e. “preferred 88% of the time.” The score is a log-odds of preference, so it has an interpretable scale.
3. It is differentiable, so we can fit \(r\) by gradient descent.

We replace the abstract strength \(r\) with a neural network \(r_\phi\) and fit \(\phi\) by maximum likelihood on the observed preferences.

Deriving the pairwise loss¶

For one comparison \((x, y_w, y_l)\), the likelihood under Bradley–Terry is \(\sigma(r_\phi(x,y_w) - r_\phi(x,y_l))\). The negative log-likelihood — the loss we minimize — is therefore

Let \(\Delta = r_\phi(x,y_w) - r_\phi(x,y_l)\) be the margin the model assigns. Then the per-example loss is \(-\log\sigma(\Delta) = \log(1 + e^{-\Delta})\), the softplus of \(-\Delta\). Note its shape: it is near zero when \(\Delta\) is large and positive (model confidently correct), grows linearly in \(-\Delta\) when the model is confidently wrong, and equals \(\log 2 \approx 0.693\) at \(\Delta = 0\). The gradient with respect to the margin is clean:

So the push on the scores is proportional to how wrong the model currently is: if it already strongly prefers the winner (\(\sigma(\Delta)\to 1\)), the gradient vanishes; if it has it backwards, the gradient saturates at magnitude \(1\). This self-limiting behavior is exactly the robustness property of logistic losses.

A subtle but critical training detail¶

InstructGPT found that if you shuffle all \(\binom{K}{2}\) pairs from many prompts together and train on them as independent examples, the RM overfits: each completion \(y_i\) appears in \(K-1\) pairs within an epoch, so a single forward pass’s value gets reused many times and the model memorizes specific completions. The fix is to put all \(\binom{K}{2}\) pairs from one prompt into a single forward/backward pass (one gradient step “sees” each of the \(K\) completions exactly once, computing all pairwise terms from those \(K\) scores). This is both more compute-efficient (each completion is encoded once, not \(K-1\) times) and a strong regularizer. Remember this for the interview.

Reward model architecture and training¶

The architecture: a transformer with a scalar head¶

A reward model is almost the same network as the policy. You take the SFT model, remove the unembedding (language-modeling) head, and bolt on a small linear layer that maps the final hidden state to a single scalar. You read off the scalar at the position of the last token of the response (or the EOS token) — that hidden state has attended over the entire prompt-and-response, so it is a reasonable summary on which to base a whole-sequence judgment.

Why initialize from the SFT model rather than from scratch or from the base? It already understands language and the task distribution; you are only teaching it a new, low-dimensional output (a ranking). InstructGPT noted that a smaller RM than the policy works fine and is far cheaper — they used a 6B RM to align a 175B policy. The RM does not need to be as capable as the policy; it only needs to recognize quality, which is easier than producing it (judging is easier than generating).

Three families of reward model coexist by 2026, and it is worth knowing which one you are building. (i) The discriminative scalar-head RM of this section: one forward pass per response, one number out. It is the only family that drops directly into a PPO inner loop, because PPO needs an absolute per-response scalar for thousands of samples per step. (ii) Pairwise / preference-ranking models put both responses in one context and predict which wins. They are typically more accurate — cross-attention between the two candidates is a real advantage — but they yield a relative verdict, so they suit best-of-N reranking and evaluation rather than PPO. (iii) Generative reward models, which write a critique against a rubric and then emit a score as text; they inherit the base model’s reasoning, can be scaled at inference time by sampling several critiques and voting (DeepSeek-GRM), and blur into LLM-as-a-Judge & Automated Evaluation. Orthogonal to all three is where the reward is attached: an outcome RM scores the finished response (this chapter), while a process reward model scores each intermediate reasoning step — the topic of RL with Verifiable Rewards (RLVR) & The Reasoning Recipe and Reward Engineering, Verifiers & Sandboxes.

The from-scratch pairwise loss¶

Here is the complete, runnable core of reward-model training — the loss every RLHF library implements. This is the central code of the chapter.

import torch
import torch.nn as nn
import torch.nn.functional as F


class RewardModel(nn.Module):
    """A transformer backbone with a scalar value head on top.

    `backbone` is any model returning per-token hidden states of shape
    (batch, seq_len, d_model) -- e.g. an SFT-initialized decoder with its
    language-modeling head removed.
    """
    def __init__(self, backbone, d_model):
        super().__init__()
        self.backbone = backbone
        # A single linear layer: d_model -> 1 scalar. Initialize small so early
        # rewards are near zero (helps optimization stability downstream).
        self.value_head = nn.Linear(d_model, 1, bias=False)
        nn.init.normal_(self.value_head.weight, std=1.0 / (d_model + 1) ** 0.5)

    def forward(self, input_ids, attention_mask):
        # hidden: (B, T, d_model)
        hidden = self.backbone(input_ids, attention_mask=attention_mask).last_hidden_state
        # Per-token scalar scores: (B, T)
        scores = self.value_head(hidden).squeeze(-1)
        # Reward is read at the LAST non-pad token of each sequence.
        # attention_mask is 1 for real tokens, 0 for padding.
        # NOTE: this index is only correct for RIGHT padding. Generation code
        # usually left-pads, so a shared tokenizer can silently make you read the
        # reward at a pad position; with left padding use `attention_mask.size(1) - 1`
        # (the true last column) instead. Always assert the padding side.
        last_idx = attention_mask.sum(dim=1) - 1            # (B,) index of final real token
        batch_idx = torch.arange(input_ids.size(0), device=input_ids.device)
        reward = scores[batch_idx, last_idx]                # (B,) one scalar per sequence
        return reward


def bradley_terry_loss(reward_chosen, reward_rejected, margin=0.0):
    """The Bradley-Terry / pairwise ranking loss.

    reward_chosen, reward_rejected : (B,) scalar reward for the preferred and
                                     dispreferred response in each pair.
    margin : optional fixed margin m. Loss becomes -log sigmoid(r_w - r_l - m),
             which only counts the pair as "solved" once r_w beats r_l by m.
             Useful when you have a graded preference strength (set m larger for
             strongly-preferred pairs). Default 0 recovers plain Bradley-Terry.

    Returns (loss, accuracy) where accuracy is the fraction of pairs the model
    currently orders correctly -- the single most important RM training metric.
    """
    delta = reward_chosen - reward_rejected - margin        # the margin Δ
    # -log σ(Δ) == softplus(-Δ), computed stably:
    loss = F.softplus(-delta).mean()
    # A pair is "correct" iff the chosen response scores higher than the rejected.
    accuracy = (reward_chosen > reward_rejected).float().mean()
    return loss, accuracy


def train_step(rm, optimizer, batch):
    """One optimization step on a batch of preference pairs.

    `batch` provides chosen/rejected token ids and masks, already tokenized as
    [prompt + response]. We run BOTH responses through the SAME reward model and
    compare their scalar scores. This shared-encoder, two-forward-pass structure
    is exactly what TRL's RewardTrainer and DPO-style code do under the hood.
    """
    rm.train()
    # Forward both completions of every pair through the same network.
    r_chosen   = rm(batch["chosen_ids"],   batch["chosen_mask"])     # (B,)
    r_rejected = rm(batch["rejected_ids"], batch["rejected_mask"])   # (B,)

    loss, acc = bradley_terry_loss(r_chosen, r_rejected)

    optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(rm.parameters(), 1.0)   # RMs train fast; clip helps
    optimizer.step()
    return {"loss": loss.item(), "pref_acc": acc.item(),
            "reward_margin": (r_chosen - r_rejected).mean().item()}

The two non-obvious lines are worth dwelling on. First, F.softplus(-delta) is the numerically stable way to compute \(-\log\sigma(\Delta)\): a naive -torch.log(torch.sigmoid(delta)) underflows to inf when delta is very negative, whereas softplus is stable everywhere (the same reason we use F.binary_cross_entropy_with_logits rather than logits-then-sigmoid-then-BCE; see Numerical Computing, Floating Point & Precision). Second, both responses pass through the same network — the reward model is a Siamese / shared-encoder architecture, and only the difference of its two outputs ever enters the loss.

Batching all-pairs-from-one-prompt¶

To implement the InstructGPT regularizer — all \(\binom{K}{2}\) pairs from a prompt in one step — you encode the \(K\) completions once and form all pairwise terms from the \(K\) scalars:

def all_pairs_bt_loss(rewards_K, chosen_better_mask):
    """rewards_K : (K,) scalar reward for the K completions of ONE prompt,
                   ordered by the human ranking (index 0 = most preferred).
       Since they are ranked, every pair (i, j) with i < j has y_i preferred.
    Computes the mean BT loss over all C(K,2) pairs, encoding each completion once.
    """
    K = rewards_K.shape[0]
    # Pairwise score differences Δ_ij = r_i - r_j for all i<j.
    diff = rewards_K.unsqueeze(1) - rewards_K.unsqueeze(0)   # (K, K), Δ_ij at [i, j]
    iu = torch.triu_indices(K, K, offset=1)                  # upper triangle: i<j
    deltas = diff[iu[0], iu[1]]                              # (C(K,2),) all preferred-minus-dispreferred
    return F.softplus(-deltas).mean()

Each of the \(K\) completions is encoded by the transformer exactly once; the \(\binom{K}{2}\) comparisons are cheap arithmetic on the resulting scalars. This is both the efficiency win and the anti-overfitting win InstructGPT reported.

Doing it for real: TRL’s RewardTrainer¶

You will not hand-roll the training loop in production. The standard implementation is HuggingFace TRL’s RewardTrainer, which is the code above plus tokenization, padding, distributed training, and evaluation. The key mapping to know: HF’s AutoModelForSequenceClassification with num_labels=1 is the scalar-head reward model — it drops the LM head, attaches a Linear(d_model, 1), and reads it at the last non-pad token, exactly like our RewardModel.

# pip install "trl>=0.12" transformers datasets accelerate
from datasets import load_dataset
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from trl import RewardTrainer, RewardConfig

model_id = "Qwen/Qwen2.5-0.5B-Instruct"      # in a real run: YOUR SFT checkpoint
tok = AutoTokenizer.from_pretrained(model_id)

rm = AutoModelForSequenceClassification.from_pretrained(model_id, num_labels=1)
rm.config.pad_token_id = tok.pad_token_id    # REQUIRED: the head must know where sequences end

# Rows carry "chosen"/"rejected" (with the prompt either separate or prepended);
# TRL applies the chat template and tokenizes both sides for you.
ds = load_dataset("HuggingFaceH4/ultrafeedback_binarized", split="train_prefs")

trainer = RewardTrainer(
    model=rm,
    args=RewardConfig(
        output_dir="rm-out",
        per_device_train_batch_size=8,
        num_train_epochs=1,                  # RMs overfit fast -- one epoch is standard
        learning_rate=1e-5,                   # ~10x lower than SFT; the head is the new part
        max_length=1024,                      # pairs longer than this are dropped, not truncated
        center_rewards_coefficient=0.01,      # auxiliary penalty on (r_w + r_l)^2, see below
    ),
    train_dataset=ds,
    processing_class=tok,                     # named `tokenizer=` in older TRL releases
)
trainer.train()

Two options deserve comment. num_train_epochs=1 is not laziness: InstructGPT reported that a second epoch overfits the reward model, and this has held up across open recipes — RM training is a short, low-learning-rate job. center_rewards_coefficient adds \(\lambda\,\mathbb{E}[(r_w + r_l)^2]\) to the loss, which does nothing to the preference likelihood (that depends only on the difference) but pins down the free additive constant we identified earlier, keeping rewards near zero so the downstream PPO advantages start on a sane scale. TRL logs train/accuracy — the pairwise preference accuracy from our bradley_terry_loss — as the metric to watch. The library-level tour of TRL, including PPOTrainer, is TRL: HuggingFace’s RL Library; for multi-node RM and PPO training the Ray + vLLM alternative is OpenRLHF, NeMo-Aligner & Ray-Based Systems.

What good RM training looks like¶

The headline metric is preference accuracy on a held-out set — the fraction of pairs the RM orders the same way the human did. For well-curated data this lands somewhere in the high-60s to high-70s percent; remember that the human–human agreement ceiling is itself around 70–75%, so an RM scoring near there is essentially at the noise floor of the labels. Other diagnostics:

• Calibration: bucket pairs by the RM’s predicted preference probability \(\sigma(\Delta)\) and check that, e.g., pairs where it predicts 80% are actually preferred ~80% of the time. A miscalibrated RM gives PPO a distorted gradient.
• Reward distribution: plot the histogram of scores on a fixed eval set across training. A healthy RM has well-separated chosen/rejected distributions; a collapsing one pushes everything to extremes.
• Score on a fixed anchor set: because rewards are only identified up to a constant, track gaps between fixed reference responses, not absolute values, to compare checkpoints.

Held-out accuracy on your own data is necessary but self-referential: it shares the biases of your annotators. The standardized external check is RewardBench (Ai2), which scores an RM on curated pairs across chat, reasoning, and safety categories, so you can see where your RM is weak rather than just that it is 72% overall. It runs as a package on any HF sequence-classification RM (pip install reward-bench, then the rewardbench CLI pointed at your model path), and RewardBench 2 is the harder best-of-4 successor on fresh unseen prompts. Treat its per-category breakdown as a smoke test before you ever start PPO — an RM that is at chance on the reasoning subset will happily reward confident wrong math.

Reward model evaluation and reward hacking¶

The reward model is a proxy. It is a learned, imperfect stand-in for “what humans actually want,” and the moment you optimize a policy against it, you create an adversary — the policy — whose entire job is to find inputs the proxy scores highly. Goodhart’s law (“when a measure becomes a target, it ceases to be a good measure”) is not a footnote here; it is the central operational risk of RLHF. We give this its own deep treatment in Reward Hacking, Over-Optimization & Alignment Failures; here we cover the mechanism and the controls baked into the pipeline.

What reward hacking looks like¶

The policy discovers responses that score high under \(r_\phi\) but are not actually better to a human. Classic patterns:

• Length exploitation. The RM learned (from biased data) that longer answers are usually preferred, so the policy learns to pad. You see reward climbing while answers get verbose and repetitive. The fix at the data level is length-balanced preference data; at the optimization level, length-penalty rewards or length-debiasing of the RM.
• Sycophancy. The RM rewards agreeing with the user’s stated position, so the policy stops pushing back even when the user is wrong.
• Format farming. The RM over-weights markdown structure, bullet points, or a confident tone; the policy emits beautifully formatted nonsense.
• Out-of-distribution exploitation. The policy finds genuinely degenerate text (gibberish, repeated tokens, special-token soup) that lands in a corner of input space the RM never trained on and happens to score high. This is the most dramatic failure: reward goes up, human quality goes to zero.

The KL leash and over-optimization¶

The reason RLHF does not immediately collapse into reward hacking is the KL penalty that anchors the policy to the frozen reference (the SFT model). Optimization maximizes a regularized objective:

The KL term penalizes the policy for drifting from the SFT model’s distribution. Why does this curb hacking? Reward-hacking outputs are usually low-probability under the SFT model (no sane assistant produces token soup), so reaching them requires a large KL move that the penalty makes expensive. The coefficient \(\beta\) is the leash length: small \(\beta\) lets the policy roam (more reward, more hacking risk); large \(\beta\) keeps it near SFT (safer, less gain).

This trade-off has a famous empirical shape. Plot true quality (measured by held-out humans or a much larger “gold” RM) against the KL distance from the reference as you optimize. Quality rises, peaks, then falls as the policy over-optimizes the proxy — the reward-model over-optimization curve characterized by Gao, Schulman & Hilton (2022). They found the gold-reward gain follows a clean functional form in \(\sqrt{\mathrm{KL}}\), with proxy reward diverging upward while true reward turns over. The practical takeaway: the KL budget is a hyperparameter you tune to sit near the peak, and “reward went up” is never sufficient evidence that the model got better.

Defenses¶

The pipeline-level defenses against reward hacking are: (1) the KL penalty and early stopping on a gold metric; (2) reward model ensembles — average several independently trained RMs so the policy must hack all of them at once, which is harder; (3) iterated RLHF — once the policy finds an exploit, collect human preferences on those exploited outputs (where humans will rank the garbage last), retrain the RM, and the hole closes; (4) uncertainty-aware rewards — penalize responses where the RM ensemble disagrees, since high disagreement flags out-of-distribution inputs. We expand all of these in Reward Hacking, Over-Optimization & Alignment Failures.

The four-model setup¶

When people say RLHF is “memory-hungry and operationally heavy,” this is what they mean. The PPO stage juggles four models simultaneously. Understanding what each one is, whether it is trained or frozen, and why it exists is the single most-tested piece of RLHF systems knowledge.

This table is not an abstraction — it is literally the constructor signature of the libraries. TRL’s PPOTrainer takes four model arguments named policy, ref_model, reward_model, and value_model; veRL calls the same four roles actor, ref, reward, and critic and gives each its own resource pool in its config (veRL: HybridFlow & The Single-Controller Architecture). If you can name the four and say which are frozen, you can read any RLHF config file.

The actor and the reference¶

The actor is the policy \(\pi_\theta\) — the model you will ship. It is initialized from the SFT model and is the only model whose weights we want to end up with. The reference \(\pi_{\text{ref}}\) is a frozen snapshot of that same SFT model. At each PPO step we need both \(\log\pi_\theta(y\mid x)\) (from the actor) and \(\log\pi_{\text{ref}}(y\mid x)\) (from the reference) to compute the per-token KL penalty. The reference never updates — that is the whole point; it is the fixed star the policy is leashed to. (Some recipes periodically refresh the reference to the current policy; the default is to freeze it for the entire run.)

The reward model¶

The reward model \(r_\phi\) from Stage 2 is frozen during PPO. It is consulted once per response, at the end, to produce the terminal scalar reward. Because the reward is terminal (only the last token “earns” \(r_\phi\)), the per-token reward fed to the optimizer is the KL penalty at every token plus the reward model’s scalar at the final token:

That is, every token pays a little KL “rent,” and the whole sequence collects the RM’s score at the end. This sparse-terminal-reward structure is exactly why a critic is so valuable here — and why critic-free methods (next section) had to find another way.

The critic¶

The critic \(V_\psi\) is the subtle one. It is a value network that, at each token position, predicts the expected total future reward of the sequence from that point on. Why do we need it? Policy-gradient methods have enormous variance if you weight each action by the raw return; subtracting a baseline — the expected return — dramatically reduces variance without adding bias. The critic is that learned baseline, and combined with the rewards it yields the advantage \(A_t\) (via Generalized Advantage Estimation, GAE) that PPO actually optimizes. The critic is trained jointly with the actor, usually initialized from the reward model (which already knows how to map sequences to scalar values). The full machinery of advantages, GAE, and the clipped PPO objective is the subject of Policy Gradients & PPO for Language Models; the infrastructure view of advantage estimation and KL control is Advantage Estimation, KL Control & Stability Tricks.

The memory bill¶

Four models is a lot of GPU. If the policy is a \(7\)B-parameter model in bf16, then in the worst case:

• Actor: trained, so it carries weights + gradients + Adam optimizer states. Roughly \(2\) bytes (bf16 weights) \(+ 2\) (grad) \(+ 8\) (fp32 Adam moments, two states) \(\approx 16\) bytes/param \(\to \sim 112\) GB before activations.
• Critic: also trained; another full set of weights+grads+optimizer states, often similar size \(\to\) on the order of another \(\sim 100\) GB (a \(7\)B critic), or less if smaller.
• Reward: frozen, inference only \(\to \sim 14\) GB (bf16 weights, no grad/optimizer).
• Reference: frozen, inference only \(\to \sim 14\) GB.

So a “\(7\)B RLHF run” can have a resident footprint on the order of \(250\)+ GB before activations and the KV cache for generation — which is why RLHF training is dominated by sharding tricks and by the generation/training split covered in The Anatomy of an RL-for-LLM System. Practical economizers: make the RM and critic smaller than the policy (InstructGPT’s 6B RM for a 175B policy); use LoRA so actor and critic share a frozen backbone and the reference is “free” (the policy minus its LoRA adapters is the reference — see PEFT I: LoRA, QLoRA, DoRA & The Adapter Family); or fold the value head onto the actor so actor and critic share a trunk.

Putting it together: a skeleton RLHF loop¶

To cement how the pieces interlock, here is the control flow of the PPO stage in pseudo-real PyTorch. The actual PPO update (clipped surrogate, GAE) is intentionally abstracted into ppo_update, which the next chapter builds in full; the goal here is to see where each of the four models is called.

import torch

def rlhf_ppo_epoch(actor, critic, reward_model, ref_model,
                   prompts, tokenizer, beta_kl=0.02, ppo_update=None):
    """One outer iteration of the RLHF/PPO loop.

    actor       : policy π_θ        (trained)   -- generates and is updated
    critic      : value V_ψ         (trained)   -- baseline for advantages
    reward_model: r_φ               (frozen)    -- terminal scalar reward
    ref_model   : π_ref (SFT copy)  (frozen)    -- KL anchor
    """
    # ---- 1. ROLLOUT: actor generates responses to a batch of prompts. ----
    #     In production this runs on a fast inference engine (vLLM/SGLang);
    #     see "The Generation–Training Loop & Rollout Engines".
    queries = tokenizer(prompts, return_tensors="pt", padding=True)
    with torch.no_grad():
        responses = actor.generate(**queries, max_new_tokens=512, do_sample=True)

    # full sequence = prompt ++ response; build a mask marking response tokens
    seq, resp_mask = build_sequences_and_mask(queries, responses)

    # ---- 2. SCORE & ANCHOR (all under no_grad; these models are not updated). ----
    with torch.no_grad():
        # Terminal scalar reward for each complete response (frozen RM).
        scores = reward_model(seq, attention_mask=(seq != tokenizer.pad_token_id))  # (B,)

        # Per-token log-probs from the FROZEN reference (for the KL penalty).
        ref_logprobs = token_logprobs(ref_model, seq, resp_mask)                    # (B, T)

    # ---- 3. The actor's own per-token log-probs and the critic's values
    #         (these DO require grad -- they define the PPO objective). ----
    actor_logprobs = token_logprobs(actor, seq, resp_mask)                          # (B, T)
    values         = critic(seq, resp_mask)                                         # (B, T)

    # ---- 4. Build the per-token reward: KL penalty everywhere + RM score at end. ----
    #     Convention: every per-token tensor here is aligned to the RESPONSE
    #     segment only (T = max response length), so column t is the t-th
    #     GENERATED token and resp_mask marks real vs right-padding. Getting this
    #     alignment wrong (leaving prompt columns in) is the #1 RLHF plumbing bug.
    kl_per_token = actor_logprobs.detach() - ref_logprobs                           # (B, T)
    rewards = -beta_kl * kl_per_token                                               # KL "rent"
    last = resp_mask.sum(dim=1) - 1                                                 # final real resp token
    rewards[torch.arange(rewards.size(0)), last] += scores                          # add terminal r_φ

    # ---- 5. PPO update: compute advantages (GAE) from (rewards, values), then
    #         take clipped policy-gradient + value-function steps. See chapter 5.6. ----
    stats = ppo_update(actor, critic,
                       logprobs=actor_logprobs, old_logprobs=actor_logprobs.detach(),
                       values=values, rewards=rewards, resp_mask=resp_mask)
    return stats

Trace the four models through it: the actor generates (step 1) and supplies differentiable log-probs (step 3); the reward model scores once (step 2); the reference supplies frozen log-probs for KL (step 2); the critic supplies values (step 3) that become the advantage baseline (step 5). Every line of an industrial RLHF library is an elaboration of this skeleton — better generation engines, sharded models, vectorized KL control, and the PPO clip — but the data flow is exactly this.

Key Takeaways¶

Further reading¶

• Ouyang, Wu, Jiang, et al., Training language models to follow instructions with human feedback (InstructGPT, 2022) — the canonical end-to-end RLHF recipe this chapter follows.
• Christiano, Leike, Brown, et al., Deep Reinforcement Learning from Human Preferences (2017) — the foundational paper that introduced learning a reward model from pairwise human comparisons.
• Bradley & Terry, Rank Analysis of Incomplete Block Designs (1952) — the original paired-comparison model underlying the reward loss; Luce/Plackett for the ranking generalization.
• Stiennon, Ouyang, Wu, et al., Learning to summarize from human feedback (2020) — the precursor that worked out reward modeling and PPO for a concrete task (summarization).
• Bai, Jones, Ndousse, et al. (Anthropic), Training a Helpful and Harmless Assistant with RLHF (2022) — the helpfulness/harmlessness split and separate reward models per axis.
• Gao, Schulman & Hilton, Scaling Laws for Reward Model Overoptimization (2022) — the proxy-vs-gold reward over-optimization curve and the \(\sqrt{\mathrm{KL}}\) functional form.
• Touvron, Martin, Stone, et al., Llama 2: Open Foundation and Fine-Tuned Chat Models (2023) — iterated RLHF, separate helpfulness/safety RMs, rejection-sampling + PPO in practice.
• HuggingFace TRL (RewardTrainer, PPOTrainer) and OpenRLHF repositories — production implementations of the reward loss and four-model loop; see TRL: HuggingFace’s RL Library.

Exercises¶

1. A colleague trains two reward models on the same preference data with different random seeds. On a fixed eval set, RM-A gives the chosen responses an average score of \(+5.2\) and RM-B gives them \(-1.8\). Your colleague concludes “RM-A is the stronger model, its rewards are higher.” Explain why this reasoning is invalid, citing the relevant property of the Bradley-Terry model. Then explain the concrete consequence this same property has for the PPO stage (why the pipeline mean-centers rewards).

2. For one comparison a reward model outputs \(r_\phi(x, y_w) = 1.2\) (chosen) and \(r_\phi(x, y_l) = 1.0\) (rejected). (a) What preference probability does the model assign to the chosen response? (b) What is the per-example Bradley-Terry loss in nats? © What is the gradient of the loss with respect to the margin \(\Delta\), and which direction does gradient descent push \(\Delta\)? (d) Now the model has a different pair backwards: \(r_\phi(x, y_w) = 0.3\), \(r_\phi(x, y_l) = 2.0\). Recompute (a)-© and comment on how the loss and gradient magnitudes changed. Use \(e^{-0.2}\approx 0.819\), \(e^{1.7}\approx 5.474\).

3. A labeler ranks \(K\) completions of a single prompt from best to worst. (a) How many pairwise preference rows \(\binom{K}{2}\) does this ranking yield for \(K = 4, 5,\) and \(9\)? (b) InstructGPT puts all pairs from one prompt into a single forward/backward pass. Compare the number of transformer encodings of completions needed per prompt under (i) the naive approach that shuffles all pairs together and encodes both completions of each pair independently, versus (ii) the batched approach that encodes each completion once. Give a general formula for the ratio and evaluate it at \(K = 9\). © Besides the compute saving, what training pathology does the batched approach prevent, and why?

4. Implement uncertainty_aware_reward(rms, seq, attention_mask, lam), a defense against reward hacking that (a) scores a batch with an ensemble of reward models and (b) penalizes each response by the ensemble’s disagreement (standard deviation), so out-of-distribution responses where the RMs disagree get a lower effective reward. Each element of rms is a RewardModel from the chapter that maps (input_ids, attention_mask) to a (B,) reward. Return the (B,) effective reward. Explain in one sentence why the standard-deviation term flags reward-hacking candidates.

import torch

def uncertainty_aware_reward(rms, seq, attention_mask, lam=0.5):
    """Ensemble reward with an uncertainty (disagreement) penalty.

    rms            : list of M frozen RewardModel instances (independently trained)
    seq, attention_mask : (B, T) token ids and mask for B full sequences
    lam            : weight on the disagreement penalty
    Returns effective reward (B,) = mean_reward - lam * std_reward.
    """
    # (M, B): each row is one RM's scalar reward for the whole batch.
    stacked = torch.stack([rm(seq, attention_mask) for rm in rms], dim=0)
    mean_r = stacked.mean(dim=0)                 # (B,) ensemble consensus
    # Unbiased=False so a single-model ensemble gives std 0 rather than NaN.
    std_r  = stacked.std(dim=0, unbiased=False)  # (B,) disagreement
    return mean_r - lam * std_r                  # (B,) penalized reward

5. In the skeleton RLHF loop, step 4 builds the per-token reward \(R_t\) from the terminal RM score plus a per-token KL penalty (the formula in the chapter). Implement build_per_token_rewards(scores, actor_logprobs, ref_logprobs, resp_mask, beta_kl) returning a (B, T) tensor where every response token pays KL “rent” and the RM’s scalar is added at each sequence’s last response token. scores is (B,); the log-prob tensors and resp_mask are (B, T) aligned to the response segment (column \(t\) is the \(t\)-th generated token; the mask is 1 on real generated tokens and 0 on right-padding). State which term of \(R_t\) must be detached from the actor’s graph and why.

import torch

def build_per_token_rewards(scores, actor_logprobs, ref_logprobs,
                            resp_mask, beta_kl=0.02):
    """Per-token reward R_t = -beta * KL_t  + r_phi * 1[t = T].

    scores         : (B,)   terminal reward-model score per sequence
    actor_logprobs : (B, T) log pi_theta(y_t | ...)   (from the actor)
    ref_logprobs   : (B, T) log pi_ref(y_t | ...)     (frozen reference)
    resp_mask      : (B, T) 1 on real generated tokens, 0 on right-padding
    All (B, T) tensors are aligned to the RESPONSE segment (t = t-th
    generated token), not to the prompt++response sequence.
    Returns rewards (B, T).
    """
    # KL estimate per token; detach the actor term -- the KL enters as a
    # REWARD, not through the policy-gradient objective, so it must not
    # backprop into the actor here (that path is handled by the PPO update).
    kl_per_token = actor_logprobs.detach() - ref_logprobs      # (B, T)
    rewards = -beta_kl * kl_per_token * resp_mask              # KL "rent" on resp tokens

    # Add the terminal RM score at each sequence's LAST response token.
    last = resp_mask.sum(dim=1) - 1                            # (B,) final resp index
    batch_idx = torch.arange(rewards.size(0), device=rewards.device)
    rewards[batch_idx, last] += scores                        # terminal reward
    return rewards

6. You are budgeting GPU memory for an RLHF run whose policy is a 13B-parameter model in bf16, using the chapter’s byte-accounting (bf16 weight = 2 bytes/param; a trained model also carries 2 bytes gradient + 8 bytes fp32 Adam moments = 16 bytes/param total; a frozen model carries only its 2 bytes/param of weights). (a) Estimate the resident footprint of the naive four-model setup, assuming the critic is also 13B. (b) Now apply two economizers from the chapter: use a small 1.5B reward model, and fold the value head onto the actor so there is no separate critic model. Re-estimate the footprint. © By what factor did the footprint shrink, and which single model dominates it now?