PPO? 强化学习
基本概念
强化学习五要素:智能体、行为、环境、状态、奖励。
先直接将五个要素映射到文本生成任务上:
智能体:生成模型。
行为: 获取当前step token的概率分布,并选取其中一个作为生成的token。
环境:一个评判标准,可以是模型,一般都是分类模型,seqgan是使用判别是否机器生成的分类器模型,chatgpt是人类偏好分类器;也可以是人为制定的评估标准,类似relu,rouge。
状态:当前step已生成的文本。
奖励:由环境结合当前的状态给出的分数。
状态价值的计算
V ( s ) = R ( s ) + γ ∑ s ′ ∈ S P ( s ′ ∣ s ) V ( s ′ ) V(s) = R(s) + \gamma \sum_{s^{'} \in S} P(s'|s) V(s') V(s)=R(s)+γs′∈S∑P(s′∣s)V(s′)
R ( s ) R(s) R(s): 当前状态的奖励
γ \gamma γ:折扣因子,因为当前的状态对后续的状态会随着步数的增加而减小。
先给出一个简单的例子:
一个生活中最常见的“吃饭”例子
好比吃饭可以带来奖励为6,折扣因子为0.8,对于吃饭状态的奖励值的计算为:
V ( 吃饭 ) = R ( 吃饭 ) + γ ∑ s ′ ∈ S P ( s ′ ∣ s ) V ( s ) = 6 + 0.8 ( 0.5 ∗ V ( 喝酒 ) + 0.5 ∗ V ( 喝茶 ) ) = 6 + 0.8 ( 0.5 ∗ ( 3 + 0.8 ( 1.0 ∗ 10 ) ) + 0.5 ∗ ( 5 + 0.8 ( 1.0 ∗ 4 ) ) ) \begin{aligned} V(吃饭) &= R(吃饭) + \gamma \sum_{s^{'} \in S} P(s^{'}|s) V(s)\\ &=6 + 0.8(0.5 * V(喝酒) + 0.5 * V(喝茶))\\ &=6 + 0.8(0.5 * (3 + 0.8(1.0 * 10)) + 0.5 * (5 + 0.8(1.0 * 4)))\\ \end{aligned} V(吃饭)=R(吃饭)+γs′∈S∑P(s′∣s)V(s)=6+0.8(0.5∗V(喝酒)+0.5∗V(喝茶))=6+0.8(0.5∗(3+0.8(1.0∗10))+0.5∗(5+0.8(1.0∗4)))
问题:1)对于文本生成而言,我们无法对每个状态给出一个奖励,我们只有在一个样本生成结束之后,才可以评判一个样本生成的好坏。
2)假设我们要生成一个20字长的文本,生成器的词表大小是20000,我们的状态空间就是20000的20次方,这是显然我们无法接受。
求解方法:
蒙特卡洛
蒙特卡洛方法,也称为统计模拟方法,就是通过大量的随机样本来估算或近似真实值,比如近似估算圆的面积。
可以通过 圆的面积/ 正方形的面积 = 圆中点的个数/正方形中点的个数 来计算圆的面积。
类似的,在文本生成中我们也可以用蒙特卡洛方法来估计一个模型的状态价值。考虑到 一个状态的价值是它的期望回报,那么如果我们用当前模型采样很多条序列,然后计算从这个状态出发的回报再求其期望就可以了,这样我们就同时解决了上述的两个问题。
例子:
输入:今天的天气怎么样?
模型采样输出:1)s1 = 很抱歉,我无法回答当前天气情况,因为我没有实时获取天气信息的功能。 7
2)s2 = 今天天气晴朗。 3
3)s3 = 今天会下冰雹。 5
…
E π ( " 今天天气怎么样 ? " ) = ∑ τ R ( τ ) P π ( τ ) E_{\pi}(^"今天天气怎么样?^") = \sum_{\tau}R(\tau)P_{\pi}(\tau) Eπ("今天天气怎么样?")=τ∑R(τ)Pπ(τ)
τ \tau τ 是我们任何可能生成的文本, π \pi π 是我们当前的生成模型,我们期望最大化 E π ( " 今天天气怎么样 ? " ) E_{\pi}(^"今天天气怎么样?^") Eπ("今天天气怎么样?"),既然要最大化,我们肯定需要对期望进行求导。
∇ E π ( " 今天天气怎么样 ? " ) = ∑ τ R ( τ ) ∇ P π ( τ ) = ∑ τ R ( τ ) P π ( τ ) ∇ P π ( τ ) P π ( τ ) = ∑ τ R ( τ ) P π ( τ ) ∇ l o g ( P π ( τ ) ) = E π ( " 今天天气怎么样 ? " ) [ R ( τ ) ∇ l o g ( P π ( τ ) ) ] \begin{aligned} \nabla E_{\pi}(^"今天天气怎么样?^") &= \sum_{\tau}R(\tau)\nabla P_{\pi}(\tau)\\ &= \sum_{\tau}R(\tau)P_{\pi}(\tau) \frac {\nabla P_{\pi}(\tau)} {P_{\pi}(\tau)}\\ &=\sum_{\tau}R(\tau)P_{\pi}(\tau)\nabla log(P_{\pi}(\tau))\\ &= E_{\pi}(^"今天天气怎么样?^")[R(\tau)\nabla log(P_{\pi}(\tau))] \end{aligned} ∇Eπ("今天天气怎么样?")=τ∑R(τ)∇Pπ(τ)=τ∑R(τ)Pπ(τ)Pπ(τ)∇Pπ(τ)=τ∑R(τ)Pπ(τ)∇log(Pπ(τ))=Eπ("今天天气怎么样?")[R(τ)∇log(Pπ(τ))]
显然 E π ( " 今天天气怎么样 ? " ) [ R ( τ ) ∇ l o g ( P π ( τ ) ) ] E_{\pi}(^"今天天气怎么样?^")[R(\tau)\nabla log(P_{\pi}(\tau))] Eπ("今天天气怎么样?")[R(τ)∇log(Pπ(τ))] 我们无法求解,这里我们就可以使用蒙特卡洛来近似求期望。
∇ E π ( " 今天天气怎么样 ? " ) ≈ 1 N ∑ n = 1 N R ( τ n ) ∇ l o g ( P π ( τ n ) ) = 1 N ∑ n = 1 N R ( " 今天天气怎么样 ? " + s n ) ∇ l o g ( P π ( s n ∣ " 今天天气怎么样 ? " ) ) \begin{aligned} \nabla E_{\pi}(^"今天天气怎么样?^") &\approx \frac 1N \sum_{n=1}^{N}R(\tau^n)\nabla log(P_{\pi}(\tau^n))\\ &= \frac 1N \sum_{n=1}^{N}R(^"今天天气怎么样?^" + s^{n})\nabla log(P_{\pi}(s^{n}|^"今天天气怎么样?^")) \end{aligned} ∇Eπ("今天天气怎么样?")≈N1n=1∑NR(τn)∇log(Pπ(τn))=N1n=1∑NR("今天天气怎么样?"+sn)∇log(Pπ(sn∣"今天天气怎么样?"))
重要性采样
对于蒙特卡洛来说,我们需要采样尽可能多的样本,才能让估计更加精准,就像上图,如果只拿两个点估计圆的面积,显然会出现很大的差错。
而对于模型来说,每一轮的更新,都会导致模型 π \pi π 的变化,从 π t → π t + 1 \pi_{t} \rightarrow \pi_{t+1} πt→πt+1 ,在上一轮采样获取的文本就无法和生成模型 π t + 1 \pi_{t+1} πt+1 生成的概率分布,因此上一轮采样的样本就无法使用了,所以在一轮都采样大量样本都是不现实的,而且对于chatgpt来说,还需把新采样获取的样本放入奖励模型获取奖励值,进一步加剧了时间的消耗,因此我们可以使用一个初始生成模型 π ′ \pi^{'} π′ ,使用 π ′ \pi^{'} π′ 一次性采样大量的样本后,然后一直使用这些样本进行训练。
重要性采样具体而言:
如果我们直接使用 π ′ \pi^{'} π′ 采样的样本进行训练,就会存在近似期望不准确的情况,毕竟 π ′ \pi^{'} π′的期望和 π \pi π的期望是不相等,因此我们希望将 π ′ \pi^{'} π′的期望来表示 π \pi π的期望:
E π = ∑ τ R ( τ ) P π ( τ ) = ∑ τ R ( τ ) P π ′ ( τ ) P π ( τ ) P π ′ ( τ ) = E π ′ P π ( τ ) P π ′ ( τ ) \begin{aligned} E_{\pi} &= \sum_{\tau}R(\tau)P_{\pi}(\tau)\\ &=\sum_{\tau}R(\tau)P_{\pi^{'}}(\tau) \frac {P_{\pi}(\tau)}{P_{\pi^{'}}(\tau)}\\ &= E_{\pi^{'}}\frac {P_{\pi}(\tau)}{P_{\pi^{'}}(\tau)} \end{aligned} Eπ=τ∑R(τ)Pπ(τ)=τ∑R(τ)Pπ′(τ)Pπ′(τ)Pπ(τ)=Eπ′Pπ′(τ)Pπ(τ)
因此梯度更新公式可以改为:
∇ E π ( " 今天天气怎么样 ? " ) ≈ 1 N ∑ n = 1 N R ( τ n ) P π ( τ n ) P π ′ ( τ n ) ∇ l o g ( P π ( τ n ) ) \begin{aligned} \nabla E_{\pi}(^"今天天气怎么样?^") &\approx \frac 1N \sum_{n=1}^{N}R(\tau^n) \frac {P_{\pi}(\tau^n)}{P_{\pi^{'}}(\tau^n)} \nabla log(P_{\pi}(\tau^n))\\ \end{aligned} ∇Eπ("今天天气怎么样?")≈N1n=1∑NR(τn)Pπ′(τn)Pπ(τn)∇log(Pπ(τn))
Advantage Actor-Critic
还是按照上述采样的生成文本作为例子:
1)s1 = 很抱歉,我无法回答当前天气情况,因为我没有实时获取天气信息的功能。 7
2)s2 = 今天天气晴朗。 3
3)s3 = 今天会下冰雹。 5
假如我们只采样到了s1和s2,没有采样到s3,由于7和3都是正向奖励,s1和s2的训练后生成的概率都会变大,且s1的概率变的更大,这看似合理,但是s3是未参与训练的,它的概率反而减小了。
所以为了避免这种情况,会对奖励增加基准线,一般使用采样获得的样本的均值作为基准线 b 奖励的计算公式转为 r ( τ ) − b r(\tau) - b r(τ)−b。
在RLHF实际代码实现方面,并没有使用该方法。但是在我以前的尝试中,该方法是一个较为有效,且很多生成+强化学习的论文中都会使用的方法。
同时为了进一步提升训练效率,有些论文就会选取最佳的一句文本进行训练来加快奖励值收敛速度。用上述的例子来说,在基线b是 (3 + 5 + 7)/ 3 = 5,而后只使用奖励值最高的s1作为训练样本进行训练。而这种方法往往在提升奖励值上会有很好的效果,但生成的样本的语义连贯性较差。
TRPO算法
重要性采样还是存在一些问题,对于增加权重 P π ( τ ) P π ′ ( τ ) \frac {P_{\pi}(\tau)}{P_{\pi^{'}}(\tau)} Pπ′(τ)Pπ(τ) ,的确可以是得两者的期望值趋于一致,但是并没有考虑方差,TRPO算法就是为了避免两者的方差差距过大而产生的。
简而言之 TRPO算法就是增加了一个KL散度限制,使用KL散度作为信任度,防止当前生成模型采样到KL散度过大的样本。
具体而言原始生成模型采样获取的一个文本,同时我们也可以获取原始生成模型对于生成该文本每个token的概率,我们也可以获取当前模型生成该文本每个token的概率。如果这两个概率分布的KL散度差距过大,我们就会认为该样本不在信任区域内,就不会使用该样本训练。
PPO算法
TRPO的问题在于把 KL 散度约束当作一个额外的约束,没有放在目标里面,导致TRPO很难计算,总之因为信任域的计算量太大了。
PPO-penalty:
直接将KL散度限制项增加到损失函数中:
∇ E π ( " 今天天气怎么样 ? " ) ≈ 1 N ∑ n = 1 N R ( τ n ) P π ( τ n ) P π ′ ( τ n ) ∇ l o g ( P π ( τ n ) ) − β K L ( P π ( τ n ) , P π ′ ( τ n ) ) \begin{aligned} \nabla E_{\pi}(^"今天天气怎么样?^") &\approx \frac 1N \sum_{n=1}^{N}R(\tau^n) \frac {P_{\pi}(\tau^n)}{P_{\pi^{'}}(\tau^n)} \nabla log(P_{\pi}(\tau^n)) -\beta KL(P_{\pi}(\tau^n),P_{\pi^{'}}(\tau^n))\\ \end{aligned} ∇Eπ("今天天气怎么样?")≈N1n=1∑NR(τn)Pπ′(τn)Pπ(τn)∇log(Pπ(τn))−βKL(Pπ(τn),Pπ′(τn))
PPO-clip:
直接使用截断 P π ( τ n ) P π ′ ( τ n ) \frac {P_{\pi}(\tau^n)}{P_{\pi^{'}}(\tau^n)} Pπ′(τn)Pπ(τn)的方法进行限制,期望重要性权重在 ( 1 − ϵ , 1 + ϵ ) (1-\epsilon, 1+\epsilon) (1−ϵ,1+ϵ) 之内,从而保证了两者KL散度不会过大。
chatgpt强化学习的实现:
首先对于RLHF模型的初始化,并没有直接使用SFT初始化,而是GPT3经过SFT数据集在+10%的pretrain数据训练两轮后作为初始化模型;奖励模型都是使用6B模型进行训练,当时也尝试使用175B进行训练,然后并不稳定。强化学习的loss如下:
L = E ( x , y ) ∼ D π ϕ R L ( r θ ( x , y ) − β l o g ( π ϕ R L ( y ∣ x ) / π S F T ( y ∣ x ) ) + γ E x ∼ D p r e t r a i n [ l o g ( π ϕ R L ( y ∣ x ) ) ] L = E_{(x,y)\sim D_{\pi^{RL}_{\phi}}}(r_{\theta}(x,y) - \beta log(\pi^{RL}_{\phi}(y|x)/\pi^{SFT}(y|x)) + \gamma E_{x\sim D_{pretrain}}[log(\pi^{RL}_{\phi}(y|x))] L=E(x,y)∼DπϕRL(rθ(x,y)−βlog(πϕRL(y∣x)/πSFT(y∣x))+γEx∼Dpretrain[log(πϕRL(y∣x))]
第二项 β l o g ( π ϕ R L ( y ∣ x ) / π S F T ( y ∣ x ) ) \beta log(\pi^{RL}_{\phi}(y|x)/\pi^{SFT}(y|x)) βlog(πϕRL(y∣x)/πSFT(y∣x)) 也就是对应了PPO-penalty算法的第二项,防止KL散度过大。
第三项和强化学习无关,只是增加了预训练任务。
第一项,也是最关键的一项, r θ ( x , y ) r_{\theta}(x,y) rθ(x,y) 具体是怎么训练的并没有提及。参考其他人实现的 RLHF :
1)首先会使用SFT模型针对每个prompt采样生成大量的样本作为强化训练语料,并过RM模型获取每个生成样本的reward。
2)使用RL模型生成SFT模型生成的相同样本 S ,并获取生成该样本的概率序列 P π R L ( S ) P_{\pi^{RL}}(S) PπRL(S) ,同时我们也有 P π S F T ( S ) P_{\pi^{SFT}}(S) PπSFT(S), 以及 r ( S ) r(S) r(S)
3)计算损失,按照上述的重要性采样和PPO算法即可计算出。具体代码实现如下。
kl_penalty = kl_div(old_action_probs, action_probs) * self.kl_div_loss_weight
rewards = rewards - kl_penalty
#PPO-penalty算法,减去KL散度,保证RL模型和SFT模型差距不会过大
ratios = (action_log_probs - old_log_probs).exp()
#前面介绍的重要性采样,这里使用指数函数替代,可以发现两者的在值为1时,指数函数和除法相差不大,且使用指数时求导更方便。
surr1 = ratios * reward
surr2 = ratios.clamp(1 - self.eps_clip, 1 + self.eps_clip) * reward
#这边使用PPO-clip算法,对loss进行截断,防止更新过大
policy_loss = - torch.min(surr1, surr2) - self.beta_s * entropies
#增加一个预训练损失。
本该到此结束,但是上述实现的时候其实是把生成的每一步的奖励都使用统一的句子级reward,但该代码其实也额外按照每个token来计算奖励值的,为了获取每个token的奖励,我们在生成模型的隐层表示上,多加一个线性层,映射到一维,作为每个状态的预测奖励值。
在上述代码实现的reward的基础上会减去一个当前token的预测reward。代码中 r e w a r d = r e w a r d − V c u r reward = reward - V_{cur} reward=reward−Vcur
由于上面的token级奖励线性层是随机初始化的,我们也需要对其训练,所使用的的损失如下所示:
L o s s v a l u e = ∣ ∣ V c u r − r e w a r d ∣ ∣ Loss_{value} = ||V_{cur} - reward|| Lossvalue=∣∣Vcur−reward∣∣