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

📄 arXiv: 2512.21852v2 📥 PDF

作者: Vedant Shah, Johan Obando-Ceron, Vineet Jain, Brian Bartoldson, Bhavya Kailkhura, Sarthak Mittal, Glen Berseth, Pablo Samuel Castro, Yoshua Bengio, Nikolay Malkin, Moksh Jain, Siddarth Venkatraman, Aaron Courville

分类: cs.LG, cs.AI

发布日期: 2025-12-26 (更新: 2026-01-06)

💡 一句话要点

研究KL散度估计器对LLM的RL训练影响，提出无偏梯度配置以提升性能。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 强化学习 大型语言模型 KL散度 梯度偏差 无偏估计 策略优化 模型训练

📋 核心要点

  1. 现有方法在LLM的RL训练中，对KL散度的估计存在偏差，导致目标函数与实际实现不符，影响模型性能。
  2. 论文分析了不同KL散度估计器配置的梯度偏差，并提出使用无偏梯度配置来优化RL训练过程。
  3. 实验表明，使用无偏梯度配置的KL估计器能提升LLM在in-domain和out-of-domain任务上的性能，并稳定训练。

📝 摘要（中文）

大型语言模型（LLM）的推理性能可以通过强化学习（RL）训练得到显著提升。LLM训练的RL目标包含一个正则化项，即训练策略与参考策略之间的反向Kullback-Leibler（KL）散度。由于精确计算KL散度是不可行的，因此实践中通常使用各种估计器从on-policy样本中估计它。尽管KL散度估计器被广泛采用，包括在多个开源库中，但目前还没有系统的研究分析将KL估计器纳入目标函数的各种方式及其对RL训练模型下游性能的影响。最近的研究表明，目前采用的KL正则化方法并没有为既定目标提供正确的梯度，从而造成目标函数与其实现之间的差异。在本文中，我们进一步分析了这些实践，并研究了几种估计器配置的梯度，揭示了设计选择如何影响梯度偏差。我们通过使用不同的配置对 exttt{Qwen2.5-7B}、 exttt{Llama-3.1-8B-Instruct}和 exttt{Qwen3-4B-Instruct-2507}进行RL微调，并在in-distribution和out-of-distribution任务上评估它们的性能，从而证实了这些发现。通过我们的分析，我们观察到，在on-policy设置中：（1）具有有偏梯度的估计器配置可能导致训练不稳定；（2）使用产生无偏梯度的估计器配置可以提高in-domain和out-of-domain任务的性能。我们还研究了off-policy设置中不同KL配置产生的结果，并观察到KL正则化可以帮助稳定异步设置产生的off-policy RL训练。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLM）通过强化学习（RL）进行训练时，KL散度正则化项的估计问题。现有的KL散度估计方法存在梯度偏差，导致实际优化目标与理论目标不一致，进而影响模型的训练稳定性和最终性能。这些偏差源于对KL散度估计器的不同配置方式，以及这些配置对梯度计算的影响。

核心思路：论文的核心思路是分析不同KL散度估计器配置的梯度，识别导致梯度偏差的因素，并提出使用无偏梯度配置来优化RL训练过程。通过确保梯度与实际优化目标一致，可以提高训练的稳定性和模型性能。这种方法强调了在RL训练中，选择合适的KL散度估计器配置的重要性。

技术框架：论文的研究框架主要包括以下几个阶段：1) 分析不同KL散度估计器配置的梯度，揭示设计选择如何影响梯度偏差；2) 通过实验验证梯度偏差对RL训练的影响，包括训练稳定性和模型性能；3) 提出并测试使用无偏梯度配置的KL估计器，以提高训练效果。该框架涵盖了理论分析、实验验证和方法改进三个方面。

关键创新：论文最重要的技术创新在于对KL散度估计器配置的梯度偏差进行了深入分析，并提出了使用无偏梯度配置来优化RL训练过程。与现有方法相比，该方法能够更准确地估计KL散度，从而提供更可靠的梯度信息，提高训练的稳定性和模型性能。这种方法强调了在RL训练中，选择合适的KL散度估计器配置的重要性。

关键设计：论文的关键设计包括：1) 对不同KL散度估计器配置的梯度进行数学推导和分析，识别导致梯度偏差的因素；2) 设计实验来验证梯度偏差对RL训练的影响，包括训练稳定性和模型性能；3) 提出并测试使用无偏梯度配置的KL估计器，例如，通过调整损失函数中的KL散度项的计算方式，确保梯度与实际优化目标一致。具体的参数设置和网络结构可能因不同的LLM和任务而异，但核心思想是确保KL散度估计的无偏性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，使用无偏梯度配置的KL估计器能够显著提高LLM在in-domain和out-of-domain任务上的性能。例如，在对 exttt{Qwen2.5-7B}、 exttt{Llama-3.1-8B-Instruct}和 exttt{Qwen3-4B-Instruct-2507}进行RL微调时，使用无偏梯度配置的KL估计器能够稳定训练过程，并提升模型在各种任务上的准确率和生成质量。

🎯 应用场景

该研究成果可应用于各种需要通过强化学习训练大型语言模型的场景，例如对话系统、文本生成、代码生成等。通过选择合适的KL散度估计器配置，可以提高模型的训练效率和性能，从而提升用户体验和应用价值。该研究还有助于推动RL在LLM训练中的应用，并为未来的研究提供指导。

📄 摘要（原文）

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.