作者：暧暧内含光，中山大学 · 核工程与技术
声明：本文只做分享，版权归原作者，侵权私信删除！
原文：https://zhuanlan.zhihu.com/p/646934197

现在有很多效果卓越的基座语言模型，但它们只会自回归式地生成句子，很多时候不能理解人类意图。 为了让语言模型与人类意图“对齐”，基于人类反馈的强化学习（RLHF）是现在比较好的方法。

InstructGPT 这篇文章：Training language models to follow instructions with human feedback^[1] 简要介绍了 RLHF，可以分成三个阶段：有监督微调、训练奖励模型、利用 PPO 算法微调 LLM。

下面就结合论文以及 TRL 库对 RLHF 的代码^[2]实现，看看 RLHF 三阶段分别是如何工作的。

prompt 收集

为什么要把它作为第一部分呢？因为这是整个 InstructGPT 构建数据集的基础。

OpenAI 首先收集了一个 prompt set，这个提示数据集包含各种各样的 instruction，包括 chat, open QA, classification, extract, generation, brainstorming 等任务。

这些 prompt 的来源分为两部分：一部分是标注工人手动写的 prompt；另一部分是用户与初代版本的 InstructGPT 交互时输入的 prompt。

这里有个先后顺序：OpenAI 先让标注工人写了一些 prompt 和相应的答案，用这些很少的数据微调了一个 Demo InstructGPT，开放 Playground 给用户。然后收集用户与模型交互的 prompt，去补充原始的 prompt set。而用户输入的这部分 prompt 其实更能反映真实的应用场景。这就是 Human in the loop，可能也是 AI 产品都要走的路。

接下来，OpenAI 在这个 prompt set 的基础上构建了三个数据集：SFT dataset, RM dataset, PPO dataset. 我们在之后讲到这三个阶段时分别介绍数据集长什么样子。

有监督微调

SFT (Supervised fine-tuning) 数据集是一些 (prompt, answer) pair，在 prompt set 里挑选出一些 prompt，让标注工人手工写答案。它包含了 13k 个 prompt，prompt 来源既有标注工人写的，也有用户真实输入的。

SFT 的技术细节没什么好讲的，就是在预训练模型上利用这些人工标注的数据进一步微调。这里想做一下概念辨析。我其实一直没有搞清楚指令微调 (Instruction fine tuning, IFT) 和有监督微调 (SFT) 的区别。

What Makes a Dialog Agent Useful? Hugging Face 的这篇博文说，IFT 的主要目的是让模型适应、听从人类的指令，比如当 prompt 出现 "summarize" 时，模型就应该知道现在的任务是总结。IFT 可以算作 SFT 的一个子集，或者说先驱步骤。

经过 IFT 后，模型学会了听从指令，但生成的内容却不一定安全可靠。为了提升帮助性、降低有害性，人们继续做 SFT。通过给模型展示无害的、有帮助性的回答，规训模型的生成内容。

SFT and IFT are very closely linked. Instruction tuning can be seen as a subset of supervised fine-tuning. In the recent literature, the SFT phase has often been utilized for safety topics, rather than instruction-specific topics, which is done after IFT. In the future, this taxonomy and delineation should mature into clearer use-cases and methodology.

说说我的理解。IFT 阶段的数据比较好收集，一般是对每个 NLP 任务编写特定的 prompt，然后把相应数据集拿过来即可。而 SFT 的回答是人工编写的，收集成本较高。另外，根据我浅薄的论文阅读量，IFT 在谷歌 Flan-T5 这篇论文中使用较多；而 SFT 在 OpenAI 的论文中被使用。

训练奖励模型

Reward model 的输入是 (prompt+answer)，输出是一个标量，衡量回答的好坏。可以把 SFT 模型最后的分类头去掉，加上回归头，作为奖励模型的初始化，在此基础上微调。但考虑到计算成本和训练稳定性，一般会选择小几号的模型（架构不变，参数更少）。对于 175B 的 SFT，OpenAI 用了 6B 的奖励模型。

训练奖励模型是一个回归任务。一般情况下，数据集应该长这样：input (prompt+answer); label (score given by human). 但一个现实的问题：给定一个 prompt 和回答，很难量化地去衡量这个回答的好坏。这个回答看起来不错，到底应该给 7 分，8 分还是 9 分呢？很难确定。

OpenAI 的做法是：对不同的回答进行排序。具体来说，从 prompt set 里拿一些 prompt 出来，对于每个 prompt，让 SFT 模型生成 K 个回答 (K=9)，人工对这些回答进行排序。有了排序，就有了回答间的相对好坏。一个 prompt 对应九个回答，那么就有 C92=36 对 (good answer, bad answer)

损失函数就定为排序中常见的 pairwise ranking loss。其中 rθ 是奖励模型的输出标量； yw 是一对回答中相对较好的那个； yl 是相对较差的那个回答。

训练过程中，InstructGPT 将每个 prompt 的 36 对回答作为一个单独的 batch。这样一来，只需要进行 9 次前向传播（计算 9 个 r(x,y) ）就可以进行 36 次参数更新，避免了重复的计算。

RM 数据集包含 33k 个 prompt，既有标注工人写的，也有用户真实输入的。注意：每个 prompt 对应 36 对回答，所以实际训练 RM 用到的数据量要比 prompt 大一个数量级。

利用 PPO 微调 LLM

PPO 数据集就是一些 prompt，全部是用户输入的，目的是尽可能贴合实际应用场景。

这一部分有蛮多实现上的细节，但 OpenAI 的论文里写得很简略。我们结合论文和 TRL 库的代码仔细理解一下。下面我会用一些强化学习的术语。

首先要搞清楚强化学习中的“状态”、“动作”在 LLM 中对应什么。把 prompt 作为初始状态，模型自回归地生成 token 时，把这些 token 加在 prompt 上就构成了当前状态。动作就对应生成 token 的过程 ，action space 就是整个词库。需要注意的是，不像很多强化学习环境，这里的状态转移没有随机性——根据上一个状态和当前动作，就可以完全确定下一个状态。

在 PPO 微调阶段，我们有四个模型：

• • SFT LM：第一阶段 SFT 过后得到的模型，它在 PPO 阶段参数冻结，作为 reference model，不进行更新。

• • RL LM：我们要微调的 LLM，初始参数与 ref LM 一致。用强化学习的术语，可以称它为 actor model

• • reward model：第二阶段得到的奖励模型，在 PPO 阶段参数冻结。

• • critic model：用于估算 PPO 中的状态值函数，是一种值函数近似（value function approximation），参数与 RL LM 一起更新

PPO 中通常的做法是让 actor model 与 critic model 共享参数，然后一起训练。此时应该把 policy gradient 的损失函数与 critic model 的损失函数（通常是 MSE）加权求和，作为总的损失。就像 PPO 论文中说的：

If using a neural network architecture that shares parameters between the policy and value function, we must use a loss function that combines the policy surrogate and a value function error term.

问题来了，actor model 与 critic model 是如何共享参数的？直接看源码：

base_model_output = self.pretrained_model(
      input_ids=input_ids,
      attention_mask=attention_mask,
      **kwargs,
)
...
# 取最后一层的隐藏状态
last_hidden_state = base_model_output.hidden_states[-1]
# v_head 是一个输出维度是 1 的全连接层，把每个 token 映射成一个标量
# value shape: (B, L)
value = self.v_head(last_hidden_state).squeeze(-1)

这样我们就了解了，critic model 就是在 actor model (RL LM) 的基础上加了一个全连接层，让每个 token 对应一个标量，作为该 token 对应的状态下，之后能获得的累积（折扣）回报。所以 critic model 的输出形状与输入是一致的，都是 (batch size, seq_len) 。

注：InstructGPT 论文中说 value function (critic model) 和奖励模型是一样的架构（6B），说明它和 actor model（175B）并不共享参数。这里我们先沿着 TRL 的实现思路继续。

PPO 中的奖励函数

PPO 的奖励函数是什么？我一开始以为就是 reward model 的输出，后来发现没那么简单。

上图是 InstructGPT 给出的目标函数，也就是 PPO 中的奖励函数。我们可以先忽略最后一项（我们之后再聊）。第一项 rθ(x,y) 就是 reward model 的输出；第二项 log(πϕRL(y∣x) / πSFT(y∣x)) 是我们额外添加的奖励，或者说惩罚。如果当前的 RL LM 模型与 reference model（也就是 SFT LM）相差过大，那么获得的奖励减少。

注1： x 是 prompt； y 是 RL LM 模型生成的回答； πϕRL(y∣x) 的形状与 y 相同，是序列长度 (L) （假设 batch size 等于 1），代表每个位置 token 的生成概率（如下图所示）。

注2：结合上面公式 (2) 中对 RL 分布的期望，第二项就是 KL 散度。但它和下图 PPO1 公式中的 KL 散度没有任何关系！下图是 PPO policy gradient 的目标函数，其中的 KL 散度是为了限制（重要性采样中）采样策略与当前策略相差不要太大。

对于每个 y ，reward model 只输出一个标量 rθ(x,y) ；但 log(πϕRL(y∣x) / πSFT(y∣x)) 是个和 y 等长的向量（在每个 token 位置上都有 KL 散度）。那么奖励函数怎么计算呢？看下图

只在 log(πϕRL(y∣x) / πSFT(y∣x)) 最后一位加上 rθ(x,y) ，然后这个向量就是总的奖励。

这就是 InstructGPT 论文里说的：

Given the prompt and response, it produces a reward determined by the reward model and ends the episode. In addition, we add a per-token KL penalty from the SFT model at each token to mitigate over optimization of the reward model.

源码求奖励函数的部分：

# 求 KL 散度
kl = self._kl_penalty(logprob, ref_logprob)
non_score_reward = -self.kl_ctl.value * kl
reward = non_score_reward.clone()
# y 中非 padding 的最后一位
last_non_masked_index = mask.nonzero()[-1]
# 最终的奖励 = KL 散度 + reward model 的输出（在 KL 散度最后一位相加）
# score 就是 reward model 的输出
reward[last_non_masked_index] += score

理解了上面的内容，再来看TRL 文档中的示意图，会更加清晰。利用 policy gradient 优化模型，我们后面会讲到，现在只需要关注奖励函数的计算。

优势函数的计算

PPO 中优势函数的计算公式：

其中 δt 就是时间差分误差 (TD error)

比较通用的计算方法是：先计算最后的 AT−1，然后有递推公式 At=δt+γλAt+1 ，倒序计算即可。

源码求优势函数的部分：

# reversed：倒序计算
for t in reversed(range(gen_len)):
   # nextvalues 就是 TD error 中的 V(s_{t+1})，即下一个状态的状态值函数
   nextvalues = values[:, t + 1] if t < gen_len - 1 else 0.0
   # 计算 TD error，这里的 rewards 数组就是上一部分求的奖励函数
   delta = rewards[:, t] + self.config.gamma * nextvalues - values[:, t]
   # 利用递推公式求当前优势函数
   lastgaelam = delta + self.config.gamma * self.config.lam * lastgaelam
   advantages_reversed.append(lastgaelam)
# 由于是倒序计算，需要把列表翻转
advantages = torch.stack(advantages_reversed[::-1]).transpose(0, 1)

优势函数的计算用到了上一部分求的奖励函数，以及 critic model 的输出。所以奖励函数是通过影响优势函数而间接影响了损失函数（梯度）的计算。

注：No discount is applied when estimating the generalized advantage —— InstructGPT InstructGPT 中，应该是令上式中 λ=1

actor model & critic model 的更新

有了优势函数，就可以计算 PPO 的目标函数了。它长这样：

其中， rt(θ) 指（由于重要性采样而引入的）当前策略与采样策略的概率比值。

前面提到，actor model 与 critic model 共享参数与损失函数，一起训练。我们有两部分的损失函数要计算：policy gradient 的损失函数与 critic model 的损失函数。

# ratio 对应当前策略与采样策略的概率比值
ratio = torch.exp(logprobs - old_logprobs)
# 由于是求损失函数，所以在目标函数前加个负号
pg_losses = -advantages * ratio
# clip 部分
pg_losses2 = -advantages * torch.clamp(ratio, 1.0 - self.config.cliprange, 1.0 + self.config.cliprange)
# 求两部分的最大值
pg_loss = masked_mean(torch.max(pg_losses, pg_losses2), mask)

上面是计算 policy gradient 损失函数的源码。

vpredclipped = clip_by_value(
     vpreds,
     values - self.config.cliprange_value,
     values + self.config.cliprange_value,
)
vf_losses1 = (vpreds - returns) ** 2
vf_losses2 = (vpredclipped - returns) ** 2
vf_loss = 0.5 * masked_mean(torch.max(vf_losses1, vf_losses2), mask)

上面是计算 critic model 损失函数的源码。整体上是 MSE 的形式，但是计算了 critic model 的预测 vpreds 以及裁剪过的预测 vpredclipped 分别与 returns 的 MSE，取最大值。这里我有一点没有明白：为什么要最小化 vpreds - returns？（这里的 returns=advantages + old_v_values，相当于我们最小化 vpreds - old_v_values - advantages） 如果有大佬明白，烦请评论区告知一下。

将两部分 loss 求和，一起反向传播，更新 actor model 与 critic model：

loss = pg_loss + self.config.vf_coef * vf_loss
self.accelerator.backward(loss)

整体流程

至此，这张图中的最后一块拼图—— Policy gradients optimize model 也讲完了。

Hugging face 的博客^[3]中有一张流程图，现在看应该会更加清晰：

Technical detail note: The above diagram makes it look like both models generate different responses for the same prompt, but what really happens is that the RL policy generates text, and that text is fed into the initial model to produce its relative probabilities for the KL penalty.

提示：上图中看起来是由 RL LM 和 STF LM 对同一个 prompt 生成分别不同的回答。但实际上是 RL LM 生成回答 y ，然后送入 SFT LM，计算概率，进而求得 KL 散度。

Deepspeed 代码实现的差异

看了一下 Deepspeed 库对于第三阶段 PPO 微调的代码^[4]，与 TRL 的实现有些区别。

上图是 Deepspeed RLHF 三阶段的流程。看起来第三阶段时开始时，actor model 拿 SFT LM 作为初始化；critic model 直接拿了 reward model 作为初始化。这就说明 critic model 与 actor model 不是共享参数，而是分开训练的。

这一段和这一段源码印证了我们的直觉：

self.actor_model.backward(actor_loss)
...
self.critic_model.backward(critic_loss)

Deepspeed 的文档中有这么一段，说他们的训练极其不稳定：

we are not able to update the actor model multiple times after generating experimental data. Therefore, in all of our successful runs, we have setper_device_train_batch_size=per_device_mini_batch_sizeandppo_epochs=generation_batch_numbers=1. This is unexpected for a standard RL training pipeline, and we have tried different methods to overcome this, but all have failed.

按照 PPO 算法，每次按照策略 πθold 收集数据，组成大小为 N 的 data buffer；然后利用 buffer 中的数据，训练 πθM 个 epochs。文档中的那段话的意思是，他们的 M 只能设置为 1，N 只能设置为 1. 也就是每次 πθ 只能更新一次，就必须重新收集数据，否则训练就崩掉了。

猜想一下，会不会是 actor model 和 critic model 分开训练的缘故呢？

我一开始猜想训练不稳定是 actor model 和 critic model 分开训练的缘故，但我后来注意到 InstructGPT 里 value function（也就是这里的 critic model）也是直接拿 reward model 作为初始化：

As previously mentioned, for all PPO models we use a 6B RM and a 6B value function, and the latter is initialized from the former.

所以 critic model 和 actor model 在 InstructGPT 里不是共享参数的。Deepspeed 的实现更贴合 InstructGPT。

Bonus: 奖励函数的另一项

InstructGPT 在 PPO 的奖励函数中还定义了关于预训练损失的一项： Ex∼Dpretrain[log(πϕRL(x))] 。这一项主要是为了保持 RL LM 在标准 NLP 任务上的性能，让它不要“忘本”。

我没有看到 TRL 关于这一项的实现，但 Deepspeed 有相关部分：

def train_unsupervised(self, inputs, unsup_coef):

    outputs = self.actor_model(**inputs, use_cache=False)
    loss = outputs.loss
    self.actor_model.backward(unsup_coef * loss)
    self.actor_model.step()

在参数更新的顺序方面：

for i, (exp_data, unsup_data) in enumerate(zip(exp_dataset, unsup_dataset)):
      actor_loss, critic_loss = trainer.train_rlhf(exp_data)
      ...
      ...
      if unsupervised_training_enabled:
         unsup_loss = trainer.train_unsupervised(unsup_data, args.unsup_coef)

思路很简单：先根据 (x,y)∼DπϕRL 更新 actor 和 critic；如果开启了无监督训练，那么继续更新 actor model.

笔者注：我觉得这个实现的思路可以借鉴。我们好像无法把无监督学习的损失放在奖励函数中，因此只能直接作为损失函数，与 PPO 损失函数是“平行”的。

因此公式 (2) 中第一个期望 （(x,y)∼DπϕRL） 是作为奖励函数，影响优势函数的计算，进而影响 PPO 损失函数；而第二个期望 （x∼Dpretrain） 直接作为损失函数，更新参数。

后来在翻看 InstructGPT 论文附录的时候，发现确实和上面的思路一样。

For each minibatch, we compute the PPO gradients and pretraining gradients in consecutive steps and accumulate them both into the gradient buffers.

Next step of RLHF...

现阶段 RLHF 的训练不稳定性（这是强化学习的通病了）是一个棘手问题。并且 RLHF 的训练数据质量要求比较高，收集成本也是个不小的费用。Anthropic 发布过大小为 160k 的公开数据集 hh-rlhf^[5] 是现在为数不多的 RLHF 训练数据。

另外，PPO 是 2017 年提出的，没有证据表明它相对于其他强化学习算法在 RLHF 中有优势。或许只是 OpenAI 对 PPO 比较熟悉，所以叫了它的作者 Schulman 一起搞 RLHF。所以在强化学习算法的选择上，是一个比较开放的领域。并且在第三阶段 PPO 微调时，reward model 并非必须要保持冻结，它可以和 actor model 交替更新；只不过这需要在微调时及时为新的回答进行排序，人工成本极高。

总之，RLHF 还是个等待各位大佬探索、发掘的领域。

引用链接

[1] Training language models to follow instructions with human feedback: https://arxiv.org/abs/2203.02155
[2] TRL 库对 RLHF 的代码: https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py
[3] Hugging face 的博客: https://huggingface.co/blog/rlhf
[4] Deepspeed 库对于第三阶段 PPO 微调的代码: https://github.com/microsoft/DeepSpeedExamples/tree/b371f8f15ce32cb88bf3aae841e53e364c12aaa5/applications/DeepSpeed-Chat/training/step3_rlhf_finetuning
[5] 公开数据集 hh-rlhf: https://huggingface.co/datasets/Anthropic/hh-rlhf