Score Approximation for Diffusion Models on Arbitrary Low-Dimensional Structures

問題提起
スコアベース拡散モデルの理論的基礎は、多くの場合、基礎となるデータ分布に対する制限的な仮定、具体的にはリプシッツ連続な密度関数や滑らかな多様体上のサポートを必要としてきました。しかし、現実世界の知覚データ（画像など）は、密度が不連続であったり、勾配が適切に定義されなかったりする特異点、鋭い境界、あるいは離散的なクラスターを伴うことがよくあります。スコア関数 $\nabla \log p_t(x)$ を近似するための既存の複雑性境界は、通常、周囲の次元 $n$ に対して指数的な依存性を抱えているか、あるいは滑らかさを欠く分布には適用できないという問題を抱えています。このギャップは、次のような根本的な理論的問いを提起します：拡散モデルは、正則性や滑らかさを仮定することなく、任意のコンパクトなデータ分布に対して理論的に成功することができるのだろうか？

手法
著者らは、すべての連続性および滑らかさの仮定を取り除く、ユニバーサルなスコア近似フレームワークを提案しています。彼らのアプローチの核心は、新しい離散混合形式と、ガウスカーネルによる平滑化効果に焦点を当てた解析的な戦略にあります。

1. 離散混合形式: 著者らは、初期分布 $p_0$ を、それぞれ小さな球 $B(y_i, r)$ にサポートを持つ有限個の劣分布 $p_i$ の混合として分解します。真のスコア関数は、各劣分布上の積分を、局所的な平均 $\tilde{y}_i$ における点質量評価に置き換えることで近似されます。これにより、スコア近似問題は、重み付きガウスカーネルの和を近似する問題へと変換されます。
2. 分割統治法による誤差解析: 離散化誤差を抑えるために、著者らはクエリ点 $x$ に対して「近い」グループと「遠い」グループに空間を分割します。
   • 近いグループ: サポートに近い点については、局所的な平均の周りでのガウスカーネルのテイラー展開により、近似誤差が小さいことが示されます。
   • 遠いグループ: サポートから遠い点については、ガウス重みが指数関数的に減衰します。著者らは、これらの領域において局所的な近似誤差が大きくなったとしても、カーネルの急速な減衰により、総誤差への寄与は無視できるほど小さいことを証明します。
3. ユニバーサルな正則性: 理論的な大きな障壁は、「悪い」点（近似が失敗する可能性のある点）の測度が無視できるものであることを保証することです。著者らは $C_1, C_2, r$-正則性条件を導入し、「ユニバーサル正則性」定理を証明しました：任意のコンパクトにサポートされた分布に対して、正則ではない点の集合は、任意に小さな測度しか持ちません。これにより、ほとんど至る所で近似が成立することが保証されます。
4. ニューラルネットワークの実装: この離散混合式は、ReLUネットワークによって実装可能であることが示されています。著者らは、二乗、指数、除算、およびSoftmaxを制御された誤差で近似するネットワークを構築し、出力がサポートの凸包内に留まることを確実にしています。

主要な貢献

• ユニバーサル・スコア近似定理: 本論文は、任意のコンパクトなサポートを持つ分布（上ミンコフ次元 $d$ によって特徴付けられる）に対して、スコア関数がReLUネットワークによって近似可能であることを証明する最初の定理を確立しました。これにより、リプシッツ、ヘルダー、またはソボレフの滑らかさの仮定が不要になります。
• 複雑性境界: 著者らは、ネットワークの複雑性（パラメータサイズ）が、周囲の次元 $n$ に対しては多項式的に成長し、固有の次元 $d$ に対してのみ指数的に成長することを証明しました。具体的には、サイズはおよそ $O(n^3 d/4 (1/\epsilon)^{d/2})$ となります。これは、既存の理論を悩ませている「周囲の次元の呪い」を打破するものです。
• 緩和された正則性条件: ユニバーサルな正則性の概念を導入することで、本研究は、近似が失敗する「悪い集合」が、確率密度関数を持たない（例：カントール測度のような）フラクタル的な測量を含む、あらゆるコンパクトな分布において無視できる測度であることを示しています。

結果
主な結果（定理1）は、直径 $R$ のコンパクトなサポートを持つ分布 $p_0$ と上ミンコフ次元 $d$ に対して、真のスコア関数 $\nabla \log p_t(x)$ は、データ予測形式 $\nabla \log p_{t,appr}(x) = s_\theta(x, t) - x/t$ によって近似可能であり、すべての $t \in [t_0, T]$ において平均二乗誤差が $7\epsilon^2$ 以下に抑えられることを述べています。

• 近似ネットワーク $s_\theta$ は、深さが $O(\log \omega \log \log \omega)$ であり、幅が $O(C_K r^{-d} (n \log \omega + \log^2 \omega))$ です（ここで $\omega$ は $n$ と $1/\epsilon$ の多項式に依存します）。
• 誤差境界はグローバルに成立し、アウトオブディストリビューションの点における近似誤差は $R/t$ であり、これは小さな確率領域において可積分です。

意義
本論文は、不規則で非滑らかなデータ構造に対する拡散モデルの経験的な成功に対し、厳密な理論的正当付けを提供すると主張しています。滑らかさの仮定なしに、任意のコンパクトな分布に対して正確なスコア推定が可能であることを証明することで、本研究は、理論的な保証と現実世界のデータ特性（ピクセルクリッピングや特異点など）との間の溝を埋めています。著者らは、彼らの分割統治の手法と正則性の結果が、高次の導関数（ヘシアンなど）を抑えるためにも適応可能であり、現在これらの項の有界性に依存している拡散モデルの将来的な理論研究を助ける可能性があると述べています。本研究は、スコア近似理論の範囲を、滑らかな多様体から一般的なコンパクトサポートを持つ分布へと拡張する、基礎的な一歩として位置付けられています。