KL Divergence Penalty › The KL-Constrained Objective

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:

$$\mathcal{J}(\theta) = \mathbb{E}_{x \sim \mathcal{D}}\left[\mathbb{E}_{y \sim \pi_\theta(\cdot|x)}[r_\phi(x, y)] - \beta \cdot D_{KL}(\pi_\theta(\cdot|x) \| \pi_{\text{ref}}(\cdot|x))\right]$$

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:

$$\mathcal{J}(\theta) = \mathbb{E}_{x, y}\left[r_\phi(x, y) - \beta \sum_{t=1}^{T} \log \frac{\pi_\theta(y_t | x, y_{<t})}{\pi_{\text{ref}}(y_t | x, y_{<t})}\right]$$

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:

$$\max_\theta \mathbb{E}[r_\phi(x, y)] \quad \text{subject to} \quad \mathbb{E}[D_{KL}(\pi_\theta \| \pi_{\text{ref}})] \leq \epsilon$$

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:

$$\mathcal{J}(\theta) = \mathbb{E}_{x, y}\left[\sum_{t=1}^{T} \gamma^t \cdot r_t\right]$$

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:

$$r_t = \begin{cases} -\beta \cdot \text{KL}_t & \text{if } t < T \\ r_\phi(x, y) - \beta \cdot \text{KL}_T & \text{if } t = T \end{cases}$$

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.