[論文レビュー] Stochastic Interpolants: A Unifying Framework for Flows and Diffusions

Michael S. Albergo, Nicholas M. Boffi+1 more|arXiv (Cornell University)|Mar 15, 2023

Gaussian Processes and Bayesian Inference被引用数 28

ひとこと要約

有限時間で任意の2つのターゲット密度を正確に橋渡しする確率的補間子を導入し、流れベースと拡散ベースの生成モデルを統一。学習可能な時間依存の速度場とスコアを二次目的関数を通じて導出。

ABSTRACT

A class of generative models that unifies flow-based and diffusion-based methods is introduced. These models extend the framework proposed in Albergo and Vanden-Eijnden (2023), enabling the use of a broad class of continuous-time stochastic processes called stochastic interpolants to bridge any two probability density functions exactly in finite time. These interpolants are built by combining data from the two prescribed densities with an additional latent variable that shapes the bridge in a flexible way. The time-dependent density function of the interpolant is shown to satisfy a transport equation as well as a family of forward and backward Fokker-Planck equations with tunable diffusion coefficient. Upon consideration of the time evolution of an individual sample, this viewpoint leads to both deterministic and stochastic generative models based on probability flow equations or stochastic differential equations with an adjustable level of noise. The drift coefficients entering these models are time-dependent velocity fields characterized as the unique minimizers of simple quadratic objective functions, one of which is a new objective for the score. We show that minimization of these quadratic objectives leads to control of the likelihood for generative models built upon stochastic dynamics, while likelihood control for deterministic dynamics is more stringent. We also construct estimators for the likelihood and the cross entropy of interpolant-based generative models, and we discuss connections with other methods such as score-based diffusion models, stochastic localization, probabilistic denoising, and rectifying flows. In addition, we demonstrate that stochastic interpolants recover the Schrödinger bridge between the two target densities when explicitly optimizing over the interpolant. Finally, algorithmic aspects are discussed and the approach is illustrated on numerical examples.

研究の動機と目的

• rho0とrho1という2つの任意の密度を、確率的補間子を用いて有限時間で正確に橋渡しする。
• 輸送とフォッケ-プランク方程式を通じて、決定論的なODEと確率論的なSDEの両方の生成モデルを生み出す、統合的な枠組みを開発する。
• 単純な二次目的関数を通じてドリフトとスコアを学習する方法を提供する。
• 補間子ベースのモデルの尤度と交差エントロピーの推定を可能にし、シュレーディンガー橋と結びつける。

提案手法

• rho0とrho1を結ぶ x_t = I(t,x0,x1) + γ(t) z という確率的補間子を定義する。
• 時間依存密度 rho(t) が、速度場 b(t,x) = E[∂t I + γ′(t) z | x_t = x] による1階の輸送方程式を満たすことを示す。
• 調整可能な拡散 ε(t) を伴う前向き・後向きのフォッケ-プランク方程式を導出し、確率的フロー（ODE）およびSDE形式への結びつきを示す。
• 一意な2次の目的関数を用いて、速度場 b およびスコア s = ∇ log ρ を特徴づける。
• η_z(t,x) を条件付き期待値 E[z | x_t = x] として導入し、これをスコアと関連づける。
• データからこれらの目的関数を推定する手法と、補間子ベースのモデルの尤度および交差エントロピーを計算する手順を提供する。

実験結果

リサーチクエスチョン

• RQ1 確率的補間子を用いて、任意の2つの密度を有限時間で正確に結ぶにはどうすればよいか。
• RQ2 補間子の時間発展密度に対する支配的輸送およびフォッケ-プランク方程式は何か、データからどのように学習できるか。
• RQ3 この枠組みの下で、決定論的（ODE）と確率論的（SDE）生成モデルをどのように設計し、尤度はどのように制御されるか。
• RQ4 確率的補間子とシュレーディンガー橋の関係は何か、スコアとデノイザーの概念は訓練とサンプリングにどう統合されるか。
• RQ5 補間子ベースのモデルの尤度と交差エントロピーの実用的な推定量をどのように構築できるか。

主な発見

• 補間子の法則は絶対連続性をもち、その密度 rho(t) は、調整可能な拡散を伴う輸送方程式および前方/後方フォッケ-プランク方程式の族を解く。
• ODE/SDE 表現におけるドリフト係数は、単純な二次目的関数の一意的最小化子であり、補間子密度の新しいスコアベースの目的関数を含む。
• η_z(t,x) = E[z | x_t = x] が存在し、スコアへの実用的なリンクを提供し、スコアベースの解釈を可能にする。
• 前方・後方フォッケ-プランクの定式化はSDEベースのモデルの尤度推定を可能にし、ODEベースのモデルにおけるFisher情報量の追加考慮を明らかにする。
• フレームワークは補間子上で最適化されたときrho0とrho1の間のシュレーディンガー橋を回復し、最適輸送の概念と結びつく。
• スコアベース拡散モデル、確率的局在化、デノイジング法、および整正化流の関連性を確立し、実用的なアルゴリズム指針と数値例を提供する。