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

要約

大規模言語モデル (LLM) をトレーニングする場合、調整は非常に重要です。
これに対処する主な戦略は、ヒューマン フィードバックからの強化学習 (RLHF) を使用することであり、PPO が事実上のアルゴリズムとして機能します。
しかし、PPO は計算効率が悪いことが知られており、このホワイトペーパーではこの課題に取り組むことを目指しています。
私たちは、RLHF タスクの 3 つの重要な特性を特定します。それは、高速シミュレーション、決定論的な遷移、および PPO では活用されない軌道レベルの報酬です。
このような観察に基づいて、私たちは ReMax と呼ばれる RLHF に合わせた新しいアルゴリズムを開発しました。
ReMax のアルゴリズム設計は、有名なアルゴリズム REINFORCE に基づいて構築されていますが、新しい分散削減技術が搭載されています。
私たちの方法には PPO に比べて 3 つの利点があります。まず、原理的にメモリ使用量が約 50% 節約されます。
その結果、8xA100-40GB GPU で Llama2 (7B) モデルを微調整すると、PPO がメモリ不足になりますが、ReMax はトレーニングに余裕ができます。
このメモリの向上は、PPO の値モデルを削除することで実現されます。
第 2 に、ReMax は実装が簡単で、スケールに敏感で調整に手間がかかる PPO の多くのハイパー パラメータを削除します。
3 番目に、GPT2 (137M) では、実測時間の観点から 2.2 倍の高速化が観察されます。
重要なのは、上記の計算上の改善によってパフォーマンスが犠牲になるわけではないということです。
これらの利点は、より大規模なモデルでも維持できると仮説を立てています。
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, 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, it 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. Second, ReMax is simple to implement and removes many hyper-parameters in PPO, which are scale-sensitive and laborious to tune. Third, on GPT2 (137M), we observe 2.2x speed-up in terms of wall-clock time. 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