KL Divergence in RLHF: Why Models Are Kept Close to the Base

Contents

The unglamorous term that stops your model from chasing the score off a cliff

Train a language model to maximize a reward score, and it will happily discover that the score and the thing you actually want are not the same. It might repeat flattering phrases, pad answers with filler, or emit strange text that happens to fool the scorer.

RLHF pipelines guard against this with one small, unglamorous term: a KL penalty. It tells the model to improve, but not to wander far from where it started. This guide explains what KL divergence is, why RLHF needs it, how implementations compute it, and how to tune it.

Implementations differ, so check the original papers and your library's documentation before you build on any detail here.

The Problem: A Proxy Reward Can't Be Trusted Forever

RLHF does not optimize true human preference. It optimizes a reward model: a learned stand-in trained on a limited set of human comparisons. That stand-in works well near the data it saw and poorly far away from it, which is why the reward models guide treats its data coverage as a boundary, not a detail.

A strong optimizer pushes the policy toward high reward-model scores, including regions where the reward model is wrong. Quality at first improves, then stalls, then falls while the proxy score keeps climbing.

Researchers call this reward overoptimization, and studies of reward model scaling show the pattern clearly: as the policy moves further from the starting model, the proxy score keeps rising while true quality eventually peaks and declines. The mechanics of that failure are covered in the reward hacking guide.

So we need a leash. The KL penalty is that leash.

What KL Divergence Measures

KL divergence (Kullback-Leibler divergence) measures how different one probability distribution is from another. For language models, the two distributions are the next-token probabilities of your trained policy and a frozen reference model, usually the supervised fine-tuned (SFT) model you started from.

Three properties matter:

• It is zero only when the distributions match exactly. Any difference makes it positive.
• It is not symmetric. KL from A to B differs from KL from B to A, which matters for how the penalty behaves.
• It grows as the policy puts probability on outputs the reference considers unlikely.

In RLHF, the penalty uses the KL from policy to reference, measured on responses the policy generates:

KL = E over responses sampled from the policy of [ log policy(response) − log reference(response) ]

In plain terms: for the responses your model now writes, how much more likely does the policy find them than the reference does?

The RLHF Objective With KL

The full objective combines reward and a penalty:

maximize: E[ reward(prompt, response) ] − beta × KL(policy || reference)

The coefficient beta sets the trade-off:

• High beta: the policy stays close to the reference. Training is safe and gains are modest.
• Low beta: the policy chases reward aggressively. Gains come faster, along with a higher risk of reward hacking and degraded text.

This objective also has an elegant closed-form solution: the optimal policy equals the reference model's probabilities, reweighted by the exponential of the reward divided by beta. That result is exactly what DPO builds on, which is why the DPO objective has a beta too.

How It Works in Practice

You cannot compute the exact KL over all possible responses, so implementations estimate it from sampled tokens.

The Simple Estimator

For each generated token, compute the difference between the policy's and the reference's log-probabilities:

def kl_per_token(logp_policy, logp_ref):
    return logp_policy - logp_ref   # shape (T,)

Sampled this way, the average gives an unbiased KL estimate. But individual token values can be negative and noisy, even though true KL is never negative.

Folding KL Into the Reward

Classic PPO-based RLHF subtracts the penalty from the reward at every token, and adds the reward model's score at the final token:

def shaped_rewards(rm_score, logp_policy, logp_ref, beta):
    rewards = -beta * (logp_policy - logp_ref)   # per-token penalty
    rewards[-1] += rm_score                      # reward model score on last token
    return rewards

The Lower-Variance Estimator

John Schulman's well-known note on approximating KL describes an alternative that stays non-negative and has lower variance. With ratio = reference probability divided by policy probability, it computes ratio − 1 − log(ratio):

import torch

def kl_k3(logp_policy, logp_ref):
    log_ratio = logp_ref - logp_policy
    return torch.exp(log_ratio) - 1 - log_ratio

GRPO-style methods commonly add a KL term like this directly to the loss, instead of folding it into the reward — the placement change that the GRPO guide describes. Implementations differ, so check which estimator and placement your library uses.

Estimator	What it computes	Behaviour
Simple (logp_policy − logp_ref)	Per-token log-probability difference	Unbiased on average, noisy, individual tokens can go negative
k3 (ratio − 1 − log(ratio))	Ratio of reference to policy probability	Always non-negative, lower variance, the usual default
Exact	Full distribution divergence	Intractable for language models; the thing the estimators approximate

Choosing and Adapting Beta

beta has no universal right value. Teams set it by experiment, and many reduce the guesswork with an adaptive controller that adjusts beta to hold KL near a target. The approach comes from early RLHF work by Ziegler and colleagues:

import numpy as np

class AdaptiveKL:
    def __init__(self, beta, target, horizon=10_000):
        self.beta, self.target, self.horizon = beta, target, horizon

    def update(self, current_kl, n_steps):
        error = np.clip(current_kl / self.target - 1, -0.2, 0.2)
        self.beta *= 1 + error * n_steps / self.horizon

When the measured KL exceeds the target, beta rises and pulls the policy back. When KL sits below target, beta falls and gives the policy more freedom.

Rules of Thumb

• Start conservative. A larger beta costs you speed and protects you from disasters.
• Think in KL budgets. Decide how far from the reference you are willing to travel, then let beta enforce it.
• Compare methods on the reward-versus-KL frontier, not on reward alone. A method that reaches the same reward at lower KL is better.

What the KL Penalty Protects

The leash does several jobs at once:

• It limits reward hacking. The policy cannot wander into regions where the reward model is unreliable.
• It preserves fluency. The reference model writes coherent text, and staying close keeps the policy coherent.
• It preserves diversity. Without it, the policy collapses toward a few high-reward outputs.
• It retains general capabilities. Large drift can erase skills the reference had, a cost sometimes called the alignment tax.
• It stabilizes training. Smaller policy changes mean steadier optimization.

Treating KL as a mere regularizer undersells it. It defines the region where you trust your reward signal.

Figure 1: KL Divergence in RLHF: Why Models Are Kept Close to the Base — beta is the leash, and you adjust it while the run is still alive

Which Direction of KL?

Because KL is not symmetric, the direction matters. RLHF uses KL from the policy to the reference, measured on the policy's own samples. This is called the reverse KL, and it has a distinctive behaviour: it punishes the policy for putting mass where the reference has little, while it tolerates the policy dropping some behaviours the reference could produce.

In plain terms, reverse KL is mode-seeking. The policy may concentrate on a subset of the reference's behaviours, which partly explains why RLHF-trained models can lose diversity. Researchers study alternative divergences, and the standard choice remains the reverse KL because it is easy to estimate from the policy's own samples.

KL in Other Methods

• DPO: the KL constraint lives inside the loss through beta and the reference model's log-probabilities. There is no separate penalty to tune, and beta plays the same role.
• GRPO and RLOO: they add a KL term to the loss. Some recent recipes for long reasoning chains reduce the coefficient or set it to zero, arguing that the policy needs room to move and that verifiable rewards are harder to hack. Others keep a small value for stability. Check the recipe you follow.
• RLVR for reasoning: verifiers resist hacking better than learned reward models, so the leash matters less for reward hacking, though it still helps against drift and instability — the argument made in the RLVR walkthrough.

What to Monitor

Track KL during every run:

• Mean KL to the reference, per token and per sequence.
• KL against reward: plot reward versus KL, and watch for the point where quality stops improving.
• Response length: rising length often accompanies rising KL.
• Entropy: a steep drop signals collapse.
• Held-out quality: evaluate with an independent judge or human review as KL grows.

Warning Signs

• KL explodes: lower the learning rate, raise beta, or check your reference model.
• KL stays near zero: beta is too high, or the learning rate is too low, so the policy is not learning.
• Reward climbs while samples get worse: classic overoptimization. Raise beta, or stop earlier.

Common Implementation Bugs

• Wrong reference model. The reference must be the frozen starting model, not a moving copy of the policy.
• Mismatched sampling settings. Compute log-probabilities with the same temperature used to sample, or your ratios will be off.
• Unmasked padding. Padding tokens corrupt KL statistics.
• Double counting. Adding KL to both the reward and the loss penalizes drift twice.
• Misreading negative estimates. A single negative token-level KL value is estimator noise, not a negative divergence.
• Precision mismatches. Different numerical precision between rollout and training adds noise to the log-ratios.

Mistakes I See Teams Make

• Setting beta once and forgetting it. Watch KL, and adjust.
• Judging runs by reward alone. Compare reward at equal KL.
• Dropping the penalty without safeguards. Verify that your reward signal resists hacking first.
• Ignoring the reference model's quality. A weak reference caps what you can reach.
• Skipping held-out evaluation. A rising KL with falling quality means you have gone too far.

Your Quick-Start Checklist

1. Freeze a reference model from your SFT checkpoint.
2. Pick an estimator and placement (reward shaping or loss term), and use it consistently.
3. Start with a conservative beta, or an adaptive controller with a KL target.
4. Log KL, reward, length, and entropy every step.
5. Plot reward versus KL, and stop near the peak of held-out quality.
6. Mask padding, and match sampling settings between rollout and training.
7. Revisit beta whenever you change the reward model or the data.

Unlock AI That Actually Works

Get lifetime access to GPT-6 Astra, Claude Fable 5.1, Gemini 3.5, Grok 4.5, and more — all in one platform. Build websites, apps, videos, content, and digital products from a single command. No monthly fees. No tool-hopping.

Click here to get GPTAstra Max now — one-time payment, lifetime access.

Frequently Asked Questions

What is KL divergence in RLHF?

KL divergence, or Kullback-Leibler divergence, measures how different one probability distribution is from another. In RLHF the two distributions are the next-token probabilities of your trained policy and a frozen reference model, usually the supervised fine-tuned checkpoint you started from. The penalty in the RLHF objective is the KL from policy to reference, measured on responses the policy itself generates, written as the expectation over policy samples of log policy minus log reference. In plain terms it asks: for the responses this model now writes, how much more likely does the policy find them than the reference does? It is zero only when the distributions match exactly, it is never negative in expectation, and it is not symmetric, which is why the direction of the measurement matters.

Why does RLHF need a KL penalty at all?

Because the reward model is a proxy, not human preference itself. It is trained on a limited set of human comparisons, so it works well near the data it saw and badly far away from it. A strong optimizer pushes the policy toward high reward-model scores, including regions where the reward model is simply wrong. Quality improves at first, then stalls, then falls while the proxy score keeps climbing, which is reward overoptimization. The KL penalty is the leash that stops the policy wandering into territory where the reward signal cannot be trusted, and it also preserves fluency, output diversity and general capabilities that large drift tends to erase.

What beta should I use and how do I tune it?

There is no universal right value; teams set beta by experiment. A high beta keeps the policy close to the reference, which is safe but gives modest gains, while a low beta chases reward aggressively with faster gains and a higher risk of reward hacking and degraded text. The practical approach is to start conservative, decide how far from the reference you are willing to travel as a KL budget, and let beta enforce that budget. Many teams use an adaptive controller that adjusts beta to hold measured KL near a target, raising beta when KL runs above target and lowering it when KL sits below. Judge runs on the reward-versus-KL frontier rather than reward alone: a method that reaches the same reward at lower KL is the better method.

Which direction of KL does RLHF use, and why does it matter?

RLHF uses the reverse KL, that is the divergence from the policy to the reference, measured on the policy's own samples. Because KL is not symmetric, this direction behaves differently from the forward version: it punishes the policy for putting probability mass where the reference has little, while it tolerates the policy dropping behaviours the reference could produce. That makes reverse KL mode-seeking, so the policy can concentrate on a subset of the reference's behaviours, which is part of the explanation for the reduced output diversity of RLHF-trained models. It is also the standard choice for practical reasons: it is easy to estimate from samples the policy already generates.

How do implementations estimate and apply the KL penalty?

The exact KL over all possible responses is intractable, so implementations estimate it from sampled tokens. The simple estimator takes the per-token difference between policy and reference log-probabilities, which is unbiased on average but noisy and can be negative for individual tokens. Schulman's k3 estimator computes the reference-to-policy ratio and returns ratio minus one minus log ratio, which stays non-negative and has lower variance. Placement then varies: classic PPO-based RLHF folds the penalty into the reward at every token and adds the reward model score at the final token, while GRPO-style methods add a KL term directly to the loss. Bugs to avoid are using a moving policy as the reference, computing log-probabilities with different sampling settings than the rollout, leaving padding tokens unmasked, and penalizing drift twice by adding KL to both the reward and the loss.

Wrapping Up

The KL penalty is the leash that keeps RLHF from wandering into territory where its reward model cannot be trusted. It preserves fluency, diversity and general ability, and it turns how far the model should move into a number you can measure and control.

Set it conservatively, watch it throughout training, and judge your results by the reward you earn per unit of KL spent. A model that improves while staying close to its roots is a model you can still understand.

Related Articles

• Reward Models Explained: Teaching a Model What "Good" Looks Like
• DPO Explained: Preference Tuning Without Reinforcement Learning
• GRPO Explained: Group Relative Policy Optimization
• Reward Hacking in LLMs: Why Models Game Their Rewards
• RLVR: Reinforcement Learning with Verifiable Rewards Explained