[论文解读] ReMax: A Simple, Effective, and Efficient Reinforcement Learning Method for Aligning Large Language Models

Ziniu Li, Xu Tian+5 more|arXiv (Cornell University)|Oct 16, 2023

Topic Modeling被引用 6

一句话总结

ReMax 是一种用于对齐大语言模型的新型强化学习方法，摒弃了 PPO 中使用的价值模型，通过利用快速模拟、确定性转移和轨迹级奖励，实现 50% 的内存占用降低、每轮训练速度提升 2 倍，并且仅需 6 行代码即可实现更简单的实现——同时在大语言模型对齐任务中保持或提升性能。

ABSTRACT

Reinforcement Learning from Human Feedback (RLHF) is key to aligning Large Language Models (LLMs), typically paired with the Proximal Policy Optimization (PPO) algorithm. While PPO is a powerful method designed for general reinforcement learning tasks, it is overly sophisticated for LLMs, leading to laborious hyper-parameter tuning and significant computation burdens. To make RLHF efficient, we present ReMax, which leverages 3 properties of RLHF: fast simulation, deterministic transitions, and trajectory-level rewards. These properties are not exploited in PPO, making it less suitable for RLHF. Building on the renowned REINFORCE algorithm, ReMax does not require training an additional value model as in PPO and is further enhanced with a new variance reduction technique. ReMax offers several benefits over PPO: it is simpler to implement, eliminates more than 4 hyper-parameters in PPO, reduces GPU memory usage, and shortens training time. ReMax can save about 46% GPU memory than PPO when training a 7B model and enables training on A800-80GB GPUs without the memory-saving offloading technique needed by PPO. Applying ReMax to a Mistral-7B model resulted in a 94.78% win rate on the AlpacaEval leaderboard and a 7.739 score on MT-bench, setting a new SOTA for open-source 7B models. These results show the effectiveness of ReMax while addressing the limitations of PPO in LLMs.

研究动机与目标

• 解决在大语言模型的基于人类反馈的强化学习（RLHF）中，近端策略优化（PPO）方法带来的高计算成本和内存开销问题。
• 识别 RLHF 任务中被低估的特性——快速模拟、确定性转移和轨迹级奖励——以期实现更简单、更高效的算法。
• 设计一种新强化学习算法，完全移除价值模型而不损失性能，从而降低复杂度和资源需求。
• 证明一种简化的、方差减少的 REINFORCE 基础方法可在效率方面超越 PPO，同时保持对齐质量。

提出的方法

• ReMax 基于 REINFORCE 算法，但引入了一种专为大语言模型设计的新方差减少技术。
• 完全移除了价值模型，从而无需训练、存储或反向传播额外的神经网络。
• 每个提示使用两次响应生成：一次用于策略梯度，一次用于方差减少，利用了标记生成的确定性特性。
• 利用了大语言模型响应生成速度快、基于历史标记的转移是确定性的，且奖励仅在序列结束时提供这一事实。
• 该算法仅用 6 行代码实现，相比 PPO 显著减少了超参数调优需求。
• 内存节省源于移除了价值模型的参数、梯度和优化器状态，内存使用量减少约 50%。

实验结果

研究问题

• RQ1 能否设计一种无需价值模型的强化学习算法用于大语言模型对齐，同时保持或提升性能？
• RQ2 在 RLHF 中，快速模拟、确定性转移和轨迹级奖励等特性能在多大程度上被利用，以简化并加速训练？
• RQ3 移除价值模型是否能带来可测量的内存效率和训练速度提升，而不损失对齐质量？
• RQ4 一种简化版、基于 REINFORCE 的方差减少方法是否能在效率和有效性方面均超越 PPO，用于大语言模型对齐？

主要发现

• 与 PPO 相比，ReMax 将内存使用量减少了约 50%，使得在 A100-80GB GPU 上训练 Llama2-7B 成为可能，而 PPO 在无显存卸载的情况下会失败。
• 在启用卸载的情况下，ReMax 在 4 张 A100-80GB GPU 上支持最大批量大小为 160，而 PPO 仅为 120，从而实现更高的吞吐量。
• 在单张 GPU 上，ReMax 在启用卸载时支持批量大小为 42，而 PPO 仅能支持 32，且两者在无卸载时均会失败。
• 通过消除卸载需求，ReMax 在 4 张 GPU 上将训练时间缩短了 2.5 倍，单轮训练时间从 10.2 小时减少至 2.8 小时。
• 每轮迭代的训练时间，ReMax 约为 PPO 的一半：在启用卸载的 4 张 GPU 上，ReMax 每轮耗时 2.8 小时，而 PPO 为 6.0 小时。
• 尽管实现简单，ReMax 仍保持了强大的对齐性能，与 PPO 相比在任务有效性上无任何下降。