テクニカル・サマリー：分布マッチング・ポリシー最適化による拡散型LLMの推論能力向上

問題提起

拡散型大規模言語モデル（dLLM）は、任意の順序での生成を可能にし、潜在的に高い推論スループットを実現できるため、自己回帰型LLM（AR-LLM）に代わる有望な選択肢となっている。しかし、dLLMをAR-LLMと同等の推論能力までスケールアップさせることは依然として大きな課題である。グループ相対ポリシー最適化（GRPO）を含む既存の強化学習（RL）ファインチューニング手法は、主に以下の2つの要因により、dLLMには適していない。

1. 計算上の不一致： dLLMは双方向の処理に依存しているため、標準的なポリシー勾配法で必要となるシーケンスの対数確率（log-probability）の推定は、ARモデルと比較して計算コストが高く、しばしば困難を極める。
2. アルゴリズムの限界： dLLM向けの現在のRLアプローチは、トークン位置間の相関を無視したヒューリスティックな近似（例：平均場分解）に依存することが多い。さらに、これらは通常、報酬最大化に焦点を当てており、これは単一の報酬モードへの最適化を招く傾向がある。この「モード探索（mode-seeking）」行動は、報酬ハッキングを引き起こし、dLLMが理論的に生成可能な多様で高品質な推論パスの探索を妨げる可能性がある。

手法：分布マッチング・ポリシー最適化（DMPO）

本論文では、確率的最適制御（SOC）に基づいた原理的なRLファインチューニング・フレームワークである**分布マッチング・ポリシー最適化（DMPO）**を提案する。期待報酬を最大化する従来のメソッドとは異なり、DMPOはモデルのポリシー分布を、報酬によって傾斜付けされた最適なターゲット分布に一致させることを目的とする。

コア・フレームワーク

目的は、誘導されるシーケンス分布 $p_\theta(o|q)$ が、以下の最適な分布 $p^*(o|q)$ と一致するようにポリシー $\pi_\theta$ を学習することである。
$$p^*(o|q) = \frac{1}{Z(q)} p_{ref}(o|q) e^{r(q,o)/\alpha}$$
ここで、$p_{ref}$ は事前学習済みポリシー、$r$ は報酬関数、$\alpha$ は温度パラメータである。

これを達成するために、著者らは問題を、ターゲットと現在のポリシー間の前方KLダイバージェンス（クロスエントロピーと同等）を最小化することとして定式化する：
$$\min_\theta \mathbb{E}_{q \sim D} \mathbb{E}_{p^*(o|q)} \left[ \log \frac{p^*(o|q)}{p_\theta(o|q)} \right]$$
この「マス・カバリング（mass-covering）」目的関数は、単一のモードへの崩壊を防ぎ、多様な高報酬軌道を探索することをモデルに促す。

主要な技術コンポーネント

1. 重み付きデノイジング・クロスエントロピー（WDCE）：
   $p^*$ は（未知の分配関数 $Z(q)$ や直接的なサンプリングの不可能性により）計算不可能であるため、DMPOは重要度サンプリング（Importance Sampling）を採用する。これは、参照ポリシー $p_v$（モデルのストップグラディエント・コピー）からのサンプルを、$p^*$ からの重み付きサンプルとして扱うものである。
   結果として得られる損失関数WDCEは、dLLM特有の前方（ノイジング）プロセスを活用する。これは、最終的なクリーンなサンプルとその関連する重要度重みのみを使用して損失を計算し、中間的なロールアウト軌跡を破棄する。これにより、WDCEはオフポリシーかつ**前方のみ（forward-only）**の目的関数となる。

2. 重みベースライン減算：
   重要な実装上の課題として、小さな学習バッチサイズを使用する場合のWDCEの不安定性が特定された。小規模なバッチでは、「悪い」レスポンス（それでも $p^*$ の下でゼロではない確率を持つもの）に対する正の重みが支配的になり、誤った勾配更新を招くことがある。
   これに対処するため、著者らは重要度重みからベースライン重み $w_{base}$ を減算する重みベースライン減算を提案する：
   $$w_{real} = w - w_{base}$$
   3つの戦略が $w_{base}$ として提案されている：

   • グループ重みベースライン： 収束を促すための定数値（1）。
   • 個別重みベースライン： 負の報酬傾斜に基づくもの。
   • モデル重みベースライン： 現在のモデル自身の重要度重みに基づくもの。
     この手法は、低報酬のサンプルに対して効果的に負の勾配を注入し、低バッチシナリオにおいてモデルが悪質なレスポンスを促進してしまうのを防ぐ。

3. 重み付き直接判別最適化（WDDO）：
   クロスエントロピーの代替案として、悪いサンプルに対する負の勾配を自然に取り入れ、分布マッチングをさらに強化するGANに着想を得た目的関数を検討している。

主な貢献

1. 新しいRLフレームワーク： 報酬最大化のパラダイムから、報酬によって傾斜付けされた全ポリシー分布のマッチングへと転換した、dLLM専用の分布マッチング・フレームワークとしてのDMPOの提案。
2. 理論的根拠： 重要度サンプリングと確率的最適制御を用いたWDCE損失の導出により、高価なシーケンス対数確率の推定を必要としない、理論的に健全なオフポリシー学習の手法を提供。
3. 実用的な解決策： 前方KL目的関数における小バッチの失敗モードの特定、および学習を安定させサンプル効率を向上させるための重みベースライン減算とWDDOの導入。
4. 効率性： WDCEが前方プロセスのみによる訓練を可能にし、ロールアウトの再利用（オフポリシー）を可能にすること、および（KVキャッシュなどの高速推論技術と組み合わせることで）dLLMの高速な推論能力と互換性があり、軌跡ベースの手法と比較して訓練を大幅に加速させることを実証。

実験結果

著者らは、LLaDA-8B-InstructおよびDream-7Bモデルを用いて、4つの推論ベンチマーク（GSM8K、MATH500、Countdown、Sudoku）でDMPOを評価した。

• 性能向上： DMPOは、ベースモデルおよび既存のRLベースライン（diffu-GRPO、cGRPO、d1を含む）を一貫して上回った。
  • Sudokuにおいて、DMPO-LLaDAは、最高の非DMPO RLベースラインに対して39.63パーセントポイント、ベースモデルに対して67.97パーセントポイントの改善を達成した。
  • Countdownでは、強力なベースラインに対して最大16.40パーセントポイントの向上が観察された。
• SFT不要： 本手法は、厳選されたデータセットを用いた教師あり微調整（SFT）なしでこれらの結果を達成しており、事前学習済みモデルから直接推論能力をスケールアップできることを示した。
• 訓練効率： そのオフポリシーな性質と前方のみの計算により、DMPOはGRPO型のメソッドと比較して2〜3倍の高いサンプル効率と、最大8倍のロールアウトサンプリング速度（高速dLLM推論と組み合わせた場合）を実現した。d1の報酬レベルに、わずか31%の訓練計算予算で到達した。
• 安定性： シーケンスレベルの重要度重みの安定性に起因する特性により、非常に古いデータ（同一のロールアウトバッチを最大24回再利用）を使用した場合でも、手法は安定していた。

意義と主張

本論文は、DMPOが推論タスクにおけるdLLMをスケールアップするための重要な一歩であることを主張している。報酬最大化のパラダイムと軌跡ベースのRLによる計算上のボトルネックから脱却することで、DMPOは拡散モデルのユニークな特性を完全に活用している：

• 多様性： 前方KL目的関数はモード崩壊を防ぎ、dLLMが多様な推論パスを探索することを可能にする。
• 効率性： WDCEの前方のみのオフポリシーな性質は、dLLMの高スループットな推論能力と合致しており、大規模モデルにとって計算可能な、スケーラブルなRLファインチューニングへの道を提供する。
• 汎用性： 本フレームワークは、LLaDAとDreamの両方のアーキテクチャでの成功に示されている通り、様々なdLLMバックボーンに適用可能な原理的なアプローチとして提示されている。

著者らは、現在の研究は特定のモデルとベンチマークに焦点を当てているものの、分布マッチング・フレームワークは、生成モデルを改善し、より効果的な重みベースライン技術を設計するための新しい方向性を切り開くものであると結論付けている。