PAPER KEY: C7VCWNCW
TITLE: Stochastic Interpolants: A Unifying Framework for Flows and Diffusions
AUTHORS: Albergo, Michael S.; Boffi, Nicholas M.; Vanden-Eijnden, Eric

Stochastic Interpolants: A Unifying Framework for Flows and Diffusions
Michael S. Albergo∗1, Nicholas M. Boffi∗2, and Eric Vanden-Eijnden2
1Center for Cosmology and Particle Physics, New York University 2Courant Institute of Mathematical Sciences, New York University
November 7, 2023
Abstract
A class of generative models that unifies flow-based and diffusion-based methods is introduced. These models extend the framework proposed in [2], enabling the use of a broad class of continuoustime stochastic processes called ‘stochastic interpolants’ to bridge any two arbitrary 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 probability density function of the stochastic interpolant is shown to satisfy a first-order 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 immediately 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 of the interpolant density. 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 processes, probabilistic denoising techniques, and rectifying flows. In addition, we demonstrate that stochastic interpolants recover the Schro ̈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.
∗Author ordering alphabetical; authors contributed equally.
1
arXiv:2303.08797v3 [cs.LG] 6 Nov 2023


Contents
1 Introduction 3 1.1 Background and motivation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 3 1.2 Main contributions and organization . . . . . . . . . . . . . . . . . . . . . . . . . . . . 4 1.3 Related work . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 5 1.4 Notation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 7
2 Stochastic interpolant framework 7 2.1 Definitions and assumptions . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 7 2.2 Transport equations, score, and quadratic objectives . . . . . . . . . . . . . . . . . . . 10 2.3 Generative models . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 14 2.4 Likelihood control . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 15 2.5 Density estimation and cross-entropy calculation . . . . . . . . . . . . . . . . . . . . . 17
3 Instantiations and extensions 20 3.1 Diffusive interpolants . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 20 3.2 One-sided interpolants for Gaussian ρ0 . . . . . . . . . . . . . . . . . . . . . . . . . . 22 3.3 Mirror interpolants . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 24 3.4 Stochastic interpolants and Schr  ̈odinger bridges . . . . . . . . . . . . . . . . . . . . . 25
4 Spatially linear interpolants 26 4.1 Factorization of the velocity field . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 26 4.2 Some specific design choices . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 27 4.3 Impact of the latent variable γ(t)z and the diffusion coefficient ε(t) . . . . . . . . . . 29 4.4 Spatially linear one-sided interpolants . . . . . . . . . . . . . . . . . . . . . . . . . . . 31
5 Connections with other methods 32 5.1 Score-based diffusion models and stochastic localization . . . . . . . . . . . . . . . . 32 5.2 Denoising methods . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 33 5.3 Rectified flows . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 34
6 Algorithmic aspects 36 6.1 Learning . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 36 6.2 Sampling . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 38
7 Numerical results 40 7.1 Deterministic versus stochastic models: 2D . . . . . . . . . . . . . . . . . . . . . . . . 40 7.2 Deterministic versus stochastic models: 128D Gaussian mixtures . . . . . . . . . . . . 41 7.3 Image generation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 44
8 Conclusion 45
A Bridging two Gaussian mixture densities 47
B Proofs 49 B.1 Proof of Theorems 2.6, 2.7, and 2.8, and Corollary 2.10. . . . . . . . . . . . . . . . . . 49 B.2 Proof of Lemma 2.19 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 54 B.3 Proofs of Lemmas 2.21 and 2.22, and Theorem 2.23. . . . . . . . . . . . . . . . . . . . 56 B.4 Proofs of Lemma 2.25 and Theorem 2.26 . . . . . . . . . . . . . . . . . . . . . . . . . . 58 B.5 Proof of Theorem 3.2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 61 B.6 Proof of Lemma 3.12 and Theorem 3.13. . . . . . . . . . . . . . . . . . . . . . . . . . . 62 B.7 Proof of Theorem 5.3 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 63 B.8 Proof of Theorem 5.5 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 64
C Experimental Specifications 64 C.1 Image Experiments . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 64
2


1 Introduction
1.1 Background and motivation
Dynamical approaches for deterministic and stochastic transport have become a central theme in contemporary generative modeling research. At the heart of progress is the idea to use ordinary or stochastic differential equations (ODEs/SDEs) to continuously transform samples from a base probability density function (PDF) ρ0 into samples from a target density ρ1 (or vice-versa), and the realization that inference over the velocity field in these equations can be formulated as an empirical risk minimization problem over a parametric class of functions [24, 58, 25, 60, 5, 2, 41, 39].
A major milestone was the introduction of score-based diffusion methods (SBDM) [60], which map an arbitrary density into a standard Gaussian by passing samples through an Ornstein-Uhlenbeck (OU) process. The key insight of SBDM is that this process can be reversed by introducing a backwards SDE whose drift coefficient depends on the score of the time-dependent density of the process. By learning this score – which can be done by minimization of a quadratic objective function known as the denoising loss [68] – the backwards SDE can be used as a generative model that maps Gaussian noise into data from the target. Though theoretically exact, the mapping takes infinite time in both directions, and hence must be truncated in practice.
While diffusion-based methods have become state-of-the-art for tasks such as image generation, there remains considerable interest in developing methods that bridge two arbitrary densities (rather than requiring one to be Gaussian), that accomplish the transport exactly, and that do so on a finite time interval. Moreover, while the highest quality results from score-based diffusion were originally obtained using SDEs [60], this has been challenged by recent works that find equivalent or better performance with ODE-based methods if the score is learned sufficiently well [32]. If made to match the performance of their stochastic counterparts, ODE-based methods exhibit a number of desirable characteristics, such as an exact, computationally tractable formula for the likelihood and the easy application of well-developed adaptive integration schemes for sampling. It is an open question of significant practical importance to understand if there exists a separation in sample quality between generative models based on deterministic dynamics and those based on stochastic dynamics.
In order to satisfy the desirable characteristics outlined in the previous paragraph, we develop a framework for generative modeling based on the method proposed in [2], which is built on the notion of a stochastic interpolant xt used to bridge two arbitrary densities ρ0 and ρ1. We will consider more general designs below, but as one example the reader can keep in mind:
xt = (1 − t)x0 + tx1 + p2t(1 − t)z, t ∈ [0, 1], (1.1)
where x0, x1, and z are random variables drawn independently from ρ0, ρ1, and the standard Gaussian density N(0, Id), respectively. The stochastic interpolant xt defined in (1.1) is a continuoustime stochastic process that, by construction, satisfies xt=0 = x0 ∼ ρ0 and xt=1 = x1 ∼ ρ1. Its paths therefore exactly bridge between samples from ρ0 at t = 0 and from ρ1 at t = 1. A key observation is that:
The law of the interpolant xt at any time t ∈ [0, 1] can be realized by many different processes, including an ODE and forward and backward SDEs whose drifts can be learned from data.
To see why this is the case, one must consider the probability distribution of the interpolant xt. As
shown below, for a large class of densities ρ0 and ρ1 supported on Rd, this distribution is absolutely continuous with respect to the Lebesgue measure. Moreover, its time-dependent density ρ(t) satisfies a first-order transport equation and a family of forward and backward Fokker-Planck equations in which the diffusion coefficient can be varied at will. Out of these equations, we can readily derive generative models that satisfy ODEs and SDEs, respectively, and whose densities at time t are given by ρ(t).
3


March 13, 2023 1
Ω(t)
v s
t i m e
Ω0
Ω1
ODE
Ω(t)
v s
t i m e
Ω0
Ω1
SDE
Figure 1: The stochastic interpolant paradigm. Example generative models based on the proposed framework, which connects two densities ρ0 and ρ1 using samples from both. The design of the time-dependent probability density ρ(t) that bridges between ρ0 and ρ1 is separated from the choice of how to sample it, which can be accomplished with deterministic or stochastic generative models. Left panel: Sampling with a deterministic (ODE) generative model known as the probability flow equation. Right panel: Sampling with a stochastic generative model given by an SDE with a tunable diffusion coefficient. The probability flow equation and the SDE have different paths, but their time-dependent density ρ(t) is the same. Moreover, the two equations rely on the same estimates for the velocity and the score.
Interestingly, the drift coefficients entering these ODEs/SDEs are the unique minimizers of quadratic objective functions that can be estimated empirically using data from ρ0, ρ1, and N(0, Id). The resulting least-squares regression problem allows us to estimate the drift coefficients of the ODE/SDEs, which can then be used to push samples from ρ0 onto new samples from ρ1 and vice-versa.
1.2 Main contributions and organization
The approach introduced here is a versatile way to build generative models that unifies and extends many existing algorithms. In Sec. 2, we develop the framework in full generality, where we emphasize the following key contributions:
• We prove that the stochastic interpolant defined in Section 2.1 has a distribution that is absolutely continuous with respect to the Lebesgue measure on Rd, and that its density ρ(t) satisfies a first-order transport equation (TE) as well as a family of forward and backward Fokker-Planck equations (FPEs) with tunable diffusion coefficients.
• We show how the stochastic interpolant can be used to learn the drift coefficients that enter the TE and the FPEs. We characterize these coefficients as the minimizers of simple quadratic objective functions given in Section 2.2. We introduce a new objective for the score ∇ log ρ(t) of the interpolant density, as well as an objective function for learning a denoiser ηz, which we relate to the score.
• In Section 2.3, we derive ordinary and stochastic differential equations associated with the TE and FPEs that lead to deterministic and stochastic generative models. In Section 2.4, we show that regressing the drift for SDE-based models controls the likelihood, but that regressing the drift alone is not sufficient for ODE-based models, which must also minimize a Fisher divergence. We show how to optimally tune the diffusion coefficient to maximize the likelihood for SDEs.
• In Section 2.5, we develop a general formula to evaluate the likelihood of SDE-based generative models that serves as a natural counterpart to the continuous change-of-variables formula commonly used to compute the likelihood of ODE-based models. In addition, we give formulas to
4


estimate the cross-entropy.
In Section 3, we discuss instantiations of the stochastic interpolant method. In Section 3.4 we first show that interpolants are equivalent to a class of stochastic bridges, but that they avoid the need for Doob’s h-transform, which is generically unknown; we show that this simplifies the construction of a broad class of generative models. In Section 3.2, we define the one-sided interpolant, which corresponds to the conventional setting in which the base ρ0 is taken to be a Gaussian. With a Gaussian base, several aspects of the interpolant simplify, and we detail the corresponding objective functions. In Section 3.3, we introduce a mirror interpolant in which the base ρ0 and the target ρ1 are identical. Finally, in Section 3.4, we show how the interpolant framework leads to a natural formulation of the Schr  ̈odinger bridge problem between two densities.
In Section 4, we discuss a special case in which the interpolant is spatially linear in x0 and x1. In this case, the velocity field can be factorized, which we show in Section 4.1 leads to a simpler learning problem. We detail specific choices of linear interpolants in Section 4.2, and in Section 4.3 we illustrate how these choices influence the performance of the resulting generative model, with a particular focus on the role of the latent variable and the diffusion coefficient. For exposition, we focus on Gaussian mixture densities, for which the drift coefficients can be computed analytically. We provide the resulting formula in Appendix A. Finally, in Section 4.4, we discuss the case of spatially linear one-sided interpolants.
In Section 5, we formalize the connection between stochastic interpolants and related classes of generative models. In Section 5.1, we show that score-based diffusion models can be re-written as one-sided interpolants after a reparameterization of time; we highlight how this approach eliminates singularities that appear when naively compressing score-based diffusion onto a finite-time interval. In Section 5.2, we show how interpolants can be used to derive the Bayes-optimal estimator for a denoiser, and we show how this approach can be iterated to create a generative model. In Section 5.3, we consider the possibility of rectifying the flow map of a learned generative model. We show that the rectification procedure does not change the underlying generative model, though it may change the time-dependent density of the interpolant.
In Section 6, we provide the details of practical algorithms associated with the mathematical results presented above. In Section 6.1, we describe how to numerically estimate the objectives given empirical datasets from the base and the target. In Section 6.2, we complement this discussion on learning with algorithms for sampling with the ODE or an SDE.
We provide numerical demonstrations in line with these recommendations in Section 7, and we conclude with some remarks in Section. 8.
1.3 Related work
Deterministic Transport and Normalizing Flows. Transport-based sampling and density estimation has its contemporary roots in Gaussianizing data via maximum entropy methods [23, 12, 64, 63]. The change of measure under such transformation is the backbone of normalizing flow models. The first neural network realizations of these methods arose through imposing clever structure on the transformation to make the change of measure tractable in discrete, sequential steps [52, 16, 50, 28, 19]. A continuous time version of this procedure was made possible by viewing the map T = Xt(x) as the solution of an ODE [11, 24], whose parametric drift defining the transport is learned via maximum likelihood estimation. Training this way is intractable at scale, as it requires simulating the ODE. Various methods have introduced regularization on the path taken between the two densities to make the ODE solves more efficient [22, 48, 65], but the fundamental difficulty remains. We also work in continuous time; however, our approach allows us to learn the drift without simulation of the dynamics, and can be formulated at sample generation time through either deterministic or
5


stochastic transport.
Stochastic Transport and Score-Based Diffusions (SBDMs). Complementary to approaches based on deterministic maps, recent works have realized that connecting a data distribution to a Gaussian density can be viewed as the evolution of an Ornstein-Ulhenbeck (OU) process which gradually degrades samples from the distribution of interest to Gaussian noise [54, 25, 58, 60]. The OU process specifies a path in the space of probability densities; this path is simple to traverse in the forward direction by addition of noise, and can be reversed if access to the score of the time-dependent density ∇ log ρ(t) is available. This score can be approximated through solution of a least-squares regression problem [29, 68], and the target can be sampled by reversing the path once the score has been learned. Interestingly, the resulting forward and backward stochastic processes have an equivalent formulation (at the distribution level) in terms of a deterministic probability flow equation, first noted by [4, 49, 33] and then applied in [44, 57, 34, 7]. The probability flow formulation is useful for density estimation and cross-entropy calculations, but it is worth noting that the probability flow and the reverse-time SDE will have densities that differ when using an approximate score. The SBDM framework, as it has been originally presented, has a number of features which are not a priori well motivated, including the dependence on mapping to a normal density, the complicated tuning of the time parameterization and noise scheduling [69, 26], and the choice of the underlying stochastic dynamics [17, 32]. While there have been some efforts to remove dependency on the OU process using stochastic bridges [51], the resulting procedures can be algorithmically complex, relying on inexact mixtures of diffusions with limited expressivity and no accessible probability flow formulation. Some of these difficulties have been lifted in follow-up works [42]. As another step in this direction, we observe that the key idea behind SBDMs – the bridging of densities via a time-dependent density whose evolution equation is available – can be generalized to a much wider class of processes in a straightforward and computationally accessible manner.
Stochastic Interpolants, Rectified Flows, and Flow matching. Variants of the stochastic interpolant method presented in [2] were also presented in [41, 39]. In [41], a linear interpolant was proposed with a focus on straight paths. This was employed as a step toward rectifying the transport