趁着周末稍微补了一下RLHF-PPO的理论和trl实现，网上有很多关于这部分的讲解：磨菇书, PPO代码讲解等。所以这篇博客还是以补充细节为主，分享在学习中可能疑惑或忽略的一些细节。
Reward Model给出的是一个scalar, PPO怎么根据这个值进行模型更新的？
总所周知，PPO会先通过pair-wise的偏好数据训一个Reward Model, 它负责提供奖励信号，具体来说，给定prompt和某个answer，它输出一个标量值表示这个答案的好坏。训练奖励模型用到了Bradley-Terry模型进行建模，也就是maximize $\log P(\frac{exp[r(x,y_w)]}{exp[r(x,y_l)] + exp[r(x,y_w)]})$， 最后可以得到$\log\sigma[r(x,y_w) - r(x,y_l)]$，这也是OpenAI给出的优化目标。
然而，这个reward是针对整个句子的。而LLM生成是token级别的，所以这个信号应该加到哪？常理来说，这个奖励是针对整个句子的，所以它同样也只能加到LLM生成最后一个token的Action上。这样一来，我们只有在LLM生成的最后一步给出信号，并没有限制中之前token的生成，可想而知训练效率会比较低。
如何给LLM生成的中间步骤也添加一个reward信号呢？我们知道PPO的优化目标有一个最小化$\pi$和$\pi_{ref}$的KL散度，这个散度自然是每个token都可以计算的($\pi$和$\pi_{ref}$的logprobs之差)，所以我们可以把每个token的KL散度作为当前生成步的reward，直观理解就是：如果每一步生成的token概率分布和Ref model的相差不大，那我就给奖励，鼓励$\pi$和$\pi_{ref}$不要相差太远。
这个操作细化了Reward的粒度，同时也去掉了PPO loss中的KL项(已经体现在reward中，所以没必要放在loss里)。
Ref Model/Reward Model/Critic Model/Actor Model这么多模型，我等凡人怎么训的动？
实际上，在trl中这四个模型只存了一份完整的weight，还有两份lora weight作为切换。
首先完整的weight肯定就是Ref Model，这个一般是经过sft后(RLHF-PPO stage1)的那个模型。
然后Actor Model也是在Ref Model上初始化，我们如果使用lora进行微调，那么它可以和Ref Model共用weight参数。
Ref Model 和 Actor Model都是负责generate response，所以他们都用的LM Head
Reward Model一般是从Pretrain Model直接训偏好数据得到的，LM Head会换成一个Score Head，而且只作用在最后一个token上。虽然它是从Pretrain Model而不是Ref Model初始化，但是大家结构是一样的，所以lora微调的话还是可以共用(每个参数的lora delta = $param_{reward} - param_{ref}$)。
最后是Critic Model，用来估计每个状态的价值(或者叫回报？)。LM Head会换成一个Value Head，输出同样是一个标量，但是它会对每个token都估计一个价值(对应不同时间步的状态)。从实现上来说，他会直接和Actor Model共享backbone，所以它只要用Actor Model那个lora weight就可以。(或者说Actor Model加一个Value Head做价值估计)
重要性采样模型又是个啥玩意？
这个涉及on-policy和off-policy的概念，on-policy指用某个策略a生成的数据直接更新策略a，而off-policy指用策略a生成的数据更新策略b，在这里a策略可以是b策略之前的某个状态，个人感觉有点像gradient-accumulation。
off-policy允许我们多次利用某些数据反复更新策略，提高了数据的利用率，当然，这里需要加一个因子$p(x)/q(x)$做转换，也是因为这个因子，Actor Loss的梯度式子中$\nabla\log P_\theta(a|s)$符号约掉了$\log$。
$$
\begin{align*}
\nabla J(\theta) &= \frac{1}{N} \sum_{n=1}^N \sum_{t=1}^{T} Adv_\theta^{GAE}(s_n^t,a_n^t) \frac{P_\theta(a_n^t|s_n^t)}{P_{\theta'}(a_n^t|s_n^t)} \nabla \log P_\theta(a_n^t|s_n^t) \\
&= \frac{1}{N} \sum_{n=1}^N \sum_{t=1}^{T} Adv_\theta^{GAE}(s_n^t,a_n^t) \frac{P_\theta(a_n^t|s_n^t)}{P_{\theta'}(a_n^t|s_n^t)} \frac{\nabla P_\theta(a_n^t|s_n^t)}{P_\theta(a_n^t|s_n^t)} \\
&= \frac{1}{N} \sum_{n=1}^N \sum_{t=1}^{T} Adv_\theta^{GAE}(s_n^t,a_n^t) \frac{\nabla P_\theta(a_n^t|s_n^t)}{P_{\theta'}(a_n^t|s_n^t)}
\end{align*}
$$
DPO为啥不要Reward Model了，Critic Model是不是也没了？
DPO其实是用训练奖励模型的目标$\log\sigma[r(x,y_w) - r(x,y_l)]$ 去训练最终的$\pi*$。
具体来说，它找到了$r*$和$\pi*$的对应关系，其中$r*$是最优的奖励模型，$\pi*$是在$r*$基础上炼出来的最优策略。然后用$\pi*$带入刚刚提到的目标中(这里假设$r$已经是$r*$，因为按照传统PPO，训练$\pi \rightarrow \pi*$的时候，Reward Model已经训练好了，也就是$r*$)，就可以得到DPO要优化的目标。这个目标里要优化的参数直接就是$\pi*$ ，因此我们不需要在训一个Reward Model了。同时，由于不需要构造奖励信号，自然也不需要之后的强化学习来学习最大化奖励期望的策略了(不需要Critic Model)。
小结
RL涉及的概念还是相当多的，这里甚至都没提Advantages/GAE/DQN等概念。这篇与其说是博客，更像是学习之余的碎碎念。因为平时没什么时间，也许以后的博客多以这种形式出现了。
评论