Build and diagnose an RLHF training pipeline

Establish the model and artifact handoffs#

RLHF connects a starting policy, preference data, a reward model and policy optimization. Stiennon and colleagues study these components in summarization. Their task-specific recipe motivates a reproducible experiment; it is not a universal configuration.

Save the supervised policy checkpoint, tokenizer, conversation template and demonstration revision. Export a tiny reference batch containing token IDs, answer masks and expected output boundaries. Load that batch through the rollout and reward paths before training. Text that looks identical can tokenize differently when separators or special tokens change.

Define four roles: the trainable policy generates answers; the frozen reference anchors behavior; the frozen reward model scores completions; the trainable value model estimates returns. They may share initialization, but their parameter updates and saved versions must remain distinguishable.

Freeze the data split and reward interpretation#

Use a family-level split for demonstrations, comparisons and evaluation. Keep all alternatives for one prompt together. Preserve label reasons and ties; decide their treatment before fitting. Reserve a development set for reward-model selection and a final policy test that neither selection process can inspect.

A usable comparison rubric might prefer the response that preserves release conditions and admits missing evidence, while separately flagging invented functionality. Test short accurate answers against verbose inventions. Export the scorer’s exact input serialization, truncation rule, scalar extraction point and score transformation. A new normalization changes the effective strength of every downstream penalty.

Run a frozen-scoring check twice on an inspected batch. Investigate different scores, incorrect answer boundaries or a score taken from padding before launching policy training. Save the held-out comparison report alongside the reward checkpoint. Pairwise accuracy alone cannot establish that the scorer remains reliable on outputs produced later by optimization.

Define reward and the reference penalty#

For prompt x and response y, consider a sequence-level objective:

The reference stays frozen. For a rollout from behavior policy πold, one sampled penalty uses:

Here s_t is the prompt plus preceding response tokens. Huang and colleagues’ reproduction details reward and reference scoring, terminal handling and implementation choices. The displayed token rewards are fixed rollout data during that batch’s update. Their expected summed log ratio estimates sequence KL under the behavior policy; one sampled sum can be negative.

Decide whether reported KL is a token mean or sequence sum. Record reward scaling and β together. Compare neither values across incompatible reductions nor differently normalized scorers. The reference penalty controls departure from the reference; it does not certify correctness or permissions.

Compute advantages and the clipped update#

Store behavior log probabilities and value estimates with every rollout. For a completed episode, use zero continuation value at termination. A token limit needs an explicit truncation policy; treating truncation as termination changes the target.

Generalized advantage estimation combines value residuals with a decay controlled by γ and λ:

Apply episode boundaries and response masks to these sums. Save advantages before the update so policy changes do not silently change that batch’s target.

PPO’s clipped objective:

Compute ρ_t by exponentiating current minus saved behavior log probability. The reference model is absent from that ratio. For an illustrative positive advantage of 2, ratio 1.4 and ε=0.2, the terms are 2.8 and 2.4; the surrogate selects 2.4. This arithmetic explains one sample’s objective, not a guaranteed limit on parameter movement.

Maximize the policy surrogate, or minimize its negative, while fitting the value model to declared return targets. Refresh rollouts after the configured update cycle. The reproduction paper also shows why scoring and optimization details need inspection rather than blind adoption.

Instrument the loop and diagnose symptoms#

Record generation, reward/reference scoring, value estimation and optimization costs separately. Log raw reward, penalized return, reference KL, update ratios, clipped fraction, value loss, response length and termination rate. State masks and reductions; changing length can change several metrics without improving quality.

These are hypotheses to test, not diagnoses supplied by a single metric. Alter one suspected cause and replay the relevant check.

Follow one fictional rollout through the system#

Use the fictional release-note corpus from the alignment guide. Generate 32 training responses with a 96-token limit, temperature 1 and no probability filtering. These are teaching settings, not hardware or model recommendations. Save the exact generation distribution used to compute behavior probabilities.

Suppose one invented response receives reward 1.8 and a summed behavior/reference log ratio of 4.0. With illustrative β=0.05, its summed penalty is 0.2 and shaped return is 1.6. Inspect the token trace to confirm that the terminal reward is added once and padding contributes nothing. Compute its advantages, check the positive-advantage example above, and save the updated checkpoint.

Now imagine development review finds longer, confident answers asserting planned features as available, despite higher reward. Pause the run. Compare matched responses using the frozen scorer and inspect training labels. If the scorer prefers the invention, repair reward data and issue a new scorer revision; if rewards were attached to the wrong token, repair the implementation. Restart from the declared baseline instead of mixing incomparable scoring regimes.

Select a checkpoint using independent evidence#

Choose candidates on development data using a written rule that first blocks critical failures, then compares quality and cost. Randomize response order, allow ties and retain adjudication reasons. Inspect unfamiliar prompt families, evidence gaps and legitimate requests that might be refused. Report sample counts and uncertainty.

After selection, evaluate once on the sealed assessment set with the same evidence and decoding settings as the baseline. The summarization study directly examines divergence between optimized reward and human judgments; a training-score maximum is insufficient. Preserve the selected checkpoint, rejected alternatives, evaluation outputs and full cost record. A failed final gate requires a new documented experiment and assessment plan, not repeated tuning against the same supposedly unseen test.