从第一性原理推导 PPO 损失函数
引言
近期关于 DPO、GRPO 和 RLVR 等强化学习方法在大语言模型(LLM)后训练中的研究成果令人瞩目。对于刚接触强化学习的研究者来说,从 Proximal Policy Optimization(PPO)入手是个不错的选择。这正是 OpenAI 在 InstructGPT 论文中展示的用于改进 LLM 对齐的算法。理解 PPO 能帮助建立策略梯度方面的正确心智模型,为后续学习基于类似思想构建的新型 LLM 专用强化学习方法打下基础。
强化学习涉及大量数学推导。本文将通过详细的数学推导从第一性原理构建 PPO 目标函数,推导过程参考了 PPO 和 InstructGPT 原始论文以及 Umar Jamil 的视频教程。
Umar Jamil 关于 RLHF 和 PPO 的视频对于建立直觉和理解 PPO 损失背后的数学原理非常有帮助。
一、强化学习核心概念
| 概念 | 通用强化学习定义 | LLM 场景(RLHF) |
|---|---|---|
| 强化学习 | 智能体在环境中学习行动以最大化期望累积奖励的学习框架 | 使用基于奖励的反馈微调语言模型,使其生成更符合人类偏好的响应 |
| 环境 | 智能体交互的外部系统,产生观察和奖励 | 提示词分布、交互循环以及来自奖励模型评估生成响应的奖励信号 |
| 智能体 | 观察状态、执行动作并接收奖励的学习者/决策者 | 逐词生成文本的语言模型 |
| 动作 (a) | 智能体做出的选择,通常基于状态 s | 在生成的每一步选择下一个词元 |
| 状态 (s) | 给定时间步智能体可获得的信息 | 提示词加上目前为止生成的响应(当前词元上下文) |
| 奖励 ® | 告诉智能体结果好坏的标量信号 | 奖励模型(基于偏好数据训练)对响应质量的评分 |
| 策略 (π) | 从状态到动作分布的随机映射 | 模型基于上下文对下一个词元的概率分布 |
| 目标 | 找到最优策略 π* 使期望累积奖励最大化 | 更新(对齐)模型,使其倾向于生成更高奖励分数的响应 |
二、RLHF 中的奖励模型
**奖励模型(RM)**是一个神经网络,输入提示词 x 和响应 y,输出标量奖励 r_φ(x,y) ∈ ℝ,表示该响应在人类偏好下的"好坏"程度。
策略梯度方法(包括 PPO)需要标量目标来更新策略参数。在标准强化学习中,环境提供这个信号。但对于文本生成,没有天然环境为"好"的响应给出奖励。让人类评估每个输出既不现实,而且基于梯度的优化需要可微的标量信号来反向传播。因此强化学习训练期间需要一个廉价、可微的人类偏好代理。学习到的奖励模型正好提供了这个功能。
奖励模型的训练方法
奖励模型的标准训练流程:
- 采样提示词 (x)
- 从基线策略(通常是 SFT 模型)生成多个候选补全 (y₁, y₂, ..., yₖ)
- 让人类比较候选项(成对比较比绝对评分更容易)
- 训练奖励模型 (r_φ) 预测这些偏好
架构上,奖励模型通常是:
- 从预训练语言模型(通常是 SFT 模型本身)初始化
- 移除最后的非嵌入层(投影到词汇表的层)
- 替换 为将最后一个词元的隐藏状态投影到单个标量输出的线性层
奖励模型损失函数
奖励模型使用 Bradley-Terry 模型进行成对比较训练。对于任意提示词 x,响应 y_w(更优)相比 y_l(较差)被偏好的概率建模为:
P ( y w ≻ y l ∣ x ) = σ ( r θ ( x , y w ) − r θ ( x , y l ) ) (II.I) P(y_w \succ y_l | x) = \sigma(r_\theta(x, y_w) - r_\theta(x, y_l)) \tag{II.I} P(yw≻yl∣x)=σ(rθ(x,yw)−rθ(x,yl))(II.I)
其中 σ 是 sigmoid 函数: σ(z) = 1/(1 + e^(-z))
负对数似然损失为:
L RM ( ϕ ) = − E ( x , y w , y l ) ∼ D [ log σ ( r ϕ ( x , y w ) − r ϕ ( x , y l ) ) ] \mathcal{L}{\text{RM}}(\phi) = -\mathbb{E}{(x, y_w, y_l) \sim \mathcal{D}} [\log \sigma(r_\phi(x, y_w) - r_\phi(x, y_l))] LRM(ϕ)=−E(x,yw,yl)∼D[logσ(rϕ(x,yw)−rϕ(x,yl))]
这个损失强制奖励模型为更优响应分配更高奖励(详见 InstructGPT 论文或 Umar Jamil 视频)。
这里有两个关键洞察:
- 不需要绝对分数,只需要奖励模型能正确排序响应
- 损失只依赖于差值 (r_φ(x,y_w) - r_φ(x,y_l)),因此对所有奖励加常数不变。这在后续讨论 PPO 损失时很有用。
奖励模型作为人类偏好的学习代理,将从每次生成获取人类反馈这个难以处理的问题转换为可处理的监督学习问题。训练完成后,它提供标量信号 r_φ(x,y) 用于使用 PPO 等强化学习算法优化策略(LLM)。
三、轨迹与回报
轨迹
轨迹(也称推出或回合)是智能体与环境交互生成的状态(s)、动作(a)和奖励®序列:
τ = ( s 0 , a 0 , r 0 , s 1 , a 1 , r 1 , ... , s T , a T , r T ) \tau = (s_0, a_0, r_0, s_1, a_1, r_1, \ldots, s_T, a_T, r_T) τ=(s0,a0,r0,s1,a1,r1,...,sT,aT,rT)
在 LLM 场景中,轨迹对应整个词元生成序列,即提示词后跟所有生成的词元直到序列结束词元。
注意状态总是随机建模的,s_{t+1} 可表示为 s_{t+1} ~ P(s_{t+1} | s_t, a_t)。给定随机策略 π_θ(a_t | s_t),轨迹 τ 的概率是以下三者的乘积:
- 初始状态分布 ρ_0(s_0)
- 随机策略 π_θ(a_t | s_t)
- 环境转移动态 P(s_{t+1} | s_t, a_t)
P ( τ ∣ π θ ) = ρ 0 ( s 0 ) ∏ t = 0 T − 1 π θ ( a t ∣ s t ) ⋅ P ( s t + 1 ∣ s t , a t ) (III.I) P(\tau | \pi_\theta) = \rho_0(s_0) \prod_{t=0}^{T-1} \pi_\theta(a_t | s_t) \cdot P(s_{t+1} | s_t, a_t) \tag{III.I} P(τ∣πθ)=ρ0(s0)t=0∏T−1πθ(at∣st)⋅P(st+1∣st,at)(III.I)
回报
回报 是完整轨迹 τ 收集的累积奖励。最简单的形式是无折扣回报:
R ( τ ) = ∑ t = 0 T r t R(\tau) = \sum_{t=0}^{T} r_t R(τ)=t=0∑Trt
更一般地,使用折扣回报:
R ( τ ) = ∑ k = 0 ∞ γ k r k = r 0 + γ r 1 + γ 2 r 2 + ⋯ (III.II) R(\tau) = \sum_{k=0}^{\infty} \gamma^k r_k = r_0 + \gamma r_1 + \gamma^2 r_2 + \cdots \tag{III.II} R(τ)=k=0∑∞γkrk=r0+γr1+γ2r2+⋯(III.II)
其中 γ ∈ [0,1] 是折扣因子。折扣因子 γ 有几个作用:
- 确保无限时域任务(T→∞)的回报有限
- 优先考虑即时奖励而非远期奖励
四、策略梯度优化与 REINFORCE 算法
强化学习的目标是找到策略 π_θ 使所有可能轨迹的期望回报最大化:
J ( θ ) = E τ ∼ π θ [ R ( τ ) ] (IV.I) \boxed{J(\theta) = \mathbb{E}{\tau \sim \pi\theta}[R(\tau)]} \tag{IV.I} J(θ)=Eτ∼πθ[R(τ)](IV.I)
这是目标函数,需要找到参数 θ* 使得:
θ ∗ = arg max θ J ( θ ) \theta^* = \arg\max_\theta J(\theta) θ∗=argθmaxJ(θ)
要使用基于梯度的方法最大化 J(θ),需要计算 ∇_θ J(θ) 并执行梯度上升:
θ k + 1 = θ k + α ∇ θ J ( π θ ) ∣ θ k (IV.II) \boxed{\theta_{k+1} = \theta_k + \alpha \left. \nabla_\theta J(\pi_\theta) \right|_{\theta_k}} \tag{IV.II} θk+1=θk+α∇θJ(πθ)∣θk(IV.II)
这个策略梯度看起来简单,但实际上难以计算。期望是对从 π_θ 采样的轨迹求的,而 π_θ 本身依赖于 θ。无法简单枚举所有可能的轨迹,对于任何合理大小的状态-动作空间(当然也不适用于 LLM)这在计算上是不可行的。
因此下一步需要推导 ∇_θ J(θ) 的某种合理且可处理的近似。使用对数导数技巧来实现:
∇ θ J ( θ ) = ∇ θ E τ ∼ π θ [ R ( τ ) ] \nabla_\theta J(\theta) = \nabla_\theta \mathbb{E}{\tau \sim \pi\theta}[R(\tau)] ∇θJ(θ)=∇θEτ∼πθ[R(τ)]
这个期望可以写成积分:
= ∇ θ ∫ τ P ( τ ∣ θ ) R ( τ ) d τ = \nabla_\theta \int_\tau P(\tau | \theta) R(\tau) d\tau =∇θ∫τP(τ∣θ)R(τ)dτ
将梯度移到积分内:
= ∫ τ ∇ θ P ( τ ∣ θ ) R ( τ ) d τ = \int_\tau \nabla_\theta P(\tau | \theta) R(\tau) d\tau =∫τ∇θP(τ∣θ)R(τ)dτ
应用对数导数技巧:
∇ θ log P ( τ ∣ θ ) = ∇ θ P ( τ ∣ θ ) P ( τ ∣ θ ) \nabla_\theta \log P(\tau | \theta) = \frac{\nabla_\theta P(\tau | \theta)}{P(\tau | \theta)} ∇θlogP(τ∣θ)=P(τ∣θ)∇θP(τ∣θ)
重新排列: ∇_θ P(τ|θ) = P(τ|θ) ∇_θ log P(τ|θ) 并代回,得到:
= ∫ τ P ( τ ∣ θ ) ∇ θ log P ( τ ∣ θ ) R ( τ ) d τ = \int_\tau P(\tau | \theta) \nabla_\theta \log P(\tau | \theta) R(\tau) d\tau =∫τP(τ∣θ)∇θlogP(τ∣θ)R(τ)dτ
也可以写成期望形式:
∇ θ J ( θ ) = E τ ∼ π θ [ ∇ θ log P ( τ ∣ θ ) ⋅ R ( τ ) ] (IV.III) \boxed{\nabla_\theta J(\theta) = \mathbb{E}{\tau \sim \pi\theta}[\nabla_\theta \log P(\tau | \theta) \cdot R(\tau)]} \tag{IV.III} ∇θJ(θ)=Eτ∼πθ[∇θlogP(τ∣θ)⋅R(τ)](IV.III)
注意这里的梯度现在是轨迹对数概率梯度的期望。使用轨迹概率表达式(III.I)可以进一步简化:
P ( τ ∣ θ ) = ρ 0 ( s 0 ) ∏ t = 0 T − 1 π θ ( a t ∣ s t ) ⋅ P ( s t + 1 ∣ s t , a t ) P(\tau | \theta) = \rho_0(s_0) \prod_{t=0}^{T-1} \pi_\theta(a_t | s_t) \cdot P(s_{t+1} | s_t, a_t) P(τ∣θ)=ρ0(s0)t=0∏T−1πθ(at∣st)⋅P(st+1∣st,at)
取对数:
log P ( τ ∣ θ ) = log ρ 0 ( s 0 ) + ∑ t = 0 T − 1 log π θ ( a t ∣ s t ) + ∑ t = 0 T − 1 log P ( s t + 1 ∣ s t , a t ) \log P(\tau | \theta) = \log \rho_0(s_0) + \sum_{t=0}^{T-1} \log \pi_\theta(a_t | s_t) + \sum_{t=0}^{T-1} \log P(s_{t+1} | s_t, a_t) logP(τ∣θ)=logρ0(s0)+t=0∑T−1logπθ(at∣st)+t=0∑T−1logP(st+1∣st,at)
求 ∇_θ 时,只有策略项依赖于 θ:
∇ θ log P ( τ ∣ θ ) = ∑ t = 0 T − 1 ∇ θ log π θ ( a t ∣ s t ) \nabla_\theta \log P(\tau | \theta) = \sum_{t=0}^{T-1} \nabla_\theta \log \pi_\theta(a_t | s_t) ∇θlogP(τ∣θ)=t=0∑T−1∇θlogπθ(at∣st)
初始状态分布和转移动态独立于 θ,其梯度为零。代回得到策略梯度定理:
∇ θ J ( θ ) = E τ ∼ π θ [ ∑ t = 0 T ∇ θ log π θ ( a t ∣ s t ) ⋅ R ( τ ) ] (IV.IV) \boxed{\nabla_\theta J(\theta) = \mathbb{E}{\tau \sim \pi\theta}\left[ \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t | s_t) \cdot R(\tau) \right]} \tag{IV.IV} ∇θJ(θ)=Eτ∼πθ[t=0∑T∇θlogπθ(at∣st)⋅R(τ)](IV.IV)
这是个了不起的结果。可以计算目标的梯度而无需对环境动态求导,只需要策略对数概率的梯度。
由于无法精确计算期望,通过采样 N 条轨迹用样本均值近似:
∇ θ J ( θ ) ≈ g ^ = 1 N ∑ i = 1 N ( ∑ t = 0 T ∇ θ log π θ ( a i , t ∣ s i , t ) ) R ( τ i ) (IV.V) \boxed{\nabla_\theta J(\theta) \approx \hat{g} = \frac{1}{N} \sum_{i=1}^{N} \left( \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_{i,t} | s_{i,t}) \right) R(\tau_i)} \tag{IV.V} ∇θJ(θ)≈g^=N1i=1∑N(t=0∑T∇θlogπθ(ai,t∣si,t))R(τi)(IV.V)
这给出了 REINFORCE 算法:
-
初始化: 从预训练或监督微调(SFT)语言模型 π_θ 开始
-
采样提示词: 从数据集抽取 N 个提示词批次 {x₁, x₂, ..., x_N}
-
生成轨迹: 对每个提示词 x_i,通过从策略 π_θ 采样词元生成响应 y_i = (a₀, a₁, ..., a_T)。每条轨迹是状态(提示词+目前生成的词元)和动作(选中的词元)的序列。
-
计算对数概率: 对每条轨迹,计算每个生成词元在其上下文下的对数概率:
log π θ ( a t ∣ s t ) for t = 0 , 1 , ... , T \log \pi_\theta(a_t | s_t) \quad \text{for } t = 0, 1, \ldots, T logπθ(at∣st)for t=0,1,...,T
-
计算奖励: 使用奖励模型对每个完整(提示词,响应)对评分: R(τ_i) = r_φ(x_i, y_i)
-
估计策略梯度: 使用(IV.V)计算梯度估计:
g ^ = 1 N ∑ i = 1 N ( ∑ t = 0 T ∇ θ log π θ ( a i , t ∣ s i , t ) ) R ( τ i ) \hat{g} = \frac{1}{N} \sum_{i=1}^{N} \left( \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_{i,t} | s_{i,t}) \right) R(\tau_i) g^=N1i=1∑N(t=0∑T∇θlogπθ(ai,t∣si,t))R(τi)
-
更新策略: 执行梯度上升步骤: θ ← θ + α ĝ
-
重复: 返回步骤 2 并迭代直到收敛
虽然 REINFORCE 提供无偏梯度估计,但存在两个使其在 LLM 训练中不实用的关键问题:
- 高方差: 梯度估计 ĝ 因采样的轨迹而有高方差。这种方差可能很大,导致梯度噪声和训练不稳定。
再看(IV.V),每个动作的梯度估计都由整条轨迹的回报 R(τ) 加权。这意味着即使某个动作很好,它也可能因为轨迹中其他动作导致差结果而收到负梯度更新(反之亦然)。在多个样本中,这种耦合引入的噪声可能很大,导致高方差。
- 在线策略约束(样本效率低) : REINFORCE 需要从当前策略 π_θ 采样的轨迹。因此每次梯度更新后,之前收集的轨迹必须丢弃,需要从更新后的策略采样新轨迹。对于 LLM,每条轨迹需要对数十亿参数模型完整前向传播,这代价极高,尤其是需要许多小梯度步骤才能有效训练时。
五、降低方差与优势函数
REINFORCE 算法提供无偏梯度估计(IV.V)。但虽然无偏,该估计器存在高方差问题。
用 Reward-to-Go 替换完整轨迹回报(利用因果性)
第一个方差降低来自注意到时刻 t 采取的动作 a_t 无法影响 t 之前接收的奖励 。这是因果性的基本结果。这些过去的奖励项只贡献噪声给梯度估计,增加方差而不贡献任何信号。因此可以移除它们,只考虑 rewards-to-go:
R ^ t = ∑ t ′ = t T r t ′ \hat{R}t = \sum{t'=t}^{T} r_{t'} R^t=t′=t∑Trt′
这给出更低方差的估计器:
∇ θ J ( θ ) ≈ 1 N ∑ i = 1 N ∑ t = 0 T ∇ θ log π θ ( a i , t ∣ s i , t ) ⋅ R ^ i , t (V.I) \boxed{\nabla_\theta J(\theta) \approx \frac{1}{N} \sum_{i=1}^{N} \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_{i,t} | s_{i,t}) \cdot \hat{R}_{i,t}} \tag{V.I} ∇θJ(θ)≈N1i=1∑Nt=0∑T∇θlogπθ(ai,t∣si,t)⋅R^i,t(V.I)
其中 Ř_{i,t} = Σ_{t'=t}^T r_{i,t'} 是从时刻 t 开始的轨迹 i 的 rewards-to-go。
减去基线
第二个互补的方差降低技术是从奖励中减去基线 b(s_t)。关键洞察是可以从奖励信号中减去任何不依赖动作的函数而不改变梯度的期望值。
因此可以从 rewards-to-go 中减去状态相关的基线 b(s_t) 得到无偏梯度估计器:
∇ θ J ( θ ) ≈ 1 N ∑ i = 1 N ∑ t = 0 T ∇ θ log π θ ( a i , t ∣ s i , t ) ⋅ ( R ^ i , t − b ( s i , t ) ) (V.II) \boxed{\nabla_\theta J(\theta) \approx \frac{1}{N} \sum_{i=1}^{N} \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_{i,t} | s_{i,t}) \cdot (\hat{R}{i,t} - b(s{i,t}))} \tag{V.II} ∇θJ(θ)≈N1i=1∑Nt=0∑T∇θlogπθ(ai,t∣si,t)⋅(R^i,t−b(si,t))(V.II)
值函数: V^π(s) 和 Q^π(s,a)
基线仍是任意函数。为了使其更系统和具体,强化学习理论中有两个基本函数。
状态值函数 : 状态值函数 V^π(s) 是智能体在状态 s 并按策略 π 行动时的期望回报:
V π ( s ) = E τ ∼ π [ ∑ t = 0 ∞ γ t r t | s 0 = s ] V^\pi(s) = \mathbb{E}{\tau \sim \pi}\left[\sum{t=0}^{\infty} \gamma^t r_t \middle| s_0 = s\right] Vπ(s)=Eτ∼π[t=0∑∞γtrt s0=s]
直观上,V^π(s) 表示**"这个状态平均有多好?"**,用作基线 b(s) = V^π(s)。
动作值函数(Q 函数) : 动作值函数 Q^π(s,a) 是从状态 s 开始采取动作 a 然后按策略 π 行动的期望回报:
Q π ( s , a ) = E τ ∼ π [ ∑ t = 0 ∞ γ t r t | s 0 = s , a 0 = a ] Q^\pi(s, a) = \mathbb{E}{\tau \sim \pi}\left[\sum{t=0}^{\infty} \gamma^t r_t \middle| s_0 = s, a_0 = a\right] Qπ(s,a)=Eτ∼π[t=0∑∞γtrt s0=s,a0=a]
直观上,Q^π(s,a) 表示**"这个特定动作在这个状态有多好?"**,在强化学习中,rewards-to-go 估计为 Q^π(s,a)。
在 LLM 场景:
- V^π(s) 估计给定提示词+部分响应的期望奖励,假设模型继续按当前策略生成
- Q^π(s,a) 估计从当前提示词+部分响应开始,如果模型生成特定下一个词元 a 然后继续按策略生成的期望奖励
优势函数
优势函数 A^π(s,a) 衡量特定动作 a 相比策略下的平均动作好(或差)多少:
A π ( s , a ) = Q π ( s , a ) − V π ( s ) (V.III) \boxed{A^\pi(s, a) = Q^\pi(s, a) - V^\pi(s)} \tag{V.III} Aπ(s,a)=Qπ(s,a)−Vπ(s)(V.III)
优势函数直接告诉:"这个特定动作相比在这个状态通常做的好多少?" 这正是策略改进需要的信号。需要增加正优势动作的概率,减少负优势动作的概率。
来自 Umar Jamil 视频:
在 LLM 场景中考虑提示词是"上海在哪里?"且模型已生成"上海是"的状态。从这个状态:
- 如果模型采样词元"在"(导向"上海是在中国"),这个动作可能有正优势,因为它比模型可能产生的平均词元更好
- 如果模型采样词元"美味"(导向不连贯的响应),这个动作可能有负优势,因为它比模型可能产生的平均词元更差
优势加权策略梯度
代入 rewards-to-go 和值函数作为基线,得到策略梯度的如下形式:
∇ θ J ( θ ) = E τ ∼ π θ [ ∑ t = 0 T ∇ θ log π θ ( a t ∣ s t ) ⋅ ( Q π ( s t , a t ) − V π ( s t ) ) ] \nabla_\theta J(\theta) = \mathbb{E}{\tau \sim \pi\theta}\left[\sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t | s_t) \cdot (Q^\pi(s_t, a_t) - V^\pi(s_t))\right] ∇θJ(θ)=Eτ∼πθ[t=0∑T∇θlogπθ(at∣st)⋅(Qπ(st,at)−Vπ(st))]
可以写成:
∇ θ J ( θ ) = E τ ∼ π θ [ ∑ t = 0 T ∇ θ log π θ ( a t ∣ s t ) ⋅ A π θ ( s t , a t ) ] (V.IV) \boxed{\nabla_\theta J(\theta) = \mathbb{E}{\tau \sim \pi\theta}\left[\sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t | s_t) \cdot A^{\pi_\theta}(s_t, a_t)\right]} \tag{V.IV} ∇θJ(θ)=Eτ∼πθ[t=0∑T∇θlogπθ(at∣st)⋅Aπθ(st,at)](V.IV)
基于样本的近似为:
∇ θ J ( θ ) ≈ 1 N ∑ i = 1 N ∑ t = 0 T ∇ θ log π θ ( a i , t ∣ s i , t ) ⋅ A ^ i , t (V.V) \boxed{\nabla_\theta J(\theta) \approx \frac{1}{N} \sum_{i=1}^{N} \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_{i,t} | s_{i,t}) \cdot \hat{A}_{i,t}} \tag{V.V} ∇θJ(θ)≈N1i=1∑Nt=0∑T∇θlogπθ(ai,t∣si,t)⋅A^i,t(V.V)
其中 Â_{i,t} 是轨迹 i 时刻 t 的优势函数估计。这是常用的策略梯度形式。
实践中,A^π(s_t,a_t) 可如下估计:
-
学习值函数: 训练神经网络 V_φ(s)(常称"评论家"或"值头")逼近 V^π(s)。在 LLM 微调中,这通常是策略使用的同一 Transformer 骨干上的线性层。
-
从样本估计 Q^π: 给定轨迹,rewards-to-go Ř_t = Σ_{t'=t}^T γ^{t'} r_{t'} 提供 Q^π(s_t,a_t) 的无偏(但高方差)估计。
-
计算优势估计: Â_t = Ř_t - V_φ(s_t)
更复杂的方法如**广义优势估计(GAE)**通过使用多步回报的加权组合在高方差低偏差估计和低方差高偏差估计之间插值。详见 GAE 论文。
六、重要性采样与离线策略梯度
注意 : 在强化学习文献中,"离线策略"通常指行为策略 (生成数据)与目标策略(被优化)差异很大的方法,比如重用数千次更新前策略的转移。本节称为"离线策略"的更准确应称为"局部离线策略"。
优势加权策略梯度(V.IV)需要从当前策略 π_θ 采样的轨迹。这造成基本的低效率,即每次梯度更新 θ → θ' 后,所有之前收集的轨迹都变"陈旧",必须丢弃这些轨迹并从更新后的策略采样新轨迹。
对于 LLM,每条轨迹需要对数十亿参数模型完整前向传播,这代价极高,尤其是需要许多小梯度步骤才能有效训练时。
需要一种方法为多次梯度更新重用相同轨迹。重要性采样提供了实现这一目标的数学机制!
重要性采样
重要性采样是使用从不同分布抽取的样本估计一个概率分布下期望的技术。考虑分布 p(x) 的期望:
E x ∼ p [ f ( x ) ] = ∫ p ( x ) f ( x ) d x \mathbb{E}_{x \sim p}[f(x)] = \int p(x) f(x) dx Ex∼p[f(x)]=∫p(x)f(x)dx
可以通过乘除另一个分布 q(x)(其中 p(x)>0 处 q(x)>0)重写:
= ∫ q ( x ) p ( x ) q ( x ) f ( x ) d x = E x ∼ q [ p ( x ) q ( x ) f ( x ) ] = \int q(x) \frac{p(x)}{q(x)} f(x) dx = \mathbb{E}_{x \sim q}\left[\frac{p(x)}{q(x)} f(x)\right] =∫q(x)q(x)p(x)f(x)dx=Ex∼q[q(x)p(x)f(x)]
比率 p(x)/q(x) 称为重要性权重。这个恒等式告诉:
E x ∼ p [ f ( x ) ] = E x ∼ q [ p ( x ) q ( x ) f ( x ) ] (VI.I) \boxed{\mathbb{E}{x \sim p}[f(x)] = \mathbb{E}{x \sim q}\left[\frac{p(x)}{q(x)} f(x)\right]} \tag{VI.I} Ex∼p[f(x)]=Ex∼q[q(x)p(x)f(x)](VI.I)
现在可以使用 q 的样本估计 p 下的期望,只要按概率比重新加权每个样本。
将重要性采样应用于策略梯度
可以将此技术应用于策略梯度设置。在线策略优势加权梯度(V.IV)是:
∇ θ J ( θ ) = E τ ∼ π θ [ ∑ t = 0 T ∇ θ log π θ ( a t ∣ s t ) ⋅ A π θ ( s t , a t ) ] \nabla_\theta J(\theta) = \mathbb{E}{\tau \sim \pi\theta}\left[\sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t | s_t) \cdot A^{\pi_\theta}(s_t, a_t)\right] ∇θJ(θ)=Eτ∼πθ[t=0∑T∇θlogπθ(at∣st)⋅Aπθ(st,at)]
为应用重要性采样,在时间步层面而非轨迹层面工作(完整轨迹重要性权重有极高方差)。对单个时间步:
∇ θ J ( θ ) = E ( s t , a t ) ∼ π θ [ ∇ θ log π θ ( a t ∣ s t ) ⋅ A π θ ( s t , a t ) ] \nabla_\theta J(\theta) = \mathbb{E}{(s_t, a_t) \sim \pi\theta}[\nabla_\theta \log \pi_\theta(a_t | s_t) \cdot A^{\pi_\theta}(s_t, a_t)] ∇θJ(θ)=E(st,at)∼πθ[∇θlogπθ(at∣st)⋅Aπθ(st,at)]
使用 π_{θ_old} 样本的重要性采样:
= E ( s t , a t ) ∼ π θ o l d [ π θ ( a t ∣ s t ) π θ o l d ( a t ∣ s t ) ∇ θ log π θ ( a t ∣ s t ) ⋅ A π θ ( s t , a t ) ] = \mathbb{E}{(s_t, a_t) \sim \pi{\theta_{old}}}\left[\frac{\pi_\theta(a_t | s_t)}{\pi_{\theta_{old}}(a_t | s_t)} \nabla_\theta \log \pi_\theta(a_t | s_t) \cdot A^{\pi_\theta}(s_t, a_t)\right] =E(st,at)∼πθold[πθold(at∣st)πθ(at∣st)∇θlogπθ(at∣st)⋅Aπθ(st,at)]
应用对数导数恒等式 ∇_θ log π_θ = (∇_θ π_θ)/π_θ,得到代理目标 L(θ),其梯度等于这个重要性加权策略梯度:
∇ θ J ( θ ) = E ( s t , a t ) ∼ π θ o l d [ ∇ θ π θ ( a t ∣ s t ) π θ o l d ( a t ∣ s t ) A π θ o l d ( s t , a t ) ] \nabla_\theta J(\theta) = \mathbb{E}{(s_t, a_t) \sim \pi{\theta_{old}}}\left[\frac{\nabla_\theta \pi_\theta(a_t | s_t)}{\pi_{\theta_{old}}(a_t | s_t)} A^{\pi_{\theta_{old}}}(s_t, a_t)\right] ∇θJ(θ)=E(st,at)∼πθold[πθold(at∣st)∇θπθ(at∣st)Aπθold(st,at)]
其中重要性加权代理目标也称**保守策略迭代(CPI)**目标:
L CPI ( θ ) = E ( s t , a t ) ∼ π θ o l d [ π θ ( a t ∣ s t ) π θ o l d ( a t ∣ s t ) A π θ o l d ( s t , a t ) ] L^{\text{CPI}}(\theta) = \mathbb{E}{(s_t, a_t) \sim \pi{\theta_{old}}}\left[\frac{\pi_\theta(a_t | s_t)}{\pi_{\theta_{old}}(a_t | s_t)} A^{\pi_{\theta_{old}}}(s_t, a_t)\right] LCPI(θ)=E(st,at)∼πθold[πθold(at∣st)πθ(at∣st)Aπθold(st,at)]
还定义概率比为:
r t ( θ ) = π θ ( a t ∣ s t ) π θ o l d ( a t ∣ s t ) (VI.II) r_t(\theta) = \frac{\pi_\theta(a_t | s_t)}{\pi_{\theta_{old}}(a_t | s_t)} \tag{VI.II} rt(θ)=πθold(at∣st)πθ(at∣st)(VI.II)
注意按构造 r_t(θ_old) = 1。因此 CPI 目标可写为:
L CPI ( θ ) = E t [ π θ ( a t ∣ s t ) π θ o l d ( a t ∣ s t ) A ^ t ] = E t [ r t ( θ ) A ^ t ] (VI.III) \boxed{L^{\text{CPI}}(\theta) = \mathbb{E}t\left[\frac{\pi\theta(a_t | s_t)}{\pi_{\theta_{old}}(a_t | s_t)} \hat{A}_t\right] = \mathbb{E}_t[r_t(\theta) \hat{A}_t]} \tag{VI.III} LCPI(θ)=Et[πθold(at∣st)πθ(at∣st)A^t]=Et[rt(θ)A^t](VI.III)
其中 Â_t 是时间步 t 的估计优势,𝔼_t[·] 表示在 π_{θ_old} 下收集的样本批次上的经验平均。
这个目标有清晰的解释:
- 如果 Â_t > 0(动作优于平均),希望增加 r_t(θ),即使新策略更可能采取这个动作
- 如果 Â_t < 0(动作劣于平均),希望减少 r_t(θ),即使新策略更不可能采取这个动作
对应的基于样本的近似为:
L CPI ( θ ) ≈ 1 N ∑ i = 1 N ∑ t = 0 T π θ ( a i , t ∣ s i , t ) π θ o l d ( a i , t ∣ s i , t ) A ^ i , t (VI.IV) \boxed{L^{\text{CPI}}(\theta) \approx \frac{1}{N} \sum_{i=1}^{N} \sum_{t=0}^{T} \frac{\pi_\theta(a_{i,t} | s_{i,t})}{\pi_{\theta_{old}}(a_{i,t} | s_{i,t})} \hat{A}_{i,t}} \tag{VI.IV} LCPI(θ)≈N1i=1∑Nt=0∑Tπθold(ai,t∣si,t)πθ(ai,t∣si,t)A^i,t(VI.IV)
离线策略学习:重用轨迹
CPI 目标使离线策略学习成为可能:可以从 π_{θ_old} 采样轨迹,存储它们,然后使用相同数据批次对 θ 执行多次梯度更新。典型工作流程变为:
- 收集: 从当前策略 π_{θ_old} 采样轨迹 {τ_i}
- 计算: 计算优势 Â_{i,t} 和对数概率 log π_{θ_old}(a_{i,t}|s_{i,t})
- 存储: 保存轨迹及其优势和旧对数概率
- 优化: 使用存储数据的小批次对 L^CPI(θ) 执行多次梯度上升步骤
- 重复: 设置 θ_old ← θ 并返回步骤 1
这极大提高了样本效率。不用在单次梯度步骤后丢弃轨迹,可以从每批昂贵的 LLM 推出中提取多次更新。
不稳定性问题
虽然 CPI 目标提高样本效率,无约束优化 L^CPI(θ) 是不稳定的。核心问题是当 π_θ 偏离 π_{θ_old} 太远时重要性采样变得不可靠:
- 极端概率比: 比率 r_t(θ) 可能变得任意大或小,破坏梯度估计的稳定性
- 陈旧优势: 估计 Â_t 是在 π_{θ_old} 下计算的,随着 π_θ 发散变得不准确。优化器可能利用这些陈旧估计,做出看似有益但实际有害的更新
实践中,无约束最大化 L^CPI(θ) 常导致过大的策略更新,造成灾难性性能崩溃。
LLM 场景(来自 Umar Jamil): 假设有一条模型生成"上海是在中国"的高优势轨迹。无约束优化可能大幅增加"中国"作为"上海是在"后下一个词元的权重------但这可能同时在其他地方造成意外的概率偏移,也许使模型在完全无关的上下文中过度倾向于说"中国",或以不可预测的方式破坏整个词汇表的概率质量分布。
需要一种机制约束 π_θ 不要偏离 π_{θ_old} 太远,保持比率 r_t(θ) 接近 1,同时仍允许有意义的策略改进。
七、信赖域策略优化(TRPO)
CPI 目标很有吸引力,因为它允许通过重要性比重用数据,但无约束优化是不稳定的。当 π_θ 偏离 π_{θ_old} 太远时,概率比 r_t(θ) 变得极端,优势估计 Â_t 变陈旧且可能被优化器利用。
信赖域策略优化(TRPO)的关键洞察是代理目标 L^CPI(θ) 只在 θ_old 的局部邻域内是真实目标的有效近似。TRPO 论文通过证明只要连续策略间的 KL 散度保持有界,策略性能就保证改进,形式化了这一点。这个理论结果促使将策略更新约束在代理目标保持可靠的"信赖域"内。详见 TRPO 论文的形式化证明。
TRPO 将这个洞察转换为约束优化问题,确保策略更新停留在代理目标保持可靠的"信赖域"内。
max θ L CPI ( θ ) = E t [ π θ ( a t ∣ s t ) π θ o l d ( a t ∣ s t ) A ^ t ] subject to E t [ D KL ( π θ o l d ( ⋅ ∣ s t ) ∥ π θ ( ⋅ ∣ s t ) ) ] ≤ δ (VII.I) \boxed{ \begin{aligned} \max_\theta \quad & L^{\text{CPI}}(\theta) = \mathbb{E}t\left[\frac{\pi\theta(a_t | s_t)}{\pi_{\theta_{old}}(a_t | s_t)} \hat{A}t\right] \\[6pt] \text{subject to} \quad & \mathbb{E}t[D{\text{KL}}(\pi{\theta_{old}}(\cdot|s_t) \| \pi_\theta(\cdot|s_t))] \leq \delta \end{aligned} } \tag{VII.I} θmaxsubject toLCPI(θ)=Et[πθold(at∣st)πθ(at∣st)A^t]Et[DKL(πθold(⋅∣st)∥πθ(⋅∣st))]≤δ(VII.I)
超参数 δ 定义信赖域大小,即连续策略间允许的最大散度。这个约束确保 r_t(θ) 保持接近 1,保持重要性加权估计的可靠性。
求解(VII.I)需要二阶优化 。TRPO 线性近似目标,二次近似 KL 约束(使用 Fisher 信息矩阵),然后通过共轭梯度算法 求解得到的问题,随后进行线搜索确保满足约束。
对于大规模 LLM 训练,这种方法不实用:
- 计算开销: 每次策略更新需要多次共轭梯度迭代和线搜索步骤,比标准梯度下降昂贵得多
- 内存需求: 计算 Fisher-向量积为数十亿参数模型增加大量内存开销
TRPO 背后的理论还建议使用 KL 惩罚而非硬约束。这更易实现且计算效率更高。
max θ E t [ r t ( θ ) A ^ t − β ⋅ D KL ( π θ o l d ( ⋅ ∣ s t ) ∥ π θ ( ⋅ ∣ s t ) ) ] (VII.II) \max_\theta \; \mathbb{E}t[r_t(\theta) \hat{A}t - \beta \cdot D{\text{KL}}(\pi{\theta_{old}}(\cdot|s_t) \| \pi_\theta(\cdot|s_t))] \tag{VII.II} θmaxEt[rt(θ)A^t−β⋅DKL(πθold(⋅∣st)∥πθ(⋅∣st))](VII.II)
然而,选择在不同问题甚至不同训练阶段都有效的惩罚系数 β 是出了名的困难。这促使了近端策略优化(PPO)的出现:一种一阶方法 ,通过裁剪代理目标而非显式约束实现 TRPO 的稳定性。
八、近端策略优化(PPO)
近端策略优化(PPO)仅使用一阶优化 实现 TRPO 的稳定性保证。PPO 不是显式约束 KL 散度,而是修改目标函数本身,通过裁剪机制阻止大的策略更新。它使用标准梯度下降隐式限制策略移动的距离,提供"软"信赖域。
裁剪代理目标
第六节的 CPI 目标和概率比:
L CPI ( θ ) = E t [ r t ( θ ) A ^ t ] where r t ( θ ) = π θ ( a t ∣ s t ) π θ o l d ( a t ∣ s t ) L^{\text{CPI}}(\theta) = \mathbb{E}t[r_t(\theta) \hat{A}t] \quad \text{where} \quad r_t(\theta) = \frac{\pi\theta(a_t | s_t)}{\pi{\theta_{old}}(a_t | s_t)} LCPI(θ)=Et[rt(θ)A^t]wherert(θ)=πθold(at∣st)πθ(at∣st)
L^CPI 的问题是没有阻止 r_t(θ) 变得任意大或小。PPO 通过裁剪概率比使其保持在 [1-ε, 1+ε] 内解决这个问题:
L CLIP ( θ ) = E t [ min ( r t ( θ ) A ^ t , clip ( r t ( θ ) , 1 − ϵ , 1 + ϵ ) ⋅ A ^ t ) ] (VIII.I) \boxed{L^{\text{CLIP}}(\theta) = \mathbb{E}_t[\min(r_t(\theta) \hat{A}_t, \; \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) \cdot \hat{A}_t)]} \tag{VIII.I} LCLIP(θ)=Et[min(rt(θ)A^t,clip(rt(θ),1−ϵ,1+ϵ)⋅A^t)](VIII.I)
其中 ε 是超参数(PPO 论文中 ε=0.2),裁剪函数定义为:
clip ( r , 1 − ϵ , 1 + ϵ ) = { 1 − ϵ if r < 1 − ϵ r if 1 − ϵ ≤ r ≤ 1 + ϵ 1 + ϵ if r > 1 + ϵ \text{clip}(r, 1-\epsilon, 1+\epsilon) = \begin{cases} 1-\epsilon & \text{if } r < 1-\epsilon \\ r & \text{if } 1-\epsilon \leq r \leq 1+\epsilon \\ 1+\epsilon & \text{if } r > 1+\epsilon \end{cases} clip(r,1−ϵ,1+ϵ)=⎩ ⎨ ⎧1−ϵr1+ϵif r<1−ϵif 1−ϵ≤r≤1+ϵif r>1+ϵ
(VIII.I)中的 min 运算符很重要。它确保取裁剪和未裁剪目标之间更悲观(更低)的估计。这根据优势的符号创建不同行为:
情况 1:正优势 (Â_t > 0)
当动作优于平均时,希望增加其概率,意味着增加 r_t(θ)。目标变为:
L t CLIP = min ( r t ( θ ) , 1 + ϵ ) ⋅ A ^ t L^{\text{CLIP}}_t = \min(r_t(\theta), 1+\epsilon) \cdot \hat{A}_t LtCLIP=min(rt(θ),1+ϵ)⋅A^t
- 如果 r_t(θ) ≤ 1+ε:目标是 r_t(θ)Â_t,所以梯度上升增加 r_t(θ)
- 如果 r_t(θ) > 1+ε:目标变为 (1+ε)Â_t
裁剪移除了将 r_t(θ) 增加到超过 1+ε 的动机。
情况 2:负优势 (Â_t < 0)
当动作劣于平均时,希望减少 其概率,意味着减少 r_t(θ)。由于 Â_t < 0,乘以更小的 r_t 使乘积更不负(更大)。目标变为:
L t CLIP = max ( r t ( θ ) , 1 − ϵ ) ⋅ A ^ t L^{\text{CLIP}}_t = \max(r_t(\theta), 1-\epsilon) \cdot \hat{A}_t LtCLIP=max(rt(θ),1−ϵ)⋅A^t
(负值的 min 在选择哪个 r_t 方面变成 max。)
- 如果 r_t(θ) ≥ 1-ε:目标是 r_t(θ)Â_t,所以梯度上升减少 r_t(θ)
- 如果 r_t(θ) < 1-ε:目标变为 (1-ε)Â_t
裁剪移除了将 r_t(θ) 减少到低于 1-ε 的动机。
这里的要点是 PPO 提供 L^CPI 的悲观下界。当更新会使事情"好得难以置信"时忽略它们。
LLM 场景(来自 Umar Jamil 视频): 在语言模型微调中,策略 π_θ(a_t|s_t) 是模型给定上下文 s_t(提示词+之前生成的词元)对词元 a_t 分配的概率。概率比 r_t(θ) 衡量微调模型相比参考策略生成特定词元的可能性增减程度。裁剪确保单次更新迭代中没有单个词元的概率能改变超过 (1±ε) 倍,防止模型对高优势词元"反应过度"。
PPO 目标
实践中,PPO 将裁剪策略目标与两个额外项结合:
L PPO ( θ ) = E t [ L t CLIP ( θ ) − c 1 L t VF ( θ ) + c 2 S [ π θ ] ( s t ) ] (VIII.II) \boxed{L^{\text{PPO}}(\theta) = \mathbb{E}_t[L^{\text{CLIP}}_t(\theta) - c_1 L^{\text{VF}}t(\theta) + c_2 S[\pi\theta](s_t)]} \tag{VIII.II} LPPO(θ)=Et[LtCLIP(θ)−c1LtVF(θ)+c2S[πθ](st)](VIII.II)
1. 值函数损失 (L^VF): 回顾第五节,需要值函数 V_φ(s) 计算优势估计。值函数训练以最小化其预测与实际回报的平方误差:
L t VF ( θ ) = ( V θ ( s t ) − V t target ) 2 L^{\text{VF}}t(\theta) = (V\theta(s_t) - V_t^{\text{target}})^2 LtVF(θ)=(Vθ(st)−Vttarget)2
其中 V_t^target 通常是折扣 return-to-go。当策略和值函数共享参数时(LLM 微调中常见,两者使用相同 Transformer 骨干),这个损失从目标中减去(因此是负号,因为最大化 L^PPO 但最小化 L^VF)。
2. 熵奖励 (S[π_θ]): 为鼓励探索并防止过早收敛到确定性策略,PPO 添加熵损失:
S [ π θ ] ( s t ) = − ∑ a π θ ( a ∣ s t ) log π θ ( a ∣ s t ) S[\pi_\theta](s_t) = -\sum_a \pi_\theta(a|s_t) \log \pi_\theta(a|s_t) S[πθ](st)=−a∑πθ(a∣st)logπθ(a∣st)
这里系数 c₁, c₂ > 0 控制正则化强度。
九、带 KL 惩罚的完整 PPO 目标
用"原版"PPO 微调 LLM 时,策略学习最大化奖励模型的奖励。但奖励模型是人类偏好的不完美代理,是在有限数据上训练的神经网络,可能被利用。没有约束,策略可能发现达到高奖励分数的对抗性输出,同时产生的文本:
- 退化为愚弄奖励模型的重复或无意义模式
- 偏离自然语言太远,失去流畅性和连贯性
- 利用奖励模型学到的虚假相关性
这种现象称为奖励黑客。策略找到"欺骗"奖励模型的方法,而非真正改进响应质量。
为防止奖励黑客,InstructGPT 论文添加 KL 散度惩罚 ,正则化策略使其接近参考模型 π_ref(通常是强化学习微调前的 SFT 模型)。
从第八节,PPO 目标(通过梯度上升最大化)包含三项:
L PPO ( θ ) = L CLIP ( θ ) ⏟ 裁剪策略目标 − c 1 L VF ( θ ) ⏟ 值函数损失 + c 2 S [ π θ ] ⏟ 熵奖励 L^{\text{PPO}}(\theta) = \underbrace{L^{\text{CLIP}}(\theta)}{\text{裁剪策略目标}} - \underbrace{c_1 L^{\text{VF}}(\theta)}{\text{值函数损失}} + \underbrace{c_2 S[\pi_\theta]}_{\text{熵奖励}} LPPO(θ)=裁剪策略目标 LCLIP(θ)−值函数损失 c1LVF(θ)+熵奖励 c2S[πθ]
现在不直接使用原始奖励模型分数,而是定义KL 惩罚奖励,正则化策略使其接近参考模型 π_ref:
r total ( s t , a t ) = r RM ( s t , a t ) − β ⋅ D KL ( π θ ( ⋅ ∣ s t ) ∥ π ref ( ⋅ ∣ s t ) ) (IX.I) \boxed{r_{\text{total}}(s_t, a_t) = r_{\text{RM}}(s_t, a_t) - \beta \cdot D_{\text{KL}}(\pi_\theta(\cdot|s_t) \| \pi_{\text{ref}}(\cdot|s_t))} \tag{IX.I} rtotal(st,at)=rRM(st,at)−β⋅DKL(πθ(⋅∣st)∥πref(⋅∣st))(IX.I)
其中:
- r_RM(s_t, a_t) 是时间步 t 的奖励信号
- β 是 KL 惩罚系数
- π_ref 是冻结的参考模型
在每个词元位置,KL 散度简化为:
D KL ( π θ ( ⋅ ∣ s t ) ∥ π ref ( ⋅ ∣ s t ) ) = E a ∼ π θ [ log π θ ( a ∣ s t ) π ref ( a ∣ s t ) ] D_{\text{KL}}(\pi_\theta(\cdot|s_t) \| \pi_{\text{ref}}(\cdot|s_t)) = \mathbb{E}{a \sim \pi\theta}\left[\log \frac{\pi_\theta(a|s_t)}{\pi_{\text{ref}}(a|s_t)}\right] DKL(πθ(⋅∣st)∥πref(⋅∣st))=Ea∼πθ[logπref(a∣st)πθ(a∣st)]
实践中用采样的词元 a_t 估计这个期望,得到:
d ^ t = log π θ ( a t ∣ s t ) π ref ( a t ∣ s t ) \hat{d}t = \log \frac{\pi\theta(a_t|s_t)}{\pi_{\text{ref}}(a_t|s_t)} d^t=logπref(at∣st)πθ(at∣st)
注意奖励模型 r_φ(x,y) 为完整响应 (x,y) 产生单个标量。这个分数只在最终词元 T 分配,而 KL 惩罚应用于每个词元。
r ~ ϕ = { − β ⋅ log π θ ( a t ∣ s t ) π ref ( a t ∣ s t ) if t < T r ϕ ( x , y ) − β ⋅ log π θ ( a T ∣ s T ) π ref ( a T ∣ s T ) if t = T \tilde{r}\phi = \begin{cases} -\beta \cdot \log \frac{\pi\theta(a_t | s_t)}{\pi_{\text{ref}}(a_t | s_t)} & \text{if } t < T \\[8pt] r_\phi(x, y) - \beta \cdot \log \frac{\pi_\theta(a_T | s_T)}{\pi_{\text{ref}}(a_T | s_T)} & \text{if } t = T \end{cases} r~ϕ=⎩ ⎨ ⎧−β⋅logπref(at∣st)πθ(at∣st)rϕ(x,y)−β⋅logπref(aT∣sT)πθ(aT∣sT)if t<Tif t=T
KL 惩罚有两个目的:
- 防止奖励黑客: 策略不能任意偏离自然语言
- 保持流畅性: 输出在分布上与训练良好的 SFT 模型保持相似
它通过修改的每词元奖励修改 PPO 中使用的优势估计 Â_t。但在数学上等价(且实现上更高效)的是直接将 KL 项添加到目标。带 KL 惩罚的 PPO 目标是:
J ( θ ) = E a ∼ π θ [ r RM ( s , a ) ] ⏟ 原版 PPO 目标 − β ⋅ D KL ( π θ ∥ π ref ) ⏟ KL 惩罚项 J(\theta) = \underbrace{\mathbb{E}{a \sim \pi\theta}[r_{\text{RM}}(s, a)]}{\text{原版 PPO 目标}} - \underbrace{\beta \cdot D{\text{KL}}(\pi_\theta \| \pi_{\text{ref}})}_{\text{KL 惩罚项}} J(θ)=原版 PPO 目标 Ea∼πθ[rRM(s,a)]−KL 惩罚项 β⋅DKL(πθ∥πref)
第一项正是原版 PPO 使用裁剪代理优化的内容。KL 惩罚项作为单独的加性组件出现,惩罚与参考模型的偏离。将第一项替换为 PPO 裁剪代理:
J c ( θ ) = E t [ min ( r t ( θ ) A ^ t , clip ( r t ( θ ) , 1 − ϵ , 1 + ϵ ) A ^ t ) ] − β ⋅ D KL ( π θ ∥ π ref ) J_{\text{c}}(\theta) = \mathbb{E}t[\min(r_t(\theta) \hat{A}t, \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) \hat{A}t)] - \beta \cdot D{\text{KL}}(\pi\theta \| \pi{\text{ref}}) Jc(θ)=Et[min(rt(θ)A^t,clip(rt(θ),1−ϵ,1+ϵ)A^t)]−β⋅DKL(πθ∥πref)
综合所有组件,带 KL 惩罚的完整 PPO 目标(要最大化)是:
L RLHF ( θ ) = L CLIP ( θ ) ⏟ 策略目标 − c 1 L VF ( θ ) ⏟ 值损失 + c 2 S [ π θ ] ⏟ 熵奖励 − β ⋅ D KL ( π θ ∥ π ref ) ⏟ KL 惩罚 (IX.II) \boxed{L^{\text{RLHF}}(\theta) = \underbrace{L^{\text{CLIP}}(\theta)}{\text{策略目标}} - \underbrace{c_1 L^{\text{VF}}(\theta)}{\text{值损失}} + \underbrace{c_2 S[\pi_\theta]}{\text{熵奖励}} - \underbrace{\beta \cdot D{\text{KL}}(\pi_\theta \| \pi_{\text{ref}})}_{\text{KL 惩罚}}} \tag{IX.II} LRLHF(θ)=策略目标 LCLIP(θ)−值损失 c1LVF(θ)+熵奖励 c2S[πθ]−KL 惩罚 β⋅DKL(πθ∥πref)(IX.II)
这里每一项都有明确目的:
| 项 | 作用 |
|---|---|
| 策略目标 L^CLIP | 改进策略,同时通过裁剪防止破坏性更新 |
| 值损失 c₁L^VF | 训练评论家进行准确优势估计(减去以最小化) |
| 熵奖励 c₂S[π_θ] | 鼓励探索,防止过早收敛 |
| KL 惩罚 βD_KL | 防止奖励黑客,保持语言质量 |
重要的是区分完整损失中两个 KL 相关机制。PPO 裁剪机制作为短期锚点 ,约束单次更新中策略能改变多少;而 KL 惩罚是长期锚点,约束整个训练过程中策略能从起点偏离多远。
终于完成了...
这就是完整推导!令人满意的是最终损失中的每一项都有特定目的,每项存在都是因为遇到特定问题并需要修复它。必须承认理解损失背后的所有数学和概念并不容易,仍然不能完全理解每个细节,但比几天前理解得好多了。
希望这篇文章有用。如果发现推导中有任何错误(肯定会有)或有建议,欢迎联系。
HF BLOG:https://huggingface.co/blog/garg-aayush/ppo-from-first-principle
参考文献
-
视频:
- Umar Jamil 关于 RLHF 和 PPO 的视频:涵盖 RLHF 和 PPO 概念的全面且必看视频
-
论文:
- Proximal Policy Optimization Algorithms:介绍裁剪代理目标的基础 PPO 论文
- Training language models to follow instructions with human feedback:展示带 KL 惩罚的 PPO 以减轻 LLM 微调中奖励黑客的 InstructGPT 论文
- Trust Region Policy Optimization:提出 PPO 中使用的信赖域约束的 TRPO 论文
- High-Dimensional Continuous Control Using Generalized Advantage Estimation:介绍用于策略梯度方差降低的指数加权优势估计器的 GAE 论文