A Comedy of Estimators: On KL Regularization in RL Training of LLMs

TL;DR

This work analyzes the practical use of reverse KL regularization in RL fine-tuning of LLMs, focusing on estimators that approximate $D_\mathrm{KL}(\pi_\theta(\cdot|x) \| \pi_\mathrm{ref}(\cdot|x))$ and how their placement (reward vs loss) affects gradient bias. It contrasts two common estimators, K1 (naive) and K3 (low-variance Schulman), deriving gradient expressions and validating biases with toy models and large-scale LLM fine-tuning (Qwen2.5-7B, Llama-3.1-8B-Instruct, Qwen3-4B-Instruct-2507). Empirically, biased gradient configurations can cause instability or collapse in on-policy RL, whereas unbiased configurations yield better in-domain and out-of-domain performance; in asynchronous off-policy settings, KL regularization helps stabilize training. The authors advocate defaulting to unbiased gradient configurations for stable, generalizable reasoning-focused RL post-training and call for careful treatment of KL gradients in public libraries.

Abstract

The reasoning performance of large language models (LLMs) can be substantially improved by training them with reinforcement learning (RL). The RL objective for LLM training involves a regularization term, which is the reverse Kullback-Leibler (KL) divergence between the trained policy and the reference policy. Since computing the KL divergence exactly is intractable, various estimators are used in practice to estimate it from on-policy samples. Despite its wide adoption, including in several open-source libraries, there is no systematic study analyzing the numerous ways of incorporating KL estimators in the objective and their effect on the downstream performance of RL-trained models. Recent works show that prevailing practices for incorporating KL regularization do not provide correct gradients for stated objectives, creating a discrepancy between the objective and its implementation. In this paper, we further analyze these practices and study the gradients of several estimators configurations, revealing how design choices shape gradient bias. We substantiate these findings with empirical observations by RL fine-tuning \texttt{Qwen2.5-7B}, \texttt{Llama-3.1-8B-Instruct} and \texttt{Qwen3-4B-Instruct-2507} with different configurations and evaluating their performance on both in- and out-of-distribution tasks. Through our analysis, we observe that, in on-policy settings: (1) estimator configurations with biased gradients can result in training instabilities; and (2) using estimator configurations resulting in unbiased gradients leads to better performance on in-domain as well as out-of-domain tasks. We also investigate the performance resulting from different KL configurations in off-policy settings and observe that KL regularization can help stabilize off-policy RL training resulting from asynchronous setups.

Figures (12)

• Figure 1: The bias and variance of expected gradients with respect to the parameters of $A$, in different configurations (logarithmic scale). While all estimators are unbiased, the expected gradients are unbiased only in the case of \ref{['eq:naive']} estimator when used in reward. \ref{['eq:lowvar']} estimator when used in reward exhibits the highest bias. While \ref{['eq:naive']} estimator when used in loss has relatively lower variance, it also suffers from high bias.
• Figure 2: Training Instabilities when using \ref{['eq:naive']} in loss. Pass@1 performance for [Left] training Qwen2.5-7B with \ref{['eq:naive']} leads to training instabilities for $\beta = 0.1$ and $1$. [Center] Training Qwen2.5-7B with 4 policy update steps per sampled batch accentuates the instabilities owing to the increased off-policyness, leading to definitive training collapse in all cases. [Right] Training Llama-3.1-8B-Instruct with \ref{['eq:naive']} in loss leads to instabilities for all $\beta$ except 0.1.
• Figure 3: Collapse in the case of adding \ref{['eq:lowvar']} to the reward. Pass@1 performance for [Left]Qwen2.5-7B trained on MATH train dataset [Right]Llama-3.1-8B-Instruct trained on MATH train dataset with \ref{['eq:lowvar']} -in-reward. The collapse maybe attributed to high bias and variance of the configuration.
• Figure 4: Pass@1 performance on MATH test set with \ref{['eq:lowvar']} -in-loss (biased; Left) and \ref{['eq:naive']}-in-reward (unbiased; Right). Although biased with respect to reverse KL, \ref{['eq:lowvar']} -in-loss yields stable training. In both cases, lower $\beta$ values lead to higher performance.
• Figure 5: Comparison of Qwen2.5-7B trained with two stable estimator configurations - \ref{['eq:naive']} in reward and \ref{['eq:lowvar']} in loss. Baseline (CT) refers to the performance of base Qwen2.5-7B when prompted with a chat template. Baseline (No-CT) represents the performance when it is prompted with a chat template. \ref{['eq:naive']} in reward (unbiased gradient performs the beats on both in-domain and out-of-domain tasks. Increasing $\beta$ consistently deteriorates performance.