KL Divergence Penalty › Effects on Training Dynamics
Language AI HandbookAll chapters KL Divergence Penalty › Effects on Training DynamicsOne section · Sections in this chapter10 / 13 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
The KL coefficient fundamentally shapes how RLHF training unfolds. Different values of 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.
Reward vs KL Tradeoff Frontier
During training, there is a Pareto frontier between reward maximization and KL minimization. Different 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.
def simulate_training_trajectory(
beta: float,
n_steps: int = 100,
reward_scale: float = 1.0,
noise_scale: float = 0.08,
) -> dict:
"""
Simulate RLHF training trajectory for a given beta.
This is a simplified model where:
- Reward increases with distance from reference (up to a point)
- KL increases proportionally with this distance
- Higher beta limits how far the policy moves
"""
kl_values = [0.0]
reward_values = [0.0]
# Stronger KL penalties produce lower equilibrium divergence.
equilibrium_kl = 18.0 / (1.0 + 8.0 * beta)
for step in range(n_steps):
current_kl = kl_values[-1]
# Optimization approaches a beta-dependent equilibrium with small noise.
kl_update = 0.08 * (equilibrium_kl - current_kl)
kl_update += noise_scale * np.random.randn()
new_kl = max(0, current_kl + kl_update)
# The learned reward keeps increasing beyond the reliable region. That
# extrapolative gain is the reward-hacking risk marked in the plot.
reward_model_score = reward_scale * (
2.0 * (1.0 - np.exp(-new_kl / 5.0)) + 0.08 * max(0, new_kl - 10.0)
)
new_reward = reward_model_score + noise_scale * np.random.randn()
kl_values.append(new_kl)
reward_values.append(new_reward)
return {
"kl": np.array(kl_values),
"reward": np.array(reward_values),
"beta": beta,
}
# Simulate for different beta values
betas = [0.01, 0.05, 0.1, 0.3, 1.0]
trajectories = [simulate_training_trajectory(b) for b in betas]
The visualization illustrates the basic tradeoff. Low- policies (warm colors) quickly accumulate KL divergence and can reach high rewards, but risk entering regions where the reward model is unreliable. High- policies (cool colors) progress slowly but stay in well-calibrated regions of the reward model's coverage. The diagonal trajectory of intermediate 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 Distribution Over Tokens
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.
# Simulate realistic per-token KL patterns
seq_len = 50
# Different patterns for different response types
# Pattern 1: High KL at beginning (new phrasing)
pattern_1 = np.exp(-np.linspace(0, 3, seq_len)) * 2 + 0.1
# Pattern 2: High KL at end (novel conclusions)
pattern_2 = np.exp(np.linspace(-3, 0, seq_len)) * 1.5 + 0.1
# Pattern 3: Spiky (specific word choices)
pattern_3 = 0.2 * np.ones(seq_len)
spike_positions = [5, 12, 23, 35, 45]
for pos in spike_positions:
pattern_3[pos] = np.random.uniform(1.5, 3.0)
# Add some noise
patterns = {
"Reformulated opening": pattern_1 + 0.1 * np.random.randn(seq_len),
"Novel conclusion": pattern_2 + 0.1 * np.random.randn(seq_len),
"Specific word choices": pattern_3
+ 0.05 * np.abs(np.random.randn(seq_len)),
}
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.
Comparison: Fixed vs Adaptive KL
Let us compare training stability between fixed and adaptive KL control to see why practitioners increasingly prefer the adaptive approach for production-scale RLHF:
def simulate_training_with_controller(
controller_type: str,
init_beta: float = 0.1,
target_kl: float = 6.0,
n_steps: int = 200,
kl_disturbances: np.ndarray | None = None,
reward_disturbances: np.ndarray | None = None,
) -> dict:
"""
Simulate training with fixed or adaptive KL control.
"""
beta = init_beta
kl_values = []
beta_values = []
reward_values = []
penalized_reward_values = []
current_kl = 0.5
if kl_disturbances is None:
kl_disturbances = np.zeros(n_steps)
if reward_disturbances is None:
reward_disturbances = np.zeros(n_steps)
for step in range(n_steps):
# Reward optimization pushes outward while the KL coefficient supplies
# a restoring force. Both controller variants see identical noise.
optimization_pressure = 0.14
constraint_force = 0.08 * beta * current_kl
current_kl = max(
0.1,
current_kl
+ optimization_pressure
- constraint_force
+ kl_disturbances[step],
)
reward = (
1.5 * (1 - np.exp(-current_kl / 5.0)) + reward_disturbances[step]
)
# Evaluate both trajectories under the same initial coefficient so the
# comparison isolates the value of controlling KL rather than changing
# the reporting scale whenever the adaptive coefficient moves.
penalized_reward = reward - init_beta * current_kl
kl_values.append(current_kl)
beta_values.append(beta)
reward_values.append(reward)
penalized_reward_values.append(penalized_reward)
# Smooth multiplicative feedback raises beta above target and lowers it
# below target without the extreme oscillations of hard doubling.
if controller_type == "adaptive":
relative_error = (current_kl - target_kl) / target_kl
beta *= np.exp(0.12 * relative_error)
beta = float(np.clip(beta, 0.02, 0.5))
return {
"kl": np.array(kl_values),
"beta": np.array(beta_values),
"reward": np.array(reward_values),
"penalized_reward": np.array(penalized_reward_values),
"controller": controller_type,
}
shared_kl_disturbances = np.random.normal(0.0, 0.08, 200)
shared_reward_disturbances = np.random.normal(0.0, 0.05, 200)
fixed_results = simulate_training_with_controller(
"fixed",
kl_disturbances=shared_kl_disturbances,
reward_disturbances=shared_reward_disturbances,
)
adaptive_results = simulate_training_with_controller(
"adaptive",
kl_disturbances=shared_kl_disturbances,
reward_disturbances=shared_reward_disturbances,
)
The adaptive controller maintains KL divergence near the target by increasing 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 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.