抱歉,您的浏览器无法访问本站
本页面需要浏览器支持(启用)JavaScript
了解详情 >

KL 散度是大模型强化学习中常见的正则项。我们希望模型提高任务奖励,同时保留参考模型的部分行为,因此通常会在奖励目标中加入对分布偏离的约束。不过,从目标函数到训练代码,还有几处需要区分的细节。同样使用当前模型与参考模型的对数概率,有的实现将它们作为奖励修正,有的则直接构造可微的损失,两者未必对应相同的参数更新。

其中一个容易忽略的原因是,KL 本身是关于概率分布的期望。模型参数变化时,样本出现的概率也随之变化,而对一批已经采到的样本直接反传,通常只计算固定样本上的函数变化。因此,估计量在数值上是否无偏,与它能否产生所需的梯度,是两个需要分别讨论的问题。进一步考虑自回归生成时,还要计入当前选择对后续前缀分布的影响。

本文从一个二元分布出发,先整理 KL 的方向和常用估计量,再推导奖励修正与直接反传各自对应的梯度,最后将这些结果推广到完整序列。这样既能看清几个常见公式之间的联系,也便于理解它们在 PPO、GRPO 等实现中的适用条件。

参数量通常是我们认识一个大模型时最先看到的数字。不过,知道模型有多少参数,还不足以判断它需要多少显存、一次生成又要做多少计算。以一个 8B 模型为例,用 BF16 保存权重大约需要 16 GB,但运行时还要容纳 KV Cache、中间激活和临时工作区,因此 16 GB 显存未必能完成推理。Mixtral 8×7B 则提供了另一个例子,它总共有约 47B 参数,每个 token 参与计算的参数却只有约 13B。

这些差别都可以从模型实际执行的运算中得到解释。参数量取决于模型包含哪些权重,计算量还与这些权重被使用的次数有关,而显存占用则要进一步考虑运行过程中同时保留了哪些张量。三者有共同的结构基础,也各自依赖不同的条件。

本文从一层 Transformer 的矩阵形状出发,逐步推导这三个量的估算方法,并用 Llama 和 Mixtral 的配置核对结果。我们关心的不只是公式本身,也包括公式中各项的来源,以及哪些实现细节会使实际结果偏离估算。文中的“B 参数”表示十亿个参数;容量使用 GiB 时,按 2302^{30} 字节计算。

搜索文章