Reformulate LLM Reinforcement Learning for Efficient Training under Black-box Discrepancy

技術要約：ブラックボックス的な不一致下での効率的な学習のためのLLM強化学習の再定式化

1. 問題提起

強化学習（RL）は、大規模言語モデル（LLM）の推論能力を向上させるための重要なポストトレーニング・パラダイムとなっています。しかし、RLによるポストトレーニングは、しばしば予測不可能な性能低下や、完全な学習の崩壊に直面します。近年の調査では、これらの失敗の原因として、広範にわたる訓練と推論の間の不一致（discrepancy）（またはミスマッチ）が挙げられています。

この不一致は、産業界の要求により、訓練時と推論時で全く同じモデル重みを使用しているにもかかわらず、異なる基盤エンジンや精度フォーマット（例：異なるカーネル、並列化戦略、あるいはBF16対FP8のような数値精度）の使用を強制されることから生じます。これにより、訓練ポリシー（$\pi_{Train}$）と推論ポリシー（$\pi_{Inf}$）の間の行動確率分布の間に、隠れた偏差が生じます。

既存の緩和策には、以下のような重大な限界があります：

• 一様な精度整合（Uniform Precision Alignment）： 単に訓練精度を推論に合わせて（例：FP16へ）引き下げることは、表現可能な値の範囲を縮小させ、勾配爆発のリスクを悪化させます。
• ヒューリスティックな介入： 高い不一致を示すトークンから勾配をマスクするような手法は、単にノイズの多いデータのインパクトを抑制するだけです。これらは根本的な原因に対処しておらず、最終的な目標である「推論デプロイメントのためのポリシー最適化」を見落としています。
• インフラストラクチャの整合： 低レベルのオペレータやハードウェアスタックを一致させようとする試みは、多大な労力を要し、無数のコーナーケースが存在するため、多くの場合実現不可能です。

核心となる課題は、報酬最大化という主要な目的を犠牲にすることなく、訓練ポリシーがその不一致を動的に感知し、能動的に無効化できるようにすることです。

2. 手法：不一致制約付きMDP (DCMDP)

著者らは、標準的なマルコフ決定過程（MDP）から、**不一致制約付きMDP（DCMDP）**へのパラダイムシフトを提案しています。このフレームワークは、訓練と推論の間の不一致をブラックボックス最適化問題として扱い、ヘテロジニアス（異種混合）な訓練および推論セットアップをシームレスに処理することを可能にします。

主要構成要素：

• 不一致指標：
  著者らは、3つのトークンレベルの不一致信号、すなわち、確率差（$\delta_{diff}$）、確率比（$\delta_{ratio}$）、および対数確率差（$\delta_{log}$）を経験的に評価しました。その結果、絶対確率差（$\delta_{diff} = |p_{Train} - p_{Inf}|$）が最も効果的なペナルティであることを特定しました。これは、比率ベースのペナルティが低確率のトークンに対して保守的になりすぎるのと異なり、デプロイメントに関連するドリフトを修正しつつ、報酬駆動型の探索を過度に制限しない「可塑性」のある正則化を提供します。

• 不一致許容領域（Discrepancy Tolerance Region）：
  重要な経験的知見として、不一致許容領域の存在が挙げられます。訓練初期に微小で無視できる不一致を積極的に罰することは、ポリシーの探索を抑制し、学習効率を低下させます。逆に、この領域の外側にある過度な不一致を削減することは、最適化の一貫性を向上させます。

  • メカニズム： DCMDPの定式化では、不一致が定義された閾値 $c$ を超えた場合にのみペナルティを適用します。この許容領域内では、ポリシーは報酬を最大化するために自由に探索することが許されます。

• ラグランジュ緩和による二重目的最適化：
  この問題は、平均トークンレベルの不一致に対する制約条件下で、期待報酬を最大化するものとして定式化されます：
  $$\max_{\theta} J(\theta) = \mathbb{E}_{\tau \sim \pi_{\theta}} \left[ \sum r(o_t) \right]$$
  $$\text{subject to } \frac{1}{|\tau|} \sum |\pi_{Train}(o_t) - \pi_{Inf}(o_t)| \leq c$$

これを解くために、著者らは**ラグランジュ緩和（Lagrangian Relaxation）**を採用しています。学習可能な乗数 $\lambda$ が、不一致ペナルティの重みを動的に調整します：
* 不一致が許容境界（$>c$）を逸脱した場合、$\lambda$ が増加し、ポリシーを安全圏に引き戻すためのペナルチーを強化します。
* 不一致が境界内にある場合、$\lambda$ はゼロに向かって減衰し、ポリシーが報酬最大化に集中できるようにします。

最終的な目的関数（Dual-GRPO）は、この適応型ペナルティをグループ正規化アドバンテージに直接注入します：
$$\hat{A}^{DCMDP}_{i,t} = \hat{A}_{i,t} - \lambda_k \cdot \delta_{i,t}$$

3. 主な貢献

1. 自律的な不一致補償： 本論文は、堅牢なブラックボックス・ペナルティ（トークンレベルの絶対確率差）を組み込むことで、訓練ポリシーに、不一致による訓練崩壊のリスクを自律的に排除する能力を付与できることを示しています。
2. DCMDPへのパラダイムシフト： 著者らは、LLMの強化学習を不一致制約付きMDPとして再定式化しました。このアプローチは、タスク報酬の最大化と訓練・推論間の不一致の最小化という二重の目的を最適にバランスさせ、単純な目的関数のスカラー化を超えた手法を提供します。
3. ヘテロジニアスな訓練パラダイム： この手法は、高忠実度でリソース集約的な環境（例：BF16精度）でモデルを最適化しつつ、低コストでリソース制約のあるデプロイメント環境（例：FP8）に向けて明示的に整合させる訓練セットアップを可能にします。これにより、インフラのヘテロジニアス性を、負債ではなく制御可能な最適化のノブへと変貌させます。

4. 実験結果

本手法は、Qwen3-8B（Dense）およびQwen3-30B-A3B（Mixture-of-Experts）モデルを用いて、6つの数学的推論ベンチマーク（AIME24/25, AMC23, MATH-500, Minerva, OlympiadBench）で評価されました。

• BF16 訓練/推論： DC-GRPOは、両方のモデルにおいてすべてのベンチマークでバニラGRPOベースラインを一貫して上回り、より高い平均スコアを達成しました（例：Qwen3-8Bは67.3から71.0に向上）。これは、不一致制約を課すことが、主要な報酬最適化を阻害するのではなく、むしろ強化することを裏付けています。
• ヘテロジニアス訓練 (BF16) + デプロイメント (FP8)： この設定では、訓練と推論のギャップが増幅されるため、バニラGRPOは大幅な性能低下または崩壊に見舞われました。DC-GRPOはこれをうまく緩和し、BF16のみのベースラインと同等、あるいはそれを上回る性能を維持しました。
• アブレーション研究： 本手法は、初期ラグランジュ乗数（$\lambda_0$）および二重学習率（$\eta_\lambda$）の変化に対して堅牢であることが証明されました。許容予算（$c$）は極めて重要であり、設定が小さすぎると過剰な正則化を招き、大きすぎると手法がバニラGRPOへと退化することが判明しました。

5. 意義と主張

本論文は、基礎となるハードウェアやソフトウェアのスタックに対して網羅的かつ特注の最適化を必要とせずに、訓練と推論の間の不安定性を緩和するための統一されたアルゴリズム的枠組みをDCMDPが提供すると主張しています。不一致をブラックボックスのフィードバック信号として扱うことで、本手法は以下を可能にします：

• 安定した訓練： RLのポストトレーニングでよく見られる急激な崩壊を防ぎます。
• 効率的な探索： 安全な不一致領域内で、ポリシーが自由に探索することを許容します。
• リソース効率： 高忠実度の訓練セットアップを使用して、制約のあるデプロイメント環境をシミュレートおよび適応させることを可能にし、それによって、より効率的な実世界のモデルデプロイメントを促進します。

著者らは、現在の経験的な範囲がBF16およびFP8の精度と特定のモデルファミリーに限定されていることを認め、他のバックエンドの組み合わせやモデルアーキテクチャへ本手法を拡張するための今後の課題があるとしています。