The KL-Regularised RL Objective

A 1.3 billion-parameter InstructGPT model is preferred by human raters over a 175x larger raw GPT-3, despite seeing no additional capability training. The difference is not scale; it is the objective the model was optimised against. That objective contains a term most practitioners gloss over: the KL divergence penalty. Understanding why it is there, what it is doing at every training step, and precisely where it breaks down separates practitioners who can tune RLHF pipelines from those who cargo-cult hyperparameters.

The Bare Reward Problem

Reinforcement learning requires a scalar reward signal. In language model post-training the reward typically comes from a learned reward model (RM) trained on human preference comparisons. Naively, the RL objective is:

maximise E_{x~D, y~pi_theta} [ r_phi(x, y) ]

where x is a prompt drawn from distribution D, y is the response sampled from the current policy pi_theta, and r_phi is the reward model parameterised by phi.

This objective is straightforward but dangerous. The reward model is an imperfect proxy for human preference. It was fitted on a finite dataset of comparison pairs, so it has a definite generalisation boundary. The policy has every incentive to find inputs that score high on the proxy while drifting far from the distribution on which the RM was evaluated. This is Goodhart's Law in closed-loop form: once a measure becomes a target, it ceases to be a good measure. Concretely, a policy trained purely against reward rapidly learns to produce degenerate text - repetitive phrases, garbled tokens, confident nonsense - that happens to exploit blind spots in the reward model.

Adding the KL Term

The standard fix, introduced in the RM fine-tuning work on summarisation (Stiennon et al., 2020) and solidified in InstructGPT (Ouyang et al., 2022), is to regularise the objective with a KL divergence between the current policy and a frozen reference policy:

maximise E_{x~D, y~pi_theta} [ r_phi(x, y) - beta * KL(pi_theta(y|x) || pi_ref(y|x)) ]

Written out token by token, the KL term is:

KL(pi_theta || pi_ref) = sum_t log( pi_theta(y_t | x, y_{<t}) / pi_ref(y_t | x, y_{<t}) )

This sum accumulates over every generated token. A single response that diverges only slightly at each step still accrues a meaningful penalty by the end of a long generation - by design.

The reference policy pi_ref is almost always the supervised fine-tuned (SFT) checkpoint that precedes RL training. It is frozen throughout RL and acts as an anchor. The scalar beta controls the trade-off:

The choice of beta is not a clean engineering decision - it is a research bet about how trustworthy the reward model is.

Why KL and Not L2 or Some Other Distance?

KL divergence on token distributions is the natural choice for three reasons.

First, language model outputs are categorical distributions over vocabularies of 50,000+ tokens. KL divergence is the canonical measure of divergence between probability distributions; it directly quantifies the extra bits needed to encode samples from pi_theta using a code optimal for pi_ref.

Second, the KL penalty is differentiable through the policy log-probabilities, which are already computed during the forward pass. Adding it costs almost nothing in compute; it reuses the same logits evaluated on the generated tokens.

Third, and most importantly, KL divergence is asymmetric. KL(pi_theta || pi_ref) penalises the policy for placing mass on tokens that the reference assigns very low probability - exactly the failure mode we want to discourage. High-reward but low-reference-probability tokens get penalised heavily. Low-reward tokens already in the reference distribution are cheap to generate and will not be punished for appearing. The asymmetry aligns with the intended regularisation direction.

Implementation: Where the Penalty Enters

In practice, the KL penalty is usually added to the per-token reward before passing it to the PPO advantage computation, not as a separate loss term:

# Pseudocode: per-token reward shaping in RLHF with PPO
for t in range(T):
    log_ratio = log_probs_policy[t] - log_probs_ref[t]   # log pi_theta / pi_ref
    kl_penalty[t] = log_ratio                              # approximates KL contribution
    shaped_reward[t] = reward_model_score * (t == T-1)    # RM score only at final token
    shaped_reward[t] -= beta * kl_penalty[t]              # subtract KL at every token

# PPO then uses shaped_reward to compute advantages

The RM score is usually assigned only at the end-of-sequence token (treating the whole response as one action), while the KL penalty is spread across every token. This creates a mixed credit assignment: dense KL penalties throughout the sequence, sparse terminal reward. The interaction between these two signals is a major source of training instability.

Some implementations instead compute the full KL as an auxiliary loss and add it to the PPO objective with its own coefficient. Both approaches have been used; per-token shaping is more common in open-source implementations such as TRL and OpenRLHF.

When It Falls Down

Beta sensitivity. There is no principled formula for beta. Too small, and the policy exploits the reward model within hundreds of steps. Too large, and RL training does nothing useful. Published work typically searches over one or two orders of magnitude and reports the best result; reproducibility across different reward models and base models is poor.

Reference policy staleness. The SFT checkpoint used as pi_ref was trained on a different data distribution from what RL later produces. If the RL distribution drifts far enough, the KL term begins to penalise genuinely good responses simply because they are unlike the SFT model. This is a form of regularisation overshoot.

Reward over-optimisation despite KL. Gao, Schulman, and Hilton (2022) showed empirically that optimising a proxy RM causes gold-reward to rise then fall as a function of the KL budget consumed, forming an inverted-U curve. The KL penalty slows this process but does not eliminate it. Even moderate beta values permit enough drift for reward hacking at scale. The tipping point depends on the number of RL steps, the capacity of the policy, and how well the RM generalises.

Does not apply to RLVR settings cleanly. More recent systems like DeepSeek-R1 use verifiable rewards (correct/incorrect) rather than a learned RM. A verifiable reward is not a proxy; it cannot be "hacked" in the same sense. In these settings, the motivation for aggressive KL regularisation is weaker. DeepSeek-R1 uses GRPO with a KL term, but the operating beta values and their interpretation differ from RLHF against a preference RM.

Group-relative methods change the accounting. GRPO (used in DeepSeek-R1) folds the KL penalty directly into the policy gradient estimator rather than shaping rewards, making the effective beta per-step harder to interpret and compare with PPO-based systems.

Frozen reference is a modelling assumption. Nothing guarantees the SFT checkpoint is the right anchor. If the SFT model itself was overtrained or had distributional issues, the reference distribution is corrupted, and all KL computations inherit that corruption.

Further Reading

• Stiennon, N. et al. (2020). "Learning to summarize from human feedback." NeurIPS 2020. https://arxiv.org/abs/2009.01325
• Ouyang, L. et al. (2022). "Training language models to follow instructions with human feedback." https://arxiv.org/abs/2203.02155
• Gao, L., Schulman, J., and Hilton, J. (2022). "Scaling laws for reward model overoptimization." https://arxiv.org/abs/2210.10760
• DeepSeek-AI (2025). "DeepSeek-R1: Incentivizing reasoning capability in LLMs via reinforcement learning." https://arxiv.org/abs/2501.12948