技術要約：人間からのフィードバックを用いた分布ロバスト強化学習

問題提起

人間からのフィードバックを用いた強化学習（RLHF）は、大規模言語モデル（LLM）を人間の意図に適合させるための標準的な手法となっています。しかし、既存のRLHFパイプラインはロバスト性に欠けています。推論時（ダウンストリームタスク）に遭遇するプロンプトの分布が、ファインチューニングに使用された選好データセットの分布と大きく異なる場合、モデルの性能はしばると低下します。この問題は、真の報酬関数を犠牲にして代理報酬を改善してしまう「過剰最適化（over-optimization）」や、環境の変化や人間の好みのドリフトによる分布シフトといった現象によって悪化します。ロバスト性に関する従来のアプローチは、特定の敵対的破損モデル、選好モデルのシフト、あるいはヒューリスティックな手法（モデル・マージなど）に焦点を当ててきましたが、これらは緩やかな分布シフトの下では機能しないことがよくあります。

手法

著者らは、RLHFのための**分布ロバスト最適化（DRO）**フレームワークを提案しています。核心となる目的は、ファインチューニングされたモデル（報酬モデルおよびポリシーの両方）が、訓練分布から$\phi$ダイバージェンス（具体的には全変動距離：Total Variation distance）の範囲内にある場合でも、性能を維持できるようにすることです。著者らは、学習者が経験的な訓練分布の周囲にある$\phi$ダイバージェンス・ボール内の最悪の分布に対して損失を最小化するという、ミニマックス最適化問題を導入することでこれを定式化しています。

本論文では、2つの主要なRLHFパイプラインに対してロバストなバリアントを開発しています。

1. ロバストな報酬推定:

   • 定式化: 報酬モデル $r$ は、$d_{TV}(D, D_{src}) \leq \rho$ を満たす分布 $D$ に対する最悪の期待損失を最小化するように訓練されます。
   • アルゴリズム: ミニバッチ勾配降下法（アルゴリズム1）が提案されています。各イテレーションにおいて、アルゴリズムはミニバッチをサンプリングし、距離 $\rho$ 内の最悪の分布シフトに対応する最適な重み $q^*_t$ を計算します。その後、報酬パラメータを更新する前に、これらの最適な重みによって勾配が再重み付けされます。
   • 収束: 線形報酬関数クラスの仮定の下で、アルゴリズムが $T = O(1/\epsilon^2)$ イテレーションおよび $n = \tilde{O}(1/\epsilon^2)$ サンプル/イテレーションで $\epsilon$-最適解に収束することを証明しています。

2. ロバストなポリシー最適化:

   • 定式化: ポリシー $\pi$ は、最悪のKL正則化値関数を最大化するように最適化されます。
   • アルゴリズム: 著者らは、ロバストな自然方策勾配（NPG）法（アルゴリズム2）を提案しています。報酬推定と同様に、これは最悪の分布のための重みを計算します。決定的な違いとして、標準的な経験的フィッシャー行列ではなく、これらの最悪の重みから導出された重み付きフィッシャー情報行列を用いてポリシーを更新します。
   • 収束: 対数線形ポリシークラスと標準的な集中条件（concentrability conditions）を仮定すると、この手法は線形収束率（$T = O(\log(1/\epsilon))$）を達成しますが、最終イテレーションの収束を保証するために、より高いサンプル複雑度 $n = \tilde{O}(1/\epsilon^4)$ を必要とします。

3. ロバストな直接選好最適化（DPO）:

   • 定式化: 明示的な報酬モデルを持たず、直接ポリシーを最適化する、分布ロバスト版のDPOが定式化されています。
   • アルゴリズム: アルゴリズム3は、最悪の分布シフトに基づく再重み付けされた勾配を用いて、ミニバッチ確率的勾配降下法を適用します。
   • 収束: 解析により、ロバストな報酬推定フェーズと同様に、イテレーション複雑度 $T = O(1/\epsilon^2)$ およびサンプル複雑度 $n = \tilde{O}(1/\epsilon^2)$ が示されています。

主な貢献

• 理論的フレームワーク: 本論文は、報酬推定およびポリシー最適化の両フェーズ、さらにはDPOに対する初の分布ロバストな定式化を確立しました。
• アルゴリズム: 報酬推定、ロバストNPG、およびロバストDPOのための実用的なミニバッチベースのアルゴリズムを導入しています。これらのアルゴリズムは、重み付け係数の計算（凸最適化により解決）と勾配の再重み付けのみを必要とするため、既存のパイプラインへの侵襲性が最小限になるよう設計されています。
• 収束保証: 著者らはこれらのアルゴリズムの収束に関する理論的証明を提供しています。線形報酬モデルおよび対数線形ポリシーに対するサンプルおよびイテレーション複雑度の境界を確立し、ロバスト性が証明可能な保証とともに達成できることを示しています。
• 実証的検証: 提案手法は、訓練にはUnified-Feedbackデータセットを、評価にはRewardBench、HHH-Alignment、MT-Benchを使用し、分布外（OOD）タスクを用いて評価されています。

実験結果

著者らは、ベースモデルとして Gemma-2B-it および Mistral-7B-Instruct-v0.2 を使用して評価を行いました。

• 報酬モデル: RewardBenchベンチマークにおいて、分布ロバスト報酬モデル（DR RM）は、特にロバスト性パラメータ $\rho$ が 0.1 または 0.2 に設定されている場合に、標準的なBradley-Terry（BTL）モデルと比較して平均精度が向上しました。特筆すべきは、**推論（Reasoning）**サブセットの性能が大幅に向上したことです（例：$\rho=0.1$ のとき、あるサブセットで +23.4%）。ただし、$\rho$ が高すぎると（例：$\rho \geq 0.6$）、性能が低下しており、ロバスト性と訓練分布への適合性の間のトレードオフが示されました。
• ポリシー最適化（PPO）: ロバストな報酬モデルを使用してPPOを通じてポリシーを訓練する場合、ロバスト・バリアント（DR PPO）は、分布内（Unified-Feedback）および分布外（HHH-Alignment, MT-Bench）の両方のデータセットにおいて、標準的なPPOと同等またはそれを上回る性能を示しました。推論タスクにおいて顕著な改善が見られ、$\rho=0.1$ で精度が +10.0% 上昇しました。
• 直接選好最適化（DPO）: ロバストDPO（DR DPO）は、標準的なDPOと比較して、OODデータセットにおける性能を一貫して向上させました。HHH-Alignmentでは $\rho=0.8$、MT-Benchでは $\rho=0.5$ で最高の性能が得られました。DR DPOは一般に、OODタスクにおいてDR PPOを上回りました。

意義と主張

本論文は、分布ロバスト最適化が、未知のタスクや分布シフトに対するLLMの汎化性能を向上させるための原理的なアプローチを提供すると主張しています。著者らは、提案手法が既存の効率的なRLHFパイプライン（Hugging Faceライブラリを使用しているものなど）への変更を最小限に抑えつつ、理論的な収束保証を提供することを強調しています。

著者らが強調する重要な発見は、分布ロバスト性を強制した際に、OODデータセットにおける推論能力に予期せぬ向上が見られたことです。著者らは、なぜDROが推論にプラスの影響を与えるのかは直ちに明らかではないと述べており、推論と分布ロバスト性の関連についてのさらなる調査が必要であることを示唆しています。

本研究は、標準的なRLHFが訓練中に見た特定の選好例を超えて汎化することに失敗しがちな限界を超え、データ分布が静的ではない動的な現実世界の環境において、人間の意図に真に適合するLLMへの一歩となるものです。