The Optimization Loop

Lesson 3 of 4 in RLHF: Learning from Preferences.

With a reward model in hand, stage three closes the loop: make the language model seek high scores. In reinforcement-learning vocabulary, the language model becomes the policy — the thing that acts. Generating a response is a sequence of actions (one per Token), and when the response is complete, the reward model hands back a score. The training question is: how do you nudge billions of Weights so that higher-scoring responses become more probable, without wrecking everything the model already knows?

The workhorse answer is PPO — Proximal Policy Optimization (Schulman et al. 2017, arXiv:1707.06347). PPO was not invented for language models; it is a general RL algorithm whose signature virtue is conservatism. Each update is clipped so the new policy cannot lurch too far from the policy that generated the current batch of responses. That built-in caution is why it became the default: RL on language models is unstable enough without giant steps.

But PPO’s clipping only limits each step. Nothing about it stops thousands of small steps from marching the model somewhere strange. For that, RLHF adds a second restraint — the one to remember.

The RLHF optimization loop — with the KL leash

  1. Batch of prompts

    Sampled from a prompt dataset each iteration; no human writes or reads anything in this loop.

  2. Policy generates responses

    The current model samples a response to every prompt in the batch.

  3. Reward model scores each response

    One forward pass per response through the frozen reward model: the proxy for human preference.

  4. KL leash: compare to the frozen SFT reference

    A frozen copy of the SFT model scores the same responses. Where the policy’s token probabilities drift from the reference’s, a KL penalty subtracts from the reward — the farther the drift, the bigger the subtraction.

  5. Combined signal: reward − β · KL

    The scalar the optimizer actually maximizes. β sets the leash length: how much reward is worth how much drift.

  6. PPO update to policy weights

    A clipped, conservative gradient step toward responses that scored well.

  7. Preference-tuned checkpoint

    Training stops on schedule or when held-out evaluation says quality has peaked — lesson four explains why stopping matters.

The KL penalty deserves its own paragraph, because it is doing two different jobs at once.

Job one: stay fluent. The reward model only judges; it does not teach language. If the optimizer’s only signal were the score, the policy could drift into strange phrasings, broken grammar, or degenerate repetition — as long as the reward model happened to score them well. The frozen SFT reference anchors the policy to a distribution that is known to produce coherent text. Drift costs reward, so the policy stays in the neighborhood of language that made sense.

Job two: don’t exploit the judge. The reward model was trained on responses from the SFT-era model. Its judgments are trustworthy near that distribution and increasingly untested far from it. The KL leash keeps the policy inside the region where the judge has actually seen evidence — which slows down (though never prevents) the policy’s discovery of scoring exploits. Loosen β toward zero and you remove the leash: the policy is free to wander into territory where the reward model’s scores are confidently wrong, and it will find that territory, because that is where unearned reward lives.

The objective, one term at a time

The quantity being maximized, for prompts x and sampled responses y:

J(θ) = E[ r_φ(x, y) − β · KL( π_θ(·|x) ‖ π_ref(·|x) ) ]

  • π_θ — the policy being trained; π_ref — the frozen SFT reference; r_φ — the frozen reward model.
  • The KL term is computed per token in practice. At each position, the penalty is proportional to log(π_θ(token) / π_ref(token)) — cheap to compute since both models already produce token probabilities, and effectively a per-token tax on deviating from the reference (this per-token formulation is how Ouyang et al. 2022, arXiv:2203.02155, apply it).
  • β trades reward for closeness. It is the price of drift, and labs tune it (or adapt it during training) against held-out evaluations.
  • PPO adds its own, separate restraint: the clipped surrogate objective, which limits how far any single update moves the policy from the version that generated the batch (Schulman et al. 2017, arXiv:1707.06347). Two brakes at two timescales — PPO clipping per update, the KL leash over the whole run. They are frequently confused; keep them apart.
  • The memory bill is real. A textbook PPO-RLHF step typically holds four models: the policy (training), the reference (frozen), the reward model (frozen), and a value model that estimates expected reward to reduce gradient variance. This footprint is a large part of why lighter-weight alternatives — DPO next module — were so eagerly adopted.
  • InstructGPT additionally mixed in gradients from the original pretraining objective (their “PPO-ptx” variant) to limit regressions on standard NLP benchmarks — a reminder that even the canonical recipe needed patches to avoid forgetting.

Interactive checkpoint quiz (1 questions) — open this page in a browser to take it.