纠正强化学习中KL散度梯度估计的常见错误，提升大模型训练稳定性。
On a few pitfalls in KL divergence gradient estimation for RL
- 错误地对KL估计值求导会得到错误梯度，影响训练效果。
- 部分实现忽略序列依赖性，仅计算不完整梯度。
- 提供正确实现方案，适用于大模型强化学习训练场景。
我们指出了在大语言模型强化学习训练中，基于KL散度梯度估计的若干实现陷阱，这些错误出现在多个开源项目和论文中。首要问题是将KL估计值作为损失函数进行反向传播以最小化KL散度，这种做法通常不正确，无法产生期望的梯度。其次，一些实现未考虑估计问题的序列特性，仅能生成部分梯度。我们通过示例性的表格实验和大模型实验验证了这些问题的影响，并给出了正确的KL梯度实现方法。
原文摘要 · Abstract (English)
We point out a few pitfalls in implementing gradient estimation for KL divergence in RL training for LLM, as seen in a number of open source projects and papers. The first major pitfall is to differentiate through the KL estimate as loss functions to minimize KL divergence. We show that such implementations are generally incorrect and do not produce the desired KL gradient. Secondly, we show that some implementations do not account for the sequential nature of the estimation problem and produce a partial gradient at best. We demonstrate the impact of such issues with illustrative tabular and LLM experiments, and show the correct way to implement the KL gradient.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。