Rethinking KL Regularization in RLHF: From Value Estimation to Gradient Optimization (2510.01555v1)

Abstract: Reinforcement Learning from Human Feedback (RLHF) leverages a Kullback-Leibler (KL) divergence loss to stabilize training and prevent overfitting. However, in methods such as GRPO, its implementation may be guided by principles from numerical value estimation-a practice that overlooks the term's functional role as an optimization loss. To analyze this issue, we establish a unified framework that connects two seemingly distinct implementation styles: using the mathematical term $k_n$ as a detached coefficient for the policy's score function ('$k_n$ in reward') or as a direct loss function through which gradients are propagated ('$k_n$ as loss'). We show that the latter can always be analyzed via an equivalent gradient coefficient in the former, unifying the two perspectives. Through this framework, we prove that the conventional '$k_1$ in reward' (like in PPO) is the principled loss for Reverse KL (RKL) regularization. We further establish a key finding: under on-policy conditions, the '$k_2$ as loss' formulation is, in fact, gradient-equivalent to '$k_1$ in reward'. This equivalence, first proven in our work, identifies both as the theoretically sound implementations of the RKL objective. In contrast, we show that the recently adopted '$k_3$ as loss' (like in GRPO) is merely a first-order, biased approximation of the principled loss. Furthermore, we argue that common off-policy implementations of '$k_n$ as loss' methods are biased due to neglected importance sampling, and we propose a principled correction. Our findings provide a comprehensive, gradient-based rationale for choosing and correctly implementing KL regularization, paving the way for more robust and effective RLHF systems.

• The paper introduces a unified gradient-centric framework that reinterprets KL regularization in RLHF from value estimation to direct gradient optimization.
• The authors prove that conventional in-reward (k1) and as-loss (k2) formulations are gradient-equivalent under on-policy sampling, ensuring effective reverse KL regularization.
• Empirical results show that principled methods yield lower variance and improved stability, while first-order approximations like k3 as loss can lead to bias and instability.

Gradient-Centric Analysis of KL Regularization in RLHF

Introduction and Motivation

The paper "Rethinking KL Regularization in RLHF: From Value Estimation to Gradient Optimization" (2510.01555) addresses a critical but under-theorized component of Reinforcement Learning from Human Feedback (RLHF): the implementation of Kullback-Leibler (KL) regularization. While KL penalties are ubiquitous in RLHF to prevent reward hacking and distributional drift, the field has largely adopted KL loss forms based on value estimation heuristics rather than a principled analysis of their optimization properties. This work systematically re-examines KL regularization from a gradient-centric perspective, providing a unified framework that clarifies the theoretical and practical implications of different KL loss formulations.

Unified Framework for KL Regularization

The authors introduce a formal framework that distinguishes two main classes of KL regularization implementations in RLHF:

1. $k_n$ in reward: The KL term is used as a detached coefficient multiplying the policy's score function, e.g., $$k_1(y|x) = \log \pi_\theta(y|x) - \log \pi_{\text{ref}}(y|x)$$.

1. $k_n$ as loss: The KL term is treated as a direct loss, with gradients propagated through it, e.g., $$k_2(y|x) = \frac{1}{2}(\log \frac{\pi_{\text{red}\theta}(y|x)}{\pi_{\text{ref}}(y|x)})^2$$ or $$k_3(y|x) = \frac{\pi_{\text{ref}}(y|x)}{\pi_{\text{red}\theta}(y|x)} - 1 - \log \frac{\pi_{\text{ref}}(y|x)}{\pi_{\text{red}\theta}(y|x)}$$.

The framework enables a direct comparison of these approaches by analyzing the scalar coefficient $c(x, y)$ that multiplies the policy gradient $$\nabla_{\text{red}\theta} \log \pi_{\text{red}\theta}(y|x)$$. This abstraction reveals that, under on-policy sampling, certain "as loss" and "in reward" forms are gradient-equivalent, while others are not.

Figure 1: Comparison of KL regularization gradient coefficients. Principled implementations yield a linear restoring force in log-probability, while first-order approximations can be loose or unstable in the tails.

Theoretical Analysis of KL Loss Forms

Counterexample: $k_1$ as Loss

The paper demonstrates that using $k_1$ as a direct loss (i.e., differentiating through $$\log \pi_{\text{red}\theta}(y|x) - \log \pi_{\text{ref}}(y|x)$$) yields a gradient that is independent of the reference policy and has zero mean. This means it provides no regularization signal, only noise, and is thus ineffective for KL regularization.

Principled KL Loss: $k_1$ in Reward and $k_2$ as Loss

A central result is the proof that, under on-policy sampling, the conventional $k_1$ in reward (as in PPO) and $k_2$ as loss are gradient-equivalent and both implement the correct reverse KL (RKL) regularization. The gradient induced by these forms is:

$$\nabla_{\text{red}\theta} \mathcal{J}_{\mathrm{RKL}}(\text{red}\theta) = \mathbb{E}_{x, y \sim \pi_\theta} \left[ \log \frac{\pi_\theta(y|x)}{\pi_{\text{ref}}(y|x)} \nabla_{\text{red}\theta} \log \pi_{\text{red}\theta}(y|x) \right]$$

This equivalence is formalized and proven in the appendix, and it provides a rigorous foundation for preferring these forms in RLHF implementations.

First-Order Approximation: $k_3$ as Loss

The $k_3$ as loss form, popularized by GRPO and justified as an "unbiased estimator" in value estimation literature, is shown to be a first-order Taylor approximation of the true RKL gradient. Its coefficient, $1 - \delta$ (where $\delta = \frac{\pi_{\text{ref}}}{\pi_\theta}$), matches $-\log \delta$ only near $\delta = 1$ but diverges in the tails, leading to:

• Bias: The update direction is not aligned with the true RKL gradient except near the reference.
• Asymmetry: The regularization is much weaker for over-coverage and can be explosively strong for under-coverage.
• Statistical Instability: The variance of the stochastic gradient is governed by the chi-squared divergence, which can be extremely large or infinite if the policy and reference distributions diverge.

Empirical Validation

The authors conduct controlled experiments on mathematical reasoning tasks, comparing the training dynamics of different KL loss forms.

Figure 2: Comparison of k_1 as loss versus no KL regularization. Training dynamics are nearly indistinguishable, confirming the theoretical prediction that k_1 as loss is ineffective as a KL regularizer.

Figure 3: Comparison of the principled k_2 as loss against its first-order surrogate k_3 as loss. k_2 as loss demonstrates superior regularization, maintaining tighter coupling to the reference policy and more stable optimization.

Key findings include:

• $k_1$ as loss: Empirically indistinguishable from no KL regularization, confirming its theoretical ineffectiveness.
• $k_2$ as loss vs. $k_3$ as loss: Both constrain the policy, but $k_2$ as loss yields lower variance, tighter coupling to the reference, and more stable training. $k_3$ as loss allows greater drift and instability, especially in large-scale models.

Implementation Guidance and Off-Policy Correction

The paper provides detailed recommendations for practitioners:

• Do not use $k_1$ as loss: It provides no regularization and increases gradient variance.
• Prefer $k_1$ in reward or $k_2$ as loss: These are theoretically sound and gradient-equivalent under on-policy sampling.
• Recognize $k_3$ as loss as a biased approximation: It may be useful for weaker late-stage constraints but is not a principled choice.
• Correct for off-policy bias: When using "as loss" forms in off-policy settings (e.g., PPO), importance sampling and PPO-style clipping must be applied to the KL term itself. The "in reward" form naturally inherits these corrections.

The authors also discuss bounded alternatives (e.g., MiniMax-01 loss) for enhanced stability, which induce gradient coefficients bounded within $[-1, 1]$.

Implications and Future Directions

This work clarifies long-standing ambiguities in RLHF regarding KL regularization. By establishing a gradient-centric framework, it provides a rigorous basis for selecting and implementing KL penalties, moving the field away from ad-hoc, value-estimation-driven heuristics. The findings have immediate practical impact: the $k_2$ as loss formulation has already been adopted in major RLHF frameworks, and the analysis of off-policy correction addresses a subtle but critical source of bias in large-scale RLHF pipelines.

Theoretically, the paper's approach suggests that future RLHF research should prioritize the analysis of optimization gradients over value estimation properties when designing regularization and auxiliary losses. Practically, the recommendations enable more stable, effective, and reproducible RLHF systems, especially as model and data scales continue to increase.

Conclusion

This paper provides a comprehensive, gradient-based analysis of KL regularization in RLHF, unifying disparate implementation styles and rigorously identifying the principled forms for reverse KL regularization. The work exposes the limitations of value-estimation-motivated losses, demonstrates the superiority of $k_1$ in reward and $k_2$ as loss, and offers actionable guidance for both on-policy and off-policy RLHF training. These contributions are expected to inform both theoretical research and practical system design in the alignment of LLMs.