背景

在LLM的强化学习方法中,KL用于限制Actor模型相对于SFT模型的变化幅度。围绕KL,本文主要讨论以下几点:

  • 要不要加KL?
  • KL加在哪里?放在Loss还是Reward,两者有什么关系,各自有什么效果?
  • KL如果放在Reward里,需不需要考虑未来token,既需不需要折扣累计?
  • KL具体实现方式,用哪种近似计算方法?

符号定义

KL放哪里

KL Penalty放在Loss里

PPO和GRPO把KL Penalty放在了Actor Loss里(这里忽略clip),PPO用Critic计算,而GRPO用采样计算

由于没办法准确计算,所以在具体实现上,PPO和GRPO采取了不同办法。

首先看下PPO的办法,它把KL的计算融入到了 其中,采用GAE表示整个序列的收益,每个token共享。

这里看到,是单独计算的,和KL没什么关系。 但在DeepSeppd-Chat、TRLX等开源PPO实现里,把放到了里,然后再计算,那么KL就会参与Advantage的计算,这样就和原始推导有区别了。那么这些改动影响如何?下一节单独讨论。

再看下GRPO的办法,它采用了k3估计来计算KL。

用k1和k2估计可以吗?

可以看到,k1的话,由于是负的且没有下界,只要一直减小它就行了,模型会学偏。

k2看起来没问题。

回过头来再看下GRPO为什么选择k3,可以看到k3也是有下界的,不会导致像k1那样的问题。至于k2和k3哪个效果好,就得实验验证一下了。

KL Penalty放在Reward里

上一节的公式可以看到,PPO把KL放在Actor Loss里,然后通过推导可以和Advantage放一起。但很多开源实现,是放在里的,这样KL项就会参与Advantage的计算,理论上和KL放在Actor Loss里不等价了。

这里的虽然形式和k1一致,但并不是指k1估计,而是公式推导得出的一项。表示停止梯度回传,因为会参与到Advantage计算中,但Actor Loss并不需要对Advantage求导,所以这里阻断了对KL中的求导,而KL放在Actor Loss里是需要对KL中求导的。

继续分析,既然的形式和k1一样,那么能不能用类似的k2和k3?

接下来,去掉Critic,然后使用计算Advantage。 其中,是第个采样的累计收益。中的已经考虑了未来收益,所以设为即可。 那么是否需要将改成。这里注意,很多论文在Loss公式里把放在了的位置,看起来像,但其实代码里还是要用GAE或其他方式来计算的,论文Loss公式表述的不清楚。所以,这里的讨论我们以代码为准。

根据Policy Gradient基础里的介绍,需要考虑未来token,那么这样看很合理。但是,的累加已经通过GAE形式的隐含进来了。如果在里也累加KL,相当于累加了两次。论文ReMax4实验证明了这两种方法效果接近,只是对于需要把调更小,这可能就是KL累加两次所以才要调小的原因。

参考

  • [1][Training language models to follow instructions with human feedback](http://arxiv.org/abs/2203.02155)
  • [2][DeepSeek-V3 Technical Report](http://arxiv.org/abs/2412.19437)
  • [3][ReMax: A Simple, Effective, and Efficient Reinforcement Learning Method for Aligning Large Language Models](http://arxiv.org/abs/2310.10505)
  • [4][REINFORCE++: A Simple and Efficient Approach for Aligning Large Language Models](http://arxiv.org/abs/2501.03262)
  • [5][DAPO: An Open-Source LLM Reinforcement Learning System at Scale](http://arxiv.org/abs/2503.14476)
  • [6][GPG: A Simple and Strong Reinforcement Learning Baseline for Model Reasoning](http://arxiv.org/abs/2504.02546)