The KL leash in RLHF

With a reward model in hand, the third stage of InstructGPT-style RLHF trains the language model to score well on it. The model samples completions, the reward model scores them, and a policy-gradient method such as PPO updates the weights. Left alone, that optimization drifts. The KL leash is the penalty that keeps the model close to where it started.

The objective

The policy maximizes the reward minus a KL term:

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

β sets the leash. A small β lets the policy chase the reward. A large β keeps the model almost unchanged from the supervised version.

• Too small: the model finds reward-model quirks and drifts into odd styles.
• Too large: almost nothing changes, and the reward barely moves.

Walking through a long-answer failure

Suppose the reward model has learned that long, formal answers score higher. With no leash, an optimizer will produce essays even for "What is 2 + 2?"

With a moderate β, the policy keeps the reward gains on hard questions. It stays close to the supervised model on easy ones, which answered briefly. Standard practice is to sweep β and plot two curves together: length and KL on one axis, and win rate against the supervised model on the other.

Stopping on the gold set

The reward model is a proxy. As its score climbs, the true human win rate can peak and then fall. This is over-optimization, Goodhart's law applied to training.

A practical guard is a frozen gold preference set: prompts labeled once by trusted annotators, and never used for training. Measure the policy's win rate against the supervised model on that set every few steps.

Common mistakes

• Tuning β by the reward curve. Check the gold win rate and the KL together.
• Using a frozen reference that is not the starting point. The leash has to pull toward the model the policy began from.
• Judging the policy only by its reward model. The reward model is the thing being gamed.

Check yourself

1. What happens to the policy if β is very large? Answer: the KL term dominates, so the policy barely moves from the reference and the reward barely improves.
2. Why keep a gold set that is never trained on? Answer: the reward model can be gamed, and only untouched human judgments reveal when the reward has stopped tracking quality.

Key takeaways

• The third RLHF stage maximizes reward minus β times the KL to a frozen reference.
• β is the leash. Sweep it, and watch length, KL and win rate together.
• Over-optimization is when the proxy rises and human win rate falls.
• A never-trained gold set is the stopping signal.

Further learning

The PPO mechanics behind this stage are in PPO and the KL leash. An offline alternative that skips the reward loop is DPO.