TLDR: A new research paper challenges the conventional implementation of KL regularization in RLHF, arguing that it has been misguided by value estimation principles rather than gradient optimization. The paper introduces a unified framework, identifies ‘k1 in reward’ and ‘k2 as loss’ as principled and gradient-equivalent methods for Reverse KL regularization, and exposes ‘k3 as loss’ (used in GRPO) as a biased approximation. It also provides crucial corrections for off-policy implementations, offering a robust, gradient-centric foundation for building more stable and effective RLHF systems.
Reinforcement Learning from Human Feedback (RLHF) has become a cornerstone in refining Large Language Models (LLMs). It helps these powerful models align with human preferences and excel in complex tasks like mathematics and code generation. A critical component of this process is KL regularization, which uses a Kullback-Leibler (KL) divergence loss to stabilize training and prevent the model from straying too far from its initial training.
However, a recent research paper highlights a significant oversight in how KL regularization has often been implemented. Many methods, including popular ones like GRPO, have historically approached KL regularization from the perspective of numerical value estimation. This means they focused on how accurately a KL term estimates a value, rather than its functional role as an optimization loss that guides the model’s learning process.
The paper, titled “Rethinking KL Regularization in RLHF: From Value Estimation to Gradient Optimization,” by Kezhao Liu, Jason Klein Liu, Mingtao Chen, and YiMing Liu, introduces a unified framework to address this issue. It connects two distinct ways KL terms are used:
Two Ways to Implement KL Regularization
The first way, termed ‘kn in reward,’ treats the KL term as a detached coefficient that weights the policy’s score function. A well-known example is PPO (Proximal Policy Optimization), which uses ‘k1 in reward.’
The second way, ‘kn as loss,’ uses the KL term directly as a loss function, through which gradients are propagated. GRPO, for instance, has adopted ‘k3 as loss’ based on its perceived properties as an unbiased value estimator.
The researchers demonstrate that the ‘kn as loss’ approach can always be analyzed through an equivalent gradient coefficient in the ‘kn in reward’ style, effectively unifying these two perspectives. This framework allows for a gradient-centric analysis, which is crucial for designing robust RLHF algorithms.
Why Some Implementations Fall Short
A key finding from the paper is that focusing solely on value estimation can lead to ineffective optimization signals. For example, ‘k1 as loss,’ despite being an unbiased estimator of the KL value, provides no meaningful regularization signal because its expected gradient is zero and independent of the reference policy. In practice, it merely adds noise to the training process.
Identifying the Principled Approaches
The paper rigorously proves that the conventional ‘k1 in reward’ formulation, widely used in methods like PPO, is indeed the principled loss for Reverse KL (RKL) regularization. More importantly, it establishes a previously unrecognized equivalence: ‘k2 as loss’ is gradient-equivalent to ‘k1 in reward’ under on-policy conditions. This means both ‘k1 in reward’ and ‘k2 as loss’ are theoretically sound choices for implementing KL regularization.
In contrast, the ‘k3 as loss’ formulation, adopted by methods like GRPO, is shown to be merely a first-order, biased approximation of the principled loss. This approximation can lead to weaker regularization, pathological asymmetry (where it saturates for over-sampled tokens and explodes for under-sampled ones), and statistical instability, making it a less reliable choice.
Also Read:
- Predicting LLM Reasoning Power: Why SFT Scores Aren’t Enough
- The Art of Reward Selection: Making AI Learn More with Less Data
Addressing Off-Policy Bias and Practical Recommendations
The research also points out a common pitfall in off-policy implementations of ‘kn as loss’ methods: they often neglect importance sampling, leading to systematic bias. The paper proposes a principled correction to address this. For practitioners, the recommendations are clear: avoid ‘k1 as loss,’ prefer ‘k1 in reward’ or ‘k2 as loss’ for their theoretical soundness, understand the limitations of ‘k3 as loss,’ and always correct for off-policy bias.
Experimental validations on mathematical reasoning tasks strongly support these theoretical distinctions. ‘k1 as loss’ proved ineffective, while ‘k2 as loss’ demonstrated superior regularization properties and training stability compared to ‘k3 as loss.’
This work provides a comprehensive, gradient-based rationale for choosing and correctly implementing KL regularization, paving the way for more robust and effective RLHF systems. The ‘k2 as loss’ formulation has already been integrated into frameworks like OpenRLHF and adopted by Reinforce++, demonstrating its immediate practical impact. You can read the full paper for more details: Rethinking KL Regularization in RLHF: From Value Estimation to Gradient Optimization.
