PPO背景知识
Published:
目测CPM的PPO策略实现与这篇论文一致，所以先看这篇论文
还会带一点Training language models to follow instructions with human feedback（以下用openAI代替）的东西，因为数据处理等实现参考了这边
Reinforcement Learning from Human Feedback
训练一共分为三步，本文主要考虑奖励模型训练和PPO训练
- 监督微调：模型通过模仿人类对话样本学习参与人类对话
- 奖励模型训练：模型根据人类反馈学习对比不同回复的偏好
- 根据奖励模型的近端策略优化（PPO）：模型根据奖励模型的反馈更新，探索发现最优策略
Reward Modeling
奖励模型阶段使用去掉了最后一层unembedding layer(embedding layer的逆矩阵，将嵌入张量转换回token id)的预训练模型，然后在最后的transformer layer上加上了一层输出为一维的线性层，从而为生成的每一个token输出一个评分，回答结束之后的token作为整个句子的评分。分数越高表示样本越好。
奖励模型训练往往使用同一个问题的好坏问答对，然后模型使用好坏问答对的分数差定义损失：
\[\mathcal{L}(\psi)=\log\sigma(r(x,y_w)-r(x,y_l))\]本文中使用模仿学习，给问答对的偏好回答引入自回归LM损失，使模型能够模仿偏好回答。在实现中给LM损失添加的因子\(\beta_{rm}\)，定义损失函数为
\[\mathcal{L}(\psi)=-\lambda\mathbb{E}_{(x,y_w,y_l)\sim D_{em}}[\log\sigma(r(x,y_w)-r(x,y_l))]+\beta_{rm}\mathbb{E}_{(x,y_w)\sim D_{rm}}[\log(r'(x,y_w))]\]但在CPM的奖励模型训练脚本中没有添加后面的自回归LM损失，就是OpenAI论文中简单的分数差损失。
在强化学习时给奖励函数添加了一个KL散度的惩罚项，用以度量学到的强化学习策略和一开始监督微调的模型的区别。一方面可以作为熵奖励，支持策略探索，避免策略过早的收敛到单一模式；另一方面也可以确保强化学习后的策略输出不会明显偏离奖励模型在训练时遇到的样本，避免出现奖励模型无法判断的情况。整体的奖励函数是（奖励模型的输出，不是损失）：
\[r_{total}=r(x,y)-\eta KL(\pi_{\phi}^{RL}(y|x),\pi^{SFT}(y|x))\]Reinforcement Learning
将RL应用于对话生成任务的主要难度来自显著的状态-动作空间。 在本文中，将人类交互考虑为环境。
- 在每个时间点\(t\)，智能体（policy模型）从环境（对话历史）中收到一个状态\(st\)，由到这个时间点为止的所有对话文本组成（人类的和智能体的都有，连问题带到目前为止的回答的token list）。
- 根据模型策略\(\pi\)（就是模型此时的参数），智能体采取动作\(a_t\)（前向计算）生成下一个token。
- 环境根据上一时刻状态和采取的动作返回一个奖励\(r(s_t,a_t)\)（利用奖励模型生成）
- 智能体转移到下一个状态\(s_{t+1}\)
RL的目标是给智能体找到一个最优的行为策略来最大化路径\(\tau=\{s_1,a_1,...,s_T,a_T\}\)上的回报（累计奖励）。一种回报是有限范围未贴现回报\(R(\tau)=\sum_{t=1}^{T'}r(s_t,a_t)\)，只计算固定数量步骤内的奖励之和。另一种回报是无限范围未贴现回报\(R(\tau)=\sum_{t=0}^\infty\gamma^tr(s_t,a_t)\)，计算智能体在其轨迹上获得的所有奖励，\(\gamma\)为衰减因子。
Policy Gradient Methods
策略梯度方法是一种直接优化智能体策略的强化学习技术，中心思想是使用梯度上升(?)方法提升策略。本质上来说，策略梯度方法沿着最大化地提升预期回报的方向调整策略参数。策略通常由参数\(\theta\)组成，记为\(\pi(a\vert s,\theta)\)，代表在状态\(s\)下采取动作\(a\)的概率。更新规则可以表示为
\[\theta\leftarrow\theta+\alpha\nabla_{\theta}J(\theta)\]其中\(J(\theta)\)是遵循策略\(\pi_{\theta}\)的期望回报。策略表现的梯度\(\nabla_{\theta}J(\theta)\)称为策略梯度，通常可以表示为
\[\nabla_{\theta}J(\theta)=\mathbb{E}_{\tau\sim\pi_{\theta}}\left[\sum_{t=0}^T\nabla_{\theta}\log\pi_{\theta}(a_t|s_t)\Phi_t\right]\]其中\(\Phi_t\)可以是任意的回报函数：\(\Phi_t=R(\tau), \Phi_t=\sum_{t'=t}^TR(s_{t'},a_{t'}), \Phi_t=\sum_{t'=t}^TR(s_{t'},a_{t'})-b(s_t)\)，这些都有相同的期望值，只是有不同的方差。
回报利用蒙特卡洛抽样计算，如果回报很受青睐，所有动作被选择的概率都会被强化。这种方法的优点在于没有偏向性，因为完全依据真实反馈。但缺点在于高方差，因为不同的路径可能会由于环境的随机性和policy自身的原因导致不同的回报。
为了降低方差，常用方法是使用优势函数(advantage function)估计来替代原始收益。优势函数\(A(s_t,a_t)\)表示在当前状态\(s_t\)的同样策略下，与动作的平均质量相比，采取指定动作\(a_t\)的收益能好多少（优势能有多大）。因此可以把策略梯度的\(\Phi_t\)表示为
\[\Phi_t=A(s_t,a_t)\]数学语言来说\(A(s_t,a_t)=Q(s_t,a_t)-V(s_t)\)，其中\(Q(s_t,a_t)\)表示在采取特定动作之后的期望回报，而\(V(s_t)\)表示此状态下的动作的平均期望回报。\(Q(\cdot,\cdot)\)只考虑特定的动作，而价值函数\(V(\cdot)\)考虑所有可能的动作的平均回报。在实际情况中使用实际事件的回报（奖励之和）来估计Q函数，但因为未来的奖励可能会非常嘈杂，这种方法会带来很大的方差。减少这种噪音的一种方法是使用价值函数估计未来收益。GAE方法在使用简单的一步时间差分回报和使用完全蒙特卡洛回报之间平衡了偏差和方差。
Generalized Advantage Estimation
TD-\(k\)回报\(\hat{R}_t^k\)由从\(t\)开始往后\(k-1\)步为止的实际奖励和第\(k\)步的预估回报组成：
\[\hat{R}_t^k=r_t+\gamma r_{t+1}+\dotsb+\gamma^{(k-1)} r_{t+(k-1)}+\gamma^k V(s_{t+k})\]考虑\(\hat{A}_t^k=\hat{R}_t^k-V(s_t)\)，给每一个\(r_t\)之后的\(\gamma^ir_{t+i}\)添加对应的价值函数项\(\gamma^iV(s_{t+i})-\gamma^iV(s_{t+i})\)，可以换算得到奖励函数
\[\begin{aligned} \hat{A}_t^k&=\hat{R}_t^k-V(s_t) \\ &=-V(s_t)+r_t+\gamma r_{t+1}+\dotsb+\gamma^{(k-1)} r_{t+(k-1)}+\gamma^k V(s_{t+k}) \\ &=-V(s_t)+r_t+\gamma r_{t+1}+\gamma V(s_{t+1})-\gamma V(s_{t+1})+\dotsb+\gamma^{(k-1)} r_{t+(k-1)}+\gamma^{k-1}V(s_{t+k-1})-\gamma^{k-1}V(s_{t+k-1})+\gamma^k V(s_{t+k}) \\ &=\sum_{l=0}^k\gamma^l\delta_{t+l} \\ \delta_t &=r_t+\gamma V(s_{t+1})-V(s_t) \end{aligned}\]将\(\delta_t\)称为TD误差。\(k\)步优势函数涉及明显的偏差-方差平衡。如果\(k\)很小，则偏差很高，因为优势估计基于较少的步骤，因此在很大程度上取决于价值函数的准确性。如果\(k\)很大，则方差可能很高，因为优势估计涉及对许多噪声奖励进行求和。
为了进行平衡，GAE将优势函数定义为\(k\)步优势的指数移动平均值
\[\begin{aligned} \hat{A}_t^{GAE(\gamma,\lambda)}&=(1-\lambda)(\hat{A}_t^{(1)}+\lambda\hat{A}_t^{(2)}+\lambda^2\hat{A}_t^{(3)}+\dotsb) \\ &=(1-\lambda)(\delta_t+\lambda(\delta_t+\gamma\delta_{t+1})+\lambda^2(\delta_t+\gamma\delta_{t+1}+\gamma^2\delta_{t+2})+\dotsb) \\ &=(1-\lambda)(\delta_t(1+\lambda+\lambda^2+\dotsb)+\gamma\delta_{t+1}(\lambda+\lambda^2+\lambda^3+\dotsb)+\dotsb) \\ &=(1-\lambda)(\delta_t(\frac{1}{1-\lambda})+\gamma\delta_{t+1}(\frac{\lambda}{1-\lambda})+\gamma^2\delta_{t+2}(\frac{\lambda^2}{1-\lambda})+\dotsb) \\ &=\sum_{l=0}^{\infty}(\gamma\lambda)^l\delta_{t+l} \end{aligned}\]\(\gamma, \lambda\)都是算法的超参数。通过控制这俩超参数，GAE可以在高偏差（\(\lambda=0\)）和高方差（\(\gamma=1\)）之间获得权衡。通过GAE，我们可以通过\(\hat{A}_t\)精确预估优势函数\(A(s_t,a_t)\)。从而将策略梯度调整为：
\[\nabla_{\theta}\hat{J}(\theta)=\frac{1}{\vert\mathcal{D}\vert}\sum_{\tau\in\mathcal{D}}\sum_{t=1}^T\nabla_{\theta}\log\pi_{\theta}(a_t\vert s_t)\hat{A}_t\]其中\(\mathcal{D}\)是一个有限的样本batch，后面会使用\(\hat{\mathbb{E}}_t\)来表示\(\frac{1}{\vert\mathcal{D}\vert}\sum_{\tau\in\mathcal{D}}\sum_{t=1}^T\)。
Proximal Policy Optimization
传统的RL方法中，策略梯度原则要求新老策略在参数空间中保持接近，但参数空间中的接近不代表表现上的相似，并且参数上的细微变动可能对策略表现产生剧烈影响。幅度很大且不受限制的更新步骤可能会导致出现策略表现上的坍塌（“坠落悬崖现象”）。PPO和TRPO是两种主要的RL训练方法，基本动机都是采用小而稳定的更新步骤逐渐向最优化方向推动策略更新。
TRPO方法采用KL散度来度量策略的更新幅度，而不是参数的相似程度，要求KL散度保持在可接受范围内：
\[\begin{aligned} maximize_{\theta}\quad&\hat{\mathbb{E}}_t\left[\frac{\pi_{\theta}(a_t\vert s_t)}{\pi_{\theta_{old}}(a_t\vert s_t)}\hat{A}_t\right], \\ s.t.\quad&\hat{\mathbb{E}}_t[KL(\pi_{old}(\cdot\vert s_t),\pi_{\theta}(\cdot\vert s_t))]\leq\delta \end{aligned}\]有两种主要的PPO算法：PPO-Penalty和PPO-Clip，CPM脚本使用了后者。
PPO-Penalty使用惩罚因子\(\beta\)将TRPO的硬性限制修改为一个基于惩罚项的无限制优化问题：
\[\mathcal{L}_{ppo-penalty}(\theta)=\hat{\mathbb{E}}_t\left[\frac{\pi_{\theta}(a_t\vert s_t)}{\pi_{\theta_{old}}(a_t\vert s_t)}\hat{A}_t\right]-\beta KL(\pi_{\theta_{old}}(\cdot\vert s_t),\pi_{\theta}(\cdot\vert s_t))\]PPO-Clip方法使用修剪版的策略比例作为目标，而不是KL散度：
\[\mathcal{L}_{ppo-clip}(\theta)=\hat{\mathbb{E}}_t\left[min\left(\frac{\pi_{\theta}(a_t\vert s_t)}{\pi_{\theta_{old}}(a_t\vert s_t)}\hat{A}_t, clip\left(\frac{\pi_{\theta}(a_t\vert s_t)}{\pi_{\theta_{old}}(a_t\vert s_t)},1-\epsilon, 1+\epsilon\right)\hat{A}_t\right)\right]\]\(\frac{\pi_{\theta}(a_t\vert s_t)}{\pi_{\theta_{old}}(a_t\vert s_t)}\)是新旧策略概率的比值，使用\(\epsilon\)控制新策略与老策略的偏离程度，clip函数将比值限定在\([1-\epsilon,1+\epsilon]\)之间。修剪函数起一个正则化的作用，避免策略更新过于剧烈。
PPO算法中critic model起到价值函数的作用，度量每个状态的期望回报。critic model的学习目标是最小化其预测价值和真实返回价值之间的差异，损失函数通常使用均方差定义：
\[\mathcal{L}_{critic}(\phi)=\hat{\mathbb{E}}_t\left[\Vert V_{\phi}(s_t)-\hat{R}_t\Vert^2\right]\]\(\hat{R}_t\)表示状态\(s_t\)下的真实回报价值，通常使用\(\hat{R}_t=\sum_{l=0}^{\infty}\gamma^lr_{t+l}\)估计。
为了缓解PPO过程中模型潜在的语言技能和知识记忆的衰减，本文将预训练数据也结合到RL阶段中，PPO-ptx方法将目标函数进行了结合：
\[\mathcal{L}_{ppo-ptx}(\theta)=\mathcal{L}_{ppo-clip}(\theta)+\lambda_{ptx}\mathbb{E}_{x\sim\mathcal{D}_{pretrain}}\left[\log(\pi_{\theta}^{RL}(x))\right]\]整个PPO算法流程如下：
以下步骤循环\(n\)次
- 通过执行策略模型收集路径集合\(\mathcal{D}_n=\{\tau_i\}\)
- 计算预期回报（rewards-to-go）\(\hat{R}_t\)
- 使用当前价值函数(critic model)计算优势函数估计\(\hat{A}_t\)
- 通过最大化PPO-penalty/clip/ptx目标函数来更新策略 \(\theta_{n+1}=\arg max_{\theta}\mathcal{L}_{ppo-clip}(\theta_n)\)
- 通过均方差回归更新价值函数(critic model) \(\phi_{n+1}=\arg min_{\phi}\mathcal{L}_{critic}(\phi_n)\)
也就是说原来一直没想明白的是有两个模型需要更新
Reward Modeling for Helpfulness and Harmless
这里是一部分奖励模型的训练细节，奖励模型是因为要求人工在每次模型更新前提供反馈不太及时。
| task | model | dataset | volume | train | test | validate | ||||
|---|---|---|---|---|---|---|---|---|---|---|
| English | LLaMA-7B | HH-RLHF | all | 168.5k | all | 160k | all | 1k | all | 7.5k |
| helpful | - | helpful | 118k | helpful | 0.7k | helpful | rest | |||
| harmless | - | harmless | 42k | harmless | 0.3k | harmless | rest | |||
| Chinese | OpenChineseLLaMa | human annotate | all | 39k | all | 30k | all | 3k | all | 6k |
| helpful | 31k | helpful | 24k | helpful | 2.4k | helpful | 4.6k | |||
| harmless | 8k | harmless | 6k | harmless | 0.6k | harmless | 1.4k |
英文模型在文中用的是LLaMA-7B，使用了HH-RLHF数据集的160k问答对数据作为训练集，分为118k的helpful样本和42k的harmless样本。然后从数据集剩下的8.5k数据中随机采样了0.7khelpful样本和0.3kharmless样本，合计1k样本作为测试集。其余的作为验证集。
中文模型用的是OpenChineseLLaMa，人工标注了139k个样本对，包含31k个helpful样本和8k个harmless样本。随机采样了24k个helpful样本和6k个harmless样本作为训练集，从剩下的数据里随机收集2.4k个helpful和0.6k个harmless样本作为测试集，其余的作为验证集。
学习率为\(5e-6\)，前\(10%\) step作为warmup。使用了动态批次方法，平衡每个batch里面的token数量使训练过程尽可能高效稳定。batch size根据batch中的token数量在\([4,128]\)之间变化。固定训练1000步，整个训练集大概跑了1.06个epoch。设置\(\beta_{rm}=1\)，即LM loss在整个训练过程中都完全参与。
Exploratio of PPO
reference model和policy model都用了7B SFT model，根据OpenChineseLLaMA在1M过滤后的指令数据上微调2个epoch得到。学习率设为\(9.5e-6\)，使用consine学习率调度器，学习率最终跌到峰值的\(10%\)。全局batch size为1024。critic model和reward model都用reward model初始化。
在人工构建的HH数据集上进行训练，数据集包含8k harmless问题和20k helpful问题。训练时固定训练步数而不是epoch数。所有实验中设定从环境中采样的batch size为128，训练policy model和critic model的batch size为32。policy model和critic model的学习率设定为\(5e-7\)和\(1.65e-6\)，前10%步作为warmup。
训练使用了zero2和gradient checkpoint技术节约显存，用了8个80g的A100，1TB存储和128CPU。