ReMax: A Simple, Effective, and Efficient Reinforcement Learning Method for Aligning Large Language Models

📄 arXiv: 2310.10505v4 📥 PDF

作者: Ziniu Li, Tian Xu, Yushun Zhang, Zhihang Lin, Yang Yu, Ruoyu Sun, Zhi-Quan Luo

分类: cs.LG

发布日期: 2023-10-16 (更新: 2024-05-16)

💡 一句话要点

提出ReMax以解决PPO在大语言模型对齐中的不足

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 强化学习 人类反馈 大型语言模型 PPO算法 方差减少 模型对齐 GPU优化 开源模型

📋 核心要点

  1. 现有的PPO算法在对齐大型语言模型时过于复杂，导致超参数调优困难和计算负担重。
  2. ReMax通过利用RLHF的快速仿真、确定性转移和轨迹级奖励等特性，简化了强化学习过程。
  3. 实验结果显示，ReMax在训练Mistral-7B模型时显著提高了性能，创造了新的开源模型基准。

📝 摘要（中文）

强化学习中的人类反馈（RLHF）是对齐大型语言模型（LLMs）的关键，通常与近端策略优化（PPO）算法配合使用。尽管PPO在一般强化学习任务中表现出色，但对于LLMs而言过于复杂，导致超参数调优困难且计算负担重。为提高RLHF的效率，本文提出ReMax，利用RLHF的快速仿真、确定性转移和轨迹级奖励等三大特性，这些特性在PPO中未被充分利用。ReMax基于著名的REINFORCE算法，不需要像PPO那样训练额外的价值模型，并通过新的方差减少技术进一步增强。ReMax在实现上更简单，消除了PPO中的4个以上超参数，减少了GPU内存使用，并缩短了训练时间。使用ReMax训练Mistral-7B模型时，在AlpacaEval排行榜上获得94.78%的胜率，并在MT-bench上得分7.739，创造了开源7B模型的新SOTA。这些结果展示了ReMax的有效性，同时解决了PPO在LLMs中的局限性。

🔬 方法详解

问题定义：本文旨在解决PPO算法在对齐大型语言模型时的复杂性和计算负担，尤其是在超参数调优和内存使用方面的不足。

核心思路：ReMax通过利用RLHF的特性，简化了强化学习的过程，避免了PPO中需要的额外价值模型训练，从而提高了效率。

技术框架：ReMax的整体架构基于REINFORCE算法，主要包括快速仿真、确定性转移和轨迹级奖励的模块，形成一个高效的训练流程。

关键创新：ReMax的核心创新在于不需要训练额外的价值模型，并引入了一种新的方差减少技术，这使得其在实现上更为简单且高效。

关键设计：ReMax消除了PPO中的多个超参数，优化了内存使用，具体在训练7B模型时节省了约46%的GPU内存，并允许在A800-80GB GPU上进行训练，无需PPO所需的内存节省技术。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

ReMax在Mistral-7B模型上的实验结果显示，获得了94.78%的胜率和7.739的MT-bench得分，显著优于PPO，且在训练过程中节省了约46%的GPU内存，展现了其在大型语言模型训练中的卓越性能。

🎯 应用场景

ReMax的研究成果在大型语言模型的训练和对齐中具有广泛的应用潜力，尤其是在需要高效计算和快速迭代的场景中。其简化的设计和显著的性能提升使其在实际应用中具备较高的价值，未来可能推动更多基于RLHF的模型开发与优化。

📄 摘要（原文）

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.