• Jian Hu’s Blogpost

    以下所有实例均可在Jian Hu开发的OpenRLHF

0 - 引言

基于人类反馈的强化学习（Reinforcement Learning from Human Feedback, RLHF）结合近端策略优化（Proximal Policy Optimization, PPO）[1]是一种用于微调大语言模型（Large Language Models, LLMs）的强大方法。该方法使用可靠且高效的PPO算法，结合人类评估者的反馈来提高模型生成响应的质量。然而，使用PPO训练LLMs存在几个挑战。这些包括维持稳定的训练过程，以及实现比直接偏好优化（Direct Preference Optimization, DPO）[2]更好的性能。因此，我们总结了RLHF与PPO的实用训练技巧，以帮助研究人员更轻松地微调LLMs，确保训练稳定性和高性能。

1 - 使用PPO训练LLM的高级技巧

我们介绍三类PPO训练技巧：1) LLM特定技巧，2) PPO特定技巧，以及3)来自最新研究的创新策略。LLM特定和PPO特定的技巧已在各种RL框架[3, 4]中实现并证明有效。然而，最近论文中提出的创新策略在特定任务上的适用性仍有待验证。

1.1 - LLM特定技巧

  • Token级KL惩罚：计算强化学习（RL）模型和监督微调（SFT）模型的响应分布之间的每个token的KL散度[11]。这种散度在训练过程中作为惩罚项被纳入奖励函数。具体来说，每个token的奖励表示如下：
$$ r(s_t, a_t) = \textbf{I}(s_t =[\text{EOS}])r(x,y)-\beta \text{KL}(t) $$$$ \text{KL}(t) = \log({\pi_{\theta_{\text{old}}}(a_t|s_t)^{\text{RL}}}/{\pi^{\text{SFT}}(a_t|s_t)} $$

​

$$t$$

是否为最后一个token的标识函数。

​ 代码链接：utils.py

  • 广义优势估计（Generalized Advantage Estimation, GAE）：GAE[10]是一种TD(λ)回报估计方法，用于估计PPO中的token级奖励。在实践中，我们通常设置$λ=1$，将GAE方法转换为蒙特卡洛估计方法。

代码链接：experience_maker.py

将GAE的λ和折扣因子$γ$都设为1，这可以减少值网络引入的偏差。

代码链接：OpenRLHF/openrlhf/cli/train_ppo_ray.py

  • 添加SFT损失：在PPO中加入额外的监督下一个token预测损失，与KL散度一起，可以保持SFT模型的预存能力。

代码链接：ppo_trainer.py

1.2 - PPO特定技巧

  • 模型初始化：在使用PPO训练LLMs时，必须初始化两个模型：演员模型（actor model）和评论家模型（critic model）[6, 7]。具体来说，使用监督微调（SFT）模型初始化演员模型，使用奖励模型初始化评论家模型可以确保PPO训练的效率。

代码链接：examples/train_ppo.py

  • Adam学习率：演员模型的Adam学习率大约是SFT模型使用的十分之一。例如，在OpenRLHF中，SFT模型的Adam学习率为5e-6，而演员模型为5e-7。此外，评论家模型的Adam学习率约为SFT模型的两倍，示例率为9e-6。

代码链接：train_ppo_llama.sh

  • Mini-batch更新：在学习阶段，PPO实现会打乱训练数据的索引，数据大小为N×M（其中N是重放缓冲区的大小，M是每个响应的长度），并将其分成小批量（mini-batches）来计算梯度和更新策略。

代码链接：ppo_trainer.py

  • 值函数损失裁剪[5]：PPO像裁剪其替代目标一样裁剪值函数[6, 7]。给定$V_{targ} = returns = advantages + values$，PPO通过最小化以下损失来拟合值网络：
$$ Loss_v = \max[(V_{\theta_t} - V_{targ})^2, (\text{clip}(V_{\theta_t}, V_{\theta_{t-1}} - \epsilon, V_{\theta_{t-1}} + \epsilon) - V_{targ})^2] $$

代码链接：loss.py

  • 奖励归一化和裁剪[5]：Bradley-Terry模型倾向于过拟合简单样本而忽视困难样本，导致简单响应对之间的奖励差距比困难对更大[6, 7]。这种变化导致梯度不平衡。奖励归一化（reward normalization）和裁剪（clipping）通过均衡所有样本的奖励分布来缓解这种效应。在实践中，我们采用Z分数归一化方法，$r = (r - \mu) / \delta$，其中$μ$是均值，$δ$是奖励训练数据集的标准差。

代码链接：utils.py

  • 优势归一化：在使用均方损失（mean square loss）训练值网络时，算法对某些大值敏感。归一化优势可以减轻这些大值的影响[5]。在实践中，我们也采用Z分数归一化方法，$r=(r-μ)/δ$，其中$μ$是均值，$δ$是批次内样本的标准差。

代码链接：ppo_trainer.py

  • 较少的训练轮数：在训练策略网络时，限制训练轮数（training epochs）很重要。这个约束可以防止更新后模型的响应分布与旧模型显著偏离，从而避免对离策略（off-policy）样本进行训练。具体来说，我们将训练轮数设置为一。

代码链接：train_ppo_llama.sh

1.3 - 创新策略

在初始训练步骤中冻结演员模型的参数：随着演员更新其策略，响应分布会偏离先前的分布。这导致优势函数（advantage function）估计不准确，可能误导后续的策略更新。因此，冻结演员模型的参数可以确保响应分布保持不变[6]，使评论家模型能够在这个分布上学习准确的估计。

代码链接：ppo_trainer.py

  • 奖励基线（Reward Baseline）[7][8]：Hou等人[7]发现在单个任务内部和不同任务之间的奖励分数存在显著变化。因此，如果不考虑其相关提示词和任务的影响，不能仅凭高奖励就可靠地判定响应质量高。为解决这些缺点，作者引入了参考基线来减轻绝对值的变异性：
$$ r(x, y) = r(x ,y) - r(x, y_{\text{ref}}) $$

其中$y_{\text{ref}}$表示由当前策略对提示词$x$生成的响应。这是一种更新的奖励归一化方法。

  • 拒绝采样微调（Rejection Sampling Fine-tuning）在PPO中的应用：拒绝采样微调包括从模型中采样K个输出，基于奖励模型选择最佳候选项，并使用这个最优候选项来微调模型。Touvron等人[9]采用拒绝采样微调和近端策略优化（PPO）的顺序组合来更新策略。在每个阶段，他们首先使用拒绝采样来更新模型，然后应用PPO。