技術要約：逆強化学習による拡散モデルからのサンプリング学習

問題提起

拡散モデルは高次元データの生成における標準となっていますが、サンプリングの手順が依然として重要なボトルネックとなっています。通常、デノイジング・ニューラルネットワークは学習後に固定されますが、サンプリングアルゴリズム自体には、ハイパーパラメータ（ノイズ・スケジュール、ガイダンス・スケール、確率性プロファイルなど）の慎重なチューニングが必要です。現在、このチューニングはコストのかかる経験的なグリッドサーチによって行われています。ImageNet-64のような複雑なモデルの場合、網羅的なハイパーパラメータ探索の計算コストは、デノイザー自体の学習コストに匹敵するか、それを上回ることがあります。さらに、既存の強化学習（RL）や制御手法は、明示的な報酬関数やデノイザーの微調整を必要とすることが多く、これは事前学習済みの生成器を変更せずにサンプラーを最適化するという目標に矛盾します。

手法

著者らは、拡散サンプリングを**逆強化学習（Inverse Reinforcement Learning: IRL）**によって解決可能な逐次決定問題として定式化することを提案しています。コアとなる革新性は、エキスパートの軌跡、明示的な報酬、またはデノイザーの再学習にアクセスすることなく、ターゲットの状態分布に一致するサンプリング・ポリシーを学習することです。

1. MDPの定式化

拡散サンプリングのプロセスは、有限ホライゾンのマルコフ決定過程（MDP）として定式化されます。

• 状態 ($s_t$): $(x_t, \sigma_t)$ として定義されます。ここで $x_t$ は現在のサンプル、$\sigma_t$ はノイズレベルです。
• 行動 ($a_t$): ガイダンス・スケールの選択、ノイズ注入レベルの選択、あるいは以前のノイズレベルへの再ノイズ化（リスタート）の決定など、サンプリングのダイナミクスに適用される離散的な制御決定です。
• ダイナミクス: 遷移カーネルは、離散化された拡散ダイナミクスと選択された制御行動によって誘導されます。プロセスは、ノイズレベルが $\Sigma_0$（吸収状態）に達したときに終了します。

2. 状態のみの逆強化学習

エキスパートの行動や軌跡を想定することが多い従来のIRLとは異なり、本手法は状態のみのレジームで動作します。「エキスパート」は、選択されたノイズレベルにおけるデータの周辺分布のみによって定義されます。

• 目的: 目標は、エキスパートの占有測度 $\mu_E$ とポリシーの占有測度 $\mu_\theta$（訪問された状態の周辺分布）を一致させるポリシー $\pi_\theta$ を見つけることです。
  $$\min_\theta D_f(\mu_E \parallel \mu_\theta)$$
  本手法では、エキスパートとポリシーの占有測度の間の $f$-ダイバージェンスを直接最小化します。具体的には、KLダイバージェンスと逆KL（rKL）ダイバージェンスを調査しており、選択する $f$-ダイバージェンスが、モード被覆（多様性）とモード探索（忠実度）のトレードオフを決定することに注目しています。

3. ポリシー勾配と最適化

著者らは、$f$-ダイバージェンスを直接最小化するためのポリシー勾配定理を導出しています。勾配は以下のように表されます：
$$\nabla_\theta L_f(\theta) = \mathbb{E}_{s_{1:T} \sim p_\theta} \left[ \frac{1}{T} \sum_{t=1}^T \nabla_\theta \log \pi_\theta(a_t|s_t) \sum_{t' \ge t} h_f\left(\frac{\mu_E(s_{t'})}{\mu_\theta(s_{t'})}\right) \right]$$
ここで、$h_f$ は $f$-ダイバージェンスの生成器から導かれる関数です。

• 密度比推定: $\mu_E$ と $\mu_\theta$ は未知であるため、エキスパートの状態（ノイズが付加されたデータ）とポリシーによって生成された状態を区別するように訓練された識別器を使用して、密度比を推定します。
• 安定化: 勾配推定の高分散に対処するため、本手法はPPO型のクリップされた目的関数、重要度サンプリング補正、およびグローバルなバッチ平均ベースラインを採用しています。

4. 評価されたサンプリング戦略

このフレームワークは、以下の3つの推論時介入に適用されています。

1. 適応型Classifier-Free Guidance: 状態依存のガイダンス・スケール $\omega(x, \sigma)$ を学習します。
2. 確率性注入: 状態依存のノイズ注入レベル $\gamma_{EDM}(x, \sigma)$ を学習します。
3. 再ノイズ化（リスタート・サンプリング）: 劣悪な生成結果から脱出するために、プロセスを以前のノイズレベルに戻すタイミングを学習します。

主な貢献

1. サンプリングのためのMDP定式化: 拡散サンプリングを、デノイザーを再学習することなく推論時にサンプラーを適応させるためのMDPとして捉えました。
2. 直接的な状態周辺分布の一致: 明示的な報酬モデリングやエキスパートの行動データを必要とせずに、状態占有測度間の $f$-ダイバージェンスを最小化するポリシー勾配目的関数を導出しました。
3. 解釈可能なトレードオフ: $f$-ダイバージェンス（KL vs rKL）の選択を通じて、サンプルの多様性（Recall）と忠実度（Precision）の制御可能なトレードオフを明らかにしました。
4. 状態依存の制御: 学習されたポリシーは、特定のノイズレベルでのみ高いガイダンスを適用したり、特定のサンプルタイプに対して介入したりするなど、推論時の制御として機能し、固定されたグローバルなスケジュールよりも高度な介入を可能にします。

実験結果

本手法は、事前学習済みのEDMモデルを用いてCIFAR-10、FFHQ、およびImageNet-64で評価されました。

• パフォーマンス: 学習されたポリシーは、FID、Precision、Recallの指標において、グリッドサーチによって見出された最良の手動チューニング・ベースラインと同等またはそれを上回る性能を示しました。
  • 適応型ガイダンス: 固定ガイダンスと比較して、品質と多様性のトレードオフを改善しました。KLダイバージェンスはRecall（多様性）を向上させ、rKLはPrecision（忠実度）を向上させました。
  • 確率性注入: 最適なEDMスケジュールの定性的パターン（軌道の終盤では低確率性）を復元し、ImageNetのような複雑なデータセットにおいて固定ベースラインを上回りました。
  • 再ノーズ化: 標準的なリスタート戦略と比較して、FIDを維持または向上させつつ、関数評価数（NFE）を大幅に削減しました。
• 計算効率:
  • 学習: 単一の学習実行が、網羅的なグリッドサーチを置き換えます。ImageNet-64において、この手法は最適化予算を最大 9倍 削減しました（3つのハイパーパラメータのグリッドサーチによる約1437 EFLOPスから、IRL学習による約161 EFLOPスへ）。
  • 推論: ポリシーによるオーバーヘッドは、デノイザー単体と比較してFLOPsで 16%から40% 程度であり、手動チューニングの排除を考慮すれば無視できる程度であると考えられます。
• アブレーション研究の知見:
  • 条件付け: フル状態 $(x, \sigma)$ で条件付けられたポリシーは、ノイズレベル $\sigma$ のみで条件付けられたものよりも大幅に優れており、クラスラベル $(x, \sigma, c)$ を用いた場合とほぼ同等の性能を示しました。これは、ポリシーが画像内容からクラス識別的な特徴を暗黙的に学習していることを示唆しています。
  • ダイバージェンスの選択: KLとrKLは、Precision-Recallのトレードオフを制御する上で最も効果的であることが示されました。一方、対称的なダイバージェンス（TV、Hellinger）は、解釈性に乏しい結果となりました。

意義と主張

本論文は、このフレームワークが拡散サンプリングにおける手動のハイパーパラメータ・チューニングに対する系統的かつ自動化された代替手段を提供すると主張しています。サンプリングをIRL問題として扱うことで、手動で設計することが困難な複雑で状態依存のサンプリング戦略を最適化することが可能になります。

著者らは、この手法が特に、試行ごとの評価コストが膨大になる大規模モデルにおいて価値が高いことを強調しています。このアプローチは生成モデル自体を変更するのではなく、推論の軌道に対する最適な「制御ポリシー」を学習するものであり、生成の再学習や網羅的な探索という計算負荷を負うことなく、サンプル品質と多様性を向上させる実用的な道筋を提供します。なお、本論文は、恩恵はデータセットに依存すること、および現在のフレームワークは離散アクション空間に依存していることを認め、今後の課題として連続アクション空間やより堅牢な初期化を挙げています。