揭示强化学习中KL正则化的正确实现方式
Rethinking KL Regularization in RLHF: From Value Estimation to Gradient Optimization
- 区分两种KL正则化实现:作为奖励系数或直接作为损失函数
- 证明在策略上条件下,两种方法梯度等价且理论正确
- 指出GRPO等方法使用的一阶近似存在偏差,需修正
强化学习从人类反馈(RLHF)中利用KL散度损失稳定训练并防止过拟合。然而,如GRPO等方法中的实现方式,可能源于数值价值估计的原则,忽视了该项作为优化损失的功能角色。我们建立统一框架,连接两种看似不同的实现形式:将数学项 $k_n$ 作为策略得分函数的独立系数('k_n in reward')或作为直接传播梯度的损失函数('k_n as loss')。我们证明后者总可转化为前者的等效梯度系数,统一两种视角。在此基础上,我们证明传统的 'k_1 in reward'(如PPO)是反向KL(RKL)正则化的合理损失。进一步发现,在策略上条件下,'k_2 as loss' 实际与 'k_1 in reward' 梯度等价,二者均为理论上正确的实现。相反,近期采用的 'k_3 as loss'(如GRPO)仅为一阶、有偏近似。此外,我们指出常见的 'k_n as loss' 的非策略实现因忽略重要性采样而存在偏差,并提出修正方案。研究为选择和正确实现KL正则化提供了基于梯度的完整依据,推动更鲁棒高效的RLHF系统发展。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。