KL Divergence Penalty › Worked Example: Step-by-Step KL Computation

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:

$$Q = \{\text{blue}: 0.50,\ \text{clear}: 0.25,\ \text{dark}: 0.15,\ \text{beautiful}: 0.10\}$$

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

$$P = \{\text{blue}: 0.30,\ \text{clear}: 0.30,\ \text{dark}: 0.05,\ \text{beautiful}: 0.35\}$$

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:

$$\begin{aligned} \text{KL}_{\text{beautiful}} &= \log P(\text{beautiful}) - \log Q(\text{beautiful}) \\ &= \log(0.35) - \log(0.10) \\ &= -1.0498 - (-2.3026) \\ &= 1.2528 \text{ nats} \end{aligned}$$

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:

$$\begin{aligned} \text{KL}_{\text{dark}} &= \log P(\text{dark}) - \log Q(\text{dark}) \\ &= \log(0.05) - \log(0.15) \\ &= -2.9957 - (-1.8971) \\ &= -1.0986 \text{ nats} \end{aligned}$$

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:

$$\begin{aligned} D_{KL}(P \| Q) &= \sum_x P(x) \log \frac{P(x)}{Q(x)} \\ &= 0.30 \cdot \log\frac{0.30}{0.50} + 0.30 \cdot \log\frac{0.30}{0.25} + 0.05 \cdot \log\frac{0.05}{0.15} + 0.35 \cdot \log\frac{0.35}{0.10} \end{aligned}$$

Computing each term:

$$\begin{aligned} &= 0.30 \cdot (-0.5108) + 0.30 \cdot (0.1823) + 0.05 \cdot (-1.0986) + 0.35 \cdot (1.2528) \\ &= -0.1532 + 0.0547 - 0.0549 + 0.4385 \\ &= 0.2851 \text{ nats} \end{aligned}$$

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.