Hugging Face TRL RLOO トレーナー リリース

Hugging FaceはTRLライブラリにRLOO（REINFORCE Leave One-Out）トレーナーを導入しました。RLOOは、近接方策最適化（PPO）に代わる、より手軽で効率的なオンラインRLHFトレーニングアルゴリズムで、GPUメモリの使用量を大幅に削減し、収束時間も短くなります。

パフォーマンスと効率の向上

RLOOはPPOと比較して、リソース利用とトレーニング速度に大幅な改善をもたらします。

  • メモリ効率: モデルサイズに応じて、RLOOはPPOに比べて約50〜70％少ないvRAMを使用します。
  • トレーニング速度: RLOOは1BパラメータモデルでPPOの2倍、6.9Bモデルで最大3倍速く動作します。
  • モデル品質: RLOOは応答勝率（GPT-4による評価）でPPOと競合し、Direct Preference Optimization（DPO）などのオフライン手法を一貫して上回ります。

技術的アーキテクチャと動機

従来のPPOは、ポリシーモデル、リファレンスポリシーモデル、報酬モデル、価値モデルの4つのモデルコピーをメモリにロードする必要があるため、リソース集約的です。RLOOは価値モデルを除外し、ポリシーモデル、リファレンスポリシーモデル、報酬モデルの3つのコピーだけで済むようにアーキテクチャを簡素化します。

主なアルゴリズム上の違い

  1. 完了レベルのアクション: PPOは各完了トークンを個別のアクションとして扱うのに対し、RLOOはモデルの全体完了を単一のアクションとして扱います。これにより、通常はEOS（文末）トークンにのみ真の報酬が与えられるスパース報酬の問題が解消され、RLOOはEOS報酬を全体の完了に割り当てます。
  2. REINFORCE損失: RLOOはREINFORCE損失を利用し、報酬とベースラインの差をアクションの対数確率で掛け合わせます。これにより、PPOで必要とされる価値モデルやGeneralized Advantage Estimation（GAE）が不要になります。
  3. Leave-One-Outベースライン: RLOOはバッチ内の他のすべてのサンプルの報酬を特定のサンプルのベースラインとして使用してベースラインを算出します。あるプロンプトに対して、ある完了のベースラインは同じプロンプトから生成された他のすべての完了の平均報酬です。

TRLでの実装

RLOOトレーナーは実験的な PPOv2Trainer を基盤としています。RLOOは概念的にはREINFORCEアルゴリズムに基づいていますが、TRLの実装ではREINFORCE損失がPPO損失の特殊ケースであるため、PPO損失が使用されています。

ベンチマーク結果

Pythia 1Bおよび6.9Bモデルで実施された実験は、RLOO実装の有効性を示しています。

  • Pythia 6.9B: GPT-4を審査員として使用し、78.7％の好まれる率（k=2）を達成し、元論文で報告された77.9％（k=4）および74.2％（k=2）を上回りました。
  • Pythia 1B: SFT（Supervised Fine-Tuning）チェックポイントの21.3％の勝率に対し、40.1％の勝率を達成しました。

bf16における数値安定性の課題

RLOOはbf16精度を使用する際に、数値安定性の問題が顕著に現れます。生成時に得られる対数確率は、トレーニングのフォワードパスで得られるものとわずかに異なることがあります。

PPOでは、クリッピング係数（通常0.2）により、比率が1.2を超えるまたは0.8未満になる個々のトークンの勾配がゼロになります。RLOOでは、全体の完了が単一のアクションとして扱われるため、これらの数値的ずれが蓄積されます。その結果、シーケンス全体の勾配がより頻繁にゼロ化されます。実証的観測では、PPOはバッチデータの約3％をゼロ化するのに対し、bf16下のRLOOは20〜40％のバッチデータをゼロ化します。

Sources