KL Divergence Penalty › Mathematical Foundation

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:

$$\begin{aligned} D_{KL}(P \| Q) &= \sum_x P(x) \log \frac{P(x)}{Q(x)} \\ &= \mathbb{E}_{x \sim P}\left[\log P(x) - \log Q(x)\right] \end{aligned}$$

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:

$$\pi(y|x) = \prod_{t=1}^{T} \pi(y_t | x, y_{<t})$$

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:

$$D_{KL}(\pi_\theta(\cdot|x) \| \pi_{\text{ref}}(\cdot|x)) = \mathbb{E}_{y \sim \pi_\theta(\cdot|x)}\left[\sum_{t=1}^{T} \log \frac{\pi_\theta(y_t | x, y_{<t})}{\pi_{\text{ref}}(y_t | x, y_{<t})}\right]$$

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:

$$\text{KL}_t = \log \pi_\theta(y_t | x, y_{<t}) - \log \pi_{\text{ref}}(y_t | x, y_{<t})$$

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:

$$\begin{aligned} \text{KL}(y) &= \sum_{t=1}^{T} \text{KL}_t \\ &= \sum_{t=1}^{T} \left[\log \pi_\theta(y_t | x, y_{<t}) - \log \pi_{\text{ref}}(y_t | x, y_{<t})\right] \end{aligned}$$

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.