KL Divergence Penalty in RLHF: Theory & Implementation

In the previous chapters, we built up the RLHF pipeline: collecting human preferences, training reward models, and applying PPO to optimize language model policies. But there is a necessary problem we have only briefly mentioned: reward hacking. Without constraints, a policy can discover bizarre outputs that score highly on the reward model while being clearly worse by any human measure. The KL divergence penalty is the mechanism that prevents this collapse, keeping the fine-tuned model anchored to the capabilities of its pre-trained foundation.

Think of the KL divergence penalty as the safety tether on a spacewalk. An astronaut tethered to the station can move freely in many directions, exploring the external surface and performing repairs, but cannot float away into open space. Remove the tether, and the same astronaut could drift indefinitely in any direction, including directions from which there is no return. The pre-trained language model is the space station: a large, carefully constructed structure representing enormous accumulated effort. The KL penalty is the tether that keeps the policy from drifting away from it while still letting real exploration of the surrounding space of behaviors.

This analogy also captures an important nuance: the length of the tether matters. A very short tether (large KL coefficient) keeps the astronaut close to the station but severely limits how far they can reach. A very long tether (small KL coefficient) allows wide exploration but risks losing contact entirely. The art of RLHF involves choosing and adapting the tether length appropriately as training progresses, tightening it when the policy shows signs of drifting dangerously and loosening it when progress stalls.

Why does reward hacking happen so readily without constraint? The reward model is a neural network trained on a finite dataset of human preferences. It approximates human judgment but does not perfectly capture it. Given enough freedom, a policy can always find inputs that the reward model scores highly but that humans would rate poorly. This is not a bug in the reward model; it is an inevitable consequence of learning from finite data. The policy is in effect performing adversarial search against the reward model, finding the holes in its coverage. The KL penalty limits how far the policy can search, keeping it in regions where the reward model is well-calibrated because those regions resemble the training data distribution.

The KL divergence penalty also serves a second, subtler purpose beyond preventing reward hacking. The pre-trained model is extraordinarily capable in ways that the reward model does not measure. It can write coherent multi-paragraph arguments, maintain consistent personas, handle edge cases gracefully, and produce text with appropriate register and tone for different audiences. None of these capabilities are directly rewarded during RLHF because they are hard to specify as preferences. By anchoring to the pre-trained model via KL divergence, we preserve these capabilities for free, without needing to explicitly specify them in the reward function. This makes the KL constraint a remarkably efficient mechanism: a single scalar coefficient does the work that would otherwise require carefully engineering dozens of auxiliary reward terms.

This chapter explores why KL divergence is the right tool for this job, how to compute it efficiently for autoregressive models, how to set and adapt the KL coefficient to balance learning against stability, and what happens when the constraint is too tight or too loose. We also examine the worked arithmetic of KL computation, the reward shaping interpretation that connects the penalty to classical RL theory, and the practical considerations that distinguish successful RLHF deployments from unstable ones.

Recall from the Reward Hacking chapter that reward models are imperfect approximations of human preferences. They are trained on a finite dataset and will inevitably assign high scores to some outputs that humans would reject. When we optimize a policy using PPO, we are searching over a large space of possible behaviors, and the optimization process will find these reward model errors if given enough freedom. This creates a tradeoff in RLHF: we must improve the model using the reward signal without trusting it unconditionally.

The basic insight is that the pre-trained language model is already a strong prior over sensible, grammatical, informative text. It was trained on trillions of tokens of human-written content, absorbing patterns of coherent discourse, factual knowledge, and linguistic conventions. By constraining our policy to stay close to this reference distribution, we inherit its linguistic competence while still allowing targeted improvements in alignment. Think of the reference model as a trusted expert whose judgment we respect: we want to make adjustments based on new feedback, but we do not want to stray so far that we lose the foundational competence that made the model useful in the first place.

What distinguishes KL divergence from other possible measures of distributional similarity? We could imagine using other metrics: total variation distance, Wasserstein distance, or simple L2 distance between parameter vectors. Each of these alternatives has shortcomings in the RLHF setting. Parameter-space constraints like L2 weight regularization penalize changes in weights, not changes in behavior. A model can change behavior sharply while keeping weights close together (via cancellations across layers) or keep behavior almost identical while changing weights materially (via symmetries and redundancies). We care about behavioral similarity, measured in terms of the outputs the model produces. Wasserstein and total variation distances operate in the right space (distributions over outputs), but they are expensive to compute for high-dimensional distributions and lack convenient decomposition properties for sequential models.

KL divergence measures how different one probability distribution is from another. For language models, this means measuring how much the policy's probability assignments over tokens differ from the reference model's. At each position in a sequence, both models assign probabilities to every token in the vocabulary. KL divergence captures the extent to which these probability distributions disagree, weighted by the probability mass the policy assigns. The key properties that make KL divergence ideal for this application are:

• Non-negativity: $D_{KL}(P \| Q) \geq 0$, with equality only when $P = Q$
• Asymmetry: it penalizes the policy for assigning probability to tokens the reference does not expect
• Decomposability: for autoregressive models, we can compute KL at each token position and sum

The non-negativity property ensures that the penalty term always works in one direction: it can only reduce the objective, never artificially inflate it. This mathematical guarantee means the optimization has a clear interpretation where reward pulls the policy toward better outputs while the KL penalty resists moving too far from the reference. The asymmetry property is particularly important for RLHF because it creates an asymmetric treatment of errors. If the policy assigns high probability to something the reference considers unlikely, the penalty is severe. But if the policy assigns low probability to something the reference considers likely, the penalty is more forgiving. This asymmetry encourages the policy to remain conservative, avoiding novel but potentially problematic outputs. Finally, decomposability allows us to compute and analyze the penalty at the level of individual tokens. This gives fine-grained insight into where and how the policy is diverging.

The decomposability property deserves special attention because it is what makes KL divergence tractable at the scale of language model training. The probability of a complete sequence is exponentially small; computing sequence-level distances directly would require summing over an astronomically large space. But because KL divergence over the full sequence decomposes into a sum of token-level KL divergences, we can compute it efficiently by running both models on the same generated sequence and comparing their log probabilities token by token. This makes the per-sequence KL computation as cheap as two forward passes, regardless of vocabulary size.

The mathematical definition of KL divergence is the basis for these methods. The formula captures a precise notion of distributional difference that translates directly into the computational procedures we use during training. Before writing down the formula, it helps to understand what we are measuring: we want to quantify how surprised the reference model would be by the choices the policy makes. Tokens that the policy prefers but the reference considers unlikely contribute heavily to this surprise. The total KL divergence aggregates this surprise across all possible tokens, weighted by how often the policy chooses them.

For two discrete probability distributions $P$ and $Q$ over the same set of outcomes, KL divergence is defined as:

where:

• $D_{KL}(P \| Q)$: the Kullback-Leibler divergence from $Q$ to $P$, representing the information lost when $Q$ is used to approximate $P$
• $P(x)$: the probability of outcome $x$ under the first distribution (typically the policy we are optimizing)
• $Q(x)$: the probability of outcome $x$ under the reference distribution (the baseline or prior)
• $\mathbb{E}_{x \sim P}$: the expectation calculated over samples drawn from $P$, meaning we average the log-ratio over the events that occur under distribution $P$

The first line presents the definition as a weighted sum over all possible outcomes, where each outcome's contribution is its probability under $P$ times the log ratio of probabilities. The second line rewrites this as an expectation, the form used in practice because we work with samples rather than complete distributions. This expectation-based form tells us something important: we only need to evaluate the log probability ratio for outcomes that occur under the policy, not for every possible outcome. This property is what makes KL divergence tractable for large vocabulary language models where the vocabulary can contain 50,000 to 100,000 tokens.

The key insight is that the expectation form turns a sum over all vocabulary tokens into an average over sampled tokens. During training, we generate responses from the policy, and for each token in each response, we look up its log probability under both the policy and the reference model. We never need to compute probabilities for tokens that were not generated. This is the basic efficiency that makes KL penalty practical in large-scale training.

In the RLHF context, $P$ is the policy distribution $\pi_\theta$ and $Q$ is the reference distribution $\pi_{\text{ref}}$ (typically the supervised fine-tuned model before RLHF). The policy is the distribution we are optimizing, and the reference is the distribution we want to stay close to. Every time we sample a token from the policy, we can compute its contribution to the KL divergence by looking at how much more or less likely that token was under the policy compared to the reference.

For autoregressive language models, the probability of a complete response $y = (y_1, y_2, \ldots, y_T)$ given prompt $x$ factorizes as:

where:

• $\pi(y|x)$: the probability of the complete sequence $y$ given prompt $x$, computed as the product of conditional probabilities
• $T$: the total length of the generated sequence (number of tokens)
• $y_t$: the token generated at time step $t$
• $y_{<t}$: the history of tokens generated before step $t$, which is the context for the next prediction

This factorization is the defining characteristic of autoregressive models. The probability of an entire sequence equals the product of probabilities for each token, where each token's probability depends on the prompt and all previously generated tokens. This chain structure means that computing the probability of a complete sequence requires $T$ forward passes through the model, one for each token position, though in practice we can compute all positions in parallel when the sequence is known by feeding the complete sequence as input and computing the logits for all positions simultaneously.

Taking the log and using the chain rule of KL divergence, the KL divergence between the policy and reference over complete responses can be written as:

where:

• $D_{KL}(\pi_\theta \| \pi_{\text{ref}})$: the KL divergence between the policy and reference distributions over complete sequences
• $\mathbb{E}_{y \sim \pi_\theta}$: the expectation calculated over response sequences $y$ sampled from the policy, since we estimate KL using the model's own generations
• $\sum_{t=1}^{T}$: the summation over all $T$ tokens, accumulating the divergence contribution from each step
• $\log \frac{\pi_\theta(\dots)}{\pi_{\text{ref}}(\dots)}$: the log probability ratio at each time step, which measures how much more (or less) likely a token is under the policy compared to the reference

This decomposition is important: we can compute KL as a sum of per-token log probability ratios, evaluated along the trajectory sampled from the policy. The log of a product becomes a sum of logs, which turns the sequence-level KL into a sum of token-level contributions. Each term in the sum measures how much the policy and reference disagree about the probability of a specific token given the context. When the policy assigns higher probability than the reference, the term is positive. This contributes to the divergence. When the policy assigns lower probability, the term is negative, reducing the divergence. The total KL divergence aggregates these token-level disagreements across the entire sequence.

Given a sampled response, the per-token KL contribution is simply:

where:

• $\text{KL}_t$: the contribution to the KL divergence at time step $t$ (positive if the policy favors the token more than the reference)
• $\pi_\theta(y_t | x, y_{<t})$: the probability assigned to the token $y_t$ by the policy model given the context
• $\pi_{\text{ref}}(y_t | x, y_{<t})$: the probability assigned to the token $y_t$ by the reference model given the context

This formula shows the simplicity of per-token KL computation. For any token in a generated sequence, we simply subtract the reference model's log probability from the policy model's log probability. No summation over the vocabulary is required because we are only interested in the sampled token. This is a direct consequence of the expectation-based formulation: since we are averaging over samples from the policy, we only need to evaluate the log ratio at the sampled outcomes.

The sign of each term is diagnostic. A positive $\text{KL}_t$ means the policy assigned more probability to this token than the reference did: the policy is more confident about or more inclined toward this choice than the pre-trained model. A negative $\text{KL}_t$ means the policy assigned less probability: the pre-trained model would have made this choice more readily than the policy would. By examining the sign and magnitude of per-token KL contributions across a response, we can identify exactly which word choices stand for the policy's learned departures from the reference distribution.

Summing these over all tokens in the response gives:

where:

• $\text{KL}(y)$: the approximate KL divergence for a single sampled response $y$, calculated as the sum of per-token divergences
• $T$: the length of the response in tokens, over which the divergence accumulates

This is an unbiased estimator of the KL divergence under the policy distribution, since we are sampling trajectories from $\pi_\theta$ itself. The unbiasedness property means that if we average this quantity over many sampled responses, we converge to the true KL divergence between the policy and reference distributions. In practice, we compute this estimator for each response in a training batch and use the batch average as our estimate of the expected KL divergence.

An important consequence of this formula is that KL divergence scales with sequence length. A response of 100 tokens accumulates up to 100 per-token contributions. This means that long responses are penalized more heavily in absolute terms than short responses, even if the per-token divergence is the same. Some practitioners address this by normalizing the total KL by sequence length, computing average per-token KL rather than total KL. Others prefer total KL because it better captures the overall behavioral departure from the reference. The choice has implications for how the penalty interacts with response length: normalizing encourages the policy to generate any length freely, while using total KL creates a cost to longer responses that diverge from the reference.

To make the mathematics concrete, let us work through a numerical example with a tiny vocabulary. Suppose we are generating the single word following the prompt "The sky is". Our vocabulary contains four words: "blue", "clear", "dark", and "beautiful". The reference model assigns these probabilities:

After RLHF training on prompts related to describing pleasant weather, our policy has shifted toward more positive words:

We can verify both distributions sum to 1.0, as required. Now suppose the policy samples the token "beautiful" for this position (which is likely given the policy's 35% probability). The per-token KL contribution for this sampled token is:

This positive value tells us that the policy assigned more than three times as much probability to "beautiful" as the reference model did, and this preference contributes 1.25 nats to the divergence penalty. If instead the policy had sampled "dark", the contribution would be:

This negative value shows the policy being less likely than the reference to produce "dark". Note something important: the total KL divergence of the full distribution (summed over all tokens, weighted by $P$) is always non-negative. But the contribution from a single sampled token can be negative. The non-negativity guarantee applies to the expectation, not to individual samples.

The full distribution KL divergence, summed over all four tokens, is:

Computing each term:

This 0.285 nats is a moderate divergence. The policy has shifted toward "beautiful" at the expense of "blue" and "dark", which the KL divergence correctly captures. If we had a KL coefficient $\beta = 0.1$, the penalty contribution from this one-token sequence would be $0.1 \times 0.2851 = 0.0285$, subtracted from whatever reward the policy earned.

Now extend this thinking to a full response of 100 tokens. If each token position contributes an average KL of about 0.3 nats (similar to our example), the total KL for the response would be around 30 nats. With $\beta = 0.1$, that produces a penalty of 3.0 reward units. This illustrates why the KL coefficient must be chosen in relation to the typical scale of the reward model scores. If rewards range from -1 to 1 and the KL penalty is 3.0, the penalty dominates and the policy cannot learn anything useful from the reward signal.

The key insight from this worked example is that KL divergence accumulates multiplicatively: each position in the sequence contributes independently based on the log probability ratio at that position. Positions where the policy makes the most different choices from the reference contribute the most. And because the contribution is proportional to $P(x) \log(P(x)/Q(x))$, tokens that the policy assigns high probability to matter more than tokens the policy considers unlikely.

Having established how to compute KL divergence, we now examine how it enters the RLHF optimization objective. The KL term acts as a regularizer, pulling against the reward signal to prevent the policy from straying too far from its trusted starting point. The combined objective balances generating responses that humans find helpful against remaining close enough to the pre-trained distribution to preserve its broad capabilities.

The RLHF objective with KL penalty is:

where:

• $\mathcal{J}(\theta)$: the objective function to be maximized, balancing reward maximization against drift from the reference
• $\mathcal{D}$: the dataset of prompts used for training
• $r_\phi(x, y)$: the learned reward model score for prompt $x$ and response $y$, giving the guidance signal
• $\beta$: the KL coefficient controlling the strength of the penalty (higher values force the policy closer to the reference)
• $\pi_{\text{ref}}$: the frozen reference model distribution, serving as the anchor

The structure of this objective captures the core tradeoff in RLHF. The first term rewards the policy for generating responses that score highly according to the reward model. The second term penalizes the policy for diverging from the reference. The coefficient $\beta$ determines the relative importance of these two competing forces. When $\beta$ is large, the penalty dominates and the policy barely moves from the reference. When $\beta$ is small, the reward signal dominates and the policy can move freely toward high-reward regions, potentially exploiting imperfections in the reward model.

This can be rewritten using the per-token KL decomposition:

where:

• $\mathbb{E}_{x, y}$: the expectation over prompts and sampled responses
• $\sum_{t=1}^{T}$: the sum of log probability ratios over the sequence, representing the total divergence for that trajectory

This formulation makes the objective directly computable from model outputs. For each sampled response, we obtain a reward from the reward model and compute the sum of log probability ratios from the policy and reference models. The combination of these quantities, averaged over the batch, gives us the objective value that we then maximize through gradient ascent.

Notice that this objective has a clean interpretation: it asks for the policy that best satisfies two objectives simultaneously. The first objective is to generate responses that humans find useful, as captured by the reward model. The second objective is to remain recognizably similar to the original pre-trained model, as measured by distributional distance. The $\beta$ parameter specifies the exchange rate between these two objectives: how much reward improvement we are willing to trade for one nat of additional divergence.

The KL-penalized objective is equivalent to the Lagrangian relaxation of a constrained optimization problem:

where:

• $\max_\theta$: the maximization over policy parameters $\theta$
• $\mathbb{E}$: the expectation over the data distribution
• $\epsilon$: the maximum allowable divergence (the "budget"), limiting how far the policy can drift

In this constrained view, we are asking for the policy that maximizes expected reward subject to a hard constraint on how much it can diverge from the reference. The constraint $\epsilon$ specifies a "KL budget" that the policy must respect. The penalized objective arises when we convert this constrained problem to an unconstrained one using Lagrangian relaxation.

The coefficient $\beta$ acts as a Lagrange multiplier. Larger $\beta$ corresponds to a tighter constraint (smaller $\epsilon$), forcing the policy to stay closer to the reference. The objective seeks the highest-reward policy among those that remain within a KL "budget" of the reference. The correspondence between $\beta$ and $\epsilon$ is not one-to-one in practice because the constraint is enforced softly rather than exactly, but the intuition remains useful. Choosing $\beta$ is analogous to choosing how much divergence we are willing to tolerate in exchange for reward improvement.

Think of the constrained formulation as having a budget for a home renovation project. You have a fixed budget (the KL constraint $\epsilon$), and you want to maximize the improvements you can make within that budget. The unconstrained version with a penalty $\beta$ is like having a project that costs extra for every dollar spent but where there is no hard limit. The two framings lead to equivalent decisions for the right choice of $\beta$, but the constrained view is more interpretable: you directly specify how much divergence you are comfortable with, and the algorithm finds the best policy within that budget.

Another way to view the KL term is as a per-token bonus added to the reward. This interpretation connects the KL penalty to the classical reinforcement learning concept of reward shaping, where the reward function is augmented with additional terms to guide learning.

We can rewrite the objective as:

where:

• $\gamma$: the discount factor (typically 1 in this context, treating all tokens as equally important)
• $r_t$: the shaped reward at time step $t$, incorporating both the environment reward and the KL penalty

The shaped reward at each step is defined as:

where:

• $\text{KL}_t$: the per-token KL penalty, acting as an instantaneous cost for deviating from the reference
• $r_\phi(x, y)$: the final reward from the reward model, typically given only at the end of the sequence ($t=T$)

This interpretation shows that the policy receives negative reward proportional to its deviation from the reference at each token, plus the final reward model score at the end of generation. The per-token penalties encourage staying on-distribution throughout the generation, not just at the final output. Every token choice that differs from what the reference would have chosen incurs a cost, creating pressure to remain conservative at every step of generation.

This reward shaping perspective has practical implications for implementation. In PPO, we need to assign credit to individual actions (tokens) for the overall outcome. The shaped reward formulation gives a natural way to do this: each token receives the KL penalty immediately, while the reward model score is attributed to the final token. This decomposition helps the value function and advantage estimation work effectively, since part of the reward signal is available at every time step rather than only at the end. Without the per-token KL shaping, the value function would receive no signal until the end of the sequence, which makes it very hard to learn credit assignment for long responses.

The reward shaping interpretation also clarifies why the KL penalty discourages reward hacking more effectively than alternative regularization methods. A weight decay penalty applied to model parameters pushes the policy back toward zero weights, which has no clear relationship to the reference distribution. L2 distance between parameters penalizes large changes uniformly, regardless of whether those changes affect behavior. The KL penalty, by contrast, directly penalizes behavioral deviations: if the policy outputs the same distribution as the reference at every token position, the penalty is exactly zero, regardless of how much the parameters have changed. This behavioral grounding makes the KL penalty much more targeted and interpretable than parameter-space regularizers.

The KL coefficient $\beta$ is perhaps the most important hyperparameter in RLHF. It controls the basic tradeoff between two competing objectives that cannot be simultaneously maximized: learning from the reward signal as aggressively as possible, and staying close to the pre-trained distribution that is accumulated linguistic competence. There is no universally correct value of $\beta$; the right choice depends on the reward model quality, the magnitude of the desired behavior change, the training data distribution, and the response length distribution.

Think of $\beta$ as the stiffness of a spring connecting the policy to the reference model. A very stiff spring (high $\beta$) means the policy can only deviate a little from the reference before the restoring force becomes overwhelming. It will converge quickly because there is not much room to move, but it may never reach the optimal alignment behavior if that behavior requires substantial departure from the pre-trained distribution. A very weak spring (low $\beta$) allows the policy to explore widely but also means it can drift into regions of the space that are far from well-calibrated reward model territory, leading to reward hacking. The art of RLHF involves finding the spring stiffness that allows real movement toward better alignment without letting pathological exploitation of reward model errors.

The effects of the coefficient choice are asymmetric in their failure modes. Setting $\beta$ too low produces obvious failures: outputs that are repetitive, formulaic, or bizarre in ways that humans immediately notice. These failures are easy to detect through human evaluation and relatively easy to fix by increasing $\beta$. Setting $\beta$ too high produces subtle failures: the model barely moves from the reference, and improvements in alignment metrics are minimal or undetectable. This failure mode is harder to diagnose because the model still looks reasonable; it just has not learned much from the reward signal.

The key insight is that $\beta$ selection is not an abstract mathematical choice but an empirical one tied to the specific properties of your training setup. The optimal $\beta$ for a 7B parameter model may differ from that for a 70B parameter model; the optimal $\beta$ for a helpfulness-focused reward model may differ from that for a safety-focused one; the optimal $\beta$ for short-form responses may differ from that for long-form generation.

Understanding the qualitative effects of different coefficient ranges helps calibrate initial choices:

When $\beta$ is too low (below 0.01):

• The policy optimizes almost purely for reward, ignoring the reference constraint
• Susceptible to reward hacking and mode collapse
• May produce repetitive, formulaic outputs that exploit reward model quirks
• Loses diversity and coherence over time as the policy converges to a narrow mode of high-reward outputs

When $\beta$ is too high (above 1.0):

• Learning becomes extremely slow because the penalty overwhelms the reward signal
• The policy barely moves from the reference, even when the reward model gives clear guidance
• May never reach real alignment improvements
• Effectively wastes computational resources on training that produces minimal behavioral change

Within the typical range ($\beta \in [0.01, 0.2]$):

• Most successful RLHF implementations use values in this range
• InstructGPT used an initial value around 0.02
• Anthropic's Constitutional AI work used values around 0.001 to 0.01
• The optimal value depends on reward model quality and desired behavior change

When selecting a KL coefficient, consider these factors:

• Reward model confidence: If your reward model was trained on limited data or shows signs of miscalibration, use a higher $\beta$ to limit exploitation of its errors. A well-calibrated reward model trained on millions of preference comparisons can tolerate lower $\beta$ because there are fewer holes to exploit.

• Magnitude of desired behavior change: For small adjustments (such as reducing a specific narrow category of harmful outputs), lower $\beta$ suffices. For significant capability changes that require the model to learn substantially new behaviors, you may need to accept higher KL divergence and should verify that training produces behavioral changes rather than just optimizing the KL-penalized objective.

• Response length: Longer responses accumulate more KL. If you are optimizing for verbose outputs, the effective per-token penalty is spread over more tokens, so you might need to account for this in choosing $\beta$ or to normalize by length.

• Training stability: If you observe reward increasing while output quality degrades subjectively (which you can detect via human evaluation or automatic metrics), increase $\beta$ to strengthen the constraint. If training seems stable but the model is not improving, consider decreasing $\beta$ or increasing the reward model signal.

Rather than fixing $\beta$ throughout training, adaptive KL methods adjust the coefficient to maintain a target KL budget. This approach, introduced in the InstructGPT paper, gives more consistent training dynamics across different stages of optimization. The core insight is that the appropriate constraint strength changes as training progresses, and a fixed coefficient cannot account for these changing conditions.

Consider what happens during a typical RLHF training run. In the very early stages, when the policy is nearly identical to the reference, even small gradient updates produce modest KL divergence. As training progresses and the policy adapts, it may start making more consistent departures from the reference distribution, accumulating larger KL with each batch. Late in training, the policy may have already diverged materially, and further updates produce only marginal additional KL. A fixed $\beta$ that was appropriate in the early stages may be far too weak to prevent reward hacking in the middle stages, or far too strong to allow real learning in the late stages.

The adaptive approach treats KL control as a feedback control problem. We specify a desired operating point (the target KL) and adjust $\beta$ to steer the system toward it. When the policy drifts too far from the reference, we tighten the constraint; when the policy stays too close, we loosen it. This automatic adjustment eliminates the need to manually tune $\beta$ schedules and gives more reliable training across different model sizes, reward models, and training data distributions.

The idea is to specify a target KL divergence $D_{\text{target}}$ that we want to maintain on average. If the current KL exceeds this target, we increase $\beta$ to pull the policy back. If KL is below target, we decrease $\beta$ to allow more exploration. This creates a feedback loop that stabilizes training by keeping the policy within a controlled distance from the reference.

The adaptation rule used in many implementations is:

where:

• $\beta_{t+1}$: the updated KL coefficient for the next step
• $\beta_t$: the current KL coefficient
• $D_{KL}$: the observed average KL divergence in the current batch (the signal used for feedback)
• $D_{\text{target}}$: the desired target KL divergence (the setpoint for the controller)

This multiplicative update creates a dead zone around the target where $\beta$ remains stable, preventing oscillation while still responding to significant deviations. The factor of 1.5 defines the boundaries of this dead zone: as long as the observed KL stays between $D_{\text{target}}/1.5$ and $1.5 \cdot D_{\text{target}}$, the coefficient remains unchanged. Only when KL drifts outside this range does the controller intervene with a multiplicative adjustment.

The choice of multiplicative rather than additive updates is deliberate and important. If $\beta$ is small (say, 0.01) and needs to increase substantially to constrain a rapidly diverging policy, an additive step of 0.1 would still leave $\beta$ much too small. A multiplicative doubling, on the other hand, always makes a proportionally significant change regardless of the current value. This scale-invariant behavior means the controller responds appropriately whether $\beta$ is 0.001 or 1.0.

Adaptive KL addresses a basic challenge: the appropriate constraint strength changes during training. Early in RLHF, when the policy is close to the reference, small rewards can produce large gradient updates that rapidly increase KL. Later, when the policy has already diverged somewhat, the same learning rate produces smaller relative changes. A fixed $\beta$ cannot account for these different regimes.

The dead-zone design deserves special attention as it prevents a failure mode called "controller chatter." If we updated $\beta$ every time the observed KL differed from the target by any amount, the coefficient would oscillate rapidly because the observed KL is noisy (it depends on the specific batch of prompts and sampled responses). By only updating when the deviation exceeds 50% of the target in either direction, we allow transient fluctuations without responding to them, while still catching real systematic deviations.

By targeting a consistent KL budget, adaptive control:

• Prevents early-training instability from sudden KL spikes by tightening constraints before the policy escapes the reference distribution
• Maintains learning signal late in training when fixed $\beta$ might over-constrain the policy by letting $\beta$ to decrease as needed
• Provides a more interpretable hyperparameter: target KL in nats rather than an abstract penalty coefficient

The interpretability benefit is substantial. When using a fixed $\beta$, it is difficult to know in advance what KL divergence will result. Different prompts, response lengths, and training stages all affect the relationship between $\beta$ and actual KL. With adaptive control, you directly specify the divergence budget you are comfortable with, and the algorithm finds the appropriate $\beta$ to reach it. A practitioner can say "I want the policy to diverge no more than 6 nats on average from the reference" and the controller will figure out how to enforce that, regardless of other training hyperparameters.

Some implementations use smoother adaptation:

where:

• $\alpha$: a small step size parameter (for example, 0.1), controlling how aggressively $\beta$ changes per step
• $\text{sign}(\cdot)$: the sign function (returns +1 if the term is positive, -1 if negative), determining the direction of the update

This approach prevents the jarring factor-of-2 jumps while still steering toward the target. The updates are smaller and more frequent, creating smoother $\beta$ trajectories. However, this also means slower response to large deviations, which can be problematic if KL suddenly spikes due to a particularly influential batch. The right choice between the InstructGPT-style dead-zone controller and the smooth controller depends on how much noise exists in the observed KL estimates.

Others use proportional control:

where:

• $\exp(\cdot)$: the exponential function, so the multiplier is always positive and $\beta$ never becomes negative
• $\frac{D_{KL} - D_{\text{target}}}{D_{\text{target}}}$: the relative error from the target, scaling the update based on the magnitude of the deviation

This scales the adjustment magnitude by how far off-target the current KL is. Small deviations produce small adjustments, while large deviations produce large adjustments. The exponential ensures that $\beta$ remains positive regardless of the error magnitude. This proportional approach offers a middle ground between the dead-zone method and the constant-step method, responding proportionally to the severity of the deviation. It also has the desirable property that $\beta$ can never become negative (due to the exponential), which would flip the penalty to a bonus and completely undermine the regularization.

Let us implement KL divergence computation and adaptive control for RLHF training. We will build components that integrate with the PPO training loop from the previous chapter. The implementation focuses on clarity and correctness. This gives the building blocks that can be optimized for production use. We proceed from the simplest building block (per-token KL) to more complete integration (the full RLHF objective) and then to the adaptive controller that ties everything together.

The core computation extracts log probabilities from both the policy and reference model, then computes their difference. This function forms the foundation of the KL penalty calculation, taking log probabilities that have already been extracted from model outputs and creating both per-token and per-sequence KL values.

This computation is straightforward because we are using the sampled trajectory estimator. The log probability ratio for generated tokens equals the KL divergence in expectation. The attention mask handles variable-length sequences by zeroing out contributions from padding tokens, so the KL computation considers only content tokens. This masking is needed when processing batches of sequences with different lengths, as padding tokens should not contribute to the divergence measure. Without masking, sequences with more padding would appear to have lower divergence simply because the padding region contributes zero per-token difference, artificially inflating how similar a batch of short sequences appears.

To compute KL, we need log probabilities for the sampled tokens. Here is how to extract them from model logits:

The function first applies log_softmax to convert raw logits into log probabilities. This operation normalizes the logits so they stand for a valid probability distribution at each position. The gather operation then selects the log probability corresponding to each actual token from the full vocabulary distribution. This is much more efficient than computing probabilities for all tokens when we only need one per position. The unsqueeze and squeeze operations handle the dimension manipulation required by PyTorch's gather function.

The use of log_softmax rather than computing softmax and then taking the log is numerically important. Computing softmax first can lead to underflow for tokens with very negative logits, as the exponential can produce values below the floating point minimum. The combined log_softmax operation is numerically stable because it subtracts the maximum logit before computing the exponential, keeping all values in a reasonable range. This numerical stability is especially important for large vocabulary models where the distribution of logits can be highly skewed.

The adaptive controller maintains the KL coefficient and adjusts it based on observed KL values. It implements the target KL approach described in the mathematical section. This gives a stateful object that can be updated after each training batch.

The controller doubles $\beta$ when KL exceeds the target, pulling the policy back toward the reference, and halves it when KL drops too low, letting more exploration. The simulation shows how the controller responds to different training phases: increasing the coefficient during the high-KL early phase, maintaining stability during the near-target phase, decreasing during the low-KL phase, and stabilizing again as KL returns to the target range.

The shaded region marks the "stable zone" where $\beta$ does not change. Outside this zone, the controller applies multiplicative updates to steer KL back toward target. The log scale on the right panel makes it easy to see the multiplicative nature of the updates: each adjustment doubles or halves the coefficient, which appears as equal-height steps on the log scale.

Here is how KL computation integrates into a simplified RLHF training step. This function combines all the pieces we have developed, taking model outputs and creating the penalized objective along with diagnostic information.

With the same underlying KL divergence, a higher $\beta$ produces a much larger penalty, materially reducing the effective reward signal available for optimization. The comparison illustrates the large impact of the coefficient choice: with low $\beta$, most of the reward signal passes through to guide learning, while with high $\beta$, the penalty dominates and the effective reward becomes negative despite positive raw reward scores.

Understanding the key parameters helps practitioners set up RLHF training correctly and diagnose problems when they occur. These parameters are conceptual knobs that control the basic behavior of the alignment process, not incidental implementation details.

The parameters for the KL divergence implementation are:

• kl_coef ($\beta$): The weight of the KL penalty term. Controls the tradeoff between reward maximization and reference adherence. Higher values enforce stricter proximity to the reference at the cost of slower learning. Starting values typically range from 0.01 to 0.2 depending on the application.

• target_kl: The desired KL divergence value (in nats) for adaptive controllers. This is the setpoint for the feedback loop. A target of 6 nats is common in practice, corresponding to a moderate level of divergence that allows real behavioral change while preserving most of the reference model's capabilities. The right target depends on response length: a target of 6 nats spread over 100 tokens means an average of 0.06 nats per token, which is relatively small distributional shifts at each individual position.

• horizon: In the smooth adaptive controller, this determines the timescale over which adaptation occurs. Longer horizons make the controller less reactive to transient fluctuations but also slower to respond to real regime changes during training.

The KL coefficient fundamentally shapes how RLHF training unfolds. Different values of $\beta$ do not just change the speed of learning; they change the trajectory through the space of possible policies, the types of outputs the policy learns to produce, and ultimately the nature of the alignment achieved. Let us examine these dynamics through simulation.

During training, there is a Pareto frontier between reward maximization and KL minimization. Different $\beta$ values trace different paths along this frontier, converging to different equilibrium points. Understanding this frontier is needed for reasoning about what RLHF can and cannot reach: no policy can simultaneously maximize reward and minimize KL divergence. Every step toward higher reward involves some cost in KL, and every step toward lower KL involves some sacrifice in reward.

The visualization illustrates the basic tradeoff. Low-$\beta$ policies (warm colors) quickly accumulate KL divergence and can reach high rewards, but risk entering regions where the reward model is unreliable. High-$\beta$ policies (cool colors) progress slowly but stay in well-calibrated regions of the reward model's coverage. The diagonal trajectory of intermediate $\beta$ values traces out the Pareto frontier between reward and divergence, with different positions on the frontier preferred depending on how much we trust the reward model and how large a behavioral change we are seeking.

KL divergence is not uniform across tokens. Some positions contribute much more than others, and the pattern of contributions reveals what kind of adaptation the policy is making. A policy that consistently has high KL at the beginning of responses is learning a different way to introduce topics. A policy with high KL at specific content words is learning to prefer different vocabulary. A policy with uniformly moderate KL throughout is making subtle but consistent shifts across all aspects of generation.

Understanding where KL accumulates helps diagnose what the policy is learning. If KL concentrates at specific positions, the policy is making targeted modifications. If KL is high throughout, the policy is developing a materially different generation style. If KL is concentrated at functionally useful words (verbs, key nouns, evaluation words), the policy is learning different semantic content, not just different phrasing. This per-token analysis is one of the most useful diagnostic tools in RLHF, giving interpretable insight into what the alignment training teaches the model.

Let us compare training stability between fixed and adaptive KL control to see why practitioners increasingly prefer the adaptive approach for production-scale RLHF:

The adaptive controller maintains KL divergence near the target by increasing $\beta$ when the policy drifts too far and relaxing it when the policy remains close to the reference. The fixed controller reaches a higher raw learned reward because it permits more drift, but its advantage disappears once both trajectories are evaluated with the same $\beta=0.1$ KL cost. This shared evaluation makes the stability tradeoff visible without conflating it with the adaptive controller's changing coefficient: degradation in the penalized reward despite an improving raw reward indicates that policy drift is overwhelming the useful learning signal.

While KL divergence is the standard constraint for RLHF, it has several real limitations worth understanding in depth. These limitations do not invalidate the approach; they are inherent to the problem of aligning language models with human preferences, and understanding them helps practitioners make better decisions about when and how to apply the KL penalty.

The reference model quality bottleneck. The KL penalty anchors the policy to a specific reference, typically the SFT model. This makes the quality of supervised fine-tuning important for RLHF success. If the SFT model has problematic behaviors, the KL penalty resists correcting them because moving away from those behaviors necessarily increases KL. For example, if the SFT model has learned to be verbose (perhaps because the supervised data over-represented detailed responses), the KL penalty will resist training the policy to be more concise, even if the reward model rewards conciseness. Conversely, if the SFT model is highly capable and well-calibrated, the constraint preserves that capability throughout RLHF. This bidirectional constraint is a feature when the reference is good and a bug when it is not, which is why so much effort goes into creating high-quality SFT data.

Distribution distance versus behavioral distance. KL divergence measures similarity between probability distributions, not between actual behaviors. Two policies could have small KL divergence but generate quite different outputs when sampled, especially for long sequences where small per-token differences compound across many conditional distributions. Conversely, policies with large KL might produce similar outputs if the probability differences are concentrated in low-probability regions that rarely get sampled. This mismatch between KL as a mathematical quantity and our actual goal (behavioral similarity) means that the KL penalty sometimes constrains the wrong thing. A policy that has learned to avoid a specific harmful phrase but otherwise generates identically to the reference might show higher KL than a policy that has not changed its harmful behavior at all but has shifted probability mass in obscure parts of the distribution.

Asymmetry creates specific biases. The KL divergence $D_{KL}(\pi_\theta \| \pi_{\text{ref}})$ penalizes the policy more heavily for assigning low probability to tokens the reference model assigns high probability to. This means the policy is discouraged from becoming more confident (more peaked distribution) than the reference, which can be desirable (maintaining diversity) or undesirable (preventing confident correct answers in factual domains). The asymmetry also means that the reverse KL $D_{KL}(\pi_{\text{ref}} \| \pi_\theta)$ would have different properties and might better suit different scenarios. Research on choosing the right divergence direction for alignment remains active.

Computational overhead. Computing KL requires running both the policy and reference model on every training example. For large models, this roughly doubles the forward pass compute. In practice, this means RLHF training is approximately twice as expensive in compute as supervised fine-tuning on the same data, before accounting for the added complexity of running the reward model. Some implementations amortize this by caching reference model log probabilities for a fixed set of prompts, but this requires careful management of the cache and may miss the per-response nature of the KL computation. Memory requirements are also approximately doubled, since both models must reside in memory simultaneously during training.

KL divergence cannot constrain all capabilities simultaneously. The penalty keeps the overall distribution of the policy close to the reference, but it cannot prevent the policy from changing its behavior on specific important subsets of inputs. A policy could decrease KL on common inputs (where it is easy to stay close to the reference) while sharply changing behavior on rare but important inputs (like edge case harmful content). This "gaming" of the aggregate KL constraint is a subtle form of reward hacking that the KL penalty does not fully prevent. Supplementary evaluations on diverse test sets are needed to detect this kind of distributional shift.

Looking ahead, methods like Direct Preference Optimization (DPO), covered in the next chapter, incorporate the KL constraint implicitly in their objective. This eliminates the need for explicit KL computation during training while achieving similar regularization effects. DPO reparameterizes the RLHF objective in terms of policy log ratios, deriving a loss function that can be minimized directly on preference data without the need for a separate reward model or explicit KL penalty term. This is an elegant alternative to the explicit penalty approach we have examined here, trading the flexibility of separate reward modeling for the simplicity and stability of a closed-form training objective.

The KL divergence penalty is needed for stable RLHF training. It prevents reward hacking by keeping the policy close to a trusted reference model, preserving the language capabilities acquired during pre-training while letting targeted improvements in alignment. Without it, RLHF optimization rapidly degenerates into adversarial exploitation of reward model errors, creating outputs that are high-scoring but meaningless or harmful.

Key takeaways from this chapter:

• KL divergence measures distribution difference: For autoregressive models, it decomposes into a sum of per-token log probability ratios, allowing efficient computation along sampled trajectories. The expectation-based form means we only need to evaluate the ratio for sampled tokens, not for all tokens in the vocabulary.

• The coefficient $\beta$ controls the reward-constraint tradeoff: Low $\beta$ allows rapid learning but risks exploitation of reward model errors. High $\beta$ maintains stability but slows improvement. Typical values range from 0.01 to 0.2, with the right value depending on reward model quality, desired behavior change magnitude, and training stability.

• Adaptive KL control maintains consistent constraints: By adjusting $\beta$ to target a specific KL budget, adaptive methods give stable training dynamics throughout optimization, preventing both early-training instability and late-training stagnation. The InstructGPT-style multiplicative dead-zone controller is a practical and reliable choice.

• Per-token KL reveals what the policy learns: Analyzing where KL accumulates in generated sequences helps diagnose whether the policy is making targeted modifications or developing a fundamentally different generation style. High KL at the beginning of responses suggests changed discourse opening strategies; high KL at specific content words suggests vocabulary shifts; uniformly distributed KL suggests broad stylistic adaptation.

• The reference model quality determines the ceiling: The KL constraint anchors the policy to the SFT model, making SFT quality a necessary determinant of RLHF outcomes. Strong SFT models let more effective alignment; weak SFT models transfer their limitations to the aligned model.

The KL penalty is one approach to constraining policy optimization. As we will see in the following chapters on DPO, there are alternative formulations that reach similar goals through different mechanisms. It offers tradeoffs in implementation complexity, computational efficiency, and training stability. Understanding the KL penalty deeply, as we have done here, gives the basis for understanding why these alternatives work and when to prefer them.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about KL divergence penalties in RLHF.