大模型 | 通俗理解RLHF基础知识以及完整流程

本文基于 huggingface 源码，对 RLHF 的实现过程做一个比较通俗的讲解。

---

一、RLHF基础知识

RLHF的核心就是4个模型之间的交互过程

• Actor model：传统的语言模型，最后一层网络是 nn.Linear(hidden_size, vocab_size)

• Reference model（不参与训练）：Actor_model的一个复制

• Reward model（不参与训练）：

  • 将传统语言模型的最后一层网络，由 nn.Linear(hidden_size, vocab_size) 替换成 nn.Linear(hidden_size, 1)，也就是说该模型输出的是当前token的得分，而不是对下一个token的预测
  • 输入是prompt + answer， 输出是answer中每个token对应的值，answer中最后一个token对应的值即为这条语料的reward

• Critic model：Reward_model 的一个复制

二、强化学习基础知识

很多NLP出身的同学，经常会因为强化学习的基础概念模糊，导致长期对 RLHF 一知半解，这里我用几个例子来做帮助大家更好的认知。

1. 大模型生成完整answer的过程，视为PPO的一次完整的交互，reward_model的打分便是这次交互的reward；

2. 大模型每生成一个token，视为PPO中的一步；

3. 假设一个汉字等价为一个token。

prompt：中国的首都是哪里？answer：首都是南京

• Reward = Reward_model(‘首都是南京’)，如果我们有一个较好的reward_model，这里大概率会输出一个负数，例如-10；

• Q(2, ‘是’) = Q(‘首都’，‘是’) ，意思是在’首都’这个state下，下一步action选择’是’这个token，所能获得的reward，显然这会是一个较大的值；

• V(4) = V(‘首都是南’)，意思是在’首都是南’这个state下，能获得的reward，显然这会是一个较小的值。

上面的例子也告诉我们，语言模型的reward，只有看到结束才能确定。有时候一个token预测错误，整个句子的reward都不会很大。

三、RLHF完整流程

有了RLHF 和 RL 的基础知识后，我们来介绍每个模型的作用：

• Reward_model 负责给 LLM 生成的句子打分

• Actor_model 就是我们要优化的 LLM

• Critic_model 负责计算Actor_model的状态动作值矩阵，也就是上面提到的Q函数（Reward模型只负责给最后一个token打分，给之前token打分的重任靠Critic_model 完成）

• Reference_model 是一个标杆，为的是让我们的Actor_model在训练时不要偏离原始模型太远，保证其不会失去原本的说话能力

RLHF的第一个环节：让模型生成答案，并对其打分

• 给定 batch_size 条 prompt

• 调用actor_model生成answer，并进行token化，得到一个 B * L 的矩阵；

• reward_model 对answer进行打分，得到一个 B * 1 的矩阵；

• critic_model 对每个token进行打分，得到一个 B * L 的矩阵；

• actor_model 和 reference_model 对生成的句子进行一遍正向传播，保存output.logits，得到两个 B * L * V 的矩阵

• 利用gather_log_probs() 函数，只保存目标token的logit值，得到两个 B * L 的矩阵

{  
    'prompts': prompts,  
    'input_ids': seq,  
    "attention_mask": attention_mask  
    'logprobs': gather_log_probs(logits[:, :-1, :], seq[:, 1:]),        # batch_size  * (seq_len - 1)  
    'ref_logprobs': gather_log_probs(logits_ref[:, :-1, :], seq[:,1:]), # batch_size  * (seq_len - 1)  
    'value': values,                                                    # batch_size * seq_len  
    'rewards': reward_score,                                            # torch.Size([batch_size])  
}  
  
def gather_log_probs(logits, labels):  
    log_probs = F.log_softmax(logits, dim=-1)  
    log_probs_labels = log_probs.gather(dim=-1, index=labels.unsqueeze(-1))  
    return log_probs_labels.squeeze(-1)

RLHF的第二个环节：修正reward

前面提到，我们不能让 actor_model 偏离 reference_model 太远，因此我们要给rewards矩阵添加一个惩罚项，compute_rewards() 函数的返回是：每个token修正后的rewards：

• 最后一个token的计算方法是 Reward_score + KL_penalty

• 前面的所有的token 的计算方法是 0 + KL_penalty (除了最后一个token，前置token的reward初始值都是0，但是要加上惩罚项）

结合代码看的时候，要始终记住这个变换：

prompts = inputs['prompts']  
log_probs = inputs['logprobs']  
ref_log_probs = inputs['ref_logprobs']  
reward_score = inputs['rewards']  
values = inputs['value']  
attention_mask = inputs['attention_mask']  
seq = inputs['input_ids']  
start = prompts.size()[-1] - 1  
action_mask = attention_mask[:, 1:]  
old_values = values  
old_rewards = self.compute_rewards(prompts, log_probs, ref_log_probs, reward_score, action_mask)  
ends = start + action_mask[:, start:].sum(1) + 1  
  
# 计算reward  
def compute_rewards(self, prompts, log_probs, ref_log_probs, reward_score, action_mask):  
    kl_divergence_estimate = -self.kl_ctl * (log_probs - ref_log_probs)  
    rewards = kl_divergence_estimate  
    start = prompts.shape[1] - 1  
    ends = start + action_mask[:, start:].sum(1) + 1  
    reward_clip = torch.clamp(reward_score, -self.clip_reward_value, self.clip_reward_value)  
    batch_size = log_probs.shape[0]  
    for j in range(batch_size):  
        rewards[j, start:ends[j]][-1] += reward_clip[j]  
  
    return rewards

RLHF的第三个环节：计算优势函数和Q函数

我们再看一下这个公式：

优势函数（Advantage Function）在强化学习中是一个非常关键的概念，通常用于评估在特定状态下采取某个动作比遵循当前策略（Policy）更好或更差的程度。优势函数的主要用途是优化策略，帮助模型明确地了解哪些动作（哪个Token）在当前状态（已生成的token）下是有利的。

get_advantages_and_returns() 函数根据第二个环节修正后的 rewards 和 values 计算优势函数，有两个返回值：

• advantages矩阵
• returns矩阵，等价于 advantages + values，也就是Q函数

PPO论文中Advantage函数的计算公式

batch = {'input_ids': seq, "attention_mask": attention_mask}  
advantages, returns = self.get_advantages_and_returns(old_values, old_rewards, start)  
  
## 优势函数的返回  
def get_advantages_and_returns(self, values, rewards, start):  
    # Adopted from https://github.com/CarperAI/trlx/blob/main/trlx/models/modeling_ppo.py#L134  
    lastgaelam = 0  
    advantages_reversed = []  
    length = rewards.size()[-1]  
    for t in reversed(range(start, length)):  
        nextvalues = values[:, t + 1] if t < length - 1 else 0.0  
        delta = rewards[:, t] + self.gamma * nextvalues - values[:, t]  
        lastgaelam = delta + self.gamma * self.lam * lastgaelam  
        advantages_reversed.append(lastgaelam)  
    advantages = torch.stack(advantages_reversed[::-1], dim=1)  
    returns = advantages + values[:, start:]  
    return advantages.detach(), returns

RLHF的第四个环节：更新Actor模型

利用最新的actor模型，重新估算一遍语言模型目标token的logits，然后利用advantages矩阵进行loss计算：

• 输入是新的actor模型的语言模型logits，旧的actor模型的语言模型logits，advantages矩阵

• 在clip_loss，和原始loss之间，选择一个最小的loss进行返回

重要性采样
PPO论文中Actor模型loss

注意上文的一句话，“利用最新的actor模型”，这里涉及到一个重要的概念：重要性采样！

简单来说，我们的 Actor_model 只要训了一条语料，就会变成一个新的模型，那也就是说：我们在第一个环节所构造的语料都无法使用了，因为现在的 actor_model 已经无法生成出之前的answer。

重要性采样的变换公式：

因此，我们是在用另外一个模型的模拟轨迹，来优化我们当前的模型。利用上述公式，我们可以完整这样的近似转化操作，这就是重要性采样的简单理解。

这里不懂也无所谓，就当是引入了一个新的系数来修正 reward 即可。log_ratio = (logprobs - old_logprobs) * mask 这一行代码对应着重要性采样的修正实现。

batch = {'input_ids': seq, "attention_mask": attention_mask}  
actor_prob = self.actor_model(**batch, use_cache=False).logits  
actor_log_prob = gather_log_probs(actor_prob[:, :-1, :], seq[:, 1:])  
actor_loss = self.actor_loss_fn(actor_log_prob[:, start:], log_probs[:, start:], advantages, action_mask[:, start:])  
self.actor_model.backward(actor_loss)  
self.actor_model.step()  
  
## loss的计算  
def actor_loss_fn(self, logprobs, old_logprobs, advantages, mask):  
    ## policy gradient loss  
    log_ratio = (logprobs - old_logprobs) * mask  
    ratio = torch.exp(log_ratio)  
    pg_loss1 = -advantages * ratio  
    pg_loss2 = -advantages * torch.clamp(ratio, 1.0 - self.cliprange, 1.0 + self.cliprange)  
    pg_loss = torch.sum(torch.max(pg_loss1, pg_loss2) * mask) / mask.sum()  
    return pg_loss

RLHF的第五个环节：更新Critic模型

同理，利用最新的critic模型，重新估算一遍Q矩阵，然后利用returns矩阵（其实就是真实的Q矩阵）进行loss计算：

• 输入是新的critic模型计算的Q矩阵，旧的critic模型计算的Q矩阵，returns矩阵
• 在clip_loss，和原始loss之间，选择一个最小的loss进行返回

value = self.critic_model.forward_value(**batch, return_value_only=True, use_cache=False)[:, :-1]  
critic_loss = self.critic_loss_fn(value[:, start:], old_values[:,start:], returns, action_mask[:, start:])  
self.critic_model.backward(critic_loss)  
self.critic_model.step()  
  
## loss的计算  
def critic_loss_fn(self, values, old_values, returns, mask):  
    values_clipped = torch.clamp(values, old_values - self.cliprange_value, old_values + self.cliprange_value)  
    if self.compute_fp32_loss:  
        values = values.float()  
        values_clipped = values_clipped.float()  
    vf_loss1 = (values - returns)**2  
    vf_loss2 = (values_clipped - returns)**2  
    vf_loss = 0.5 * torch.sum(torch.max(vf_loss1, vf_loss2) * mask) / mask.sum()  
    return vf_loss

---

最后分享

AI大模型作为人工智能领域的重要技术突破，正成为推动各行各业创新和转型的关键力量。抓住AI大模型的风口，掌握AI大模型的知识和技能将变得越来越重要。

学习AI大模型是一个系统的过程，需要从基础开始，逐步深入到更高级的技术。

这里给大家精心整理了一份全面的AI大模型学习资源，包括：AI大模型全套学习路线图（从入门到实战）、精品AI大模型学习书籍手册、视频教程、实战学习、面试题等，资料免费分享！

1. 成长路线图&学习规划

要学习一门新的技术，作为新手一定要先学习成长路线图，方向不对，努力白费。

这里，我们为新手和想要进一步提升的专业人士准备了一份详细的学习成长路线图和规划。可以说是最科学最系统的学习成长路线。

2. 大模型经典PDF书籍

书籍和学习文档资料是学习大模型过程中必不可少的，我们精选了一系列深入探讨大模型技术的书籍和学习文档，它们由领域内的顶尖专家撰写，内容全面、深入、详尽，为你学习大模型提供坚实的理论基础。（书籍含电子版PDF）

3. 大模型视频教程

对于很多自学或者没有基础的同学来说，书籍这些纯文字类的学习教材会觉得比较晦涩难以理解，因此，我们提供了丰富的大模型视频教程，以动态、形象的方式展示技术概念，帮助你更快、更轻松地掌握核心知识。

4. 大模型项目实战

学以致用 ，当你的理论知识积累到一定程度，就需要通过项目实战，在实际操作中检验和巩固你所学到的知识，同时为你找工作和职业发展打下坚实的基础。

5. 大模型面试题

面试，不仅是技术的较量，更需要充分的准备。在你已经掌握了大模型技术之后，就需要开始准备面试，我们将提供精心整理的大模型面试题库，涵盖当前面试中可能遇到的各种技术问题，让你在面试中游刃有余。

全套的AI大模型学习资源已经整理打包，有需要的小伙伴可以微信扫描下方CSDN官方认证二维码，免费领取【保证100%免费】

如有侵权，请联系删除。