16. KL-Divergence Penalty for RLHF-Style Policy Training

Problem

RLHF-style post-training walks a fine line: push the policy toward higher-reward completions without letting it drift so far from the frozen reference (pre-RLHF) model that it collapses into degenerate, reward-hacking text. The standard fix is to shrink the reward itself by a per-token KL-divergence penalty between the current policy's next-token distribution and the reference model's next-token distribution, scaled by a coefficient beta.

The bigger the policy diverges from the reference at a given position, the more that position's reward gets discounted. Implement this as a full categorical KL divergence computed over the vocabulary dimension, not a sampled-token approximation.

Source: src/16_rlhf_kl_penalty.py

def kl_penalized_reward(
    policy_logits: torch.Tensor,
    ref_logits: torch.Tensor,
    reward: torch.Tensor,
    beta: float = 0.1,
) -> torch.Tensor: ...

>>> logits = torch.tensor([[[2.0, 0.0, 0.0]]])
>>> reward = torch.tensor([[1.0]])
>>> kl_penalized_reward(logits, logits, reward, beta=0.1)
tensor([[1.]])

Step-by-Step Approach

1. Convert both policy_logits and ref_logits to log-probabilities with F.log_softmax along the vocab axis, rather than calling log(softmax(...)) separately.
2. Exponentiate the policy's log-probabilities to recover the policy's probability distribution per token.
3. Compute the per-token KL divergence as sum_v policy(v) * (log policy(v) - log ref(v)), summing over the vocab axis (the last dimension).
4. Subtract beta * kl from the given per-token reward tensor and return the result, which has the same shape as reward.
5. Sanity check: when policy_logits == ref_logits, KL is exactly 0 everywhere, so the output should equal the input reward unchanged.

The key insight is that log_softmax followed by .exp() is the numerically stable way to get both the log-probabilities and probabilities you need — calling softmax and then .log() separately risks taking the log of an underflowed zero for very negative logits.

Reference solution

import torch
import torch.nn.functional as F


def kl_penalized_reward(
    policy_logits: torch.Tensor,
    ref_logits: torch.Tensor,
    reward: torch.Tensor,
    beta: float = 0.1,
) -> torch.Tensor:
    # log_softmax (not log(softmax(...))) for numerical stability -- avoids
    # taking log of a value that may have underflowed to exactly 0
    policy_log_probs = F.log_softmax(policy_logits, dim=-1)  # (batch, seq_len, vocab)
    ref_log_probs = F.log_softmax(ref_logits, dim=-1)  # (batch, seq_len, vocab)
    policy_probs = policy_log_probs.exp()
    # full categorical KL(policy || ref) per token, summed over the vocab axis
    kl = (policy_probs * (policy_log_probs - ref_log_probs)).sum(dim=-1)  # (batch, seq_len)
    return reward - beta * kl

Key Functions & Tricks

• F.log_softmax(logits, dim=-1) — numerically stable log-probabilities via the max-subtraction trick, safe even for large-magnitude logits.
• log_probs.exp() — recovers probabilities from log-probabilities without a separate, less-stable softmax call.
• (p * (log_p - log_q)).sum(dim=-1) — the categorical KL divergence formula, vectorized across the vocab axis for every batch and sequence position at once.
• Broadcasting reward - beta * kl — both tensors already share shape (batch, seq_len), so this is a plain elementwise subtraction, no reshaping needed.

How to Recognize This Pattern