Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

PPO for Language Models › The Critical Role of the KL Penalty

Michael BrenndoerferDecember 28, 20259 min read

Language AI HandbookAll chapters

Part 37 of 67Alignment and RLHF

Chapter 279PPO for Language Models › The Critical Role of the KL Penalty

Reading section 6 of 15 · The Critical Role of the KL Penalty · 9 minSection 6 of 15 · 9 min

Section 6 of 15 · 9 min read

PPO for Language Models › The Critical Role of the KL PenaltyOne section · Sections in this chapter6 / 15 ↓

Sections in this chapter15

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 most important adaptation for language model PPO is the KL divergence penalty. This constraint measures how far the policy drifts from its starting point and is an important regularizer. Without this constraint, the optimized policy can diverge substantially from the original model, often finding degenerate solutions that maximize reward without producing useful responses.

Think of the KL penalty as a tether connecting the evolving policy to its starting point. As PPO updates the model's parameters to increase reward, the tether pulls back, preventing the model from wandering too far into unknown territory. The reference model, frozen at initialization, represents a baseline of linguistic competence built through pretraining and supervised fine-tuning. The KL penalty says: "you can move to improve reward, but you must not move so far that you forget what you already know."

Reward Hacking Revisited

As we discussed in the chapter on Reward Hacking, optimizing a proxy reward (the learned reward model) rather than true human preferences creates opportunities for exploitation. The reward model is an imperfect approximation of human preferences, and any learned approximation has exploitable gaps. Without the KL penalty, the policy is free to find responses that score highly according to the reward model while being less useful, repetitive, or even nonsensical. The KL penalty is our primary defense against this failure mode, but it is not perfect. The appropriate strength of the penalty depends on how well the reward model generalizes and must be tuned carefully.

Without the KL penalty, the policy often finds degenerate solutions. The policy is free to move anywhere in the space of possible token distributions. If the reward model has exploitable shortcuts that yield high scores without corresponding quality, the unconstrained policy will find them. It might learn to generate repetitive phrases that the reward model scores highly. It might produce outputs that superficially resemble good responses while lacking substance. It might drift so far from natural language that it generates text no human would write. One classic failure mode is mode collapse, where the policy learns to always generate one particular type of response that scores well on average, losing the diversity and adaptability that make a language model useful. The KL penalty prevents these failure modes by keeping the policy anchored to the reference distribution.

The KL penalty modifies the reward at each timestep, converting the sparse terminal reward into a dense per-token signal. The shaped reward at step $t$ is:

\tilde{r}_t = r_t - \beta \cdot \text{KL}_t

where:

$\tilde{r}_t$ : the shaped reward at step $t$ used for PPO training
$r_t$ : the original sparse reward (nonzero only at the final step)
$\beta$ : the KL penalty coefficient, controlling the strength of the regularization
$\text{KL}_t$ : the KL divergence contribution at this step

The KL contribution at each step measures the divergence between the current policy and the reference for the specific token generated. Rather than summing over the full vocabulary (which would require querying both policies over all 50,000+ tokens), we use the log ratio for the chosen action as a tractable per-token approximation:

\text{KL}_t = \log \frac{\pi_\theta(y_t | s_t)}{\pi_{\text{ref}}(y_t | s_t)}

where:

$\pi_\theta(y_t | s_t)$ : the probability of the chosen token under the current policy
$\pi_{\text{ref}}(y_t | s_t)$ : the probability of the chosen token under the reference policy
$y_t$ : the token generated at step $t$
$s_t$ : the context (state) at step $t$

This formula is easy to interpret. When the current policy assigns higher probability to a token than the reference did, the log ratio is positive, and subtracting it from the reward penalizes this divergence. When the current policy assigns lower probability than the reference, the log ratio is negative, and subtracting it provides a bonus (the policy is moving closer to the reference for that action). The net effect is that the policy is discouraged from making dramatic probability changes in either direction. It can shift probabilities to improve reward, but only within bounds set by $\beta$ .

The key insight is that this per-token KL penalty serves double duty. First, it constrains policy drift globally, preventing the model from drifting too far from its original behavior. Second, it provides a dense reward signal at every token, not just at the end of generation. This density matters for credit assignment: instead of all tokens receiving the same sparse terminal reward, each token now receives a small bonus or penalty based on how much the policy changed from the reference for that particular choice.

The reference policy $\pi_{\text{ref}}$ is typically the model after supervised fine-tuning (SFT) but before any RL training. This anchor prevents the policy from drifting into regions of token space that the original model considered highly unlikely. The SFT model represents our best current understanding of how to generate helpful, coherent text. By constraining the RL policy to stay near this baseline, we ensure that the optimized model retains the linguistic competence learned during pretraining and SFT.

The coefficient $\beta$ controls the strength of this constraint. Higher values keep the policy closer to the reference but limit learning; the policy cannot deviate much even when doing so would improve reward. Lower values allow more exploration but risk instability and reward hacking; the policy might find degenerate solutions that maximize reward while creating poor text. Success in RLHF depends on finding the right balance. We will explore the mathematical properties and tuning of this penalty in detail in the upcoming chapter on KL Divergence Penalty.

In[13]:

Code

def compute_kl_penalty(log_probs_policy, log_probs_ref):
    """
    Compute per-token KL divergence between policy and reference.

    This is a simplified approximation using only the chosen actions,
    not the full distribution over vocabulary.

    Args:
        log_probs_policy: Log probabilities under current policy
        log_probs_ref: Log probabilities under reference policy

    Returns:
        Per-token KL divergence estimates
    """
    # KL(policy || ref) approximated as log(policy(a)) - log(ref(a)) for chosen action a
    # This is an approximation; true KL sums over all actions
    kl = log_probs_policy - log_probs_ref
    return kl


def apply_kl_reward_shaping(rewards, log_probs_policy, log_probs_ref, beta=0.1):
    """
    Modify rewards with KL penalty.

    Args:
        rewards: Original per-token rewards (typically sparse)
        log_probs_policy: Log probs of chosen tokens under policy
        log_probs_ref: Log probs of chosen tokens under reference
        beta: KL penalty coefficient

    Returns:
        Modified rewards with KL penalty applied
    """
    kl = compute_kl_penalty(log_probs_policy, log_probs_ref)
    shaped_rewards = rewards - beta * kl
    return shaped_rewards, kl

In[14]:

Code

# Demonstrate KL penalty effect
np.random.seed(42)
seq_len = 6

# Simulate log probabilities: gaps chosen so per-token KL penalty is
# clearly visible (~0.25–0.75 units) relative to the terminal reward
log_probs_policy = np.array([-0.8, -1.6, -0.6, -1.8, -0.7, -1.5])
log_probs_ref = np.array([-1.5, -0.9, -1.8, -0.9, -1.9, -0.8])

# Original sparse reward
original_rewards = assign_rewards(seq_len, 0.8)

# Apply KL shaping with beta=0.5 so the per-token penalty is easily readable
shaped_rewards, kl_values = apply_kl_reward_shaping(
    original_rewards, log_probs_policy, log_probs_ref, beta=0.5
)

total_kl = kl_values.sum()
total_original = original_rewards.sum()
total_shaped = shaped_rewards.sum()

Out[15]:

Console

Per-token analysis:
------------------------------------------------------------
Token 0: original_r=0.00, KL=0.700, shaped_r=-0.350
Token 1: original_r=0.00, KL=-0.700, shaped_r=0.350
Token 2: original_r=0.00, KL=1.200, shaped_r=-0.600
Token 3: original_r=0.00, KL=-0.900, shaped_r=0.450
Token 4: original_r=0.00, KL=1.200, shaped_r=-0.600
Token 5: original_r=0.80, KL=-0.700, shaped_r=1.150
------------------------------------------------------------
Total KL penalty: 0.800
Total original reward: 0.80
Total shaped reward: 0.400

Out[16]:

Visualization

Bar chart of original per-token rewards. Five bars at zero and one final bar at 0.8. — Original sparse reward: only the final token receives a non-zero value from the reward model, leaving all intermediate tokens without direct feedback.

Bar chart of negative beta-weighted KL penalty at each token position, alternating between positive and negative values. — Per-token KL contribution: the negative beta-weighted KL penalty applied at each position. Red bars indicate the policy diverged toward higher probability than the reference (penalized), while green bars indicate the policy moved toward lower probability (rewarded).

Bar chart of KL-shaped reward combining terminal reward and per-token KL corrections. — KL-shaped reward: the combination of the sparse terminal reward and the dense per-token KL penalty. Every token now receives feedback, converting the sparse reward problem into a dense one.

Notice how the KL penalty transforms the reward signal. The original sparse reward only provides signal at the final token; all intermediate tokens receive zero reward. After KL shaping, every token receives a reward component based on how much the policy diverges from the reference. Tokens where the policy assigns higher probability than the reference receive penalties (the KL term is positive, so it subtracts from the reward). Tokens where the policy assigns lower probability than the reference receive bonuses (the KL term is negative, so subtracting it adds to the reward). This transformation converts the sparse terminal reward into a dense per-token signal that guides the policy at every step.