人类对齐与强化学习—RLHF 全流程与 PPO 原理源码解读

从 deepspeed-chat 开源实现入手讲解 RLHF：强化学习中智能体与环境的交互框架、NLP 场景下的强化学习，以及 Actor、Reference、Critic、Reward 四个模型在 PPO 训练流程中的分工与公式推导。 > [原文链接：图解大模型RLHF系列之：人人都能看懂的PPO原理与源码解读](https://zhuanlan.zhihu.com/p/677607581) **本文直接从一个RLHF开源项目源码入手（deepspeed-chat），根据源码的实现细节，给出尽可能丰富的训练流程图，并对所有的公式给出直观的解释。希望可以帮助大家更具象地感受RLHF的训练流程**。关于RLHF，各家的开源代码间都会有一些差异，同时也不止PPO一种RLHF方式。 ## **一、强化学习概述** ### **1.1 强化学习整体流程**  强化学习的两个实体：**智能体（Agent）与环境（Environment）** 强化学习中两个实体的交互： - **状态空间S**：S即为State，指环境中所有可能状态的集合 - **动作空间A**：A即为Action，指智能体所有可能动作的集合 - **奖励R**：R即为Reward，指智能体在环境的某一状态下所获得的奖励。 以上图为例，智能体与环境的交互过程如下： - 在 $t$ 时刻，环境的状态为 $S_{t}$ ，达到这一状态所获得的奖励为 $R_{t}$ - 智能体观测到 $S_{t}$ 与 $R_{t}$ ，采取相应动作 $A_{t}$ - 智能体采取 $A_{t}$ 后，环境状态变为 $S_{t+1}$ ，得到相应的奖励 $R_{t+1}$ 智能体在这个过程中学习，它的最终目标是：**找到一个策略，这个策略根据当前观测到的环境状态和奖励反馈，来选择最佳的动作。** ### **1.2 价值函数** 在1.1中，谈到了奖励值 $R_{t}$ ，它表示环境进入状态 $S_{t}$ 下的**即时奖励**。**但如果只考虑即时奖励，目光似乎太短浅了**：当下的状态和动作会影响到未来的状态和动作，进而影响到未来的整体收益。所以，一种更好的设计方式是：**t时刻状态s的总收益 = 身处状态s能带来的****即时收益**** + 从状态s出发后能带来的****未来收益**。写成表达式就是 $$ V_{t} = R_{t} + \gamma V_{t+1}    $$ 其中： - $V_{t}$ ：$t$ 时刻的总收益，注意这个收益蕴涵了“即时”和“未来”的概念 - $R_{t}$：$t$ 时刻的即时收益 - $V_{t+1}$：$t+1$ 时刻的总收益，注意这个收益蕴涵了“即时”和“未来”的概念。而 $V_{t+1}$ 对 $V_{t}$ 来说就是“未来”。 - $\gamma$：折扣因子。它决定了我们在多大程度上考虑将“未来收益”纳入“当下收益”。 注：在这里，不展开讨论RL中关于价值函数的一系列假设与推导，而是直接给出一个便于理解的简化结果，方便没有RL背景的朋友能倾注更多在“PPO策略具体怎么做”及“对PPO的直觉理解”上。 ## **二、NLP中的强化学习** 在第一部分介绍了通用强化学习的流程，那么要怎么把这个流程对应到NLP任务中呢？**换句话说，NLP任务中的智能体、环境、状态、动作等等，都是指什么呢？**  回想一下对NLP任务做强化学习（RLHF）的目的：**希望给模型一个prompt，让模型能生成符合人类喜好的response**。再回想一下GPT模型做推理的过程：**每个时刻** $t$ **只产生一个token，即token是一个一个蹦出来的，先有上一个token，再有下一个token**。 复习了这两点，现在可以更好解读上面这张图了： - 先喂给模型一个prompt，期望它能产出符合人类喜好的response - 在 $t$ 时刻，模型根据上文，产出一个token，**这个token即对应着强化学习中的动作，记为** $A_{t}$。因此不难理解，在NLP语境下，强化学习任务的动作空间就对应着词表。 - 在 $t$ 时刻，模型产出token $A_{t}$ 对应着的即时收益为 $R_{t}$，总收益为 $V_{t}$。这个收益即可以理解为“对人类喜好的衡量”。此刻，模型的状态从 $S_{t}$ 变为 $S_{t+1}$，也就是从“上文”变成“上文 + 新产出的token” - 在NLP语境下，智能体是语言模型本身，环境则对应着它产出的语料 这样，就大致解释了NLP语境下的强化学习框架，不过针对上面这张图，可能还有以下问题： **（1）问题1：**图中的下标是不是写得不太对？例如根据第一部分的介绍，$A_{t}$ 应该对应着 $R_{t+1}$ ， $A_{t+1}$ 应该对应着 $R_{t+2}$ ，以此类推？ > 答：说的对。但这里不用太纠结下标的问题，只需要记住在对应的response token位置，会产生相应的即时奖励和总收益即可。之所以用图中这样的下标，是更方便后续理解代码。 **（2）问题2：** 知道$A_{t}$肯定是由语言模型产生的，那么$R_t$，$ V_{t} $ 是怎么来的呢，也是语言模型产生的吗 > 答：先直接说结论， $ A_{t} $ 是由我们的语言模型产生的， $R_t$，$ V_{t} $ 则分别由另外两个模型来产生，在后文中会细说。 **（3）问题3：** 语言模型的参数在什么时候更新？是观测到一个$R_t$，$ V_{t} $ ，就更新一次参数，然后再去产生 $A_{t+1}$ 吗？ > 答：当然不是。只看到某个时刻的收益，就急着用它更新模型，这也太莽撞了。肯定是要等有足够的观测数据了（例如等模型把完整的response生成完），再去更新它的参数。 **（4）问题4：** 再谈谈 $R_t$，$ V_{t} $  吧，在NLP的语境下我还是不太理解它们 - 首先，“收益”的含义是“对人类喜好的衡量” - $R_t$ ：即时收益，指语言模型当下产生token  $A_t$ 带来的收益 - $V_t$： 实际期望总收益（即时+未来），指对语言模型“当下产生token  ，一直到整个response生产结束”后的期收益预估。因为当下语言模型还没产出  后的token，所以只是对它之后一系列动作的收益做了估计，因而称为“期望总收益”。 ## **三、RLHF中的四个重要角色** 本节中，在第二部分的基础上更进一步：更详细理清NLP语境下RLHF的运作流程。 从第二部分中已经知道：生成token $A_t$ 和对应收益 $R_t$，$ V_{t} $ 的并不是一个模型。那么在RLHF中到底有几个模型？他们是怎么配合做训练的？而我们最终要的是哪个模型？  如上图，**在RLHF-PPO阶段，一共有四个主要模型**，分别是： - **Actor Model**：演员模型，这就是**想要训练的目标语言模型** - **Critic Model**：评论家模型，它的作用是预**估总收益 ****$V_{t}$** - **Reward Model**：奖励模型，它的作用是**计算即时收益 **$R_{t}$ - **Reference Model**：参考模型，它的作用是**在RLHF阶段给语言模型增加一些“约束”，防止语言模型训歪**（朝不受控制的方向更新，效果可能越来越差） 其中: - **Actor/Critic Model**在RLHF阶段是**需要训练**的（图中给这两个模型加了粗边，就是表示这个含义）；而**Reward/Reference Model**是**参数冻结**的。 - Critic/Reward/Reference Model共同组成了一个“奖励-loss”计算体系（自己命名的，为了方便理解），综合它们的结果计算loss，用于更新Actor和Critic Model 我们把这四个部分展开说说。 ### **3.1 Actor Model (演员模型)** 正如前文所说，**Actor就是想要训练的目标语言模型。一般用SFT阶段产出的SFT模型来对它做初始化。**  最终目的是**让Actor模型能产生符合人类喜好的response**。所以策略是，先喂给Actor一条prompt （这里假设`batch_size = 1`，所以是1条prompt），让它生成对应的response。然后，再将“prompt + response"送入我们的“奖励-loss”计算体系中去算得最后的loss，用于更新actor。 ### **3.2 Reference Model（参考模型）** **Reference Model（以下简称Ref模型）一般也用SFT阶段得到的SFT模型做初始化，在训练过程中，它的参数是冻结的。** Ref模型的主要作用是防止Actor“训歪”，那么它具体是怎么做到这一点的呢？  “防止模型训歪”换一个更详细的解释是：**希望训练出来的Actor模型既能达到符合人类喜好的目的，又尽量让它和SFT模型不要差异太大**。简言之，**希望两个模型的输出分布尽量相似**。那什么指标能用来衡量输出分布的相似度呢？自然而然想到了**KL散度**。 如图所示： - **对Actor模型**，喂给它一个prompt，它正常输出对应的response。那么response中每一个token肯定有它对应的log_prob结果，把这样的结果记为**log_probs** - **对Ref模型**，把Actor生成的"prompt + response"喂给它，那么它同样能给出每个token的log_prob结果，我们记其为**ref_log_probs** - 那么这两个模型的输出分布相似度就可以用 **ref_log_probs - log_probs **来衡量，可以从两个方面来理解这个公式： 注：可能已经注意到，按照KL散度的定义，这里写成 **log_probs - ref_log_probs **更合适一些。但是如果你看过一些RLHF相关的论文的话，可能记得在计算损失函数时，有一项 Rt−KL散度 （对这个有疑惑不要紧，我们马上在后文细说），即KL散度前带了负号，所以这里我写成 **ref_log_probs - log_probs **这样的形式，更方便大家从直觉上理解这个公式。 现在，已经知道**怎么利用Ref模型和KL散度来防止Actor训歪了**。**KL散度将在后续被用于loss的计算**。 ### **3.3 Critic Model（评论家模型）** **Critic Model用于预测期望总收益 ****$V_{t}$**** ，和Actor模型一样，它需要做参数更新**。实践中，Critic Model的设计和初始化方式也有很多种，例如和Actor共享部分参数、从RW阶段的Reward Model初始化而来等等。我们讲解时，和deepspeed-chat的实现保持一致：从RW阶段的Reward Model初始化而来。 **你可能想问：训练Actor模型我能理解，但我还是不明白，为什么要单独训练一个Critic模型用于预测收益呢？** > 这是因为，当我们在前文讨论总收益 $V_t$ （即时 + 未来）时，我们是站在上帝视角的，也就是这个  $V_t$  就是客观存在的、真正的总收益。但是在训练模型时，就没有这个上帝视角加成了，也就是在 t 时刻，给不出客观存在的总收益  $V_t$  ，只能训练一个模型去预测它。 **所以总结来说，在RLHF中，不仅要训练模型生成符合人类喜好的内容的能力（Actor），也要提升模型对人类喜好量化判断的能力（Critic）**。这就是Critic模型存在的意义。来看看它的大致架构：  deepspeed-chat采用了**Reward模型作为它的初始化**，所以这里也按Reward模型的架构来简单画画它。你可以简单理解成，Reward/Critic模型和Actor模型的架构是很相似的（毕竟输入都一样），同时，它在最后一层增加了一个Value Head层，该层是个简单的线形层，用于将原始输出结果映射成单一的 $V_t$ 值。 在图中，$V_t$ 表示Critic模型对 $t$ 时刻及未来（response完成）的收益预估。 ### **3.4 Reward Model（奖励模型）** Reward Model用于**计算生成token At 的即时收益**，它就是RW阶段所训练的奖励模型，在RLHF过程中，它的**参数是冻结的**。 **你可能想问：为什么Critic模型要参与训练，而同样是和收益相关的Reward模型的参数就可以冻结呢？** 这是因为，Reward模型是站在上帝视角的。这个上帝视角有两层含义： - 第一点，Reward模型是经过和“估算收益”相关的训练的，因此在RLHF阶段它可以直接被当作一个能产生客观值的模型。 - 第二点，Reward模型代表的含义就是“**即时收益**”，你的token  $A_t$ 已经产生，因此即时收益自然可以立刻算出。 **你还可能想问：已经用Critic预测出** $V_t$ **了，而这个** $V_t$ **包含了“即时”和“未来”的概念，那还需要代表“即时”的** $R_t$ **做什么呢？直接用** $V_t$ **不就好了吗？** 为了解答这个问题，先回顾下1.2部分中给出的价值函数： $ V_{t} = R_{t} + \gamma V_{t+1} $ 这个函数告诉我们，当前可以用两个结果来表示 t 时刻的总收益： - 结果1：Critic模型预测的 $V_t$ - 结果2：Reward模型预测的  $R_t$ 和critic模型预测的 $V_{t+1}$ 那么哪一个结果更靠近上帝视角给出的客观值呢？当然是结果2，因为结果1全靠预测，而结果2中的 $R_t$ 是事实数据。我们知道Critic模型也是参与参数更新的，可以用`MSE(上帝视角的客观收益-Critic模型预测的收益)`来衡量它的loss。**但是上帝视角的客观收益是不知道的，只能用已知事实数据去逼近它，所以我们就用** $ R_{t} + \gamma * V_{t+1} $**来做近似。** 这就是 $ R_{t}, V_{t} $ 同时存在的意义。 Reward模型和critic模型非常相似，这里就只给出架构图，不再做过多的说明。关于Reward模型的训练过程，后续有时间也会出个原理和代码解析。  ## **四、RLHF中的loss计算** 到目前为止，已经基本了解了RLHF的训练框架，以及其中的四个重要角色（训练一个RLHF，有4个模型在硬件上跑，可想而知对存储的压力）。在本节中，一起来解读RLHF的loss计算方式。在解读中，会再一次理一遍RLHF的整体训练过程，填补相关细节。在这之后，就可以来看代码解析了。 在第三部分的讲解中，我们知道**Actor和Critic模型都会做参数更新**，所以loss也分成2个： - **Actor loss：** 用于**评估Actor是否产生了符合人类喜好的结果**，将作用于Actor的BWD上。 - **Critic loss**：用于**评估Critic是否正确预测了人类的喜好**，将作用于Critic的BWD上。 ### **4.1 Actor loss** **1）直观设计** 先来看一个直观的loss设计方式： - Actor接收到当前上文 $S_{t}$ ，产出token $A_{t}$（$P(A_{t} | S_{t})$ ） - Critic根据 $S_{t}$, $A_{t}$ ，产出对总收益的预测 $V_{t}$ - 那么Actor loss可以设计为： $actor\_loss =- \sum_{t \in { response\_timestep }} V_{t} \log P (A_{t} | S_{t})$ 求和符号表示只考虑response部分所有token的loss，为了表达简便，先把这个求和符号略去（下文也是同理），也就是说： $$ actor\_loss =-V_{t} \log P\left(A_{t} \mid S_{t}\right) $$ 我们希望minimize这个`actor_loss`。 **这个设计的直观解释是：** - 当 $V_{t}>0$ 时，意味着Critic对Actor当前采取的动作给了正向反馈，因此就需要在训练迭代中提高 $ P(A_{t} | S_{t}) $，这样就能达到减小loss的作用。 - 当 $V_{t}