Vol. XVIII · Learning From Preferences Plate 5 of 11KL divergence, KL penalty, beta, reference policy, trust region, entropy, diversityLE

Staying Close.

Every nat of KL is paid for in reward. β is the price.

The KL divergence KL(π ‖ π_ref) measures, in nats, how differently the policy distributes its probability from the reference. In RLHF it does two jobs. It is a trust region: the reward model was trained on responses from the reference, and its scores are reliable only near them, so moving far away means trusting it where it was never tested. And it preserves diversity: the optimum π_ref · e^{r/β} keeps some probability on everything the reference could say, where pure reward maximisation would put all of it on a single response.

The preference world of this volume: every prompt is a sum such as 7 + 8, half of them with the user’s wrong guess (“is it 16?”), and a response is one of nine types — the plain answer, an explanation, flattery, hedging, wrong answers, agreement with the guess. Each has a true quality fixed by the course: the plain correct answer is best, padding and flattery cost a little, hedging and wrong answers more. Simulated raters prefer by the Bradley–Terry rule on that quality plus a bias towards long, flattering and agreeable answers. The workshop computes the exact optimum for each β, using a reward model fitted to 1 000 of the raters’ comparisons, and traces reward, true quality, KL and the policy’s entropy as β varies.

Below: left, the reward-model score, true quality and KL of the optimal policy against β. Right, the probability of each response type at the chosen β beside the reference.

What you will learn

Define the KL divergence between policy and reference.
Explain the trust-region role of the KL penalty.
Show how β trades reward for distance.
Relate the penalty to the policy’s diversity.

Workshop

Exact optimum π ∝ π_ref · e^{r/β} for each β, with r the reward model fitted to 1 000 comparisons from raters of bias 1; quantities averaged over prompts with and without a guess. Entropy of the policy in nats.

A price on drift

At β ≈ 0.3 the optimum sits 0.70 nats from the reference, with entropy 1.47 nats against the reference’s 2.03: it has moved most of its probability to correct answers but still uses several of them. At β ≈ 3 it barely moves (0.02 nats, quality 1.38). At β ≈ 0.02 it goes 2.94 nats away, entropy 0.00: nearly all its probability on the single response the reward model scores highest, the flattering explanation, 20.0 words long, with reward 1.88 but true quality 1.70.

β is therefore not a detail but the main dial of preference training, and it is usually set by watching KL rather than reward: a run is stopped, or β raised, when the policy has drifted a few nats, because beyond that the reward model is being asked about text unlike anything it was trained on. The penalty also has a cost the plot makes plain — reward that could have been collected is left on the table — which is the price of not trusting a proxy further than it was tested.

The KL here is the reverse KL, KL(π ‖ π_ref): it is large when the policy puts probability where the reference puts little, and indifferent when the policy drops responses the reference liked. Optimising it makes policies mode-seeking — they narrow, rather than spread.

Terms

KL divergence: A measure, in nats, of how one distribution differs from another.
Trust region: The neighbourhood of the reference where the reward model is reliable.
β: The KL coefficient: the exchange rate between reward and distance.
Entropy: A measure of how spread a policy’s probability is.

The KL penalty keeps the policy where the reward model can be trusted and keeps its variety; β sets how much reward drift must buy.

Sources

Jaques et al., “Way Off-Policy Batch Deep Reinforcement Learning of Implicit Human Preferences in Dialog”, 2019 (KL control).
Ziegler et al., “Fine-Tuning Language Models from Human Preferences”, arXiv:1909.08593, 2019.
The course’s preference world: js/post.js.