ART for Diffusion Sampling: A Reinforcement Learning Approach to Timestep Schedule

技術的サマリー：拡散サンプリングにおける ART

問題定義
スコアベースの拡散モデルは、学習された逆時間ダイナミクスを積分することでサンプルを生成する。この過程における重要な数値的選択は、時間範囲を時間ステップのグリッドに離散化することである。標準的なアプローチは、通常、一様グリッドまたは手動設計されたスケジュール（例：EDM、DPM-Solver）を採用する。著者らは、固定された関数評価回数（NFE）の制約下では、これらの固定スケジュールがしばしば最適ではないと主張する。一様グリッドは、逆ダイナミクスが軌道上で著しく変化する場合でも、計算リソースを均等に割り当てる。一方、手動設計されたグリッドは、特定のデータ分布やスコア関数の複雑さに適応しない可能性がある。目標は、サンプリング軌道に沿って計算リソースを体系的に再配分し、集合的なオイラー離散化誤差を最小化することである。

手法
本論文は、最適な時間ステップスケジュールのための制御理論的枠組みである**適応的再パラメータ化時間（ART）**を導入する。

1. 制御定式化（ART）:

   • グリッド点を直接選択するのではなく、ART は拡散プロセスの「クロック速度」を制御変数として扱う。
   • 連続的な写像 $\psi(t)$ を介して、再パラメータ化された時間変数 $t$ を元の拡散時間 $\tau$ にマッピングする。
   • 制御 $\theta(t) = \dot{\psi}(t)$ は、拡散時間の瞬間的な変化率を表す。
   • 目的は、総時間予算が保存されるという制約（$\int_0^T \theta(t) dt = T$）の下で、集合的な局所オイラー離散化誤差を最小化するように $\theta(t)$ を選択することである。
   • 局所誤差の代理指標は、$|Q(x, \psi)| \theta(t)^2$ に比例すると導出される。ここで $Q$ はスコア関数とその微分に依存する。

2. 強化学習による解法（ART-RL）:

   • 決定論的な ART 問題は、高次元の状態空間とハミルトン・ヤコビ・ベルマン（HJB）方程式の閉形式解の欠如により、直接解くことが困難である。
   • 著者らは、連続時間強化学習（CTRL）アプローチであるART-RLを提案する。
   • 制御 $\theta$ をガウス方策 $\pi(\cdot | t, x, \psi) = \mathcal{N}(\mu, \lambda/|Q|)$ からサンプリングするランダム化された補助問題を導入する。
   • この定式化により、方策平均 $\mu$ と価値関数を確率近似を通じて学習するために、アクター・クリティックアルゴリズム（特に Jia と Zhou [14] の枠組みを適用）を使用することが可能となる。

3. 理論的架け橋:

   • 核心的な理論的貢献は、決定論的 ART 最適解と ART-RL 補助問題における最適ガウス方策との間の双方向の架け橋の確立である。
   • 十分性: 決定論的最適制御 $\mu^*$ は、その平均がまさに $\mu^*$ となる最適ガウス方策へと持ち上げられる。
   • 必要性: 補助問題に対する任意の最適ガウス方策は、その平均が決定論的 ART 制御 $\mu^*$ を回復しなければならない。
   • この結果は、アクターをガウス平均としてパラメータ化し、トレーニング後に決定論的に展開することを正当化し、学習されたスケジュールが元の制御問題に対する真の最適解であるという厳密な保証を提供する。

4. 蒸留:

   • オフライントレーニング後、確率的方策は、時間のみを考慮する決定論的グリッドへと蒸留される。これにより、推論中のステップごとのアクターネットワーク評価や分散計算の必要性が排除され、標準的なスケジュール手法と比較して追加の計算コストがゼロであることを保証する。

主な貢献

• 定式化: 適応的時間再パラメータ化による拡散時間ステップスケジュールのための、初の制御理論的枠組みである ART を提案する。
• 理論: 決定論的最適制御問題とガウス方策を用いた連続時間アクター・クリティック学習手順との間の双方向の等価性を確立する。これにより、ヒューリスティックな RL 応用を、決定論的最適解を見つけるための原理的な経路へと変換する。
• アルゴリズム: データ駆動で適応的時間スケジュールを学習するアクター・クリティックアルゴリズム ART-RL を開発する。
• 汎化: 学習されたスケジュールが、再トレーニングなしで異なるデータセットや時間ステップ数に対して蒸留・転送可能であることを実証する。

実験結果
著者らは、公式の EDM パイプライン内で CIFAR-10、AFHQv2、FFHQ、ImageNet において ART-RL を評価した。

• CIFAR-10: ART-RL は、広範な NFE バジェットにおいて、一様、DPM-Solver log-SNR、および標準的 EDM スケジュールを常に上回る。
  • 低バジェット（例：NFE=3）において、ART-RL は認識可能な物体を生成するのに対し、ベースラインはノイズに近い状態のままとなる。
  • NFE=35 において、ART-RL は FID 1.82 を達成し、EDM（1.85）および DPM-logSNR（1.89）を上回る。
  • 複数のシードによる統計的有意性が確認され、ART-RL の最悪の実行結果が EDM の最良の実行結果を上回った。
• 汎化:
  • 補間/外挿: CIFAR-10 上で $K=18$ ステップで学習されたスケジュールは、他のステップ数（$K \in \{4, 9, 12, 15, 20\}$）に補間または外挿された場合でも、その性能優位性を維持する。
  • クロスデータセット転送: CIFAR-10 で蒸留されたスケジュールは、再トレーニングなしで AFHQv2、FFHQ、ImageNet に適用される。特に低および中程度の NFE において、EDM および DPM-logSNR に対して FID を一貫して改善する。
• 1 次元解析: 解析的スコア関数を持つ設定において、ART-RL は手動設計されたスケジュールと比較してワッサーシュタイン距離（$W_2$）を著しく低減し、改善がスコア推定誤差ではなく、より良い時間割り当てに起因することを確認する。

意義と主張
本論文は、適応的時間再パラメータ化を介して拡散時間ステップスケジュールを制御問題として定式化した最初の研究であり、決定論的最適解とアクター・クリティック学習手順との間の厳密な双方向の架け橋を提供した最初の研究であると主張する。

著者らは、その手法がデータ駆動かつアモルタイズ（償却）可能であることを強調する。

• データ駆動: スケジュールは、手動設計されたものではなく、目的関数とデータダイナミクスから学習される。
• アモルタイズド転送: オフライントレーニングコスト（CIFAR-10 については Colab T4 で約 1〜2 時間）は一回限りの経費である。得られた蒸留された決定論的グリッドは、追加の推論コストなしで異なるデータセットやステップ数で再利用可能であり、一様や EDM などの標準的スケジュール手法と同様の実行時間を達成する。

この研究は、適応的時間再パラメータ化を、生成拡散モデルにおける体系的なスケジュール設計のための萌芽的だが有望な方向性として位置づけ、将来の研究のための制御理論的な出発点を提供する。