极市导读
采用 RLOO 可以持续性地超越 PPO、DPO、RAFT等一系列主流 RLHF 方法,并且对于噪声和 KL 超参数更加鲁棒。 >>加入极市CV技术交流群,走在计算机视觉的最前沿
背景介绍
原始的 RLHF 流程包含三个阶段:
-
SFT 阶段:对一个预训练好的语言模型进行指令微调,即给定指令 prompt 和(通常情况下)人类撰写的回复进行监督训练,SFT 阶段训好的模型会用于 Reward model 和 RLHF policy 的初始化。 -
Reward Model 阶段:采集一份人类偏好数据集,包含正样本和负样本,对应一个给定的 prompt 下人类喜欢和不喜欢的回复,利用这份标注可以训练一个二分类器,loss 如下:
即,让模型输出正样本的 logits 越高于负样本越好。
-
RL 阶段:Reward Model 用于提供在线的 feedback 来优化以下目标:
PPO 之殇
对于研究大模型的同学来说,RLHF(Reinforcement Learning from Human Feedback) 应该算是最熟悉的内容之一了,而 RLHF 最常见的途径之一,则是流传自 OpenAI 自家的 PPO(Proximal Policy Optimization)。
其实很多人也想到过,RL 相较于 SFT 最大的区别主要是在训练中能提供负梯度,那么很自然地可以想到,PPO 并不是 RLHF 的唯一方式,而 OpenAI 之所以选择 PPO 很大可能只是因为用起来最顺手罢了。
然而在实践中,对于非 RL 背景的研究者来说,PPO 并不是一个容易使用的算法,众所周知的缺点有:
-
巨大的计算开销
PPO 需要同时加载 4 个模型:
-
the generator -
the reference (for KL estimation) -
the critic -
the reward model
其中 generator 和 critic 的训练是交错进行的,且训练难度会随着模型规模变大而显著提升。
-
优化难度
作为 online RL 优化算法,训练相当不稳定且敏感,导致训练 PPO 需要相当高的专业知识和经验才能把它调好。
DPO 过墙梯
既然 PPO 使用起来困难重重,很自然地涌现出了一系列研究试图饶过这些困难,这也就是所谓的“RL-free”方法,其中最有名的工作要数 DPO(Direct Preference Optimization)。
“RL-free”方法的暴力之处在于,直接把全部的 RL 组件都扔掉了:既然 PPO 不好训,那我干脆直接不做 RL 不就行了。DPO 跳过了原本的对抗学习,通过公式转换把 RLHF 转化为了 SFT,直接通过收集人类偏好标注来训练一个 Reward model 然后进行优化。
这种捷径走得不可说不妙,但也带来了很多副作用,有研究显示 DPO 等算法的泛化能力是会弱于 RLHF 的,并且独立训练 Reward model 也意味着对收集到的人类偏好标注的数据质量和数量都有了更高的要求。
所以 Cohere 团队的这篇工作试图不走捷径,提问:我们能不能在避开计算复杂度和优化难度的情况,依旧保持 PPO 的性能?
RLHF 不必硬套 PPO
作者团队认为,传统的 Deep-RL 算法的一些设定和假设,与当前 LLM 的人类偏好学习之间,存在着一些差异。
在传统 RL 中,模型往往是从随机初始化的参数开始训练的,训练样本产生的梯度方差往往非常大,因此 PPO 设计的重点在于如何在梯度更新非常大的情况下保持训练稳定。
相比之下,RLHF 中的 LLM 已经经过了相当充分的 pretraining 和 SFT,因此并不存在传统 RL 训练中的这种特性。尽管理论上每一步的预测空间是整个词表范围,但由于预训练和 SFT 的影响,语言模型的输出已经很大程度上不是胡言乱语了,实际上真正的搜索空间已经被局限在了有限的一个 token 子集中了,换言之,RLHF 所需要面对的环境远远比传统 RL 要和谐友善得多,如果说传统 RL 是从蛮荒开天辟地开始荒野求生的话,RLHF 大概是现代人出去参加极限挑战了。
作者的消融实验显示,PPO 中 GAE 的超参数 (variance最大)时,能获得最大的 training reward(左图),而 PPO 中对当前 policy 的更新幅度进行 Clipping 的操作并没有明显的收益(右图):
更进一步来讲,人类偏好学习放在 RL 的语境下,生成下一个 token 被看成一个 action,而从 prompt 起始的、未完成的 sequence,被视为一个 state。
但在 RLHF 中,对于不完整的 sequence 进行建模是不必要的,我们只需要对完整的生成结果进行 reward,而不需要对中间结果进行 reward。因此 PPO 中采用传统的 RL 建模假设就有点在做无用功了,更正确且有效的建模思路是,把生成一条完整序列视为一个 action,初始状态由 prompt 决定。
基于以上观察,作者尝试使用 REINFORCE estimator 和其 multi-sample 扩展版本,RLOO(REINFORCE Leave-One-Out)来优化 sequence-level 的目标。
这样一种 sequence-level 的优化可以直接优化完整的轨迹(即完整 sequence),因此是无偏的,而不像 PPO 为代表的 actor-critic 算法会存在中间误差累积的问题。
作者在一系列的开源模型和数据集上进行了验证,总结了以下结论:
-
PPO 不是 RLHF 场景下做 RL 的正确选择:只使用最基础的策略梯度算法也能持续性超越 PPO 性能,胜率提升在 3.2% 到 20.3% 不等。 -
RLOO 比当下的主流方法都强:采用 RLOO 可以持续性地超越 PPO、DPO、RAFT等一系列主流 RLHF 方法,并且对于噪声和 KL 超参数更加鲁棒。 -
建模中间过程是不必要的:对于 LLM 的偏好学习中,建模中间 token 是不必要的,直接建模完整序列可以保持性能不变,同时大大降低复杂度、加速训练。
RLOO & ReMax
在 PPO 框架下,初始状态由 prompt 决定,每生成一个 token 被视为一个 action,而已生成的 sequence 被视为状态,只有当生成 <EOS> token 时才会获得 reward model输出的奖励,这个奖励与 KL 惩罚相结合,而对于词表中的其他 token,只有 KL 部分是非零的。
这样看来,将生成完成的整段 sequence 视为一个 action 可能是更加合理的,而生成一条 sequence 然后获取 reward 这一过程可以视为一个多臂老虎机问题(bandit formulation),因此可以直接套用 REINFORCE 框架:
然后更自然地,可以采用 REINFORCE 估计的改进版本,即 REINFORCE with baseline 来进一步提升学习效率:
其中的核心思想很简单,REINFORCE 估计面临着梯度方差过大的问题,因此可以通过减去一个 baselin 值,使得方差减小,但期望不变。在原论文中,baseline 的一个简单获取方法是对一定步数内的 reward 取滑动平均:
然而 REINFORCE 已经是 1992 年提出的算法了,在 2019 年又有了一个改进版本,也就是 RLOO(REINFORCE Leave-One-Out):
其核心思想也很简单,对当前模型进行 k 次在线采样,然后取其中一个结果作为本次的输出,剩下的 k-1 个结果 reward 的均值作为 baseline。
此处还可以补充本工作的一个同期工作 ReMax,整体思路与本文如出一辙,差别主要在于 baseline 的选取上,ReMax 方法把贪心采样(temperature=0)得到的 sequence 作为 baseline。
实验显示,采用 REINFORCE 估计的性能在各个数据集上均可以超越已有的方法。

公众号后台回复“数据集”获取100+深度学习各方向资源整理
极市干货

点击阅读原文进入CV社区
收获更多技术干货

