KL Divergence Penalty › Adaptive KL Control

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:

$$\beta_{t+1} = \begin{cases} 2 \cdot \beta_t & \text{if } D_{KL} > 1.5 \cdot D_{\text{target}} \\ \beta_t / 2 & \text{if } D_{KL} < D_{\text{target}} / 1.5 \\ \beta_t & \text{otherwise} \end{cases}$$

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:

$$\beta_{t+1} = \beta_t \cdot \left(1 + \alpha \cdot \text{sign}(D_{KL} - D_{\text{target}})\right)$$

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:

$$\beta_{t+1} = \beta_t \cdot \exp\left(\alpha \cdot \frac{D_{KL} - D_{\text{target}}}{D_{\text{target}}}\right)$$

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.