ReMax: 一种对齐语言模型的高效强化学习算法

简介

       基于人类反馈的强化学习（RLHF）是对齐大型语言模型（LLM）的关键，通常与近端策略优化（PPO）算法配对使用。虽然PPO是为一般强化学习（RL）任务设计的强大方法，但它对于LLM来说过于复杂，导致了显著的内存和计算成本。为了提高RLHF的效率，我们提出了一种名为ReMax的定制算法。

ReMax的创新之处

       特别是，ReMax利用了RLHF的三个特性：快速模拟、确定性转移和轨迹级别的奖励，而这些在PPO中并未得到充分利用。基于著名的REINFORCE算法，ReMax通过以下方式优化了RLHF流程：
       1.   无需额外的价值模型：ReMax不需要像PPO那样训练额外的价值模型，从而减少了计算复杂度。
       2.   新的方差减少技术：ReMax引入了一种新的方差减少技术，进一步提高了算法的稳定性和效率。

ReMax的优势

       ReMax相比PPO具有多个显著优势：
       •   实现简单：ReMax的实现更加简洁，减少了工程上的复杂性。
       •   减少超参数：ReMax消除了PPO中的四个超参数，使得模型调优更为直接。
       •   内存效率：使用ReMax训练一个7B模型时，可以比PPO节省约46%的GPU内存。
       •   加快训练速度：ReMax缩短了训练时间，无需像PPO那样依赖内存节省的卸载技术，后者的训练速度也慢了1.6倍。

实验结果

       我们用4xA800-80GB的机器进行实验。我们将ReMax应用于Mistral-7B模型，取得了非常好的结果：

       •   AlpacaEval排行榜：在AlpacaEval排行榜上，ReMax实现了94.78%的胜率，远超当前主流方法。
       •   MT-bench评分：ReMax在MT-bench上取得了7.739分，创下了开源7B模型的新SOTA。

       这些结果展示了ReMax的有效性，同时解决了PPO在LLM中的局限性，推动了RLHF技术的发展。

结论

       ReMax通过优化RLHF的关键流程，显著降低了内存和计算开销，同时提高了模型的训练速度和性能表现。我们的实验结果表明，ReMax在处理大型语言模型时，具备较PPO更优的效率和效果，为开源社区提供了一个更为实用的选择。