跳转至

KL 散度约束

详细推导 RLHF 中带 KL 惩罚的修改奖励函数：R_total = R_RM - β * KL(π_θ || π_ref)

1.1 从原始优化目标谈起

RLHF 的终极目标是训练一个语言模型策略 πθπθ，使其能生成符合人类偏好的回答。最直接的想法是最大化奖励模型的期望评分：

image.png

image.png

1.3 转化为无约束的拉格朗日对偶问题

使用拉格朗日乘子法，我们将约束优化转化为无约束的极大化问题。引入拉格朗日乘子 β>0，得到：

image.png

由于 δ 是常数，在优化时不影响梯度，因此等价于最大化：

image.png

这里的 β 就扮演了 KL 惩罚系数的角色。它控制着约束的强度：ββ 越大，惩罚越强，策略越不敢偏离参考模型；β 越小，惩罚越弱，策略越倾向于最大化奖励而可能远离参考模型。

1.4 展开到 token 级别的强化学习

image.png

但在实际 PPO 训练中，由于我们已经在采样阶段获得了具体的 token 序列，我们可以使用该序列下每个 token 的对数概率来近似 KL 散度。于是，对于每个 token at，定义即时 KL 惩罚：

image.png

注意这里是对数比率，不是真正的 KL 散度（KL 需要加期望），但在给定样本下，它提供了无偏的梯度估计。因此，token 级别的总即时奖励变为：

image.png

1.5 推导总结

因此，最终在 PPO 中使用的修改奖励函数为：

image.png

其中 KL 项在实际计算中由所有生成 token 的对数概率比之和近似，除以序列长度或保持为总和，这取决于实现。这个目标完美地平衡了“追求高奖励”与“保持语言流畅性”之间的矛盾。

为什么 KL 惩罚能防止策略崩塌？从信息论和策略梯度的角度分析。

2.1 策略崩塌的本质

策略崩塌（Policy Collapse）是指语言模型在 RLHF 中为了片面追求奖励模型的高分，逐渐丧失其原有的语言能力，输出变得重复、无意义或充满模板化用语。从概率分布的角度看，策略崩塌表现为策略熵的急剧下降，概率质量高度集中在少数几个能获得高奖励的 token 或 token 序列上。从行为上看，它像是一种严重的“过拟合”，过度适应了奖励模型的表面特征。

2.2 信息论视角：KL 散度作为信息约束

image.png

  • 最小描述长度原则：KL 惩罚确保了策略在寻找高奖励输出时，仍然倾向于选择那些在参考模型下“自然”的路径。这符合奥卡姆剃刀原则：如果两个解释（回复）都能获得同样高的奖励，那么选择那个对参考模型来说更“常见”的（即 KL 散度更小的）解释。这避免了策略走向怪异、不自然的解空间。

2.3 策略梯度视角：KL 惩罚稳定了更新方向

image.png

image.png

  • 降低方差：KL 惩罚作为每个 token 的即时成本，为 Critic 提供了密集的、相对稳定的训练信号。这使得价值函数的学习变得更容易，优势估计的方差更低。没有它，Critic 就必须完全依赖那个稀疏且可能有噪声的最终奖励来学习价值函数，其方差会极高，导致训练极其不稳定。

因此，KL 惩罚既是一个防止遗忘语言先验的信息瓶颈，也是一个持续将策略锚定在安全流形的引力场，还是降低 RL 训练方差的稳定器。三者共同作用，有力地防止了策略崩塌。

如何计算序列级别的 KL 散度？它是所有 token KL 的平均还是求和？

3.1 序列级别 KL 散度的定义

image.png

根据定义，序列级别的 KL 散度是两个序列分布之间的 KL 散度：

image.png

这个期望在实际计算中无法穷举所有序列，因此必须通过蒙特卡洛采样来近似。

3.2 基于采样的近似：求和还是平均？

给定我们采样得到的一条具体序列 y，我们可以计算该序列的对数似然比：

image.png

这个和就是对序列级别 KL 散度的一个无偏但高方差的单样本估计。在 PPO 实现中，为了惩罚这个偏离，我们通常会使用这个和，但在聚合多个序列的损失时，会进行平均处理。

具体来说，实际代码中计算的 KL 惩罚损失项是：

image.png

image.png

3.3 为什么实践中倾向于取平均？

  • 超参数鲁棒性：如果直接使用求和作为惩罚项，长序列的惩罚会远大于短序列，这迫使 β 超参数必须随着序列长度的变化而调整，极难调优。取平均后，每个 token 的贡献被归一化，β 就可以设定为“每个 token 允许的偏离成本”，不受序列长度影响。

  • 与损失函数的一致性：PPO 的 Actor 损失通常也是取所有有效 token 损失的平均值。KL 惩罚作为损失项的一部分，取平均可以保持整体梯度量级的一致性。

  • 更加公平的序列间比较：在不同长度的序列间，平均 KL 提供了一个公平的比较基准。一个 10 步的序列和 100 步的序列，如果每步偏离程度相同，其平均 KL 应该相等，这反映了真实的偏离密度。

因此，尽管序列级别 KL 散度的数学形式是求和，但在 RLHF 的 PPO 实践里，我们几乎总是使用 token 级别的平均对数比率来作为 KL 惩罚的有效近似和优化目标。

β 参数（KL 系数）通常的取值范围是多少？过大或过小会有什么现象？

4.1 典型取值范围

在 RLHF 的实践中，β 的常见范围在 0.01 到 0.2 之间。具体取值高度依赖于模型的规模、奖励模型的评分尺度以及 PPO 的其他超参数设置。例如，InstructGPT 论文中使用了 β≈0.02，而一些开源复现实验可能使用 0.05 或 0.1。这个参数的敏感性很高，找到一个合适的 β 是 RLHF 调参中最关键的步骤之一。

4.2 β 过小的现象——奖励黑客与语言崩溃

当 β 过小，KL 惩罚几乎形同虚设，策略优化的重心完全倒向奖励模型。

  • 奖励黑客泛滥：策略会迅速发现并利用奖励模型的各种表面偏见，如偏好超长回答、特定礼貌用语、固定格式等。模型开始生成“驴唇不对马嘴”但能骗得高分的文本。人工抽查时会发现回复变得啰嗦、模板化，或者开始输出无意义的重复词句。

  • 语言能力不可逆退化：因为没有了“保持正常语言分布”的约束力，模型参数会快速偏离 SFT 模型的流形，产生灾难性遗忘。语法错误开始出现，多样性急剧下降，最终可能退化为完全不可读的乱码。

  • 训练曲线异常：奖励得分可能迅速飙升，但 KL 散度也呈指数级增长。这种背离是崩溃的典型前兆。

4.3 β 过大的现象——优化停滞与对齐不足

当 β 过大，KL 惩罚过于严厉，策略被牢牢锁死在参考模型附近。

  • 优化停滞：Actor 几乎不敢越雷池一步，任何试图提升奖励的微小尝试都会因为那一点点偏离而受到严厉惩罚。模型的输出与原始的 SFT 模型几乎没有差异，RLHF 带来的对齐效果（如更安全、更有用）完全体现不出来。

  • 奖励提升极其缓慢：即使奖励模型提示有更好的方向，Actor 也因“不敢动”而无法接近。训练曲线表现为奖励和 KL 散度都长时间保持平坦。

  • 对齐税过高：模型保留了 SFT 的全部知识，但也保留了 SFT 的全部缺陷，如容易产生幻觉、有时会拒绝回答本可回答的问题等。RLHF 的优势荡然无存。

因此，β 的选择是一场精妙的平衡艺术：既要在初期快速拉升奖励，又要在后期稳定时防止奖励黑客。自适应调整 β 就是为了解决这个调参痛点而生的。

如何在训练过程中动态调整 β？基于 KL 目标值的自适应调整方法是什么？

5.1 静态 β 的痛点

使用固定 β 的最大问题在于，RLHF 训练的不同阶段对约束强度的需求完全不同。初期，Actor 距参考模型很近，RL 需要相对宽松的约束来允许有益的探索；中期，模型开始找到提升奖励的方向，但风险也开始累积，需要加强监控；后期，模型接近收敛，需要精细的平衡来避免过优化。一个固定的 ββ 无法适应这种动态变化。

5.2 基于 KL 目标值的自适应调整

image.png

  1. 平滑与边界：为了防止 β 剧烈抖动，可以对其应用指数移动平均进行平滑。同时设定 β 的最小值和最大值，防止它在极端情况下失控。

5.3 其他动态调整策略

  • 线性退火：在训练初期使用一个较小的 β，让其线性增长到一个较大的终值。这相当于先鼓励探索，后收紧约束。

  • 基于奖励变化的自适应：当检测到奖励在数个步内异常飙升时，立即大幅增加 β，这是一种紧急避险机制。

  • 基于困惑度（PPL）的监控：如果生成的文本在通用语言模型上的困惑度异常升高，说明语言能力在退化，此时应立刻增大 β。

通过这些动态调整，我们就能在不需要人工干预的情况下，让 PPO 训练在“奖励提升”和“策略安全”之间自动找到一个动态平衡。

如果 KL 散度过大，除了增大 β，还可以采取什么措施？（如减少 PPO 更新轮数）

当 KL 散度超出安全范围，说明策略正在发生危险的偏移。增大 β 是最直接的“加强引力”手段，但它往往有滞后性。一个完整的防御体系应当是多层次的。

6.1 减少 PPO 更新轮数（K 值）

PPO 允许对同一批采样数据重复利用多次（例如 4 个 epoch）。每次更新，策略都在远离旧策略。如果 KL 散度飙升，首要措施之一就是立刻减少甚至暂停本批数据的重复利用。将训练 epoch 从 4 降为 2 或 1，甚至提前终止当前批次的更新。这直接从源头切断了导致偏离的过拟合过程，是代价最小的紧急措施。

6.2 降低 Actor 的学习率

高学习率是导致策略快速漂移的元凶之一。当检测到 KL 过大时，立即将 Actor 的学习率按比例缩小（例如减半），可以让参数更新的步长变小，将策略稳定在当前位置附近。

6.3 缩小 PPO 裁剪范围 ε

将 PPO 的裁剪率 ε 从 0.2 调整为 0.1 或 0.05，进一步收紧信任区域。这确保即使在同一个 batch 内继续更新，每一次的策略变化也不会过于剧烈。

6.4 回滚到上一个检查点

如果 KL 散度已经大到了出现 NaN 或奖励、文本质量开始崩坏的程度，最稳妥的做法是立刻停止训练，将 Actor 回滚到上一个 KL 散度正常的检查点。重新评估问题原因（例如是否是 RM 漏洞、是否 β 设置过低）后再启动。

6.5 调整奖励模型的评分尺度

如果 RM 给出的奖励在数值上远大于 KL 惩罚项，即使 β 很大也可能不足以平衡。此时可以对 RM 的输出进行 Z-score 标准化，或者对最终奖励进行裁剪，使其数值范围与 KL 惩罚的尺度相匹配，恢复 β 的控制力。

6.6 增强参考模型的先验强度

在极端情况下，可以考虑更换一个更强大或更保守的参考模型，例如使用一个更大规模或经过更强正则化的 SFT 模型，从根上提高“正常语言分布”的引力。

6.7 对 Actor 进行重锚定微调

如果策略已经漂移，但尚未崩溃，可以在正式继续 RL 之前，将 Actor 在通用 SFT 数据上微调几个 epoch，把它先“拉回”正常的语言流形上，然后再调整参数重新开始 RL。

结论：KL 过大是一个综合性信号，需要从优化器、数据利用、模型初始化、奖励尺度等多个维度协同应对，构建“探测-减速-回滚-加固”的纵深防御体系。

KL 惩罚和裁剪目标在 PPO 中都起到约束更新的作用，它们有什么区别和联系？

7.1 两者在优化目标中的位置不同

  • KL 惩罚：它是被直接注入到奖励信号中的。它修改了 RL 问题的目标本身，告诉 Actor：“从此刻起，偏离参考模型的每一步都要付出代价。”因此，它是从目标定义层面施加约束。

  • 裁剪目标：它是 PPO 算法核心的梯度更新规则。它没有改变奖励，而是在使用旧策略数据计算梯度时，通过限制重要性采样比率的范围，防止单次参数更新过大。因此，它是从优化过程层面施加约束。

7.2 约束的维度不同——“绝对”与“相对”

  • KL 惩罚提供绝对锚点：它以固定的 πrefπref 为基准，惩罚任何偏离。无论 Actor 是否更新、更新了多少步，只要它离 πrefπref 太远，就会持续受到惩罚。这确保了长期优化中，策略不会漂移到语言荒原。

  • 裁剪目标提供相对约束：它以生成这批数据时的旧策略 πoldπold 为基准，限制新策略相对这批数据的“变化幅度”。它只关心在这一次小批量更新中，你别跑太远了。下一批数据到来时，新的 πoldπold 又变成了刚更新完的策略。因此，裁剪目标本质上是限制短期变化率的。如果长期来看策略一直在缓慢地朝一个坏方向漂移，裁剪目标自己可能无法感知，因为每步的变化都在允许范围内。

7.3 联系——双重保险

KL 惩罚和裁剪目标构成了一套完美的双保险机制：

  • 裁剪目标保证了平滑、稳定的短期更新。它让我们能安全地多次复用同一批数据，而不用担心一次大跃进就造成崩溃。它是“防冲顶”的。

  • KL 惩罚保证了长期、全局的安全性。它防止了由无数次平滑更新累积而成的、不易察觉的“温水煮青蛙”式漂移。它是“防漂移”的。

没有裁剪，单次更新就可能崩溃；没有 KL 惩罚，策略会在多次平滑更新中逐渐遗忘语言能力，最终变成奖励黑客。两者缺一不可。

是否可以使用反向 KL（KL(π_ref || π_θ)）作为惩罚？效果会有什么不同？

8.1 反向 KL 的定义与性质

image.png

8.3 在 RLHF 中的灾难性后果

如果在 RLHF 中使用反向 KL：

  • 多样性彻底丧失：策略将极度保守，只敢生成参考模型最最确定、最最主流的那一类回答，任何一点创新、任何稍微冷门的表达，即使质量很高，只要参考模型自己没怎么用过，就会被反向 KL 极度严厉地惩罚。

  • 语言能力退化：为了最小化反向 KL，策略可能把所有概率都集中在少数几个“万能”模板词上，导致输出极度单一、重复，这本身就是一种策略崩塌。

  • RL 优化瘫痪：奖励模型试图引导策略走向一些高质量但略微非主流的回答，但这个方向与反向 KL 的惩罚方向严重冲突，导致策略优化陷入僵局，完全无法提升。

结论：正向 KL 鼓励多样性，它允许策略在安全的语言区域内探索；反向 KL 扼杀多样性，它迫使策略向少数模式坍缩。因此，RLHF 必须使用正向 KL 来作为惩罚，以在安全与探索之间保持平衡。

在计算 KL 时，是逐 token 计算还是只在某些特殊位置计算？为什么通常使用 token-level 平均 KL？

9.1 逐 token 计算是绝对必要的

RLHF 中的 KL 惩罚必须逐 token 计算，覆盖模型生成的每一个 token。这有几个根本原因：

  • 细粒度的信用分配：语言生成是一个顺序决策过程，任何一个中间步骤的偏离，都可能导致后续文本走向奇怪的方向。如果只在句子末尾等特殊位置计算 KL，模型就会在生成本文过程中毫无约束，等到最后才被惩罚，但为时已晚，中间的错误决策已经无法纠正。

image.png

  • 与架构的一致性：无论是 Actor、Reference 还是 Critic，它们都是 Transformer 模型，天然支持对输入序列的每个位置输出相应的 logits 或 value。因此逐 token 计算在工程上也是最高效、最自然的。

9.2 为什么是平均而不是求和？

在聚合损失时，我们计算的是序列中所有有效 token 的 KL 的平均值。理由如下：

  • 长度无关的惩罚强度：如果采用求和，长序列的惩罚量级是短序列的数十倍。这会导致模型过度关注如何在长回复上压低 KL，而可能忽视短回复的优化。采用平均，确保了模型每个生成 token 的“偏离成本”是统一的，不受总长度影响。

  • 超参数 β 的普适性：β 作为每个 token 的惩罚系数，在取平均的场景下，它天然具有“每单位偏离所支付的代价”的物理意义，这使得 β 在不同长度、不同 batch size 下都能保持稳定，极大方便了调参。

  • 数值稳定性：平均后的损失值量级更稳定，有益于梯度更新，减少爆炸风险。

9.3 是否有只在特殊位置计算的场景？

在某些分析或调试场景中，可能会统计 EOS 位置的 KL，或者观察 KL 在生成过程中的动态变化曲线，但这仅用于监控和分析，而非作为训练的目标。训练时的 KL 惩罚，永远是基于所有生成 token 的平均值。

如果完全没有 KL 惩罚，完全依靠 PPO 的裁剪，能否阻止模型远离原始策略？为什么通常不够？

image.png

10.2 为什么裁剪自身不足以防止长期漂移

裁剪机制存在一个根本性的盲区：它没有绝对坐标。 想象一艘在雾中航行的船，裁剪机制就像一条铁链，确保每次划桨的幅度都不能太大。但它没有指南针来告诉你原本安全的港口（参考模型）在哪里。

  • 累积性漂移：在 RLHF 中，PPO 会进行成千上万次更新。每一次更新都是基于当时的旧策略采样数据，裁剪只限制相对于那个旧策略的变化。但经过 N 次小步更新后，策略可能已经无声无息地漂移到了一个离最初 SFT 模型非常遥远的地方。裁剪机制本身无法感知这个长期的、累积的偏移，因为它在每一步看来都是“合规的”。

  • 奖励模型的引力陷阱：如果奖励模型对某个畸形的模式（如反复说“当然！”）有系统性高估，那么每一次 PPO 更新，裁剪都会允许策略朝这个方向迈出一小步。无数步之后，策略就会深深陷入这个陷阱。裁剪全程都“同意”了这些微小的变化，因为它只看相对距离，不看绝对距离。

  • 信任区域的原点漂移：每一次更新后，新的策略就成为下一次的“旧策略”。所以，PPO 信任区域的原点是在不断漂移的。如果这个原点自身在奖励模型的引导下走向了不好的区域，信任区域也会跟着它一起漂走。它就像一个没有外部参照的惯性导航系统，误差会随时间无限累积。

10.3 KL 惩罚是如何填补这个空白的

KL 惩罚给系统提供了绝对坐标和外部参照。它以固定的 πref为锚点，在任何时候，任何更新步之后，都能准确度量策略离安全港口有多远。它不关心你是如何一步步走到这里的，它只关心你最终的位置是否正确。当策略试图长期漂移时，KL 惩罚会持续发出警告和惩罚，就像一个一直指向安全港口的指南针，提供了裁剪机制所缺乏的长期、全局的稳定性保障。

因此，裁剪是防止短期崩溃的“减震器”，而 KL 惩罚是防止长期漂移的“导航仪”。两者协同工作，缺一不可，才能保证 RLHF 这艘大船在波涛汹涌的奖励海洋中安全航行。