技術的概要：拡散モデルのためのノイズスケジューリング設計：最適制御の視点

問題定義

拡散モデルは、近似サンプリングのための主要な枠組みとして浮上し、生成を学習されたスコア関数によって駆動される時間反転型確率微分方程式（SDE）のシミュレーションとして定式化している。訓練目的（スコアマッチング）はよく理解されているが、ノイズスケジューリング（前方過程と逆過程を支配する関数 $f(t)$ および $g(t)$）の設計は、依然として主に経験則に依存している。

既存の理論的研究は、特定の（しばしば一定の）ノイズスケジューリングに対するサンプリング誤差 bound（クラウバック・ライブラー発散で測定）を確立している。しかし、これらの bound は通常、データ次元 $d$ および問題定数に対して非効率的にスケーリングし、$\tilde{O}(d^2/n)$ や $\tilde{O}(\max\{J^\star, d\}/n)$ として現れることが多い。ここで $J^\star$ はデータ分布のフィッシャー情報、$n$ は離散化ステップ数である。$J^\star$ は $d$ よりも著しく大きくなり得るため、これらの保証は問題依存の深刻なペナルティに苦しむ。一方、$\tilde{O}(d/n)$ の最先端の理論的 bound は存在するが、これらは実務でほとんど展開されない特定のノイズスケジューリングに依存しており、拡散モデルの理論と実証的な実装の間にギャップが生じている。

手法

著者らは、ノイズスケジューリングの設計を最適制御問題として再定式化する原理的な枠組みを提案する。中核的な革新は、確率密度の進化（フォッカー・プランク偏微分方程式によって支配される）の扱いにくい最適化を回避し、フィッシャー情報の軌道に焦点を当てることにある。

1. 最適制御の定式化:

   • 状態: 拡散過程のフィッシャー情報 $J_t = J(p_t)$。これはボフネルの公式を通じて導出された 1 次元常微分方程式（ODE）に従って進化し、次式で表される：
     $$\dot{J}_t = 2f(t)J_t - g(t)\|\nabla s_t\|_{L^2}^2$$
   • 制御: ノイズスケジューリング関数 $(f, g)$。
   • 目的: KL サンプリング誤差の上界となるフィッシャー情報を含む汎関数を最小化すること。著者らは、ギルサノフの定理、ボフネルの公式、エントロピーパワの凹性、およびブラクマン・スタムの不等式を含む確率論および情報理論のツールを用いて、この bound を導出した。

2. 閉形式スケジューリングの導出:

   • 対象データ分布が**強対数凹性（SLC）**であるか、**ガウス混合モデル（GMM）**であるという追加の仮定の下、著者らは最適ノイズスケジューリングの閉形式の特性を導出する。
   • 彼らは、$f(t)$ と $g(t)$ がアフィン関係 $f(t) = \theta g(t) + \omega$ を満たす**アフィン結合スケジューリング（ACS）**と呼ばれるクラスを導入する。
   • 得られた ODE を解くことで、$(\theta, \omega, g_0, \lambda^*)$ でパラメータ化された $f(t)$ と $g(t)$ の明示的な式が得られる。

3. 適応的パラメータ化:

   • この枠組みは、導出されたパラメータを計算予算（ニューラル関数評価、NFE、$n$ と表記）に基づいて調整可能なハイパーパラメータとして扱う。
   • 最適減衰率 $\lambda^*$ はランベルト W 関数を介して $n$ と動的にスケーリングすることが示されており、これによりスケジューリングは初期化誤差と離散化誤差を自動的にバランスさせることができる。

主要な貢献

1. 最適制御枠組み: 本論文は、状態をフィッシャー情報の軌道とする最適制御問題としてノイズスケジューリング設計を定式化する。これにより、扱いにくい PDE 制約が扱いやすい ODE に置き換えられ、最適誤差 bound に対する十分条件の導出が可能となる。
2. $\tilde{O}(d/n)$ 誤差 bound: 著者らは、最適制御問題の解が KL 発散において $\tilde{O}(d/n)$ のサンプリング誤差 bound を達成することを証明する。これは、実務で使用されているものを含む広範なノイズスケジューリングのクラスに、最先端の理論的保証を拡張するものである。
3. アフィン結合スケジューリング（ACS）: 彼らは、標準的な経験則的スケジューリングを一般化する新しいパラメトリックなノイズスケジューリングのクラス（ACS）を導出した：
   • VE-Exponential: パラメータ $(\theta, \omega) = (0, 0)$ を設定することで復元される。
   • VP-Sigmoid: $(\theta, \omega) = (1/2, 0)$ を設定することで復元される。
   • ACS 枠組みは、これらの固定された経験則を超えて体系的な最適化を可能にする追加の調整可能パラメータ（$g_0, c$）を導入する。
4. 実証的検証: CIFAR-10 および ImageNet（64x64）における実験は、ACS ハイパーパラメータを体系的に調整することが、様々な NFE 予算において標準的なベースライン（VP-Linear, VP-Cosine, VP-Sigmoid）よりも優れたフレッチェ・インセプション距離（FID）スコアをもたらすことを示している。

結果

• 理論的: 本論文は、SLC および GMM 分布に対して、提案された ACS 枠組みが最適な $\tilde{O}(d/n)$ 収束率を達成することを確立する。さらに、実務で一般的な一定または線形スケジューリングが、特にデータ分布が高度に異方的である場合（$J^\star \gg d$）、非効率的なトレードオフに苦しむ理由を分析する。
• 実証的: 無条件 CIFAR-10 およびクラス条件付き ImageNet において、適応的 ACS 枠組みは標準的なベースラインを一貫して上回る。例えば、ImageNet において NFE が 50 の場合、ACS は FID 5.30 を達成し、VP-Linear（7.46）、VP-Cosine（5.67）、VP-Sigmoid（6.77）を上回った。これらの結果は、ACS が計算予算に効果的に適応して曲率と分散の軌道を調整するのに対し、固定されたスケジューリングはステップサイズの変化に伴い品質が低下することを示している。

重要性

本論文は、ノイズスケジューリング設計のための原理的な方法を提供することで、拡散モデルの理論と実践の間のギャップを埋めると主張している。経験則的な掃引や固定された経験則に依存するのではなく、著者らは特定のスケジューリングが機能する理由についての理論的正当性と、新しく最適化されたスケジューリングを生成するメカニズムを提供する。

その重要性は以下の点にある：

1. 一般化: 標準的な経験則的スケジューリング（Exponential, Sigmoid）が、より広範で理論的に最適なクラス（ACS）の特殊なケースであることを示すこと。
2. 適応性: 静的なスケジューリングを使用するのではなく、利用可能な計算予算（$n$）と内在的なデータ特性（フィッシャー情報）に基づいてスケジューリングを動的に調整する方法を提供すること。
3. 性能: この理論的に裏付けられたアプローチが、標準的なベンチマークにおける生成品質（FID スコア）の目に見える改善をもたらすことを示し、生成モデルにおける最適制御の視点の有効性を検証すること。

著者らは、閉形式の導出が SLC または GMM 分布を仮定しているものの、得られたスケジューリングは複雑な実世界のデータでも良好に機能することを指摘しており、これは理論的洞察が導出のために作られた特定の仮定を超えて一般化されることを示唆している。今後の課題として、これらの特性記述をより広範な複雑なデータ分布のクラスに拡張することが挙げられている。