关于 LLMs 中强化学习步骤里 PPO 算法的具体实现

LLMs 场景下强化学习的整体流程

对于一般的强化学习而言，会有一个做动作的策略网络 $\pi$，它根据自己观测的状态 $s_i$，做出动作 $a_i$ 跟环境进行交互，然后环境会返回一个反馈 $r_i$，同时进入到下一个状态 $s_{i+1}$；策略网络再继续观测状态 $s_{i+1}$ 做出下一个动作 $a_{i+1}$ …直到达到最终状态。这样，策略网络和环境的一系列互动后最终会得到一个轨迹(trajectory)，一般在一个轨迹上训练完叫做一个episode:
$$
\tau={s_1,a_1,r_1,s_2,a_2,r_2,\cdots,s_T,a_T,r_T}
$$
在 LLMs 训练的场景下，策略网络就是待微调的 LLMs，它做的事情就是接受 context（状态），执行后返回一系列文本（或者只是文本上的概率分布）（动作）。该策略的动作空间是与语言模型的词汇表相对应的所有 tokens，观察空间是可能的输入token 序列的分布，该分布也十分巨大，奖励函数是偏好模型和policy shift 约束的结合。

PPO 算法的具体实现

形象地来说，就像是学生和老师一样，老师出题，学生根据题目进行答题，老师判断学生回答的正确与否，然后学生根据评判结果重新修改自己对知识的理解。具体来说，LLMs 中 PPO 算法的大体流程如下（伪代码）：

policy_model = load_model()
for k in range(20000):
    # 采样（生成答案）
    prompts = sample_prompt()
    data = respond(policy_model, prompts)
    # 反馈（计算奖励）
    rewards = reward_func(reward_model, data)
    # 学习（更新参数）
    for epoch in range(4):
        policy_model = train(policy_model, prompts, data, rewards)

接下来便分成这三个阶段对该算法的实现进行学习。

采样

采样就是 LLMs 根据提示（prompt）输出回答（response）的过程，或者说是模型自行生产 RLHF 训练数据的过程。

先明确一个概念——策略（policy），它就是 RLHF 中的“学生”，我们则扮演着老师的角色，给出有趣的问题，而模型则会像小学生一样，不断尝试给出答案。模型会对着黑板写下它的答案，有时候是正确的，有时候会有错误。我们会仔细检查每一个答案，如果它表现得好，就会给予它高声赞扬；如果它表现不佳，我们则会给予它耐心的指导和反馈，帮助它不断改进，直到达到令人满意的水平。

policy 由两个模型组成，一个叫做演员模型（Actor），另一个叫做评论家模型（Critic）。它们就像是学生大脑中的两种意识，一个负责决策，一个负责总结得失。其中演员就是我们想要训练出来的 LLMs。在用 PPO 训练它之前，它就是训练好的SFT（Supervised Fine-Tuning）model。输入一段上下文，它将输出下一个 token 的概率分布。评论家是强化学习的辅助模型，输入一段上下文，它将输出下一个 token 的“收益”，也就是从下一个 token 开始，模型能够获得的总奖励（浮点数标量）。从实现上说，评论家就是将演员模型的倒数第二层连接到一个新的全连接层上。除了这个全连接层之外，演员和评论家的参数都是共享的。

下图是 PPO 算法采样阶段的示意图，采样指的是 old_policy 从 prompt 池中抽出 M 个 prompt 后，对每个 prompt 进行语言模型生成的 token 采样：

计算 response 的第1个 token 的概率分布，然后从概率分布中采样出第1个 token
根据第1个 token，计算 response 的第2个 token 的概率分布，然后从概率分布中采样出第2个 token
……
根据前 N-1 个 token，计算 response 的第 N 个 token 的概率分布，然后从概率分布中采样出第 N 个 token

然后基于给定的 prompt，我们可以得到三个输出，假设对每个 prompt ，policy 生成的 token 的个数为 N，那么这三个输出分别是：

response：M 个字符串，每个字符串包含 N 个 token
old_log_probs：演员输出的 $M\times N$ 的张量，包含了 response 中 token 的对数概率 $log⁡(P(token|context))$
old_values：评论家输出的 $M\times N$ 的张量，包含了每次生成 token 时评论家预估的收益

得到这三个输出后，采样阶段就就结束了。这三个输出都是后续阶段重要的输入数据。至此，学生答题部分结束。

反馈

反馈就是老师检查答案的过程，是奖励模型（Reward Model）给 response 打分的过程。下图是 PPO 中 reward model 的使用方法：

从图中我们可以看出，左上角的绿色矩形 reward model 拿到 prompt 和 response，然后输出了分数 score。实际上发生的事情是，prompt 和 response 被拼接成一个字符串，接着被送入到 reward model 中，最后 reward model 计算出了匹配分数。显然，在图中，score 并不是最终的奖励，它和最终的奖励 rewards 之间还隔着一个 reward function 函数。这是因为score只能衡量结果的对错，不能衡量过程的合理性。怎么衡量过程的合理性呢？一种简单粗暴的方法是：循规蹈矩，即为合理。

在一个已有的相对完备的理论系统中，突然出现一个与理论相矛盾的假说，那大家就有必要基于现有的经验给予适当的质疑。因为并非每一次对现有真理的质疑都是正确的（事实上大多数都是错误的），而如果这个假说真的得到验证，那么就是给予肯定和荣誉的时候了。语言模型也是一样，在我们给予最终奖励之前，最好也对它的“标新立异”给予少量的惩罚（即刚刚说的质疑）。

那么我们该怎么做呢？那就是给它立一个规矩，只要它按照这个规矩来，就能获得少量奖励。而这个规矩就是我们在 SFT 阶段已经训练好的语言模型 ref_policy（图中右下角的绿色矩形），或者说是完全还没经过强化学习训练的语言模型。过程合理性奖励的计算方式是这样的：ref_policy 拿到 prompt，然后给 old_policy 生成的 response 的每个 token 计算对数概率，得到一个张量 ref_log_prob。现在假设 old_policy 的演员模型生成了第 i 个 token，此时它应该获得的奖励为 $ref\_log\_prob[i]-old\_log\_prob[i]$。来理解一下这个式子：

ref_log_prob[i] 越高，ref_policy 越认可 old_policy 的输出，说明 old_policy 更守规矩，因此应该获得更高的奖励
old_log_prob[i] 越高，old_policy 获得的奖励反而更低，old_log_prob[i] 作为正则项，可以保证概率分布的多样性

最终，我们将过程合理性奖励和结果正确性奖励合并起来，就得到了最终奖励的计算方式：
$$
\text{reward[i]}=\begin{cases}\text{ref_}\log_\text{prob}[i]-\text{old_}\log_\text{prob}[i],&i<N\\text{ref_}\log_\text{prob}[N]-\text{old_}\log_\text{prob}[N]+\textit{score},&i=N\end{cases}
$$
值得指出的是，结果正确性奖励只有在最后一个 token 上才会使用，这就是上图中 reward function 的计算方法，此外，我们只对 response 计算奖励，在整个反馈阶段，reward_model 和 ref_policy 是不更新参数的。至此，老师反馈阶段结束。

学习

“学习”就是学生根据反馈总结得失并自我改进的过程，或者说是强化优势动作的过程。

如果说前两步的采样和反馈分别是在收集数据X，以及给数据打上标签Y。那么这一步就是在利用数据(X, Y)训练模型。”强化优势动作”是PPO学习阶段的焦点。在深入探讨之前，我们首先要明确一个关键概念——优势。此处，我们将优势定义为“实际获得的收益超出预期的程度”。

在 PPO 算法中计算优势的方法就是实际收益 - 预期收益。对于语言模型而言，生成第 i 个 token 的实际收益就是：从生成第 i 个 token 开始到生成第 N 个 token 为止，所能获得的所有奖励的总和。我们用 return 来表示实际收益，它的计算方式为：$return[i]=\sum_{j=i}^N reward[j]$ 。对于预期收益来说，我们在“采样”阶段提到过，policy 包含演员模型和评论家模型，其中后者是用来预估收益的。其实，当时说的收益 old_values 就是现在我们想要计算的预期收益。评论家会为 response 中的每个 token 计算一个预期收益，第 i 个预期收益记为 old_values[i]。现在，我们可以这样计算生成第 i 个 token 的优势 a 了：$a[i]=return[i]-old\_values[i]$。接下来，我们就要使用优势 a 来进行“强化优势动作”的操作。

所谓“强化优势动作”，即强化那些展现出显著优势的动作。在语言模型中，根据上下文生成一个 token 就是所谓的“动作”。”强化优势动作”表示：如果在上下文（context）中生成了某个 token，并且这个动作的优势很高，那么我们应该增加生成该 token 的概率，也就是说我们可以给演员模型设计一个损失函数，通过优化损失函数来实现“强化优势动作”的目的，这个损失函数定义如下：
$$
\text{actor_loss}=-\dfrac{1}{N}\sum_{i=1}^{N}a[i]\times\frac{p\text{(token[i] | context)}}{p_{old}\text{(token[i] | context)}}
$$
其中：

当优势大于0时，概率越大，loss越小；因此优化器会通过增大概率（即强化优势动作）来减小loss
当优势小于0时，概率越小，loss越小；因此优化器会通过减小概率（即弱化劣势动作）来减小loss
PPO 中，每份数据在每次迭代的学习阶段中都会使用四次，P_old 指的是同一次迭代的还未进入学习阶段的演员模型，因为它不接收梯度回传，因此可以视为一个常量
分母的作用在于当生成某个 token 的概率已经很大了的时候，即便这个动作的优势很大，也不要再使劲增大概率了，也就是步子不要迈得太大
分子部分可以使用当前的演员模型计算得到，分母部分就是采样阶段的 old_log_probs

上面描述的都是演员模型的优化，下面就要介绍评论家的优化了：

我们提到过，评论家会为 response 中的每个 token 计算一个预期收益 values[i]，估计的目标是奖励模型 reward[i] 的累加 return[i]，于是，很自然地，PPO 算法设计了一个损失函数来衡量评论家预期收益和真实收益之间的差距，这个损失函数就是均方差损失：
$$
\text{critic_loss}=\frac{1}{2N}\sum_{i=1}^{N}(\text{values}[i]-\text{return}[i])^2
$$
最终优化 policy 时用的 loss 是演员和评论家的 loss 的加权和：
$$
\text{loss}=\text{actor_loss}+0.1\times\text{critic_loss}
$$
学习阶段的整体流程图如下所示：

完整伪代码

在梳理完全部流程之后我们便能补全最开始给出的伪代码：

policy_model = load_model()
ref_policy_model = policy_model.copy()

for k in range(20000):
    # 采样
    prompts = sample_prompt()
    responses, old_log_probs, old_values = respond(policy_model, prompts)

    # 反馈
    scores = reward_model(prompts, responses)
    ref_log_probs, _ = analyze_responses(ref_policy_model, prompts, responses)
    rewards = reward_func(reward_model, scores, old_log_probs, ref_log_probs)
    
    # 学习
    for epoch in range(4):
        log_probs, values = analyze_responses(policy_model, prompts, responses)
        advantages = advantage_func(rewards, old_values)
        actor_loss = actor_loss_func(advantages, old_log_probs, log_probs)
        critic_loss = critic_loss_func(rewards, values)
        loss = actor_loss + 0.1 * critic_loss
        train(loss, policy_model.parameters())

如何更加高效地训练有偏好的 LLMs

前面总结了 RLHF 中 PPO 算法的具体实现，但是现有的强化学习算法却仍有一些不足：

人工产生的偏好数据集成本较高，很难量产
三个阶段的训练（ SFT -> RM -> PPO ）过程较长，更新迭代较慢
PPO 的训练过程同时存在4个模型（2训练，2推理），对计算资源的要求较高。

因此，当前的改进优化思路主要有以下几个方向：

方式一：使用AI替代人类

这一类方法使用AI模型来替换人工标注偏好数据，或者指导模型训练。

RLAIF（RL from AI Feedback）

该方法可分为监督训练阶段和强化学习阶段两部分。

监督训练阶段，此阶段包括以下步骤：
1. 获得 Helpful 模型对 red teaming提示的响应。因此，在这些情况下，模型的响应可能是有害的。
2. 在提供了一套应该遵守的原则，让 Helpful 模型据此评论自己的响应。
3. 要求 Helpful 模型根据其提供的评论修改其先前的响应。
4. 重复步骤 2 和 3 进行 n 次迭代。
5. 针对来自所有有害提示的响应的所有修订版本微调预训练的 LLM，还包括有用的提示和响应的组合，以确保微调后的模型仍然有用，此模型即 Supervised Learning Constitutional AI (SL-CAI) 模型。
强化学习阶段，此阶段包括以下步骤：
1. 使用在上一步训练得到的 SL-CAI 模型生成针对有害提示的响应对。
2. 使用具有一个原则和一对响应的反馈模型，去选择更无害的响应。
3. 反馈模型的归一化对数概率用于训练偏好模型/奖励模型。
4. 最后，利用上一步训练的偏好模型作为奖励函数，以 RLHF 方式训练 SL-CAI 模型，得到 Reinforcement Learning Constitutional AI (RL-CAI) 模型。

RRHF（Rank Response to align Human Feedback）

该方法首先从多种来源收集某个 prompt 的响应（模型，chatgpt，人为撰写等），接着基于对数概率对响应进行评分，然后通过排名损失来讲这些分数与人类偏好奖励模型或人类偏好标签的分数进行匹配排序。下图中对比了 PPO 算法和 RRHF 算法的区别：

相比 PPO 算法，RRHF的优点有以下几个方面：

仅需要1到2个模型，而PPO需要4个模型，因此 RRHF 算法更加简单高效。
监督微调（SFT）可以被看作是 RRHF 算法的一种特殊形式。
RRHF 算法可以同时被用作语言模型和奖励模型。
RRHF 算法可以在较低的训练难度下拟合奖励模型的偏好，达到 PPO 算法的效果，并且避免了 PPO 算法中的复杂性和不稳定性问题。

方式二：优化微调数据

该类方法的核心在于仅仅通过优质数据集的获取和产生，以训练得到一个效果较好的 SFT 模型，而无需进行 RM 和 PPO 的训练。

LIMA（Less Is More for Alignment）

浅层对齐假说，即一个模型的知识和能力几乎完全是在预训练中学习的，而对齐则是教会它与用户交互时如何选择子分布。如果假说正确，对齐主要有关于学习方式，那么该假说的一个推论是，人们可以用相当少的样本充分调整预训练的语言模型。因此，该工作假设，对齐可以是一个简单的过程，模型学习与用户互动的风格或格式，以揭示在预训练中已经获得的知识和能力。消融实验显示，当扩大数据量而不同时扩大提示多样性时，收益会大大减少，而在优化数据质量时，收益会大大增加。 此外，尽管没有对话实例，LIMA可以进行连贯的多轮对话，而且这种能力可以通过向训练集添加30条手工制作的多轮对话数据而得到极大的提高。

MAYBE ONLY 0.5% DATA IS NEEDED

这篇工作的目的是通过从现有数据中识别出最有价值的核心样本来帮助模型获取下游任务的知识，并仅用少量数据来实现不相上下甚至更好的性能。

方法流程下图所示，潜在空间用三个矩形表示，每个任务代表其中一个颜色系列。具有相同色系但不同色调的点，对应于来自同一任务但来自不同数据集的数据，如 NLI 任务有 5 个数据集，因此有 5 种不同的色调。主要分为以下几步：

将每个句子编码成embedding向量，并进行均值池化和 L2 归一化的预处理。
在潜在空间中，将所有样本点聚类成几个类别。
从这些聚类样本中进行采样，找到原始分布中的核心样本。
使用这些检索到的样本来指导微调 LLM 并进行评估。

方式三：优化训练流程

该类方法通常通过改造模型的训练方式（如只保留 SFT 和 RM），以提高训练效率并减少训练成本。

RAFT（Reward rAnked FineTuning）

该算法的伪代码如下图所示，整个流程分为三个阶段：数据收集、数据排序、模型微调；这三个阶段可以单独实施和执行。因此，只要计算资源和显存允许在某些特定模型上进行 SFT，对齐过程就可以使用 RAFT 完成。在数据收集阶段，需要从 prompt 集合中采样一批样本，然后对于每个 prompt，让大模型分别生成响应，并计算这些样本的 reward。在数据排序阶段，要对这些样本排序，并选择指定百分比的具有最高奖励的样本作为训练样本。最后，在模型微调阶段，要使用这些筛选过的样本对模型进行微调。

DPO（Direct Preference Optimization）

该方法提出了一种使用二进制交叉熵目标来精确优化LLM的方法，以替代基于 RL HF 的优化目标，从而大大简化偏好学习 pipeline。也就是说，完全可以直接优化语言模型以实现人类的偏好，而不需要明确的奖励模型或强化学习。

与现有的算法一样，DPO 也依赖于理论上的偏好模型（如 Bradley-Terry 模型），以此衡量给定的奖励函数与经验偏好数据的吻合程度。然而，现有的方法使用偏好模型定义偏好损失来训练奖励模型，然后训练优化所学奖励模型的策略，而 DPO 使用变量的变化来直接定义偏好损失作为策略的一个函数。鉴于人类对模型响应的偏好数据集，DPO 因此可以使用一个简单的二进制交叉熵目标来优化策略，而不需要明确地学习奖励函数或在训练期间从策略中采样。