[論文レビュー] Sampling, Diffusions, and Stochastic Localization

Andrea Montanari|arXiv (Cornell University)|May 18, 2023

Markov Chains and Monte Carlo Methods被引用数 4

ひとこと要約

本稿は、ノイズ除去拡散モデルと確率的局在化の間の理論的関係を確立し、標準的なノイズ除去拡散が確率的局在化プロセスの特殊ケースであることを示している。また、確率的局在化をより広範な設定に一般化し、この視点がモデル学習を単純化することを示し、ガウスモデルにおける線形推定近似のもとで、目的分布への収束を証明している。

ABSTRACT

Diffusions are a successful technique to sample from high-dimensional distributions. The target distribution can be either explicitly given or learnt from a collection of samples. They implement a diffusion process whose endpoint is a sample from the target distribution. The drift of the diffusion process is typically represented as a neural network. Stochastic localization is a successful technique to prove mixing of Markov Chains and other functional inequalities in high dimension. An algorithmic version of stochastic localization was recently proposed in order to sample from certain statistical mechanics models. This expository article has three objectives: $(i)$~Generalize the algorithmic construction to other stochastic localization processes. This construction is both simple and broadly applicable; $(ii)$~Clarify the connection between diffusions and stochastic localization. This allows to derive several known sampling schemes in a unified fashion; $(iii)$~Describe the insights that follow from this unified viewpoint.

研究の動機と目的

• EAMS22 で提示されたアルゴリズム的確率的局在化フレームワークを、より広範な確率的過程のクラスへ一般化すること。
• 拡散ベースのサンプリングと確率的局在化の間の理論的関係を明確にし、標準的なノイズ除去拡散が後者の特定の例であることを示すこと。
• 確率的局在化の視点が、拡散モデルにおけるドリフト関数の学習をどのように簡素化できるかを示すこと。
• 確率的局在化のレンズを通じて、サンプリングアルゴリズムの設計と分析に新たな知見を提供すること。
• 高次元ガウスモデルにおける線形推定子に基づく手法の収束保証を確立すること。

提案手法

• 元のフレームワークを超えた時間非一様的かつ非可逆的ダイナミクスを許容する、確率的局在化プロセスの一般化を提案する。
• 既知の初期分布から目的測度へのサンプリング経路を構築するため、拡散過程の時間反転スデを導出する。
• 現在の状態を条件とした、目的変数の線形変換の条件付き期待値がドリフトとなる連続時間確率過程を導入する。
• 循環行列構造を活用して計算効率を高めるために、線形かつシフト不変な推定子構造を用いて最適ドリフトを近似する。
• 有限ウィンドウ畳み込み近似を最適推定子に適用し、複雑さを低減しつつ漸近的一貫性を維持する。
• ドリフト成分の学習を導くために、カルバック・ライバラー距離を損失関数として採用する。

実験結果

リサーチクエスチョン

• RQ1 確率的局在化フレームワークは、元の構築を超えて、より広いクラスの拡散過程を含めるように一般化可能か？
• RQ2 標準的なノイズ除去拡散モデルと確率的局在化プロセスの関係は何か？
• RQ3 確率的局在化の視点は、学習ドリフト関数の複雑さを軽減することで、スコアベース生成モデルの学習を簡素化できるか？
• RQ4 高次元ガウスモデルにおける線形推定子に基づく確率的局在化プロセスの収束挙動はいかなるものか？
• RQ5 有限ウィンドウ畳み込みで最適条件付き期待値を近似することは、目的分布への漸近的収束を保持するか？

主な発見

• 本稿は、標準的なノイズ除去拡散が確率的局在化プロセスの特殊ケースであることを証明し、2つの主要なサンプリングパラダイムの間の形式的な理論的関係を確立している。
• 有限ウィンドウ畳み込み近似を用いた線形かつシフト不変な推定子は、高次元ガウスモデルにおいて目的分布への漸近的収束を達成する。
• 時間 t が無限大に近づくにつれて、目的分布と生成分布の間の制限 Wasserstein 距離および全変動距離の両方がゼロに収束する。
• ガウス設定において、ドリフト成分の推定器は O(1/n) 項までベイズ最適であることが示され、それが原理的近似としての妥当性を裏付けている。
• 最適ドリフトが循環構造を持つ単純な線形関数によって近似可能であるため、複雑なニューラルネットワークアーキテクチャの必要性が低減され、学習タスクが簡素化される。
• 拡散と局在化の視点を統合することで、より解釈可能で解析可能なサンプリングアルゴリズムの設計が可能になる。