Manifold Drift in Flow Preference Optimization: A Root Cause of Reward Hacking

技術要約：フロー・プリファレンス最適化における多様体ドリフト（Manifold Drift）

問題定義：多様体ドリフト（Manifold Drift）
本論文は、連続時間生成モデル、特にフロー・マッチング（Flow Matching; FM）や拡散モデルに対して直接的選好最適化（Direct Preference Optimization; DPO）を適用する際の根本的な失敗モードを特定している。DPOは、出力確率を再重み付けすることで、LLMのような離散モデルに対しては効果的である。しかし、連続モデルは学習された軌道に沿ってノイズをデータ多様体へと輸送することによってサンプルを生成する。

著者らは、標準的な選好更新（例：FlowDPO）が、事前学習されたデータ多様体に対して本質的に制約を設けない方法でこれらの輸送ダイナミクスを修正してしまうと主張している。これにより、**多様体ドリフト（Manifold Drift）**が生じる。これは、ファインチューニングされたモデルの終端サンプルが、事前学習された分布のサポートから外れてしまう現象（$\text{supp}(\mu_\theta) \not\subseteq M_0$）である。

• 理論的原因： 理論的には、最適なフロー・マッチングは終端のデータ分布を復元する。しかし、選好更新は、誘導された終端変位が事前学習された多様体に対して垂直な成分を持つような勾配ステップを導入する。この「多様体外」への変位により、モデルは特定の報酬関数（例：OCRの正確性）において高いスコアを得る一方で、視覚的なアーティファクト、意味的な歪み、あるいは忠実度の低下を示すサンプルを生成するようになる。
• 報酬ハッキング（Reward Hacking）： 本論文は、このドリフトが報酬ハッキングの根本原因であると仮定している。モデルは、「多様体を認識していない（manifold-unaware）」報酬（テキストの正しさなど）に対して高いスコアを得るために、幾何学的に無効であったり、分布外であったりする画像を生成することで、有効なデータサポートを離脱し、実質的に報酬関数を回避することがある。

手法：THERMODPO
多様体ドリフトに対処するため、著者らは、事前学習された多様体に選好最適化を明示的に繋ぎ止める温度制御された目的関数、THERMODPOを提案している。

1. 熱力学的定式化： この手法は、アライメントを、報酬の最大化と多様体の保持との間のバランスを取る行為として捉え、それを時間依存の温度関数 $\tau(t)$ によって制御する。
2. 三状態システム： 目的関数は、以下の3つの状態を持つシステム内での、選好されるサンプルの負の対数確率を最小化する：
   • 選好状態 ($E_w$)
   • 拒絶状態 ($E_l$)
   • リファレンス／ベースライン状態 ($E_b$)
     損失関数は次のように定式化される：
     $$L_{\text{ThermoDPO}}(\theta) = \mathbb{E}\left[ -\tau(t) \cdot t^2 \cdot \log \frac{e^{E_w}}{e^{E_w} + e^{E_l} + e^{E_b}} \right]$$
3. 理論的特性：
   • RFTへの帰着： 温度 $\tau \to 0$ のとき、この目的関数は、選好されるサンプルの多様体上での再構成を厳格に強制する、拒絶サンプリング・ファインチューニング（Rejection Sampling Fine-Tuning; RFT）へと収束する。
   • アンカー付きFlowDPO： $\tau > 0$ の場合、この目的関数は、温度スケールされたFlowDPO項と、非負の**勝者側アンカー項（winner-side anchoring term）**に分解される。このアンカーは、選好されるサンプルに関する事前学習された多様体からの逸脱を罰する。
   • 幾何学的境界： 理論的には、この損失は、再構成された選好サンプルとデータ多様体との間の二乗距離の上界を与えるものであり、ドリフトに対する点ごとの保証を提供する。
4. 実践的な実装（THERMODPO-weighted）： 標準的な定式化には、終端エンド付近（$t=0$）で多様体アンカーが弱まる $t^2$ という因子が含まれている。これは、保存が最も重要となる場所において、保存機能が弱まることを意味する。これを解決するため、著者らはグローバルな $t^2$ を $(1-t)^2$ に置き換えたTHERMODPO-weightedを導入した。これにより、コアとなる目的関数の理論的保証を損なうことなく、終端時刻付近でアンカーを動的に活性化させ、堅牢な保存を実現している。

主な貢献

• 多様体ドリフトの定式化： 本論文は、ファインチューニングされた軌道が事前学習された終端サポートを離れる失敗モードとして多様体ドリフトを正式に定義し、更新における非ゼロの垂直成分がこのドリフトの十分条件であることを特定した。
• THERMODPOフレームワーク： RFTとFlowDPOを橋渡しし、報酬の最大化と多様体の保存の間のトレードオフを明示的に制御する統一された目的関数を導入した。
• 理論的保証： THERMODPOが低温度極限においてRFTに帰着すること、およびFlowDPOと非負のアンカー項に分解されることの証明を提供した。また、選好されるサンプルの多様体距離に対する上界も確立した。
• 実証的検証： 合成3D多様体を用いたトイ実験および実画像ベンチマークを通じて、THERMODPO-weightedが既存の手法と比較して、選好のアライメントとサンプル品質の優れたバランスを実現することを実証した。

結果

• トイ・ベンチマーク： 合成3D多様体において、THERMODPO-weightedは StrictScore 0.899 を達成し、FlowDPO (0.629) および FlowDPO+RFT (0.857) を大幅に上回った。一方、バニラのFlowDPOは高い「Win」率を達成したものの、「OnManifold（多様体上）」の妥当性が0%となり、ドリフト仮説が裏付けられた。
• 実画像生成（SD3.5-M）： OCR選好データセットを用いたStable Diffusion 3.5-Mによる評価：
  • THERMODPO-weightedは、ベースラインと比較してOCRスコアを 47.5% 向上させた（CFG=4.5時）。
  • 4つの指標（GenEval, OCR, HPSv3.0, UniReward）において、平均 16.0% の改善を達成した。
  • 極めて重要なことに、品質の顕著な低下と保持指標（GenEval, HPSv3.0）でのスコア低下が見られたFlowDPOのバリアントとは異なり、THERMODPO-weightedは視覚的な忠実度を維持し、直接最適化されていない指標においても競争力のある性能を示した。
• FLUX.2-klein-base-4B： 異なるモデルアーキテクチャにおいても同様の傾向が観察され、フローベースの生成器に対する本手法の有効性が強化された。

意義と主張
本論文は、連続時間モデルにおける選好最適化は、単なる出力の選択ではなく、輸送ダイナミクスを根本的に変化させるものであると主張している。多様体ドリフトを定式化することで、著者らは、この構造的な問題を無視することが報酬ハッキングとサンプル品質の低下につながると論じている。THERMODPOは、事前学習中に学習された生成的な事前分布を犠牲にすることなく、選好アライメントを可能にする実用的な解決策として提示されている。著者らは、これを画像およびビデオ生成における信頼性の高いアライメントのための必要なステップとして位置づけており、特に「多様体を認識していない」報酬（テキスト生成など）が、本来は分布外の振る舞いを助長する可能性があるタスクにおいて重要であるとしている。本研究は、制限のない展開に対する即時の安全性ではなく、制御されたオフライン研究設定のための技術的貢献を目指したものとして位置づけられている。