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

要約

大規模言語モデル (LLM) をトレーニングする場合、調整は非常に重要です。
これに対処する主な戦略は、ヒューマン フィードバックからの強化学習 (RLHF) を使用することであり、PPO が事実上のアルゴリズムとして機能します。
しかし、PPO は計算効率が低いことが知られており、このホワイト ペーパーではこれに対処することを目的としています。
私たちは、RLHF タスクの 3 つの重要な特性を特定します。それは、高速シミュレーション、決定論的な遷移、および PPO では活用されない軌道レベルの報酬です。
このような観察に基づいて、私たちは ReMax と呼ばれる RLHF に合わせた新しいアルゴリズムを開発しました。
ReMax のアルゴリズム設計は、有名なアルゴリズム REINFORCE に基づいて構築されていますが、新しい分散削減技術が搭載されています。
私たちの方法には PPO に比べて 3 つの利点があります。まず、ReMax は実装が簡単で、スケールに敏感で調整に手間がかかる PPO の多くのハイパーパラメータを削除します。
次に、ReMax は原則としてメモリ使用量を約 50% 節約します。
その結果、8xA100-40GB GPU で Llama2 (7B) モデルを微調整すると、PPO がメモリ不足になりますが、ReMax はトレーニングに余裕ができます。
このメモリの向上は、PPO の値モデルを削除することで実現されます。
第三に、私たちの計算に基づくと、PPO が Llama2 (7B) のトレーニングに耐えられると仮定しても、それでも ReMax よりも約 2 倍遅く実行されることがわかります。
これは、ReMax には存在しない価値モデルの計算オーバーヘッドによるものです。
重要なのは、上記の計算上の改善によってパフォーマンスが犠牲になるわけではないということです。
これらの利点は、より大規模なモデルでも維持できると仮説を立てています。
ReMax の実装は https://github.com/liziniu/ReMax で入手できます。

要約(オリジナル)

Alignment is of critical importance for training large language models (LLMs). The predominant strategy to address this is through Reinforcement Learning from Human Feedback (RLHF), where PPO serves as the de-facto algorithm. Yet, PPO is known to suffer from computational inefficiency, which is a challenge that this paper aims to address. We identify three important properties in RLHF tasks: fast simulation, deterministic transitions, and trajectory-level rewards, which are not leveraged in PPO. Based on such observations, we develop a new algorithm tailored for RLHF, called ReMax. The algorithm design of ReMax is built on a celebrated algorithm REINFORCE but is equipped with a new variance-reduction technique. Our method has three-fold advantages over PPO: first, ReMax is simple to implement and removes many hyper-parameters in PPO, which are scale-sensitive and laborious to tune. Second, ReMax saves about 50% memory usage in principle. As a result, PPO runs out-of-memory when fine-tuning a Llama2 (7B) model on 8xA100-40GB GPUs, whereas ReMax can afford training. This memory improvement is achieved by removing the value model in PPO. Third, based on our calculations, we find that even assuming PPO can afford the training of Llama2 (7B), it would still run about 2x slower than ReMax. This is due to the computational overhead of the value model, which does not exist in ReMax. Importantly, the above computational improvements do not sacrifice the performance. We hypothesize these advantages can be maintained in larger-scaled models. Our implementation of ReMax is available at https://github.com/liziniu/ReMax

arxiv情報

提供元, 利用サービス

arxiv.jp, Google