强化学习 (reinforcement learning)人类反馈 (RLHF) 中的近端策略优化 (PPO) 阶段,主要目标是调整语言模型的策略πθ,使其生成能最大化已学习奖励模型(RM)所给预期奖励的响应。然而,仅针对RM分数进行优化会带来很大风险。策略可能迅速进入策略空间中,这些区域根据RM会产生高奖励,但输出内容却可能是无意义的、重复的,或在风格上与初始监督微调 (fine-tuning)(SFT)阶段建立的期望行为不符。这种现象可以看作是策略对奖励模型“过拟合 (overfitting)”,可能借助其不准确性或局限性(一种奖励欺骗),或者简单地忘记了在此阶段之前它所具备的基本语言生成能力。
为减轻这种情况,PPO引入了一个基于库尔巴克-莱布勒(KL)散度的惩罚项。KL散度,表示为 DKL(πθ∣∣πref),衡量两个概率分布之间的差异。在RLHF的背景下,它量化 (quantization)了当前策略 πθ 偏离 参考策略 πref 的程度。通常,这个参考策略是SFT阶段之后获得的模型,我们称之为 πSFT。
当前策略与SFT策略在给定状态(提示)s 和动作(标记 (token))a 上的KL散度计算如下:
DKL(πθ(⋅∣s)∣∣πSFT(⋅∣s))=a∑πθ(a∣s)logπSFT(a∣s)πθ(a∣s)
低KL散度表明当前策略的输出分布与SFT策略的分布相似,而高KL散度则表明有明显偏离。
核心思想是增加PPO目标函数。并非单纯地最大化预期优势(这与奖励相关),我们最大化一个修改后的目标,其中包含与KL散度成比例的惩罚项:
目标≈Et[奖励t]−βDKL(πθ(⋅∣st)∣∣πSFT(⋅∣st))
此处说明:
- Et[奖励t] 表示策略 πθ 获得的预期奖励,通常在实际的PPO损失函数 (loss function)中,使用优势估计 At 进行近似。这一项促使策略生成奖励模型偏好的内容。
- DKL(πθ(⋅∣st)∣∣πSFT(⋅∣st)) 是在时间步 t 生成的标记上,当前策略与冻结的SFT策略之间的KL散度。
- β 是一个控制KL惩罚强度的超参数 (parameter) (hyperparameter)。
这个KL惩罚项作为正则化 (regularization)项。它阻止策略 πθ 在优化过程中过度偏离SFT策略 πSFT。通过惩罚每个标记输出概率分布的大范围变化,它有助于确保模型保留在SFT阶段获得的通用语言流畅性、知识和风格特征,即使它在适应以最大化奖励信号。
实现细节
实际上,在PPO训练循环中:
- 采样一批提示。
- 当前策略 πθ 对这些提示生成响应。
- 对于每个生成的标记 (token),当前策略 πθ 和冻结的参考策略 πSFT 都计算它们在词汇表 (vocabulary)上的各自概率分布(或logits)。
- 计算这两个分布之间每个标记的KL散度。
- 每个标记的KL散度通常在生成的序列和批次上进行平均。
- 这个按 β 缩放的平均KL散度将从奖励信号中减去(或与策略损失和价值损失项一起纳入PPO损失计算中)。
参考策略 πSFT 在整个PPO训练过程中保持固定;它的权重 (weight)不会更新。它作为一个稳定的锚点,代表从初始监督数据集中学习到的行为。
KL系数(β)
KL系数 β 的选择对于平衡探索与稳定性很重要。
- 低 β: 惩罚较弱。策略 πθ 有更大的自由度偏离 πSFT 以最大化奖励。这可能导致更快地向奖励信号优化,但如果策略偏离过远,会增加策略崩溃、不稳定或生成不良文本的风险。
- 高 β: 惩罚较强。策略 πθ 被严格限制,以保持接近 πSFT。这有助于提升稳定性并保留SFT模型的特性,但可能会阻碍策略充分适应奖励信号,可能导致相比SFT模型仅有微小改进。
PPO训练中的策略更新。πSFT 是初始策略。低 β 允许向高奖励区域迈出更大的步长,但有不稳定的风险。高 β 限制步长,保持与 πSFT 的接近。
寻找 β 的合适值通常需要进行实验。此外,自适应KL控制器也常被使用。这些控制器根据每个批次中观察到的KL散度值在训练期间动态调整 β。目标是将实际的 DKL(πθ∣∣πSFT) 保持在预设的目标范围内(例如,保持平均KL为6纳特)。如果观察到的KL超过目标,则增加 β 以加强惩罚;如果低于目标,则减小 β 以允许更多的优化。像Hugging Face的TRL这样的库提供了此类自适应KL控制器的实现。
总之,KL散度惩罚是PPO在RLHF中的一种机制。它防止语言模型策略在优化奖励模型中编码的人类偏好时,偏离初始SFT模型过远。这有助于提升训练稳定性,保留基础模型的理想特征,并提供了一种在奖励最大化和策略约束之间可调节的平衡。