大模型RLHF训练中的PPO算法细节

虽然了解大模型训练中的RLHF训练，但是都是有点不够深刻，特别是PPO算法的细节。

看到一篇好文章，转载并重新编辑，加入个人理解，以便日后查阅。有兴趣可以参考原文

强化学习

强化学习属于机器学习的一个分支，区别于有监督学习。关键点在于：

无监督，没有标签，通过试错和奖励来优化行为和策略。
与环境有交互。（抽象概念，不要深究）
没有明确的反馈（无标签），反馈是通过奖励信号传递的，可以是延迟的，需要考虑长期回报。

强化学习的简化图：

强化学习的两个实体：智能体（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}$。

智能体在这个过程中学习，它的最终目标是：找到一个策略，这个策略根据当前观测到的环境状态和奖励反馈，来选择最佳的动作。

价值函数

奖励$R$是一个标量，但是在实际问题中，一个动作，既有即时奖励，也要考虑长期回报。为了解决这个问题，引入了价值函数的概念。

$V_t = R_t + \gamma V_{t+1}$

其中：

$V_t$：表示在$t$时刻的状态$S_t$下的价值（包含了即时和未来的奖励）。
$R_t$：$t$时刻的即时收益。
$\gamma$：是折扣因子，用于平衡当前奖励和未来奖励。

这里最需要注意的是：$V_{t+1}$同样包含了现在和未来的奖励，但是对于$V_{t}$来说，它就相当于未来潜在收益。

NLP和强化学习

这里的NLP是特质生成模型。

生成模型的推理执行过程：给模型一个prompt，让模型能生成符合人类喜好的response。再回想一下gpt模型做推理的过程：每个时刻$t$只产生一个token，即token是一个一个蹦出来的，先有上一个token，再有下一个token。

结合上面的图，分解一下这个过程：

智能体就是生成模型。
在$t$时刻，有上下文context（$S_t$），模型产出一个token，对应RL中的动作，记为$A_t$。动作空间就是词表。
在$t$时刻，有了$A_t$动作，即时收益为$R_t$，总收益为$V_t$（注意二者不一样）。对于生成模型，收益是什么？人类喜好。
状态变化，$S_{t+1}$变为$S_t$和新生成的token。
忽略图中的下表，主要理解过程和对应的东西。

$A_t$是产出一个新token，$S_t$是词表空间，$R_t$和$V_t$是什么？答案是通过模型产生的分数，这里不要在意命名，你叫评价模型，叫奖励模型都行，只不过是两个打分模型而已。

记住，到此已经有了3个模型了啊，$A_t$模型表示智能体的动作，$R_t$和$V_t$是两个打分模型，分别表示即时奖励和未来长期奖励。

还有一个重要的点：不是生成一个token，也就是有一个动作，我们就要计算奖励、打分，可以等生成模型回答完毕（也就是eos token）再打分。

RLHF中的4个模型

OpenAI的示意图：

RLHF中使用的模型示意图：

现在大家都知道PPO有4个模型，上面我们说了3个，还有1个，这里将4个模型都列出来：

Actor Model：PPO训练的模型，也是我们最终要用于应用的模型。
Critic Model：懒得翻译了，我以前就是太纠结这个名字，一直理解不上去。实际上就是一个整体收益的打分模型，也是上面的$V_t$。
Reward Model：即时收益$R_t$。
Reference Model：这个模型是额外增加的，主要是在RLHF阶段给语言模型增加一些“约束”，防止语言模型训歪（朝不受控制的方向更新，效果可能越来越差）。这个看Loss就能明白了。

哪些模型需要更新参数？

Actor Model 和 Critic Model。Actor肯定是很好理解的，所以不多说了，Critic Model 为什么也要更新？主要是这里存在一个难点：怎么评估总体收益呢？我们自己随口一说评估总体收益，但是这个是很难的，因为没有真的标签（有监督）。所以我们需要通过一个模型来判断，而且更重要的是，这个判断模型的能力，要不断提升能力，才能做好这件事。

Actor Model

我们的最终目的是让Actor模型能产生符合人类喜好的response。所以我们的策略是，先喂给Actor一条prompt （这里假设batch_size = 1，所以是1条prompt），让它生成对应的response。然后，我们再将“prompt + response”送入我们的“奖励-loss”计算体系中去算得最后的loss，用于更新actor。

Reference Model

Reference Model一般也用SFT阶段得到的SFT模型做初始化，在训练过程中，它的参数是冻结的。Ref模型的主要作用是防止Actor”训歪”，那么它具体是怎么做到这一点的呢？

“防止模型训歪”换一个更详细的解释是：我们希望训练出来的Actor模型既能达到符合人类喜好的目的，又尽量让它和SFT模型不要差异太大。简言之，我们希望两个模型的输出分布尽量相似。那什么指标能用来衡量输出分布的相似度呢？我们自然而然想到了KL散度。

简单来说就是防止模型“高分低能”，过拟合到乱七八糟但是得分高的回答上。

关于KL散度和ref模型的计算，这里不需要展开，网上资料特别多。

Critic Model

前面已经讲了这个模型的作用以及为什么要更新参数。简单讲一下这个模型怎么训练的：一般都是采用了Reward模型作为它的初始化，所以这里我们也按Reward模型的架构来简单画画它。你可以简单理解成，Reward/Critic模型和Actor模型的架构是很相似的（毕竟输入都一样），同时，它在最后一层增加了一个Value Head层，该层是个简单的线形层，用于将原始输出结果映射成单一的$V_t$值。

Reward Model

计算即时收益，也没啥好讲的，提前训练好的（RLHF第二阶段做的事情），这里重点讲一下为啥reward模型不需要更新参数呢？

其实我觉得不要深入去纠结，我感觉PPO这么做的原因就是为了引入一个客观的、绝对的标准。这个模型最重要的区别在于，它只关心当前这个response的好坏。Critic隐含了综合考虑所有response的好坏的含义（需要才需要更新参数）。

RLHF的Loss计算

我已经看过太多次了，所以不想重新写这个了，直接看原文或者网上搜一下就知道了。注意每一项对应ref、critic、reward模型，结合前面讲解的各个模型的作用，应该能很好地理解这个Loss的含义。