4 models in PPO. RLHF 基于 SFT 模型继续进行训练,首先需要使用 SFT 模型针对同一问题输出多个不同回答交给标注员进行排序或打分,得到的人工标注的偏好数据可以用于训练 reward model,其中 reward model 直接用 SFT 模型初始化,训练的 reward model 用于在 RL 训练中提供 reward 判断 LLM 生成回答的好坏;在 PPO 训练中,我们还需要 actor model 和 critic model,其中 actor model 就是我们想要对齐偏好的 LLM,可以使用 SFT 模型初始化,critic model 则是 PPO 中估计 value function 的模型,可以使用 reward model 初始化 (也可以和 actor model 共享部分权重);此外,为了防止对齐后的模型输出显著偏离 SFT 模型,通常还会引入由 SFT 模型初始化的 reference model,将 actor model 和 reference model 的 KL 散度损失作为正则项,使得对齐后的模型输出不要过度偏离 SFT 模型;总的来说,PPO 里需要用到 4 个不同的 LLM!其中两个需要训练,两个参数冻结,总体需要巨量的内存开销!
(2) 计算 reward. reward model 可以针对 actor 的生成结果
s
T
s_T
sT 给出 reward
R
(
s
T
)
R(s_T)
R(sT),相当于我们只能在 episode 结束时得到 reward,只能反映出生成结果的好坏,那我们怎么衡量中间生成过程的合理性呢?一种简单粗暴的方法是:循规蹈矩,只要 actor 生成的 token 和 reference model 比较一致,那么就可以额外得到少量的 reward. 相当于是在 reward 里加上
log
P
r
e
f
(
a
t
∣
s
t
)
−
log
P
(
a
t
∣
s
t
)
\log P_{r e f}(a_t \mid s_t)-\log P(a_t \mid s_t)
logPref(at∣st)−logP(at∣st),
log
P
r
e
f
(
a
t
∣
s
t
)
\log P_{r e f}(a_t \mid s_t)
logPref(at∣st) 越高,说明 actor 和 reference model 越一致,所获得的 reward 也越高,而
−
log
P
(
a
t
∣
s
t
)
-\log P(a_t \mid s_t)
−logP(at∣st) 则作为正则项,保证了概率分布的多样性,使得 actor 和 reference model 的输出内容又不那么一样;其实这项 reward 可以看作是 KL 散度
KL
(
P
∥
P
r
e
f
)
\text{KL}(P\|P_{ref})
KL(P∥Pref) 的简化版本
{
r
t
=
−
k
⋅
log
P
(
a
t
∣
s
t
)
P
r
e
f
(
a
t
∣
s
t
)
,
t
≠
T
r
t
=
−
k
⋅
log
P
(
a
t
∣
s
t
)
P
r
e
f
(
a
t
∣
s
t
)
+
R
(
s
t
)
,
t
=
T
\left\{\begin{array}{l} r_t=-k \cdot\log \frac{P(a_t \mid s_t)}{P_{r e f}(a_t \mid s_t)}, \quad \quad \quad \quad\ \ \ t \neq T \\ r_t=-k \cdot\log \frac{P(a_t \mid s_t)}{P_{r e f}(a_t \mid s_t)}+R(s_t), \quad t=T \end{array}\right.
{rt=−k⋅logPref(at∣st)P(at∣st),t=Trt=−k⋅logPref(at∣st)P(at∣st)+R(st),t=T其中,
k
=
0.1
k=0.1
k=0.1 为超参
(3) 计算 advantage. critic 预测得到
V
θ
(
s
t
)
V^{\theta}(s_t)
Vθ(st),结合
r
t
r_t
rt 可以计算 advantage
A
θ
(
s
t
,
a
t
)
=
r
t
+
γ
V
θ
(
s
t
+
1
)
−
V
θ
(
s
t
)
A^{\theta}(s_t,a_t)=r_t+\gamma V^{\theta}(s_{t+1})-V^{\theta}(s_t)
Aθ(st,at)=rt+γVθ(st+1)−Vθ(st)(Generalized Advantage Estimation, GAE). 上述 advantage 代表着
t
t
t 时刻的即时优势,我们也可以引入未来优势,重新定义 advantage 为
A
θ
(
s
t
,
a
t
)
=
(
r
t
+
γ
V
θ
(
s
t
+
1
)
−
V
θ
(
s
t
)
)
+
λ
⋅
γ
A
θ
(
s
t
+
1
,
a
t
+
1
)
A^{\theta}(s_t,a_t)=(r_t+\gamma V^\theta(s_{t+1})-V^\theta(s_t))+\lambda\cdot \gamma A^{\theta}(s_{t+1},a_{t+1})
Aθ(st,at)=(rt+γVθ(st+1)−Vθ(st))+λ⋅γAθ(st+1,at+1)其中,
A
(
s
t
,
a
t
)
A(s_{t},a_{t})
A(st,at) 的计算可以采用动态规划的方法从
T
T
T 时刻往前倒推
(4) PPO 训练. 使用采样数据
(
s
t
,
a
t
)
∼
π
θ
old
(s_t,a_t)\sim\pi_{\theta_{\text{old}}}
(st,at)∼πθold 训练多个 epochs 的 actor 和 critic,actor loss 为
L
actor
=
−
min
(
p
θ
(
a
t
∣
s
t
)
p
θ
old
(
a
t
∣
s
t
)
A
θ
old
(
s
t
,
a
t
)
,
clip
(
p
θ
(
a
t
∣
s
t
)
p
θ
old
(
a
t
∣
s
t
)
,
1
−
ε
,
1
+
ε
)
A
θ
old
(
s
t
,
a
t
)
)
\begin{aligned} L_{\text{actor}}=-\min(&\frac{p_{\theta}(a_t|s_t)}{p_{\theta_{\text{old}}}(a_t|s_t)}A^{\theta_{\text{old}}}(s_t,a_t), \\&\text{clip}(\frac{p_{\theta}(a_t|s_t)}{p_{\theta_{\text{old}}}(a_t|s_t)},1-\varepsilon,1+\varepsilon)A^{\theta_{\text{old}}}(s_t,a_t)) \end{aligned}
Lactor=−min(pθold(at∣st)pθ(at∣st)Aθold(st,at),clip(pθold(at∣st)pθ(at∣st),1−ε,1+ε)Aθold(st,at))critic loss 为 “实际回报” 和 “预估回报” 的 MSE loss,其中 “实际回报” 为
r
t
+
γ
V
θ
(
s
t
+
1
)
r_t+\gamma V^\theta(s_{t+1})
rt+γVθ(st+1),由于我们还引入了未来优势,因此可以将 “实际回报” 进一步优化为
A
θ
(
s
t
,
a
t
)
+
V
θ
(
s
t
)
A^{\theta}(s_t,a_t)+V^\theta(s_{t})
Aθ(st,at)+Vθ(st);“预估回报” 为
V
θ
old
(
s
t
)
V^{\theta_{\text{old}}}(s_t)
Vθold(st),也就是 old critic model 的预估 value function;因此 critic loss 为
L
critic
=
(
A
θ
(
s
t
,
a
t
)
+
V
θ
(
s
t
)
−
V
θ
old
(
s
t
)
)
2
\begin{aligned} L_{\text{critic}}=(A^{\theta}(s_t,a_t)+V^\theta(s_{t})-V^{\theta_{\text{old}}}(s_t))^2 \end{aligned}
Lcritic=(Aθ(st,at)+Vθ(st)−Vθold(st))2总的 loss 为
L
=
L
actor
+
0.1
⋅
L
critic
\begin{aligned} L=L_{\text{actor}}+0.1\cdot L_{\text{critic}} \end{aligned}
L=Lactor+0.1⋅Lcritic
DPO 的损失函数. 在 RLHF 中,我们的目标函数为
max
π
θ
E
x
∼
D
,
y
∼
π
θ
(
y
∣
x
)
[
r
ϕ
(
x
,
y
)
]
−
β
D
K
L
[
π
θ
(
y
∣
x
)
∣
∣
π
r
e
f
(
y
∣
x
)
]
\max_{\pi_\theta}\mathbb{E}_{x\sim\mathcal{D},y\sim\pi_\theta(y|x)}\bigl[r_\phi(x,y)\bigr]-\beta\mathbb{D}_{\mathbf{KL}}\bigl[\pi_\theta(y\mid x)\mid\mid\pi_{\mathbf{ref}}(y\mid x)\bigr]
πθmaxEx∼D,y∼πθ(y∣x)[rϕ(x,y)]−βDKL[πθ(y∣x)∣∣πref(y∣x)]将 KL 散度展开可得
max
π
E
x
∼
D
,
y
∼
π
(
y
∣
x
)
[
r
(
x
,
y
)
]
−
β
D
K
L
[
π
(
y
∣
x
)
∣
∣
π
r
e
f
(
y
∣
x
)
]
=
max
π
E
x
∼
D
E
y
∼
π
(
y
∣
x
)
[
r
(
x
,
y
)
−
β
log
π
(
y
∣
x
)
π
r
e
f
(
y
∣
x
)
]
=
min
π
E
x
∼
D
E
y
∼
π
(
y
∣
x
)
[
log
π
(
y
∣
x
)
π
r
e
f
(
y
∣
x
)
−
1
β
r
(
x
,
y
)
]
=
min
π
E
x
∼
D
E
y
∼
π
(
y
∣
x
)
[
log
π
(
y
∣
x
)
1
Z
(
x
)
π
r
e
f
(
y
∣
x
)
exp
(
1
β
r
(
x
,
y
)
)
−
log
Z
(
x
)
]
\begin{aligned} &\max_{\pi}\mathbb{E}_{x\sim\mathcal{D},y\sim\pi(y|x)}\bigl[r(x,y)\bigr]-\beta\mathbb{D}_{\mathbf{KL}}\bigl[\pi(y\mid x)\mid\mid\pi_{\mathbf{ref}}(y\mid x)\bigr] \\ =&\max_\pi\mathbb{E}_{x\sim\mathcal{D}}\mathbb{E}_{y\sim\pi(y|x)}\left[r(x,y)-\beta\log\frac{\pi(y|x)}{\pi_{\mathrm{ref}}(y|x)}\right] \\ =&\min_{\pi}\mathbb{E}_{x\sim\mathcal{D}}\mathbb{E}_{y\sim\pi(y|x)}\left[\log\frac{\pi(y|x)}{\pi_{\mathrm{ref}}(y|x)}-\frac1\beta r(x,y)\right] \\ =&\min_\pi\mathbb{E}_{x\sim\mathcal{D}}\mathbb{E}_{y\sim\pi(y|x)}\left[\log\frac{\pi(y|x)}{\frac1{Z(x)}\pi_{\mathrm{ref}}(y|x)\exp\left(\frac1\beta r(x,y)\right)}-\log Z(x)\right] \end{aligned}
===πmaxEx∼D,y∼π(y∣x)[r(x,y)]−βDKL[π(y∣x)∣∣πref(y∣x)]πmaxEx∼DEy∼π(y∣x)[r(x,y)−βlogπref(y∣x)π(y∣x)]πminEx∼DEy∼π(y∣x)[logπref(y∣x)π(y∣x)−β1r(x,y)]πminEx∼DEy∼π(y∣x)logZ(x)1πref(y∣x)exp(β1r(x,y))π(y∣x)−logZ(x)由此可以定义一个新的概率分布
π
∗
(
y
∣
x
)
=
1
Z
(
x
)
π
ref
(
y
∣
x
)
exp
(
1
β
r
(
x
,
y
)
)
\pi^*(y|x)=\frac1{Z(x)}\pi_\text{ref}{(y|x)}\exp\left(\frac1\beta r(x,y)\right)
π∗(y∣x)=Z(x)1πref(y∣x)exp(β1r(x,y))其中,
Z
(
x
)
=
∑
y
π
r
e
f
(
y
∣
x
)
exp
(
1
β
r
(
x
,
y
)
)
Z(x)=\sum_y\pi_{\mathrm{ref}}(y|x)\exp\left(\frac1\beta r(x,y)\right)
Z(x)=∑yπref(y∣x)exp(β1r(x,y)) 为归一化因子,继续代入目标函数可得
min
π
E
x
∼
D
[
E
y
∼
π
(
y
∣
x
)
[
log
π
(
y
∣
x
)
π
∗
(
y
∣
x
)
]
−
log
Z
(
x
)
]
=
min
π
E
x
∼
D
[
D
K
L
(
π
(
y
∣
x
)
∥
π
∗
(
y
∣
x
)
)
−
log
Z
(
x
)
]
\begin{aligned}&\min_{\pi}\mathbb{E}_{x\sim\mathcal{D}}\left[\mathbb{E}_{y\sim\pi(y|x)}\left[\log\frac{\pi(y|x)}{\pi^*(y|x)}\right]-\log Z(x)\right]\\=&\min_{\pi}\mathbb{E}_{x\sim\mathcal{D}}\left[\mathbb{D}_{\mathrm{KL}}(\pi(y|x)\parallel\pi^*(y|x))-\log Z(x)\right]\end{aligned}
=πminEx∼D[Ey∼π(y∣x)[logπ∗(y∣x)π(y∣x)]−logZ(x)]πminEx∼D[DKL(π(y∣x)∥π∗(y∣x))−logZ(x)]由于
Z
(
x
)
Z(x)
Z(x) 和 LLM 无关,因此优化这个式子只需要优化前面的 KL 项,可见最优的
π
\pi
π 即为
π
∗
\pi^*
π∗,但
π
∗
\pi^*
π∗ 的计算仍然依赖于 reference model 和 reward model,我们想要进一步摆脱 reward model. 为此,我们可以反过来根据
π
∗
\pi^*
π∗ 的表达式去表示
r
(
x
,
y
)
r(x,y)
r(x,y),得到
r
(
x
,
y
)
=
β
log
π
r
(
y
∣
x
)
π
r
e
f
(
y
∣
x
)
+
β
log
Z
(
x
)
r(x,y)=\beta\log\frac{\pi_r(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}+\beta\log Z(x)
r(x,y)=βlogπref(y∣x)πr(y∣x)+βlogZ(x)在 RLHF 中我们希望 reward 能通过 Bradley–Terry model 来预测人类偏好,即
p
∗
(
y
w
≻
y
l
∣
x
)
=
1
1
+
exp
(
−
(
r
ϕ
(
x
,
y
w
)
−
r
ϕ
(
x
,
y
l
)
)
)
p^*(y_w\succ y_l\mid x)=\frac{1}{1+\exp\bigg(-\big(r_\phi(x,y_w)-r_\phi(x,y_l)\big)\bigg)}
p∗(yw≻yl∣x)=1+exp(−(rϕ(x,yw)−rϕ(x,yl)))1基于最大似然估计可以给出 reward model 的损失函数
L
R
(
r
ϕ
,
D
)
=
−
E
(
x
,
y
w
,
y
l
)
∼
D
[
log
σ
(
r
ϕ
(
x
,
y
w
)
−
r
ϕ
(
x
,
y
l
)
)
]
\mathcal{L}_R(r_\phi,\mathcal{D})=-\mathbb{E}_{(x,y_w,y_l)\sim\mathcal{D}}\left[\log\sigma(r_\phi(x,y_w)-r_\phi(x,y_l))\right]
LR(rϕ,D)=−E(x,yw,yl)∼D[logσ(rϕ(x,yw)−rϕ(x,yl))]我们将得到的
r
(
x
,
y
)
r(x,y)
r(x,y) 带入上述损失函数即可得到 DPO 的损失函数
L
D
P
O
(
π
θ
;
π
r
e
f
)
=
−
E
(
x
,
y
w
,
y
l
)
∼
D
[
log
σ
(
β
log
π
θ
(
y
w
∣
x
)
π
r
e
f
(
y
w
∣
x
)
−
β
log
π
θ
(
y
l
∣
x
)
π
r
e
f
(
y
l
∣
x
)
)
]
\mathcal{L}_{\mathrm{DPO}}(\pi_{\theta};\pi_{\mathrm{ref}})=-\mathbb{E}_{(x,y_{w},y_{l})\sim\mathcal{D}}\left[\log\sigma\left(\beta\log\frac{\pi_{\theta}(y_{w}\mid x)}{\pi_{\mathrm{ref}}(y_{w}\mid x)}-\beta\log\frac{\pi_{\theta}(y_{l}\mid x)}{\pi_{\mathrm{ref}}(y_{l}\mid x)}\right)\right]
LDPO(πθ;πref)=−E(x,yw,yl)∼D[logσ(βlogπref(yw∣x)πθ(yw∣x)−βlogπref(yl∣x)πθ(yl∣x))]其中,
r
^
(
x
,
y
)
=
β
log
π
θ
(
y
∣
x
)
π
r
e
f
(
y
∣
x
)
\hat r(x,y)=\beta\log\frac{\pi_{\theta}(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}
r^(x,y)=βlogπref(y∣x)πθ(y∣x) 即为 LLM 中隐含的 reward model,也就是论文标题中所说的 “Your language model is secretly a reward model”
Simple Preference Optimization (SimPO)
DPO 在优化时虽然相比 PPO 简化了许多,但仍需要同时加载 actor 和 reference model,训练开销仍然比较大;此外,DPO 的隐式 reward model
r
^
(
x
,
y
)
=
β
log
π
θ
(
y
∣
x
)
π
r
e
f
(
y
∣
x
)
\hat r(x,y)=\beta\log\frac{\pi_{\theta}(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}
r^(x,y)=βlogπref(y∣x)πθ(y∣x) 在计算 reward 时没有对生成文本长度做归一化,因此 DPO 在优化时模型会倾向于生成更长的文本来获取更大的 reward,尤其是当数据量比较少的时候;DPO 优化时还有训练和推理目标不完全对齐的问题,我们推理时想要的是
π
θ
(
y
w
∣
x
)
>
π
θ
(
y
l
∣
x
)
\pi_{\theta}(y_w\mid x)>\pi_{\theta}(y_l\mid x)
πθ(yw∣x)>πθ(yl∣x),但 DPO 的训练目标是
r
^
(
x
,
y
w
)
>
r
^
(
x
,
y
l
)
\hat r(x,y_w)>\hat r(x,y_l)
r^(x,yw)>r^(x,yl),也就是说即使
r
^
(
x
,
y
w
)
>
r
^
(
x
,
y
l
)
\hat r(x,y_w)>\hat r(x,y_l)
r^(x,yw)>r^(x,yl),我们也并不能保证 LLM 在推理时仍然能生成
y
w
y_w
yw
SimPO 针对上述问题进行了改进,无需 reference model 即可训练,并通过文本长度归一化和训推目标的一致性使得 LLM 在生成较短长度的前提下依然能够保证生成质量;至此,SimPO 的损失函数似乎和对比学习差不多了,RLHF 的过程逐渐由繁化简 (当然实际效果如何还需要实践出真知)
具体来说,SimPO 为了让训推目标对齐,直接把 DPO reward model 里的 reference model 拿掉了,此外还加入了长度归一化项
∣
y
∣
|y|
∣y∣,实验表明这个长度归一化项非常重要,由于 SimPO 中不存在 reference model 做隐式的长度归一化,因此去掉它之后模型就会输出冗长的低质量回答
r
S
i
m
P
O
(
x
,
y
)
=
β
∣
y
∣
log
π
θ
(
y
∣
x
)
=
β
∣
y
∣
∑
i
=
1
∣
y
∣
log
π
θ
(
y
i
∣
x
,
y
<
i
)
r_{\mathrm{SimPO}}(x,y)=\frac\beta{|y|}\log\pi_\theta(y\mid x)=\frac\beta{|y|}\sum_{i=1}^{|y|}\log\pi_\theta(y_i\mid x,y_{<i})
rSimPO(x,y)=∣y∣βlogπθ(y∣x)=∣y∣βi=1∑∣y∣logπθ(yi∣x,y<i)此外,SimPO 还引入了 reward margin
γ
\gamma
γ,要求正负样本对的 reward 差值大于
γ
\gamma
γ
p
(
y
w
≻
y
l
∣
x
)
=
σ
(
r
(
x
,
y
w
)
−
r
(
x
,
y
l
)
−
γ
)
p(y_w\succ y_l\mid x)=\sigma\left(r(x,y_w)-r(x,y_l)-\gamma\right)
p(yw≻yl∣x)=σ(r(x,yw)−r(x,yl)−γ)
最终,SimPO 的损失函数为
L
S
i
m
P
O
(
π
θ
)
=
−
E
(
x
,
y
w
,
y
l
)
∼
D
[
log
σ
(
β
∣
y
w
∣
l
o
g
π
θ
(
y
w
∣
x
)
−
β
∣
y
l
∣
l
o
g
π
θ
(
y
l
∣
x
)
−
γ
)
]
\mathcal{L}_{\mathrm{SimPO}}(\pi_\theta)=-\mathbb{E}_{(x,y_w,y_l)\sim\mathcal{D}}\left[\log\sigma\left(\frac\beta{|y_w|}\mathrm{log}\pi_\theta(y_w|x)-\frac\beta{|y_l|}\mathrm{log}\pi_\theta(y_l|x)-\gamma\right)\right]
LSimPO(πθ)=−E(x,yw,yl)∼D[logσ(∣yw∣βlogπθ(yw∣x)−∣yl∣βlogπθ(yl∣x)−γ)]