[論文レビュー] 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%削減、1イテレーションあたりの学習速度を2倍にし、実装も単純化（わずか6行のコード）する。一方で、LLMアライメントタスクにおける性能は維持または向上させる。

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）において、Proximal Policy Optimization（PPO）が示す高い計算コストとメモリオーバーヘッドに対処すること。
• RLHFタスクにおいて未利用の特性—高速なシミュレーション、決定論的な遷移、トレイジェクトリーレベルの報酬—を特定し、それらを活用することで、より単純で効率的なアルゴリズムを実現すること。
• 性能に妥協を来さずに価値モデルを完全に排除する新しい強化学習アルゴリズムを設計し、複雑さとリソース要件を低減すること。
• 単純化された分散低減REINFORCEベースの手法が、PPOを上回る効率性を示しつつ、アライメント品質を維持できることを実証すること。

提案手法

• ReMaxはREINFORCEアルゴリズムに基づいているが、大規模言語モデルに特化した新たな分散低減技術を組み込んでいる。
• 価値モデルを完全に排除しており、追加のニューラルネットワークの学習・保存・逆伝播の必要がなくなる。
• 1つのプロンプトに対して2つの応答を生成する：1つは方策勾配の計算に、もう1つは分散低減に使用。生成の決定論的性質を活用している。
• 大規模言語モデルの応答生成が高速であること、過去のトークンに基づいて遷移が決定的であること、報酬がシーケンスの終了時にのみ提供されることを活用している。
• アルゴリズムはわずか6行のコードで実装可能であり、PPOと比較してハイパーパrameterチューニングが著しく簡素化されている。
• メモリの節約は、価値モデルのパラメータ、勾配、最適化ステートを排除することで達成され、約50%のメモリ使用量削減が実現している。

実験結果

リサーチクエスチョン

• RQ1 価値モデルを排除した強化学習アルゴリズムを、LLMアライメントタスクにおいて性能を維持または向上させながら設計できるか？
• RQ2 RLHFタスクにおける高速なシミュレーション、決定論的な遷移、トレイジェクトリーレベルの報酬といった特性を、どの程度活用して学習の単純化と高速化を実現できるか？
• RQ3 価値モデルを排除することで、アライメント品質に影響を与えることなく、メモリ効率性と学習速度に顕著な改善が得られるか？
• RQ4 分散低減を施した単純化されたREINFORCEベースの手法が、PPOを上回る効率性と有効性を示せるか？

主な発見

• ReMaxはPPOと比較して約50%のメモリ使用量削減を達成し、PPOがオフロードなしでは失敗するLlama2-7BのA100-80GB GPU上での学習を可能にした。
• オフロードを適用した場合、4枚のA100-80GB GPUでReMaxは最大バッチサイズ160をサポートするが、PPOは120にとどまるため、より高いスループットが実現できる。
• 1枚のGPUでオフロードを適用した場合、ReMaxはバッチサイズ42をサポートするが、PPOは32にとどまり、両者ともオフロードなしでは失敗する。
• 4GPU環境で、オフロードの必要がなくなるため、ReMaxは学習時間を2.8時間（1エポック）にまで短縮し、PPOの10.2時間から2.5倍の高速化を達成した。
• オフロードを適用した4GPU環境下で、ReMaxの1イテレーションあたりの学習時間はPPOの半分で、ReMaxは1エポックあたり2.8時間、PPOは6.0時間である。
• 単純さにもかかわらず、ReMaxはPPOと同等の高いアライメント性能を維持しており、タスクの有効性に劣化は見られない。