技術的サマリー：サンプルベースの変分推論による償却を用いた少数ステップ生成モデルの整合化

問題定義

報酬関数や人間の好みといった特定の目的に、少数ステップの生成モデル（GAN、整合モデル、フローマップ、ドリフトモデルなど）を整合させることは、重大な課題です。既存の整合フレームワークは、多くの場合、少数ステップ生成器がしばしば違反する制限的な構造仮定に依存しています：

1. 扱いやすい尤度： 多くの手法は、去噪過程の軌道における明示的な尤度評価を必要としますが、GAN やワンステップマッピングのような暗黙的モデルでは利用できません。
2. 特定のソルバー： 確率的最適制御（SOC）または強化学習（RL）に基づくアプローチは、しばしば特定の常微分方程式（ODE）または確率微分方程式（SDE）ソルバーと逐次的な去噪プロセスを仮定します。
3. モデル固有性： 定式化は頻繁に特定のモデルファミリー（拡散モデルなど）に合わせて調整されており、VAE やフローマップのような多様な少数ステップアーキテクチャへの汎化に失敗します。

その結果、少数ステップ生成器の多様性を微調整しつつ、その中核的な利点である高速な単一ステップ推論を維持できる、一般的な整合フレームワークが存在しません。

手法：FAV

著者は、整合を生成器ファミリーおよびサンプリングダイナミクスから切り離す一般的なフレームワークであるFAV（サンプルベース変分推論による少数ステップ生成モデルの整合化）を提案します。FAV は、整合を、参照分布にアンカーされた報酬傾斜分布からのサンプリングとして定式化します：
$$q^*(x) \propto p_{\text{ref}}(x) \exp(\beta \cdot r(x))$$
ここで、$p_{\text{ref}}$ は参照分布（事前学習された生成器や経験的データなど）であり、$r(x)$ は目的関数です。

この手法は、以下の 3 つの主要な段階で進行します：

1. SVGD によるサンプルベース近似

FAV は、**Stein 変分勾配降下（SVGD）**を用いて目的分布 $q^*$ を近似します。密度評価を必要とする代わりに、SVGD は分布を粒子で表現し、それらを反復的に目的に向かって輸送します。最適な Stein 速度場 $\phi^*$ は、KL 発散の減少を最大化するように導出されます：
$$\phi^*_{q_\theta, q^*}(x) = \mathbb{E}_{x' \sim q_\theta} \left[ k(x', x) \nabla_{x'} \log p_{\text{ref}}(x') + \beta \cdot k(x', x) \nabla_{x'} r(x') + \nabla_{x'} k(x', x) \right]$$
この場は 3 つの成分から構成されます：

• 事前整合： サンプルを参照分布 $p_{\text{ref}}$ へと駆動します。
• 報酬ガイダンス： サンプルを高報酬領域へと押し出します。
• 多様性の強制： カーネル勾配項を通じてモード崩壊を防止します。

2. 非パラメトリックスコア推定（KDE）

重要な課題として、経験的データや暗黙的生成器に対して $\nabla \log p_{\text{ref}}$ というスコアが扱いにくいことが挙げられます。FAV は、**カーネル密度推定（KDE）**を用いてスコアを非パラメトリックに推定することでこれを解決します。参照分布をガウスカーネルで畳み込むことで、スコアは参照サンプルと現在の粒子との差の加重平均として近似され、明示的な尤度を必要とせずに実効的に回避されます。

3. 固定点回帰による償却

少数ステップモデルの高速な推論速度を維持するため、FAV は推論時に反復的な SVGD 更新を実行しません。代わりに、固定点回帰を通じて輸送を生成器パラメータ $\theta$ に償却します。生成器は、その出力と「輸送された」ターゲットとの距離を最小化するように訓練されます：
$$\mathcal{L}(\theta) = \mathbb{E}_{\epsilon} \left[ \| \text{stopgrad}(f_\theta(\epsilon) + \hat{\phi}^*_{q_\theta, q^*_\sigma}(f_\theta(\epsilon))) - f_\theta(\epsilon) \|_2^2 \right]$$
ここで、輸送されたサンプルは安定した回帰ターゲットとして機能します。これにより、生成器は単一のフォワードパスで整合されたサンプルを生成することを学習できます。

ブラックボックス報酬への拡張： 微分不可能な報酬に対して、FAV はゼロ次勾配推定（ガウス平滑化）を用いて $\nabla r(x)$ を近似し、報酬関数がブラックボックスであっても整合を可能にします。

主要な貢献

1. 一般的な整合フレームワーク： FAV は、扱いやすい尤度、特定の ODE/SDE ソルバー、または多ステップ軌道を必要とせず、多様な少数ステップ生成器（VAE、GAN、整合モデル、フローマップを含む）を整合させる最初のフレームワークです。
2. サンプルベース VI 定式化： SVGD と KDE を用いて報酬傾斜分布からのサンプリングとして整合を定式化することで、FAV は整合手順を生成器の内部ダイナミクスから切り離します。
3. 効率性の維持： 固定点回帰による償却を通じて、FAV はベース生成器の単一ステップ推論効率を維持し、反復サンプリングまたは蒸留プロセスの計算オーバーヘッドを回避します。
4. ブラックボックス互換性： ゼロ次勾配推定の統合により、FAV は任意のブラックボックス報酬の最適化を可能にします。

実験結果

著者は、FAV をロボティクス操作と画像生成器の整合という 2 つの主要なドメインで評価しました。

1. ロボティクス操作（オフラインおよびオフラインからオンラインへの RL）

• ベンチマーク： 56 のオフライン RL タスク（OGBench および D4RL）と 30 のオフラインからオンラインへの RL タスク。
• パフォーマンス： FAV は、再パラメータ化された方策勾配、随伴マッチング、蒸留された生成方策（FQL、IFQL など）を含む 10 のベースラインを上回り、すべてのタスクで最先端のパフォーマンスを達成しました。
• 主要な発見： FAV は、補助的な生成モデルや多ステップ去噪がなくても、優れた方策抽出能力とオフラインからオンライン学習への遷移パフォーマンスを実証しました。

2. 画像生成器の整合

• スケール： 実験は ImageNet-256（256x256）から高解像度のテキストから画像への合成（1024x1024）まで範囲を広げました。
• モデル： StyleGAN-XL、ドリフトモデル、iMeanFlow、SANA-Sprint でテストされました。
• 指標： 審美的スコア、人間の好みスコア（HPS）、安全性検出器（NSFW）に対して評価されました。
• パフォーマンス：
  • FAV は、報酬最適化とサンプルの多様性/品質の間のパレートフロンティアにおいて一貫して最良の結果を達成し、DRaFT や随伴マッチングなどのベースラインよりも報酬の過剰最適化を軽減しました。
  • 推論効率： FAV は、競合する画像品質を維持しながら、推論時整合手法（Best-of-N、ReNO など）よりも約180 倍から 280 倍高速に整合されたサンプルを生成しました。
  • スケーラビリティ： この手法は 256x256 から 1024x1024 の解像度まで正常にスケールし、異なるモデルアーキテクチャ間で汎化しました。

意義と主張

本論文は、FAV が少数ステップ生成モデルの整合に対する統一された、モデルに依存しないソリューションを提供すると主張しています。その意義は以下の点にあります：

• 構造障壁の除去： 扱いやすい尤度や特定のサンプリング軌道の必要性を排除し、より広範なクラスの暗黙的生成モデルへの整合を可能にします。
• 効率性と品質のバランス： 高報酬生成と少数ステップ推論の計算効率性の間のトレードオフを成功裏に調整します。これは、以前の手法がしばしばパフォーマンスを低下させたり、複雑な蒸留を必要としたりした課題です。
• 頑健性： このフレームワークは、低次元制御から高解像度画像合成まで多様なタスクにわたって頑健性を示し、微分可能な報酬関数とブラックボックス報酬関数の両方を処理します。

著者は、FAV が RL における方策抽出および画像生成における報酬整合のための強力で汎用的なアプローチを表しており、より複雑で仮定に依存した手法に対する実用的な代替手段を提供すると結論付けています。