Reading

RLHF in LLM

引言

大语言模型(LLMs)在近年来取得了显著进展,展现出上下文学习、指令跟随和逐步推理等突出特性。然而,由于这些模型是在包含高质量和低质量数据的预训练语料库上训练的,它们可能会表现出编造事实、生成有偏见或有毒文本等意外行为。因此,将LLMs与人类价值观对齐变得至关重要,特别是在帮助性、诚实性和无害性(3H)方面。

基于人类反馈的强化学习(RLHF)已被验证为有效的对齐方法,但训练过程复杂且不稳定。本文深入分析了RLHF框架,特别是PPO算法的内部工作原理,并提出了PPO-max算法,以提高策略模型训练的稳定性和效果。

RLHF的基本框架

RLHF训练过程包括三个主要阶段:

  1. 监督微调(SFT):模型通过模仿人类标注的对话示例来学习一般的人类对话方式, 优化模型的指令跟随能力
  2. 奖励模型(RM)训练:模型学习基于人类反馈比较不同回复的偏好
  3. 近端策略优化(PPO):模型基于奖励模型的反馈进行更新,通过探索和利用来发现优化的策略

奖励建模(Reward Model)

奖励模型使用预训练的基于Transformer的语言模型,移除最后的反嵌入层,并在最终Transformer层添加一个额外的线性层。给定任何文本,奖励模型会为最后一个标记分配一个标量奖励值,奖励值越大,样本越好。

奖励模型的训练损失函数为:

\[L(\psi) = \log \sigma(r(x, y_w) - r(x, y_l))\]

其中 \(\sigma\) 是sigmoid函数,\(r\) 表示参数为 \(\psi\) 的奖励模型,\(r(x, y)\) 是输入提示 \(x\) 和回复 \(y\) 的预测奖励。\(y_w\) 和 \(y_l\) 分别是较好和较差的模型response

此外,还引入了自回归LM损失,使模型模仿每个比较的pair对中的首选回复

\[L(\psi) = -\lambda E_{(x,y_w,y_l)\sim D_{rm}}[\log \sigma(r(x, y_w) - r(x, y_l))] + \beta_{rm}E_{(x,y_w)\sim D_{rm}}[\log(r'(x, y_w)]\]

其中\(D_{rm}\) 是训练集的经验分布,\(r'\)\(r\) 相同,只是顶部线性层的维度对应于词汇表大小。

奖励函数还包括基于学习的 RL 策略 \(\pi_{RL}^{\phi}\) 和初始监督模型 \(\pi^{SFT}\)之间的KL散度的惩罚项:

\[r_{total} = r(x, y) - \eta KL(\pi_{RL}^{\phi}(y|x), \pi^{SFT}(y|x))\]

其中 \(\eta\) 是KL奖励系数,控制KL惩罚的强度。这里 KL项起着两个重要作用。首先,它起到熵奖励的作用,促进对策略环境的探索,并防止策略过早地收敛到单一模式。其次,它致力于确保 RL 策略的输出不会与奖励模型在训练阶段遇到的样本大幅偏差。

为什么这里的KL散度项可以促进多样性和创新性

KL散度的方向性很重要

首先需要明确的是,在RLHF中使用的KL散度通常是惩罚项,这里的关键是KL散度的方向:我们计算的是从当前RL策略 \(\pi^{RL}\) 到初始SFT模型 \(\pi^{SFT}\) 的散度。SFT模型(通过监督微调得到的模型)通常保留了预训练语言模型的多样性特性。当我们限制RL策略不要偏离SFT模型太远时,我们实际上是在保留这种多样性。从信息论角度看,KL散度 \(KL(P||Q)\) 衡量的是使用分布Q来编码来自分布P的样本时的额外信息量。当我们最小化 \(KL(\pi^{RL}||\pi^{SFT})\) 时:
  1. 我们确保RL策略 \(\pi^{RL}\) 不会在SFT模型 \(\pi^{SFT}\) 认为可能性低的区域分配过高概率
  2. 但这不会阻止RL策略在SFT模型认为可能性高的区域中分配较低概率

这种不对称性意味着:

  • RL策略被允许在SFT模型的高概率区域内"移动"
  • 但不能探索SFT模型认为不太可能的区域

这里有一个反直觉的点:KL散度项本身并不直接增加多样性,而是通过限制偏离原始多样性分布的程度来间接保持多样性。
更准确地说:

  • KL散度项不是作为传统意义上的"熵奖励"工作的
  • 它是通过防止策略崩溃到少数高奖励模式来维持多样性
  • 它保护了SFT模型中已有的多样性和生成能力

因此,KL散度项确实能帮助模型保持多样性和创新性,但这是通过防止过度优化和保留原始模型的多样性特性实现的,而非直接鼓励探索未知区域。

LLM 强化学习

在对话生成中应用RL面临着巨大的状态-动作空间挑战。在这种情况下,我们将人类交互视为"环境"。在每个时间步 \(t\),代理(即AI助手)从环境(即对话历史)接收状态\(s_t\),基于其策略 \(\pi\),代理的动作 \(a_t\) 是生成下一个标记。环境返回奖励 \(r(s_t, a_t)\),代理然后转换到下一个状态 \(s_{t+1}\)

RL的目标是找到代理的最优行为策略,以最大化轨迹\(\tau = \{s_1, a_1, ..., s_T, a_T\}\)上的累积奖励(即回报)。

策略梯度方法

策略梯度方法是一类直接优化代理策略(状态到动作的映射)的RL技术,而不是学习价值函数。这些方法的核心思想是使用梯度上升算法来改进策略。

策略 \(\pi\) 通常由 \(\theta\) 参数化,表示为\(\pi(a|s, \theta)\),即在状态 \(s\) 下采取动作 \(a\) 的概率。策略梯度的更新规则为:

\[\theta \leftarrow \theta + \alpha\nabla_{\theta}J(\theta)\]

其中\(\alpha\)是学习率,\(J(\theta)\) 表示遵循策略 \(\pi_{\theta}\) 时的期望回报。

策略梯度的一般形式可以表示为:

\[\nabla_{\theta}J(\theta) = E_{\tau \sim\pi_{\theta}}\left[\sum_{t=0}^{T}\nabla_{\theta}\log\pi_{\theta}(a_t|s_t)\Phi_t\right]\]

其中 \(\Phi_t\) 可以是 \(\Phi_t = R(\tau)\)\(\Phi_t = \sum_{t'=t}^{T}R(s_{t'}, a_{t'})\)\(\Phi_t = \sum_{t'=t}^{T}R(s_{t'}, a_{t'}) - b(s_t)\),带有基线 \(b\)

为了减少方差,常见的策略是在策略梯度更新规则中使用优势函数估计代替原始回报:

\[\Phi_t = A(s_t, a_t)\]

数学上,\(A(s_t, a_t) = Q(s_t, a_t) - V(s_t)\),其中 \(Q(s_t, a_t)\) 是动作价值函数,\(V(s_t)\)是价值函数。

广义优势估计(GAE)

广义优势估计(Generalized Advantage Estimation, GAE)是一种在强化学习中计算优势函数的方法,旨在平衡估计的偏差和方差。在理解GAE之前,我们需要先回顾几个关键概念。

  • 时序差分(TD)误差

TD误差是强化学习中的一个基本概念,定义为:

\[\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)\]

其中:

    • \(r_t\) 是在时间步 \(t\) 获得的即时奖励
    • \(V(s_t)\)是状态 \(s_t\) 的估计价值
    • \(\gamma\)是折扣因子,通常接近但小于1

TD误差可以被视为对优势函数 \(A(s_t, a_t)\) 的一种估计。

  • \(n\)步回报估计

\(n\)步回报是将实际观察到的奖励与估计价值结合的方法:

\[\hat{R}_t^{(n)} = r_t + \gamma r_{t+1} + \gamma^2 r_{t+2} + ... + \gamma^{n-1} r_{t+n-1} + \gamma^n V(s_{t+n})\]

这种估计结合了实际奖励(前n步)和价值函数估计(第n步之后)。

  • \(n\)步优势估计

基于\(n\)步回报,我们可以定义n步优势估计:

\[\hat{A}_t^{(n)} = \hat{R}_t^{(n)} - V(s_t)\]

展开后得到:

\[\hat{A}_t^{(n)} = r_t + \gamma r_{t+1} + ... + \gamma^{n-1} r_{t+n-1} + \gamma^n V(s_{t+n}) - V(s_t)\]

我们可以用TD误差重写\(n\)步优势。将TD误差的公式代入\(n\)步优势公式,我们可以得到:

\[\begin{aligned} \hat{A}_t^{(n)} &= r_t + \gamma r_{t+1} + ... + \gamma^{n-1} r_{t+n-1} + \gamma^n V(s_{t+n}) - V(s_t) \\ &= [\delta_t + V(s_t) - \gamma V(s_{t+1})] + \gamma [\delta_{t+1} + V(s_{t+1}) - \gamma V(s_{t+2})] + ... \\ &\quad + \gamma^{n-1} [\delta_{t+n-1} + V(s_{t+n-1}) - \gamma V(s_{t+n})] + \gamma^n V(s_{t+n}) - V(s_t)\\ &= \delta_t + \gamma \delta_{t+1} + \gamma^2 \delta_{t+2} + ... + \gamma^{n-1} \delta_{t+n-1} \end{aligned}\]

这表明n步优势可以表示为一系列折扣TD误差的和。

GAE的定义与推导

GAE将不同长度的n步优势进行加权平均,权重为\((1-\lambda)\lambda^{k-1}\)

\[\begin{aligned} \hat{A}_t^{GAE(\gamma, \lambda)} &= (1-\lambda)(\hat{A}_t^{(1)} + \lambda \hat{A}_t^{(2)} + \lambda^2 \hat{A}_t^{(3)} + ...) \\ &= (1-\lambda)(\delta_t + \lambda(\delta_t + \gamma \delta_{t+1}) + \lambda^2(\delta_t + \gamma \delta_{t+1} + \gamma^2 \delta_{t+2}) + ...)\\ &= (1-\lambda)[\delta_t(1 + \lambda + \lambda^2 + ...) + \gamma \delta_{t+1}(\lambda + \lambda^2 + ...) + \gamma^2 \delta_{t+2}(\lambda^2 + ...) + ...] \\ &= (1-\lambda)[\delta_t \frac{1}{1-\lambda} + \gamma \delta_{t+1} \frac{\lambda}{1-\lambda} + \gamma^2 \delta_{t+2} \frac{\lambda^2}{1-\lambda} + ...] \end{aligned}\]

注意到几何级数 \(1 + \lambda + \lambda^2 + ... = \frac{1}{1-\lambda}\)(当\(|\lambda| < 1\)),以及 \(\lambda + \lambda^2 + ... = \frac{\lambda}{1-\lambda}\),等等。

简化后得到:

\[\begin{aligned} \hat{A}_t^{GAE(\gamma, \lambda)} &= \delta_t + \lambda \gamma \delta_{t+1} + (\lambda \gamma)^2 \delta_{t+2} + ... \\ &= \sum_{l=0}^{\infty} (\lambda \gamma)^l \delta_{t+l} \end{aligned}\]

这就是GAE的最终表达式。

GAE有两个重要的边界情况:

  1. \(\lambda = 0\) 时:
\[\hat{A}_t^{GAE(\gamma, 0)} = \delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)\]

这等同于一步TD误差,具有低方差但可能有高偏差。

  1. \(\lambda = 1\)时:
\[\hat{A}_t^{GAE(\gamma, 1)} = \sum_{l=0}^{\infty} \gamma^l \delta_{t+l} = \sum_{l=0}^{\infty} \gamma^l r_{t+l} - V(s_t)\]

这等同于蒙特卡洛回报减去基线,具有零偏差但高方差。

实际应用

在实际应用中,我们通常在有限长度的轨迹上计算GAE。假设轨迹长度为T,则:

\[\hat{A}_t^{GAE(\gamma, \lambda)} = \sum_{l=0}^{T-t-1} (\lambda \gamma)^l \delta_{t+l}\]

这可以通过一个高效的反向迭代过程计算:

  1. 初始化 \(\hat{A}_T = 0\)
  2. 对于 \(t = T-1, T-2, ..., 0\)

在PPO算法中,GAE用于计算策略梯度中的优势估计:

\[\nabla_{\theta}J(\theta) \approx \frac{1}{N} \sum_{i=1}^{N} \sum_{t=0}^{T} \nabla_{\theta} \log \pi_{\theta}(a_t^i|s_t^i) \hat{A}_t^{i, GAE(\gamma, \lambda)}\]

这里:

  • \(\pi_{\theta}\) 是参数化策略
  • \(\hat{A}_t^{i, GAE(\gamma, \lambda)}\) 是使用GAE计算的优势估计
  • \(N\) 是收集的轨迹数量

GAE在实践中表现出色的原因有以下几点:

  1. 平滑的偏差-方差权衡:通过调整 \(\lambda\) 参数,可以在不同任务上找到最佳平衡点
  2. 稳定的学习:减少了优势估计的方差,使策略梯度更加稳定
  3. 适应不同时间尺度:能够有效处理短期和长期奖励信号
  4. 与价值函数学习协同:随着价值函数估计改进,GAE的性能也会提高

数值示例

假设我们有以下轨迹数据:

  • 状态价值:\(V(s_1) = 0.5, V(s_2) = 0.6, V(s_3) = 0.7, V(s_4) = 0.8\)
  • 奖励:\(r_1 = 0.1, r_2 = 0.2, r_3 = 0.3\)
  • 参数:\(\gamma = 0.9, \lambda = 0.8\)

计算TD误差:

  • \(\delta_1 = r_1 + \gamma V(s_2) - V(s_1) = 0.1 + 0.9 \times 0.6 - 0.5 = 0.14\)
  • \(\delta_2 = r_2 + \gamma V(s_3) - V(s_2) = 0.2 + 0.9 \times 0.7 - 0.6 = 0.23\)
  • \(\delta_3 = r_3 + \gamma V(s_4) - V(s_3) = 0.3 + 0.9 \times 0.8 - 0.7 = 0.32\)

使用反向迭代计算GAE:

  • \(\hat{A}_3 = \delta_3 = 0.32\)
  • \(\hat{A}_2 = \delta_2 + \lambda \gamma \hat{A}_3 = 0.23 + 0.8 \times 0.9 \times 0.32 \approx 0.46\)
  • \(\hat{A}_1 = \delta_1 + \lambda \gamma \hat{A}_2 = 0.14 + 0.8 \times 0.9 \times 0.46 \approx 0.47\)

这样,我们得到了状态1、2、3的GAE优势估计值。

近端策略优化(PPO)

近端策略优化(Proximal Policy Optimization, PPO)是一种策略梯度算法,旨在解决策略优化中的两个关键问题:

  1. 如何确定适当的步长以避免性能崩溃
  2. 如何有效利用采样数据进行多次参数更新

PPO的核心思想是在保证性能单调改进的同时,允许较大的策略更新步长。

TRPO的约束优化问题

信任区域策略优化(Trust Region Policy Optimization, TRPO)提出了以下约束优化问题:

\[\begin{aligned} \max_{\theta} \mathbb{E}_{s,a \sim \pi_{\theta_{old}}} \left[ \frac{\pi_{\theta}(a|s)}{\pi_{\theta_{old}}(a|s)} A^{\pi_{\theta_{old}}}(s,a) \right] \\ \text{subject to } D_{KL}(\pi_{\theta_{old}} || \pi_{\theta}) \leq \delta \end{aligned}\]

其中:

  • \(\pi_{\theta}\) 是当前策略
  • \(\pi_{\theta_{old}}\) 是旧策略
  • \(A^{\pi_{\theta_{old}}}(s,a)\) 是在旧策略下的优势函数
  • \(D_{KL}\) 是KL散度
  • \(\delta\) 是信任区域大小

TRPO的主要缺点是计算复杂,难以与架构(如共享参数或循环网络)集成。

PPO的简化方法

首先,我们定义策略比率:

\[r_t(\theta) = \frac{\pi_{\theta}(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)}\]

PPO通过以下两种方式简化了TRPO:

  1. PPO-惩罚:将KL约束作为惩罚项添加到目标函数中
\[L^{CLIP}(\theta) = \mathbb{E}_t \left[ \min(r_t(\theta) A_t, \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) A_t) \right]\]

其中:

    • \(\epsilon\) 是裁剪参数(通常为0.1或0.2)
    • \(A_t\) 是优势估计(通常使用GAE计算)
    • \(\text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon)\)\(r_t(\theta)\) 限制在 \([1-\epsilon, 1+\epsilon]\) 范围内

这种设计确保了策略更新不会偏离太远,提供了类似于TRPO的信任区域约束,但计算效率更高。

  1. PPO-裁剪:直接修改目标函数,限制新旧策略比率的变化范围
\[{L}^{KLPEN}\left( \theta \right) = {\hat{\mathbb{E}}}_{t}\left\lbrack {\frac{{\pi }_{\theta }\left( {{a}_{t} \mid {s}_{t}}\right) }{{\pi }_{{\theta }_{\text{old }}}\left( {{a}_{t} \mid {s}_{t}}\right) }{\hat{A}}_{t} - \beta \operatorname{KL}\left\lbrack {{\pi }_{{\theta }_{\text{old }}}\left( {\cdot \mid {s}_{t}}\right) ,{\pi }_{\theta }\left( {\cdot \mid {s}_{t}}\right) }\right\rbrack }\right\rbrack\]

PPO的完整目标函数

实际中,PPO通常结合价值函数学习和熵正则化,完整的目标函数为:

\[L^{TOTAL}(\theta) = \mathbb{E}t \left[ L^{CLIP}(\theta) - c_1 L^{VF}(\theta) + c_2 S\pi{\theta} \right]\]

其中:

  • \(L^{VF}(\theta) = (V_{\theta}(s_t) - V_t^{target})^2\) 是价值函数的均方误差损失
  • \(S\pi_{\theta} = -\sum_a \pi_{\theta}(a|s_t) \log \pi_{\theta}(a|s_t)\) 是策略的熵
  • \(c_1, c_2\) 是权衡不同目标的系数

PPO算法的伪代码如下:

image.png

PPO在RLHF中的应用

框架

整体的流程图如下所示

image.png

在RLHF中,PPO的目标函数通常为:

\[\begin{aligned} L^{RLHF}(\phi) &= \mathbb{E}{(x,y) \sim \mathcal{D}} \left[ \min(r{\phi}(y|x) A(x,y), \text{clip}(r_{\phi}(y|x), 1-\epsilon, 1+\epsilon) A(x,y)) \right] \\&- \beta KL(\pi_{\phi}(y|x) || \pi_{SFT}(y|x)) \end{aligned}\]

其中:

  • \(r_{\phi}(y|x) = \frac{\pi_{\phi}(y|x)}{\pi_{SFT}(y|x)}\) 是策略比率
  • \(A(x,y)\) 是优势估计,通常基于奖励模型和值函数
  • \(\beta\) 是KL散度的权重系数

注意图中最终用于优化的奖励信号是由两部分组成的

\[r'(x,y) = r(x,y) - β·KL[π^{RL}_θ || π^{SFT}]\]

将KL散度作为奖励的一部分(而非单独的损失项)有几个关键优势:

  1. 强化学习框架的纯粹性
    • 保持RL算法的标准形式,只优化单一奖励信号
    • 避免多目标优化带来的复杂性
  2. 统一的优势估计
    • 优势函数 \(Â(s_t,a_t)\) 基于修正后的奖励 \(r'\) 计算
    • 使KL约束影响直接体现在优势估计中
  3. 理论上的等价性
    • 在数学上,这种形式与将KL作为单独损失项是等价的
    • 但在实现和训练动态上有显著差异

模型

在RLHF-PPO阶段,一共有四个主要模型,分别是:

  • Actor Model(Policy Model):演员模型,这就是我们想要训练的目标语言模型(用于训练策略函数), 从SFT model 参数初始化进行训练;
  • Critic Model(Value Model):评论家模型,它的作用是预估总收益 \(V_t\),从训练的Reward Model参数初始化进行训练
  • Reward Model:奖励模型,它的作用是计算即时收益 \(R_t\),参数冻结
  • Reference Model(SFT Model):参考模型,它的作用是在RLHF阶段给语言模型增加一些“约束”,防止语言模型训歪(朝不受控制的方向更新,效果可能越来越差),

其中:

  • Actor/Critic Model在RLHF阶段是需要训练的;而Reward/Reference Model参数冻结的。
  • Critic/Reward/Reference Model共同组成了一个“奖励-loss”计算体系,我们综合它们的结果计算loss,用于更新Actor和Critic Model

Actor Model (Policy Model)

正如前文所说,Actor就是我们想要训练的目标语言模型。我们一般用SFT阶段产出的SFT模型来对它做初始化。

image.png

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

Reference Model(SFT Model)

Reference Model(以下简称Ref模型)一般也用SFT阶段得到的SFT模型做初始化,在训练过程中,它的参数是冻结的。

image.png

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

Critic Model(Value Model)

Critic Model用于预测期望总收益 \(V_t\) ,和Actor模型一样,它需要做参数更新。实践中,Critic Model的设计和初始化方式也有很多种,例如和Actor共享部分参数、从RW阶段的Reward Model初始化而来等等。

当我们在前文讨论总收益 \(V_t\) (即时 + 未来)时,我们是站在上帝视角的,也就是这个 \(V_t\) 就是客观存在的、真正的总收益。但是我们在训练模型时,就没有这个上帝视角加成了,也就是在 t 时刻,我们给不出客观存在的总收益 \(V_t\) ,我们只能训练一个模型去预测它。

所以总结来说,在RLHF中,我们不仅要训练模型生成符合人类喜好的内容的能力(Actor),也要提升模型对人类喜好量化判断的能力(Critic)。这就是Critic模型存在的意义。我们来看看它的大致架构:

image.png

这里我们按Reward模型的架构来简单画画它。你可以简单理解成,Reward/Critic模型和Actor模型的架构是很相似的(毕竟输入都一样),同时,它在最后一层增加了一个Value Head层,该层是个简单的线形层,用于将原始输出结果映射成单一的 \(V_t\) 值。

在图中, \(V_t\) 表示Critic模型对 \(t\) 时刻及未来(response完成)的收益预估。

Reward Model(奖励模型)

Reward Model用于计算生成token \(A_t\) 的即时收益,它就是RW阶段所训练的奖励模型,在RLHF过程中,它的参数是冻结的。

image.png

我们知道Critic模型也是参与参数更新的,我们可以用MSE(上帝视角的客观收益-Critic模型预测的收益)来衡量它的loss。但是上帝视角的客观收益我们是不知道的,只能用已知事实数据去逼近它,所以我们就用 \(R_{t} + \gamma * V_{t+1}\) 来做近似作为Critic Model的gt。

Reference

图解大模型RLHF系列之:人人都能看懂的PPO原理与源码解读