KL penalty in RLHF: KL-regularised RL, the optimal policy and how to choose beta

The KL penalty in RLHF is a term that charges the policy for drifting away from a reference model, usually the supervised fine-tuned model it started from. The training objective becomes expected reward minus $\beta$ times the KL divergence to that reference, and its exact solution is the reference distribution reweighted by $\exp (r/\beta )$. The coefficient $\beta$ sets how far the policy may travel to chase reward.

This page derives that closed-form optimum, shows on a small toy you can drive with a slider how $\beta$ trades reward for KL, explains how the sequence-level penalty becomes a per-token one in practice, compares the three common ways of estimating KL from samples, and runs the adaptive controller that many RLHF pipelines use to pick $\beta$. It ends with what the penalty does not protect against. Every number comes from seeded numpy runs on toy problems; nothing here involves a real language model.

Why does RLHF need a KL penalty?#

Take one prompt and 40 possible responses. Each has a reward between 0.03 and 0.91, and the reference model, a stand-in for the fine-tuned language model, spreads its probability over all of them with an average reward of 0.471. The best response is one the reference model almost never writes: probability 0.00089.

Ask an optimiser to maximise reward with no penalty and the answer is trivial: put all probability on the single best response. Expected reward rises to 0.907, entropy falls from 3.34 nats to zero, and the policy has moved $\log (1/0.00089)=7.02$ nats from where it started. On a real model this is the familiar failure: a policy that repeats one high-scoring pattern for every prompt, loses fluency on everything the reward model never looked at, and finds the reward model's blind spots. The reward model was trained on responses near the reference distribution, so its scores are least trustworthy exactly where an unconstrained optimiser goes.

The KL penalty is the standard fix. It appears in the first applications of human-preference RL to language models [1], in the summarisation work that followed [9] and in instruction-following RLHF [10], and it predates them: Jaques et al. used KL-control towards a pretrained sequence model to keep a fine-tuned generator close to its data prior [2]. Three reasons it earns its place:

• It limits how far reward model errors can be exploited. Gold quality typically rises and then falls as the policy moves away from the reference; the penalty chooses how far to go [11]. Our research note on reward overoptimization measures that hump in detail.
• It keeps what the reference model already does well. Fluency, formatting and knowledge that the reward does not mention are preserved because changing them costs KL and earns nothing.
• It turns an ill-posed problem into a well-posed one. Pure reward maximisation has a degenerate, deterministic solution; the regularised problem has a unique, smooth one that we can write down.

The KL-regularised objective#

For a prompt $x$, the objective over policies $\pi$ is

Here $y$ is a complete response, $r(x,y)$ its reward, ${\pi }_{\text{ref}}$ the frozen reference policy and $\beta >0$ the KL coefficient. In words: you earn reward, and you pay $\beta$ nats-worth of reward for every nat of divergence. The divergence is the reverse KL, measured under the policy's own samples, which is what makes it cheap to estimate during training: you already have samples from $\pi$ and can score them under ${\pi }_{\text{ref}}$.

What is the optimal policy? The closed form#

Define the tilted distribution and its normaliser

$Z(x)$ is the partition function: the reference-weighted average of $\exp (r/\beta )$. Taking logs, $r(x,y)=\beta \log \frac{{\pi }^{\ast }(y\mid x)}{{\pi }_{\text{ref}}(y\mid x)}+\beta \log Z(x)$. Substitute that into $J$ for any policy $\pi$ and collect terms:

The first term does not depend on $\pi$, and a KL divergence is zero only when the two distributions match. So ${\pi }^{\ast }$ is the unique maximiser, and the optimal value is $V^{\ast }(x)=\beta \log {\mathbb{E}}_{{\pi }_{\text{ref}}}[\exp (r/\beta )]$, a "soft maximum" of the reward. This is the Gibbs variational principle, and it is why KL-regularised RL is often described as inference: ${\pi }^{\ast }$ is a posterior with the reference as prior and $\exp (r/\beta )$ as likelihood [3, 4]. Maximum-entropy RL is the special case with a uniform reference [14], and the same exponential family shows up in maximum-entropy inverse RL.

A worked example with three responses#

Let the reference put probabilities 0.6, 0.3 and 0.1 on three responses with rewards 0.2, 0.5 and 0.9, and set $\beta =0.2$. The unnormalised weights are $0.6e^1=1.631$, $0.3e^{2.5}=3.655$ and $0.1e^{4.5}=9.002$, so $Z=14.29$ and ${\pi }^{\ast }=(0.114,0.256,0.630)$. Expected reward rises from 0.360 to 0.718, the KL is 0.929 nats, and $J=0.718-0.2\times 0.929=0.532=0.2\log 14.29$, as the formula promises. At $\beta =0.5$ the same calculation gives ${\pi }^{\ast }=(0.387,0.352,0.261)$ and KL 0.137; at $\beta =0.1$ it gives $(0.005,0.052,0.943)$ and KL 2.00.

Four properties worth remembering#

• Only $r/\beta$ matters. Doubling the reward and doubling $\beta$ give the same policy, so a value of $\beta$ means nothing without the reward's scale. Reward normalisation changes the effective $\beta$.
• Constant shifts do nothing. Adding a constant to every response's reward for a prompt cancels in $Z$ (we checked: the largest probability change was $8\times 10^{-16}$). Per-prompt reward offsets are harmless; per-response ones are not.
• No new support. If ${\pi }_{\text{ref}}(y)=0$ then ${\pi }^{\ast }(y)=0$ for every $\beta$. The optimum can only reweight what the reference already does.
• Ratios among equal-reward responses are frozen. Two responses with the same reward keep the reference's probability ratio. If a reward cannot tell a good response from a bad one, the KL-regularised optimum keeps whatever mix the reference had. Our pillar article on RL for LLMs shows this with a verifier that checks only final answers.

The same identity is what direct preference methods exploit: solve it for $r$, plug into a Bradley-Terry model, and $\log Z$ cancels between the two responses of a pair, leaving a loss on $\beta \log (\pi /{\pi }_{\text{ref}})$ alone [5]. The RLHF explainer places that in the full pipeline.

How does beta trade reward against KL?#

Back to the 40-response toy. The figure computes ${\pi }^{\ast }$ for any $\beta$ you choose and plots where it lands on the reward-versus-KL plane. The dark line marks the reference probabilities; the bars are the tilted policy.

Three things are visible. First, the frontier is concave: the first nat of KL buys about 0.2 of reward, the last two nats buy 0.02. Second, its slope at any point equals $\beta$ exactly. That follows from the derivation (the frontier is a Legendre transform) and we confirmed it numerically: finite differences gave slopes of 0.300, 0.050 and 0.010 at $\beta =0.3$, 0.05 and 0.01. So $\beta$ is literally the exchange rate between reward and KL at the operating point. Third, most of the entropy collapse happens in a narrow band of $\beta$, between about 0.1 and 0.01 here, which is why small changes to the coefficient can make large qualitative differences in diversity.

Best-of-n sits just below the frontier#

Best-of-n sampling (draw $n$ responses, keep the highest-reward one) is the cheapest alternative to RL, and it is close to KL-optimal. In the toy, best-of-16 reaches reward 0.721 at a KL of 1.83 nats, while ${\pi }^{\ast }$ at the same KL reaches 0.742; best-of-256 reaches 0.868 at 4.47 nats against 0.870 on the frontier. The familiar formula $\log n-(n-1)/n$ slightly overstates best-of-n KL (4.55 versus 4.47 at $n=256$), a gap that Beirami et al. analyse and that grows when the reference puts a lot of mass on a few responses [13].

From sequences to tokens: how the penalty is implemented#

A language model does not sample whole responses as single actions; it samples tokens. The KL between two autoregressive models decomposes by the chain rule:

$y_t$ is the token at position $t$ and $y_{<t}$ the tokens before it. Sequence KL is the expected sum of per-token log-ratios. So RLHF implementations fold the penalty into a per-token reward: every token receives $-\beta \log (\pi /{\pi }_{\text{ref}})$ for that token, and the final token also receives the reward model's score [1, 10]. A standard policy-gradient learner such as PPO then sees an ordinary episodic RL problem with a dense shaping term.

Does that per-token scheme really converge to the sequence-level ${\pi }^{\ast }$? We checked on a tabular autoregressive policy with a 4-token vocabulary and 3-token responses (64 sequences), random reference, random rewards, $\beta =0.1$. Training with the exact policy gradient of the per-token-penalised return, the largest gap between the learned sequence distribution and the closed form was 0.0029 after 1,000 steps and $3.7\times 10^{-5}$ after 20,000, and the objective converged to 0.59596 against a closed-form 0.59596. The per-token KL sum and the sequence KL agreed to every printed digit throughout, as the chain rule says they must. A nice consequence, developed by Rafailov et al., is that under this token-level view the log-ratio $\beta \log ({\pi }^{\ast }/{\pi }_{\text{ref}})$ of each token plays the role of an optimal soft advantage [6].

Estimating KL from samples: k1, k2 and k3#

During training you never sum over all responses; you have a batch of samples from $\pi$ and their log-probabilities under both models. Write $\rho ={\pi }_{\text{ref}}(y)/\pi (y)$ for one sample. Three per-sample estimators are in common use, named in Schulman's note on the topic [7]:

$k_1$ is unbiased but can be negative for individual samples. $k_2$ is always non-negative but biased. $k_3$ adds $\rho -1$, which has expectation zero under $\pi$, so it stays unbiased while being non-negative for every sample. The GRPO recipe adds the KL to the loss with $k_3$ rather than to the reward [8]. How good each is depends on how far apart the two models are. We drew 20,000 batches of 64 samples from ${\pi }^{\ast }$ in the 40-response toy:

Near the reference, $k_3$ is the best of the three: unbiased with less than half the spread of $k_1$. Far from it, $k_3$ becomes almost useless, because $\rho$ is occasionally enormous (a sample the reference likes far more than the policy does) and a batch mean of 4.48 with a standard deviation of 22.7 is noise; its sampling distribution is so skewed that even the 20,000-batch average misses by 0.5. $k_2$ more than doubles the true KL at $\beta =0.05$. $k_1$ stays unbiased and tight throughout. If your run is meant to stay within a nat or so of the reference, $k_3$ is a good choice; if you log KL to monitor drift, log $k_1$ as well, because the estimator you penalise with is not always the one you should trust.

Choosing beta: fixed or adaptive#

Ziegler et al. observed that runs with the same $\beta$ but different seeds can end at quite different KLs, which makes them hard to compare, and proposed targeting a KL instead [1]. Their log-space proportional controller updates

${\mathrm{KL}}_t$ is the current KL, $e_t$ the clipped relative error and $K_{\beta }$ the gain. If KL is too high, $\beta$ grows by up to 2 percent per step; too low, it shrinks. We ran it on the 64-sequence tabular model with sampled REINFORCE (batch 64, per-token penalty, batch-mean baseline, learning rate 0.5), a 0.5-nat target, 1,500 steps and 8 seeds, feeding the controller the batch's $k_1$ estimate.

Both starts converge to the same place: final $\beta$ of 0.153 (from 0.02) and 0.153 (from 0.5), seed-to-seed SD about 0.015, with true KL at 0.51. The path is not smooth. Early on the KL is near zero, the controller relaxes $\beta$ to 0.002 and 0.009, the policy sprints, and KL overshoots to 1.71 and 0.91 nats before the penalty catches up. In a real run that overshoot is where reward hacks get learned; it is worth starting $\beta$ near a sensible value even with a controller.

One subtlety: the closed-form policy with KL 0.5 has $\beta =0.178$, not 0.153. The controller holds KL at the target while the optimiser is still moving, so the $\beta$ it settles on reflects the optimiser's speed as well as the reward. Fixed-$\beta$ runs show the same thing from the other side: after 1,500 steps at $\beta =0.1$ the true KL was 0.75 against a closed-form 1.22. Training stops before the optimum, so the KL you get is a joint property of $\beta$, learning rate and step count. Treat the target KL, not $\beta$, as the quantity you tune and report.

Forward versus reverse KL#

The penalty uses $\mathrm{KL}(\pi \Vert {\pi }_{\text{ref}})$, with the policy first. That direction is mode-seeking: it heavily penalises putting mass where the reference has little, and barely penalises dropping some of the reference's modes. The table above shows the consequence: at $\beta =0.02$ the policy has abandoned most of the 40 responses and kept a handful. The other direction, $\mathrm{KL}({\pi }_{\text{ref}}\Vert \pi )$, is what supervised fine-tuning on reference samples minimises, and it is mass-covering. Neither is "the" right regulariser. Reverse KL is the one with the clean closed form and the cheap on-policy estimator, and it explains the loss of diversity that RL-tuned models often show relative to their supervised starting points.

What the KL penalty does not do#

It does not fix a wrong reward. The penalty picks a point on the path between the reference and the reward's maximiser. If gold quality peaks somewhere along that path, a well-tuned $\beta$ stops near the peak, which is valuable, but no $\beta$ moves you off the path. Gao et al. found that in their setting the KL penalty acted much like early stopping rather than improving the gold-reward-versus-KL frontier [11].

It does not bound how much a rare behaviour changes. KL is an average. Raising one response from probability 0.0001 to 0.05 costs only 0.26 nats; raising it to 0.01 costs 0.036. A "small" KL budget can contain a fifty-fold increase in some rare, possibly harmful output. For safety-relevant behaviours, measure them directly with targeted evaluations instead of inferring them from KL.

It may not even exist under heavy-tailed reward error. If the reward model's errors have heavy enough tails, ${\mathbb{E}}_{{\pi }_{\text{ref}}}[\exp (r/\beta )]$ diverges and the KL-regularised optimum is undefined; Kwa et al. show that policies can then gain arbitrarily large proxy reward with no true improvement even with the penalty in place [12].

It cannot remove what the reference does when the reward is blind to it. Because ratios among equal-reward responses are frozen, a reward that cannot distinguish a correct derivation from a lucky guess leaves the reference's share of lucky guesses intact at the optimum. That is a problem for the reward, the subject of our reward hacking article and our note on reward hacking with verifiable rewards, not one the regulariser can solve.

In practice#

The closed form and the estimators fit in a few lines of numpy:

import numpy as np

def tilt(pi_ref, r, beta):
    """Closed-form optimum of E[r] - beta * KL(pi || pi_ref) over a finite set of responses."""
    logits = np.log(pi_ref) + r / beta
    logits -= logits.max()                  # numerical stability
    p = np.exp(logits)
    return p / p.sum()

def kl_estimates(logp, logp_ref):
    """KL(pi || pi_ref) from samples drawn from pi, given their log-probs under both models."""
    lr = logp - logp_ref                    # log(pi / pi_ref) per sample
    k1 = lr
    k2 = 0.5 * lr**2
    k3 = np.expm1(-lr) + lr                 # (ratio - 1) - log(ratio), ratio = pi_ref / pi
    return k1.mean(), k2.mean(), k3.mean()

rng = np.random.default_rng(0)
pi_ref = rng.dirichlet(np.ones(40))
r = rng.uniform(0, 1, 40)
p = tilt(pi_ref, r, beta=0.1)
y = rng.choice(40, size=64, p=p)
print("exact KL:", (p * np.log(p / pi_ref)).sum())                     # 1.218
print("k1, k2, k3:", kl_estimates(np.log(p[y]), np.log(pi_ref[y])))   # 1.252, 1.161, 1.209

The same objective also appears outside language models. Offline RL methods that constrain a learned policy towards the behaviour policy use the same idea to stay where the data is (see offline reinforcement learning), and Jaques et al. used KL-control for exactly that reason when learning dialogue policies from fixed human data [15]. For the policy-gradient machinery underneath, start with the policy gradient theorem.

What the toys do not show#

The 40-response example has a fixed, known reward and a reference we invented; real reward models are noisy functions of text and real references assign probabilities to astronomically many responses, which is why the estimators matter. The controller run uses tabular REINFORCE, not PPO with clipping and value functions, so the size of the overshoot will differ in practice. Everything here is simulation; no language model was trained.

Frequently asked questions#

Sources#

1. Daniel M. Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B. Brown, Alec Radford, Dario Amodei, Paul Christiano, Geoffrey Irving. Fine-Tuning Language Models from Human Preferences, 2019. arxiv.org
2. Natasha Jaques, Shixiang Gu, Dzmitry Bahdanau, José Miguel Hernández-Lobato, Richard E. Turner, Douglas Eck. Sequence Tutor: Conservative Fine-Tuning of Sequence Generation Models with KL-control, 2016. arxiv.org
3. Tomasz Korbak, Ethan Perez, Christopher L. Buckley. RL with KL penalties is better viewed as Bayesian inference, 2022. arxiv.org
4. Sergey Levine. Reinforcement Learning and Control as Probabilistic Inference: Tutorial and Review, 2018. arxiv.org
5. Rafael Rafailov, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D. Manning, Chelsea Finn. Direct Preference Optimization: Your Language Model is Secretly a Reward Model, 2023. arxiv.org
6. Rafael Rafailov, Joey Hejna, Ryan Park, Chelsea Finn. From r to Q*: Your Language Model is Secretly a Q-Function, 2024. arxiv.org
7. John Schulman. Approximating KL Divergence, 2020. joschu.net
8. Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, et al.. DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models, 2024. arxiv.org
9. Nisan Stiennon, Long Ouyang, Jeff Wu, Daniel M. Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, Paul Christiano. Learning to summarize from human feedback, 2020. arxiv.org
10. Long Ouyang, Jeff Wu, Xu Jiang, et al.. Training language models to follow instructions with human feedback, 2022. arxiv.org
11. Leo Gao, John Schulman, Jacob Hilton. Scaling Laws for Reward Model Overoptimization, 2022. arxiv.org
12. Thomas Kwa, Drake Thomas, Adrià Garriga-Alonso. Catastrophic Goodhart: regularizing RLHF with KL divergence does not mitigate heavy-tailed reward misspecification, 2024. arxiv.org
13. Ahmad Beirami, Alekh Agarwal, Jonathan Berant, Alexander D'Amour, Jacob Eisenstein, Chirag Nagpal, Ananda Theertha Suresh. Theoretical guarantees on the best-of-n alignment policy, 2024. arxiv.org
14. Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, Sergey Levine. Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor, 2018. arxiv.org
15. Natasha Jaques, Asma Ghandeharioun, Judy Hanwen Shen, Craig Ferguson, Agata Lapedriza, Noah Jones, Shixiang Gu, Rosalind Picard. Way Off-Policy Batch Deep Reinforcement Learning of Implicit Human Preferences in Dialog, 2019. arxiv.org