PPOを用いたRLHFの実装詳細 – Hugging Face Blog 要約

このHugging Faceのブログ記事は、OpenAIの2019年のRLHFコードベースを再現し、スタイルに関するタスクにおいてその学習曲線と一致させ、再現性に影響を与える実装の詳細を列挙しています。

学習曲線の再現

再現実験の結果、感情分析および記述性タスクにおいて、OpenAIのオリジナルコードベースとほぼ同一の学習曲線が得られました。著者らは、ベースラインの指標を得るためにAWS p3dn.24xlargeインスタンス上でオリジナルのTensorFlow 1.xリポジトリを実行し、それらの指標を使用してPyTorchベースの再現結果を検証しました。

一般的な実装の詳細

報酬モデル（reward model）とポリシー価値ヘッド（policy value head）は、連結されたクエリとレスポンスを入力として受け取ります。シーケンスは特別なパディングトークンを使用して固定長にパディングされ、長すぎる場合は切り捨てられます。また、logitの計算中にはパディングトークンに対して位置インデックスが調整されます。レスポンス生成では、EOSトークンで停止せずに固定長の出力をサンプリングし、報酬モデルとポリシーのトレーニングの両方において学習率はゼロに向かってアニーリングされます。探索を促進するために、GPUプロセスごとに異なる乱数シードが使用されます。

報酬モデルの実装詳細

報酬モデルは、連結されたクエリとレスポンスのペアの最後のトークンにおける値のみを出力します。報酬ヘッドの重みは、分散 $1/(\sqrt{d_{model}}+1)$ の正規分布から初期化され、バイアスはゼロに設定されます。報酬の正規化は、固定の参照ポリシーからの報酬の経験的な平均と標準偏差を計算し、単位分散とゼロ平均にスケーリングすることで、トレーニングの前後で行われます。

ポリシーのトレーニング実装詳細

対数確率を計算する前に、logitsはサンプリング温度で除算されます。価値ヘッドの重みは、平均ゼロ、分散ゼロで初期化されます。ポリシーのトレーニング中はドロップアウトは無効化されます。棄却サンプリング（Rejection sampling）は、トークン16から24の間に期間を強制し、それ以外の場合は固定の低い報酬を割り当てます。割引率 $\gamma$ は1に設定されます。トレーニングには、勾配蓄積のためにバッチ、ミニバッチ、およびマイクロバッチが使用されます。トークンごとのKLペナルティが報酬に加算され、ミニバッチごとの報酬とアドバンテージのホワイトニング（アドバンテージのみ平均シフト）が適用されます。価値関数は標準的なPPOと同様にクリッピングされ、KLペナルティ係数 $\beta$ は観測されたKLダイバージェンスに基づいて適応されます。

RLHFにおけるPyTorch Adamオプティマイザの数値的問題

PyTorchのAdam実装は、分母内のepsilonの配置がTensorFlowとは異なっており、これによりトレーニングの初期段階で実効的な正規化項が小さくなり、結果としてよりアグレッシブな勾配更新が行われます。実験では、PyTorchのAdamは、TensorFlowスタイルのAdamと比較して、特にgpt2-xlのような大規模なモデルにおいて、より高い対数確率分散、より大きな比率の極値、および増加した近似KLとクリップ比率をもたらすことが示されています。

制限事項

計算コストの高さとトレーニングの不安定さのため、CNN/DMまたはTL;DRからの要約タスクの再現は試みられていません。オリジナルのコードには8× V100 32GB GPUが必要であり、再現実験ではGPU利用率が低かった（約30%）ため、実行コストが高く不安定でした。

結論

OpenAIのRLHFコードベースを再現し、その実装の詳細を文書化することで、著者らはRLHFのエンジニアリング手法を理解するためのリファレンスを提供し、トレーニングのダイナミクスに大きな影響を与える可能性のあるオプティマイザの微妙な違いを明らかにしています。