KL Divergence Penalty › Code Implementation
Language AI HandbookAll chapters KL Divergence Penalty › Code ImplementationOne section · Sections in this chapter8 / 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
Let us implement KL divergence computation and adaptive control for RLHF training. We will build components that integrate with the PPO training loop from the previous chapter. The implementation focuses on clarity and correctness. This gives the building blocks that can be optimized for production use. We proceed from the simplest building block (per-token KL) to more complete integration (the full RLHF objective) and then to the adaptive controller that ties everything together.
Computing Per-Token KL Divergence
The core computation extracts log probabilities from both the policy and reference model, then computes their difference. This function forms the foundation of the KL penalty calculation, taking log probabilities that have already been extracted from model outputs and creating both per-token and per-sequence KL values.
from typing import Optional, Tuple
import torch
def compute_token_kl(
policy_logprobs: torch.Tensor,
reference_logprobs: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Compute per-token and per-sequence KL divergence.
Args:
policy_logprobs: Log probs from policy for selected tokens (batch, seq_len)
reference_logprobs: Log probs from reference for same tokens (batch, seq_len)
attention_mask: Mask showing valid tokens (batch, seq_len)
Returns:
per_token_kl: KL at each position (batch, seq_len)
sequence_kl: Total KL per sequence (batch,)
"""
# Per-token KL is simply the log ratio
per_token_kl = policy_logprobs - reference_logprobs
if attention_mask is not None:
# Zero out KL for padding tokens
per_token_kl = per_token_kl * attention_mask
# Sum over valid tokens only
sequence_kl = per_token_kl.sum(dim=-1)
else:
sequence_kl = per_token_kl.sum(dim=-1)
return per_token_kl, sequence_klThis computation is straightforward because we are using the sampled trajectory estimator. The log probability ratio for generated tokens equals the KL divergence in expectation. The attention mask handles variable-length sequences by zeroing out contributions from padding tokens, so the KL computation considers only content tokens. This masking is needed when processing batches of sequences with different lengths, as padding tokens should not contribute to the divergence measure. Without masking, sequences with more padding would appear to have lower divergence simply because the padding region contributes zero per-token difference, artificially inflating how similar a batch of short sequences appears.
Extracting Log Probabilities from Model Outputs
To compute KL, we need log probabilities for the sampled tokens. Here is how to extract them from model logits:
import torch.nn.functional as F
def get_logprobs_for_tokens(
logits: torch.Tensor, tokens: torch.Tensor
) -> torch.Tensor:
"""
Extract log probabilities for specific tokens from logits.
Args:
logits: Model output logits (batch, seq_len, vocab_size)
tokens: Token indices to get probs for (batch, seq_len)
Returns:
logprobs: Log probabilities for specified tokens (batch, seq_len)
"""
# Convert logits to log probabilities
log_probs = F.log_softmax(logits, dim=-1)
# Gather log probs for the actual tokens
# Shape: (batch, seq_len, 1) -> (batch, seq_len)
token_logprobs = torch.gather(
log_probs, dim=-1, index=tokens.unsqueeze(-1)
).squeeze(-1)
return token_logprobsThe function first applies log_softmax to convert raw logits into log probabilities. This operation normalizes the logits so they stand for a valid probability distribution at each position. The gather operation then selects the log probability corresponding to each actual token from the full vocabulary distribution. This is much more efficient than computing probabilities for all tokens when we only need one per position. The unsqueeze and squeeze operations handle the dimension manipulation required by PyTorch's gather function.
The use of log_softmax rather than computing softmax and then taking the log is numerically important. Computing softmax first can lead to underflow for tokens with very negative logits, as the exponential can produce values below the floating point minimum. The combined log_softmax operation is numerically stable because it subtracts the maximum logit before computing the exponential, keeping all values in a reasonable range. This numerical stability is especially important for large vocabulary models where the distribution of logits can be highly skewed.
Adaptive KL Controller
The adaptive controller maintains the KL coefficient and adjusts it based on observed KL values. It implements the target KL approach described in the mathematical section. This gives a stateful object that can be updated after each training batch.
class AdaptiveKLController:
"""
Adaptive KL coefficient controller that targets a specific KL value.
Based on the approach from InstructGPT (Ouyang et al., 2022).
"""
def __init__(
self,
init_kl_coef: float = 0.1,
target_kl: float = 6.0,
horizon: int = 10000,
):
"""
Args:
init_kl_coef: Initial value of beta
target_kl: Target KL divergence per response
horizon: Number of steps over which to adapt (for smooth version)
"""
self.kl_coef = init_kl_coef
self.target_kl = target_kl
self.horizon = horizon
# Track history for monitoring
self.kl_history = []
self.coef_history = []
def update(self, current_kl: float) -> float:
"""
Update KL coefficient based on current KL divergence.
Args:
current_kl: Mean KL divergence from current batch
Returns:
Updated KL coefficient
"""
self.kl_history.append(current_kl)
# Multiplicative update (InstructGPT style)
if current_kl > 1.5 * self.target_kl:
self.kl_coef *= 2.0
elif current_kl < self.target_kl / 1.5:
self.kl_coef /= 2.0
# Clamp to reasonable range
self.kl_coef = max(min(self.kl_coef, 10.0), 0.001)
self.coef_history.append(self.kl_coef)
return self.kl_coef
def get_coef(self) -> float:
"""Return current KL coefficient."""
return self.kl_coefThe controller maintains a history of observed KL values and coefficient updates, which is useful for monitoring training dynamics and diagnosing issues. The clamping bounds prevent the coefficient from reaching extreme values that could either halt learning entirely (too high) or give no regularization (too low). These bounds can be adjusted based on the specific application and model scale. The history tracking is important for debugging: if the controller shows consistently increasing, it indicates the policy is consistently exceeding the target KL, which may suggest the reward model signal is too strong or the target KL is set too low for the task at hand.
Let us test the adaptive controller with simulated KL values:
# Simulate training dynamics
initial_beta = 0.1
controller = AdaptiveKLController(init_kl_coef=initial_beta, target_kl=6.0)
# Simulate KL values that start high, then stabilize
simulated_kl = np.concatenate(
[
np.random.uniform(8, 12, 20), # Early: KL too high
np.random.uniform(5, 7, 30), # Middle: near target
np.random.uniform(2, 4, 20), # Late: KL dropping
np.random.uniform(5, 7, 30), # Stabilized near target
]
)
for kl in simulated_kl:
controller.update(kl)Initial KL coefficient: 0.1 Target KL: 6.0 KL coefficient evolution: After high-KL phase (step 20): 10.0000 After stable phase (step 50): 10.0000 After low-KL phase (step 70): 0.0010 Final value (step 100): 0.0010
The controller doubles when KL exceeds the target, pulling the policy back toward the reference, and halves it when KL drops too low, letting more exploration. The simulation shows how the controller responds to different training phases: increasing the coefficient during the high-KL early phase, maintaining stability during the near-target phase, decreasing during the low-KL phase, and stabilizing again as KL returns to the target range.
Visualizing Adaptive KL Dynamics
The shaded region marks the "stable zone" where does not change. Outside this zone, the controller applies multiplicative updates to steer KL back toward target. The log scale on the right panel makes it easy to see the multiplicative nature of the updates: each adjustment doubles or halves the coefficient, which appears as equal-height steps on the log scale.
Complete KL Penalty Integration
Here is how KL computation integrates into a simplified RLHF training step. This function combines all the pieces we have developed, taking model outputs and creating the penalized objective along with diagnostic information.
def compute_rlhf_objective(
policy_logits: torch.Tensor,
reference_logits: torch.Tensor,
response_tokens: torch.Tensor,
rewards: torch.Tensor,
attention_mask: torch.Tensor,
kl_coef: float,
) -> dict:
"""
Compute the RLHF objective with KL penalty.
Args:
policy_logits: Logits from current policy (batch, seq_len, vocab)
reference_logits: Logits from reference model (batch, seq_len, vocab)
response_tokens: Generated token indices (batch, seq_len)
rewards: Reward model scores (batch,)
attention_mask: Valid token mask (batch, seq_len)
kl_coef: KL penalty coefficient (beta)
Returns:
Dictionary with objective value and diagnostics
"""
# Get log probabilities for generated tokens
policy_logprobs = get_logprobs_for_tokens(policy_logits, response_tokens)
reference_logprobs = get_logprobs_for_tokens(
reference_logits, response_tokens
)
# Compute KL divergence
per_token_kl, sequence_kl = compute_token_kl(
policy_logprobs, reference_logprobs, attention_mask
)
# KL-penalized reward
penalized_rewards = rewards - kl_coef * sequence_kl
# Mean objective (to be maximized)
objective = penalized_rewards.mean()
return {
"objective": objective,
"mean_reward": rewards.mean().item(),
"mean_kl": sequence_kl.mean().item(),
"mean_penalized_reward": penalized_rewards.mean().item(),
"kl_penalty": (kl_coef * sequence_kl.mean()).item(),
"per_token_kl": per_token_kl,
}The function returns a dictionary containing both the objective value for optimization and diagnostic quantities for monitoring. The separation between raw reward and penalized reward helps track whether improvements come from higher reward model scores or reduced divergence. The per-token KL tensor is included for detailed analysis of where the policy diverges most from the reference. Logging these diagnostic quantities during training gives the situational awareness needed to diagnose problems: if mean_kl is rising but mean_reward is flat, the policy is drifting without improving, suggesting should be increased.
Let us show with synthetic data:
# Create synthetic data
batch_size, seq_len, vocab_size = 4, 20, 1000
torch.manual_seed(42)
# Synthetic logits (policy slightly different from reference)
reference_logits = torch.randn(batch_size, seq_len, vocab_size)
policy_logits = reference_logits + 0.5 * torch.randn_like(reference_logits)
# Synthetic tokens and rewards
response_tokens = torch.randint(0, vocab_size, (batch_size, seq_len))
rewards = torch.tensor([0.8, 0.3, -0.2, 0.5])
attention_mask = torch.ones(batch_size, seq_len)
# Compute objective with different KL coefficients
low_beta = 0.01
results_low_beta = compute_rlhf_objective(
policy_logits,
reference_logits,
response_tokens,
rewards,
attention_mask,
kl_coef=low_beta,
)
high_beta = 0.5
results_high_beta = compute_rlhf_objective(
policy_logits,
reference_logits,
response_tokens,
rewards,
attention_mask,
kl_coef=high_beta,
)Effect of KL Coefficient on Objective ============================================= Low beta (0.01): Mean reward: 0.3500 Mean KL: -2.2833 KL penalty: -0.0228 Penalized reward: 0.3728 High beta (0.5): Mean reward: 0.3500 Mean KL: -2.2833 KL penalty: -1.1416 Penalized reward: 1.4916
With the same underlying KL divergence, a higher produces a much larger penalty, materially reducing the effective reward signal available for optimization. The comparison illustrates the large impact of the coefficient choice: with low , most of the reward signal passes through to guide learning, while with high , the penalty dominates and the effective reward becomes negative despite positive raw reward scores.