[论文解读] 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

一句话总结

引入在有限时间内能将任意两个目标密度精确连接的随机插值器，统一基于流的和扩散的生成模型，具有可学习的随时间变化的速度场与得分函数，通过二次目标函数实现。

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.

研究动机与目标

• Bridge two arbitrary densities rho0 and rho1 exactly in finite time using stochastic interpolants.
• Develop a unified framework that yields both deterministic (ODE) and stochastic (SDE) generative models via transport and Fokker-Planck equations.
• Provide methods to learn drift and score through simple quadratic objective functionals.
• Enable likelihood and cross-entropy estimations for interpolant-based models and connect to Schrödinger bridges.

提出的方法

• Define stochastic interpolants x_t = I(t,x0,x1) + γ(t) z that connect rho0 and rho1.
• Show that the time-dependent density rho(t) satisfies a first-order transport equation with velocity b(t,x) = E[∂t I + γ′(t) z | x_t = x].
• Derive forward and backward Fokker-Planck equations with a tunable diffusion ε(t) and link to probability flow (ODE) and SDE formulations.
• Characterize the velocity field b and the score s = ∇ log ρ via unique quadratic objective functionals.
• Introduce a denoiser η_z(t,x) as the conditional expectation E[z | x_t = x] and relate it to the score.
• Provide procedures to estimate these objectives from data and to compute likelihoods and cross-entropy for interpolant-based models.

实验结果

研究问题

• RQ1 How can two arbitrary densities be connected exactly in finite time using stochastic interpolants?
• RQ2 What are the governing transport and Fokker-Planck equations for the interpolant's time-evolving density, and how can they be learned from data?
• RQ3 How can one design deterministic (ODE) and stochastic (SDE) generative models within this framework, and how is likelihood controlled?
• RQ4 What is the relationship between stochastic interpolants and Schrödinger bridges, and how do score and denoiser concepts integrate into training and sampling?
• RQ5 How can practical estimators for likelihood and cross-entropy be constructed for interpolant-based models?

主要发现

• The interpolant's law is absolutely continuous and its density rho(t) solves a transport equation and a family of forward/backward Fokker-Planck equations with tunable diffusion.
• The drift coefficients in the ODE/SDE representations are the unique minimizers of simple quadratic objective functions, including a new score-based objective for the interpolant density.
• A denoiser η_z(t,x) = E[z | x_t = x] exists and provides a practical link to the score, enabling score-based interpretations.
• Forward and backward Fokker-Planck formulations enable likelihood estimation for SDE-based models and reveal additional considerations for Fisher divergences in ODE-based models.
• The framework recovers the Schrödinger bridge between rho0 and rho1 when optimized over the interpolant, linking to optimal transport concepts.
• Connections are established with score-based diffusion models, stochastic localization, denoising approaches, and rectified flows, with practical algorithmic guidance and numerical illustrations.