KL散度加到Loss里？过时了！Bengio团队重磅研究：别再拿KL散度当Loss了，这才是无偏最优解- 大数跨境

首页

KL散度加到Loss里？过时了！Bengio团队重磅研究：别再拿KL散度当Loss了，这才是无偏最优解

AI前沿速递

2026-01-06

导读：KL散度加到Loss里？过时了！Bengio团队重磅研究：别再拿KL散度当Loss了，这才是无偏最优解

随着 DeepSeek-R1 的爆火，RLVR 几乎成为了大模型进阶的必修课。从 PPO 到 GRPO，所有算法都在讲同一个故事：在 Maximize Reward 和 Minimize KL Divergence 之间寻找完美的平衡点。

但在落地实操中，我们往往因为“大家都这么做”而忽略了一个关键细节：KL 惩罚项的放置位置。

目前，绝大多数主流开源框架（包括 VeRL, OpenRLHF, SkyRL）为了图省事，默认将 KL 估算器直接写进了 Loss 函数里。这种做法看似殊途同归，实则暗藏玄机。

Mila 实验室（Bengio 团队）最新发表的论文《A Comedy of Estimators》，给这场集体无意识的工程实践泼了一盆冷水。研究指出，这种主流的“KL-in-Loss”写法，在数学原理上会导致梯度估计出现系统性偏差。这种偏差的代价是昂贵的——它会让模型在训练中更容易崩盘，且严重损害模型举一反三的泛化能力。

修复方案简单到令人难以置信：放弃花哨的技巧，回归最原始的逻辑——使用朴素 K1 估算器，并将其放回 Reward 中。

这一简单的“复古”操作，竟然能带来近 20% 的 OOD 性能跃升。本文将为你拆解这项颠覆性发现，帮你纠正潜伏在 RL 训练流程中最隐蔽的工程误区。

扫码添加小助手回复“C798”

免费获取 AAAI 2026 论文合集

一、论文基本信息

• 论文标题：A Comedy of Estimators: On KL Regularization in RL Training of LLMs

• 作者姓名与单位：Vedant Shah, Johan Obando-Ceron 等（主要来自 Mila - 魁北克人工智能研究所，以及蒙特利尔大学、麦吉尔大学等）

• 论文链接：https://arxiv.org/abs/2512.21852

二、主要贡献与创新

1. 理论揭示梯度偏差：深入剖析了 K1 和 K3 两种 KL 估计器在放入“奖励”或“损失”时产生的梯度性质，证明了目前流行的“K3 放损失”配置虽然稳定但梯度有偏。
2. 确立最优配置：通过数学推导和实证发现，最朴素的 K1 估计器放入奖励函数 是唯一能产生稳定且无偏梯度估计的配置，优于复杂的 K3 估计器。
3. 广泛的实验验证：在 Qwen2.5 和 Llama-3.1 等模型上进行了详尽的数学推理任务（MATH）实验，证明无偏估计配置在分布外（OOD）任务上提升尤为明显。
4. 异步训练的稳定性：研究发现 KL 正则化对于稳定大规模异步（Asynchronous）强化学习训练至关重要，能有效防止训练坍塌。

三、研究方法与原理

该论文的核心思路是回归数学本质，推导不同 KL 散度估计公式（K1 和 K3）对参数求导后的梯度表达式，对比其与真实反向 KL 散度梯度的差异，从而找出真正“无偏”的训练配置。

【模型核心原理分析】

1. 强化学习目标函数
LLM 的强化学习通常优化以下目标，即最大化奖励同时最小化与参考模型的反向 KL 散度：

其中是正则化系数。

2. 两个关键的 KL 估计器
由于直接计算序列级别的 KL 很困难，通常分解为 token 级别的估计：

• K1 估计器（Naïve）：最直接的对数似然比。
• K3 估计器（Schulman）：方差较小，具有无偏性（指期望无偏），被认为更优。

3. 两种放置位置与梯度分析
论文详细推导了将上述估计器放入 奖励（Reward） 或 损失（Loss） 时产生的梯度：

• 配置一：K1 放入奖励 (K1-in-Reward)
梯度推导结果为：

结论：这是 无偏（Unbiased） 的梯度估计。它正确地优化了反向 KL 散度目标。
• 配置二：K1 放入损失 (K1-in-Loss)
直接将 K1 加到 Loss 中求导，其期望梯度为 0。
结论：会导致训练不稳定，因为它在期望上等于没有加 KL 正则化。
• 配置三：K3 放入奖励 (K3-in-Reward)
梯度推导显示包含额外的偏差项。
结论：实验显示会导致训练完全 坍塌（Collapse）。
• 配置四：K3 放入损失 (K3-in-Loss) —— 这是 GRPO 等算法常用的配置
其梯度推导结果为：

结论：这是 有偏（Biased） 的。论文指出，这个梯度实际上更接近于优化 前向 KL 散度（Forward KL）或类似蒸馏的目标，而不是原定的反向 KL。虽然它训练稳定，但并不是在优化预期的目标。

四、实验设计与结果分析

1. 实验设置

• 模型：Qwen2.5-7B, Llama-3.1-8B-Instruct, Qwen3-4B-Instruct。
• 数据集：训练集使用 Hendrycks MATH（7500题）；评估集包括 MATH500, MATH （分布内），以及 MMLU 的物理、化学、生物子集（分布外/OOD）。
• 评测指标：Pass@1 准确率，以及 Mean@32 准确率。

2. 训练稳定性对比

论文首先测试了不同配置在训练过程中的稳定性（Pass@1 随训练步数的变化）。

• K1-in-Loss：表现出明显的训练不稳定性，甚至在稍有 Off-policy 的情况下导致模型坍塌。
• K3-in-Reward：导致训练迅速坍塌，准确率跌至 0。
• K3-in-Loss (GRPO常用)：训练曲线平滑，表现稳定。
• K1-in-Reward (本文推荐)：训练曲线平滑，表现稳定。

3. 下游性能对比 (In-domain & Out-of-domain)

论文重点对比了两个“稳定”的配置：K3-in-Loss（有偏） vs K1-in-Reward（无偏）。

评估任务 (Accuracy)	K3-in-Loss (Biased)	K1-in-Reward (Unbiased)	结论
MATH500 (分布内)	0.613	0.640	无偏估计更强
MATH (分布内)	0.337	0.364	无偏估计更强
MMLU Physics (分布外)	0.429	0.508	差距显著扩大
MMLU Chemistry (分布外)	0.373	0.431	差距显著扩大
MMLU Biology (分布外)	0.525	0.647	差距显著扩大

注：以上数据基于 Qwen2.5-7B, 的实验结果。

分析：使用无偏的 K1-in-Reward 在所有任务上都击败了流行的 K3-in-Loss。特别是在**分布外（OOD）**任务上，相对提升高达 19.06%，这说明正确的 KL 正则化有助于保持模型的通用能力。

4. 异步训练（Asynchronous）的影响

在高度异步（Asynchrony level = 10）的设置下（模拟大规模并行训练）：

• 不加 KL：训练不稳定。
• K1-in-Reward：能够有效稳定异步带来的 Off-policy 影响，并保持高性能。

五、论文结论与评价

本论文通过严谨的数学推导和实验，揭穿了 LLM 强化学习中关于 KL 散度的一个“误解”。

主要结论：

1. 流行未必正确：目前社区中广泛使用的将 K3 估计器放入 Loss（如 GRPO 的实现）虽然能跑通且稳定，但它在数学上是有偏的，实际上是在做类似 Logit 蒸馏的前向 KL 优化，而非反向 KL。
2. 简单即是最好：最简单的 K1 估计器（直接算 Log 差值）作为奖励项扣除，能提供无偏的梯度，从而让模型在数学推理和通用知识保持上都取得最佳效果。
3. 泛化能力关键：使用正确的无偏估计器，对模型在未见过任务上的泛化能力保护效果最好。

对实际应用的启示：
对于正在开发 RLHF 或 RLVR（如 Math, Code）系统的开发者，建议检查代码库中 KL 正则化的实现。不要盲目照搬开源库（如 TRL 或某些 GRPO 实现）中将 KL 放入 Loss 的做法，而是应该尝试将 KL 计算为奖励的一部分（Reward Shaping），这可能以零成本带来显著的性能提升。

优缺点分析：

• 优点：理论分析非常扎实，直击当前热门算法（如 DeepSeek-R1 相关的 GRPO）的痛点；实验覆盖了分布内和分布外，说服力强。
• 缺点：主要聚焦于带有验证器（Verifiable Rewards）的推理任务，对于纯偏好对齐（Chat/Preference）任务的讨论相对较少（虽然理论通用）；未详细讨论计算开销的差异（虽然两者差别很小）。