Stochastic Interpolants: A Unifying Framework for Flows and Diffusions

Introduction

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) $\rho_{0}$ into samples from a target density $\rho_{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 .

A major milestone was the introduction of score-based diffusion methods (SBDM) , 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 – 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 , this has been challenged by recent works that find equivalent or better performance with ODE-based methods if the score is learned sufficiently well . 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 , which is built on the notion of a stochastic interpolant $x_{t}$ used to bridge two arbitrary densities $\rho_{0}$ and $\rho_{1}$. We will consider more general designs below, but as one example the reader can keep in mind:

where $x_{0}$, $x_{1}$, and $z$ are random variables drawn independently from $\rho_{0}$, $\rho_{1}$, and the standard Gaussian density $\mathsf{N}(0,\text{\it Id})$, respectively. The stochastic interpolant $x_{t}$ defined in (1.1) is a continuous-time stochastic process that, by construction, satisfies $x_{t=0}=x_{0}\sim\rho_{0}$ and $x_{t=1}=x_{1}\sim\rho_{1}$. Its paths therefore exactly bridge between samples from $\rho_{0}$ at $t=0$ and from $\rho_{1}$ at $t=1$. A key observation is that:

The law of the interpolant $x_{t}$ at any time $t\in$ can be realized by many different processes, including an ODE and forward and backward SDEs whose drifts can be learned from data.

Interestingly, the drift coefficients entering these ODEs/SDEs are the unique minimizers of quadratic objective functions that can be estimated empirically using data from $\rho_{0}$, $\rho_{1}$, and $\mathsf{N}(0,\text{\it 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 $\rho_{0}$ onto new samples from $\rho_{1}$ and vice-versa.

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 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 $\nabla\log\rho(t)$ of the interpolant density, as well as an objective function for learning a denoiser $\eta_{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 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 $\rho_{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 $\rho_{0}$ and the target $\rho_{1}$ are identical. Finally, in Section 3.4, we show how the interpolant framework leads to a natural formulation of the Schrödinger bridge problem between two densities.

In Section 4, we discuss a special case in which the interpolant is spatially linear in $x_{0}$ and $x_{1}$. 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.

3 Related work

Transport-based sampling and density estimation has its contemporary roots in Gaussianizing data via maximum entropy methods . 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 . A continuous time version of this procedure was made possible by viewing the map $T=X_{t}(x)$ as the solution of an ODE , 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 , 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 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 . 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 $\nabla\log\rho(t)$ is available. This score can be approximated through solution of a least-squares regression problem , 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 and then applied in . 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 , and the choice of the underlying stochastic dynamics . While there have been some efforts to remove dependency on the OU process using stochastic bridges , 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 . 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 were also presented in . In , a linear interpolant was proposed with a focus on straight paths. This was employed as a step toward rectifying the transport paths through a procedure that improves sampling efficiency but introduces a bias. In Section 5.3, we present an alternative form of rectification that is bias-free. In , the interpolant picture was assembled from the perspective of conditional probability paths connecting to a Gaussian, where a noise convolution was used to improve the learning at the cost of biasing the method. Extensions of were presented in that generalize the method beyond the Gaussian base density. In the method proposed here, we introduce an unbiased means to incorporate noise into the process, both via the introduction of a latent variable into the stochastic interpolant and the inclusion of a tunable diffusion coefficient in the asociated stochastic generative models. We provide theoretical and practical motivation for the presence of these noise terms.

Optimal Transport and Schrödinger Bridges.

There is both theoretical and practical interest in minimizing the transport cost of connecting $\rho_{0}$ and $\rho_{1}$, which. In the case of deterministic maps, this is characterized by the optimal transport problem, and in the case of diffusive maps, by the Schrödinger Bridge problem . Formally, these two problems can be related by viewing the Schrödinger Bridge as an entropy-regularized optimal transport. Optimal transport has primarily been employed as a means to regularize flow-based methods by imposing either a path length penalty or structure on the parameterization itself . A variety of recent works have formulated the Schrödinger problem in the context of a learnable diffusion . In the interpolant framework, all propose optimal transport extensions to the learning procedure. The method proposed in allows one to sequentially lower the transport cost through rectification, at the cost of introducing a bias unless the velocity field is perfectly learned. The method proposed in is an unbiased framework at the cost of solving an additional optimization problem over the interpolant function. The statement of optimal transport in only applies to Gaussians, but is shown to be practically useful in experimental demonstrations.

In the method proposed below, we provide two approaches for optimizing the transport under a stochastic dynamics. Our primary approach, based on the scheme introduced in , is presented in Section 3.4. It offers an alternative route to solve the Schrödinger bridge problem under the Benamou-Brenier hydrodynamic formulation of transport by maximizing over the interpolant . However, we stress that this additional optimization step is not necessary in practice, as our approach leads to bias-free generative models for any fixed interpolant. In addition, Section 5.3 discusses an unbiased variant of the rectification scheme proposed in .

Convergence bounds.

Inspired by the successes of score-based diffusion, significant recent research effort has been expended to understand the control that can be obtained on suitable distances between the distribution of the generative model and the target data distribution, such as $\mathsf{KL}$, $W_{2}$, or $\mathsf{TV}$. Perhaps the first line of work in this direction is , which showed that standard score-based diffusion training techniques bound the likelihood of the resulting SDE model. Importantly, as we show here, the likelihood of the corresponding probability flow is not bounded in general by this technique, as first highlighted in the context of SBDM by . Control for SBDM-based techniques was later quantified more rigorously under the assumption of functional inequalities in a discretized setting by , which were removed by and via Girsanov-based techniques. Most relevant to the PDE-based methods considered here is , which applies similar techniques to our own in the SBDM context to obtain sharp guarantees with minimal assumptions.

4 Notation

Stochastic interpolant framework

We begin by defining the stochastic processes that are central to our approach:

The pair $(x_{0},x_{1})$ is drawn from a probability measure $\nu$ that marginalizes on $\rho_{0}$ and $\rho_{1}$, i.e.

$z$ is a Gaussian random variable independent of $(x_{0},x_{1})$, i.e. $z\sim{\sf N}(0,\text{\it Id})$ and $z\perp(x_{0},x_{1})$.

Eq. (2.2) states that $I(t,x_{0},x_{1})$ does not move too fast along the way from $x_{0}$ at $t=0$ to $x_{1}$ at $t=1$, and as a result does not wander too far from either endpoint – this assumption is made for convenience but is not necessary for most arguments below. Later, we will find it useful to consider choices for $I$ that are spatially nonlinear, which we show can recover the solution to the Schrödinger bridge problem. Nevertheless, a simple example that serves as a valid $I$ in the sense of Definition 2.1 is given in (1.1). The measure $\nu$ allows for a coupling between the two densities $\rho_{0}$ and $\rho_{1}$, which affects the properties of the stochastic interpolant, but a simple choice is to take the product measure $\nu(dx_{0},dx_{1})=\rho_{0}(x_{0})\rho_{1}(x_{1})dx_{0}dx_{1}$, in which case $x_{0}$ and $x_{1}$ are independent. In Section 6 we discuss how to design the stochastic interpolant in (2.1) and state some properties of the corresponding process $x_{t}$. Examples of stochastic interpolants are also shown in Figure 2 for various choices of $I$ and $\gamma$.

The main difference between the stochastic interpolant defined in (2.1) and the one originally introduced in is the inclusion of the latent variable $\gamma(t)z$. Many of the results below also hold when we set $\gamma(t)z=0$, but the objective of the present paper is to elucidate the advantages that this additional term provides when neither of the endpoints are Gaussian. We note that we could generalize the construction by making $\gamma(t)$ a tensor; here we focus on the scalar case for simplicity. Another difference is the possibility to couple $\rho_{0}$ and $\rho_{1}$ via $\nu$.

The stochastic interpolant $x_{t}$ in (2.1) is a continuous-time stochastic process whose realizations are samples from $\rho_{0}$ at time $t=0$ and from $\rho_{1}$ at time $t=1$ by construction. As a result, it offers a way to bridge $\rho_{0}$ and $\rho_{1}$ – we are interested in characterizing the law of $x_{t}$ over the full interval $$, as it will allow us to design generative models. Mathematically, we want to characterize the properties of the time-dependent probability distribution$\mu(t,dx)$ such that

where $\mu(t,dx)$ is the time-dependent distribution of $x_{t}$ defined by (2.3), and the expectation on the right-hand side is taken independently over $(x_{0},x_{1})\sim\nu$, and $z\sim{\sf N}(0,\text{\it Id})$.

Another seemingly more general way to define the stochastic interpolant is via

To proceed, we will make the following assumption on the densities $\rho_{0}$, $\rho_{1}$, and the interplay between the measure $\nu$ to the function $I$:

The measure $\nu$ and the function $I$ are such that

where the expectation is taken over $(x_{0},x_{1})\sim\nu$.

Note that for the interpolant (1.1), Assumption 2.5 holds if $\rho_{0}$ and $\rho_{1}$ both have finite fourth moments.

2 Transport equations, score, and quadratic objectives

We now state a result that specifies some important properties of the probability distribution of the stochastic interpolant $x_{t}$:

Note that this theorem means that we can write (2.4) as

The transport equation (2.9) can be solved either forward in time from the initial condition $\rho(0)=\rho_{0}$, in which case $\rho(1)=\rho_{1}$, or backward in time from the final condition $\rho(1)=\rho_{1}$, in which case $\rho(0)=\rho_{0}$.

The proof of Theorem 2.2 is given in Appendix B.1; it mostly relies on manipulations involving the characteristic function of the stochastic interpolant $x_{t}$. The transport equation (2.9) for $\rho$ lead to methods for generative modeling and density estimation, as explained in Secs. 2.3 and 2.5, provided that we can estimate the velocity $b$. This velocity is explicitly available only in special cases, for example when $\rho_{0}$ and $\rho_{1}$ are both Gaussian mixture densities: this case is treated in Appendix A. In general $b$ must be calculated numerically, which can be performed via empirical risk minimization of a quadratic objective function, as characterized by our next result:

where $x_{t}$ is defined in (2.1) and the expectation is taken independently over $(x_{0},x_{1})\sim\nu$ and $z\sim{\sf N}(0,\text{\it Id}).$

The proof of Theorem 2.2 is given in Appendix B.1: it relies on the definitions of $b$ in (2.10), as well as the definition of $\rho$ in (2.12) and some elementary properties of the conditional expectation. We discuss how to estimate the objective function (2.13) in practice in Section 6. Interestingly, we also have access to the score of the probability density, as shown by our next result:

where $x_{t}$ is defined in (2.1) and the expectation is taken independently over $(x_{0},x_{1})\sim\nu$ and $z\sim{\sf N}(0,\text{\it Id})$

The proof of Theorem 2.2 is given in Appendix B.1. We stress that the objective function is well defined despite the fact that $\gamma(0)=\gamma(1)=0$: see Section 6 for more details about how to evaluate this objective in practice.

will be referred as the denoiser, for reasons that will be made clear in Section 5.2. By (2.14), this quantity gives access to the score on $t\in(0,1)$ (where $\gamma(t)>0$) since, from (2.14),

This denoiser is the minimizer of an equivalent expression to (2.16),

The denoiser $\eta_{z}$ is useful for numerical realizations. In particular, the objective in (2.19) is easier to use than the one in (2.16) because it does not contain the factor $\gamma^{-1}(t)$, which needs careful handling as $t$ approaches 0 and 1.

Having access to the score immediately allows us to rewrite the TE (2.9) as forward and backward Fokker-Planck equations, which we state as:

[Fokker-Planck equations]corollaryinterpolationfpe For any $\epsilon\in C^{0}()$ with $\epsilon(t)\geq 0$ for all $t\in$, the probability density $\rho$ specified in Theorem 2.2 satisfies:

Equation (2.20) is well-posed when solved forward in time from $t=0$ to $t=1$, and its solution for the initial condition $\rho(t=0)=\rho_{0}$ satisfies $\rho(t=1)=\rho_{1}$.

Equation (2.22) is well-posed when solved backward in time from $t=1$ to $t=0$, and its solution for the final condition $\rho(1)=\rho_{1}$ satisfies $\rho(0)=\rho_{0}$.

In Section 2.3 we will use the results of this theorem to design generative models based on forward and backward stochastic differential equations. Note that we can replace the diffusion coefficient $\epsilon(t)$ by a positive semi-definite tensor; also note that if we define $\rho_{\mathsf{B}}(t_{\mathsf{B}},x)=\rho(1-t_{\mathsf{B}},x)$, the reversed FPE (2.22) can be written as

which is now well-posed forward in (reversed) time $t_{\mathsf{B}}$. So as to have only one definition of time $t$, it is more convenient to work with (2.22).

Let us make a few remarks about the statements made so far:

If we set $\gamma(t)=0$ in $x_{t}$ (i.e, if we remove the latent variable), the stochastic interpolant (2.1) reduces to the one originally considered in . In this setup, the results above formally stand except that we cannot guarantee the spatial regularity of $b(t,x)$ and $s(t,x)$, since it relies on the presence of the latent variable (as shown in the proof of Theorem 2.2). Hence, we expect the introduction of the latent variable $\gamma(t)z$ to help for generative modeling, where the solution to the corresponding ODEs/SDEs will be better behaved, and for statistical approximation, since the targets $b$ and $s$ will be more regular. We will see in Section 6 that it also gives us much greater flexibility in the way we can bridge $\rho_{0}$ and $\rho_{1}$, which will enable us to design generative models with appealing properties.

We will see in Section 2.4 that the forward and backward FPE in (2.20) and (2.22) are more robust than the TE in (2.9) against approximation errors in the velocity $b$ and the score $s$, which has practical implications for generative models based on these equations.

We could also obtain $b(t,\cdot)$ at any $t\in$ by minimizing

and $s(t,\cdot)$ at any $t\in(0,1)$ by minimizing

where $s$ is the score given in (2.14) and we defined the velocity field

Learning this velocity and the score separately may be useful in practice.

The objectives in (2.13) and (2.16) (as well as the ones in(2.19) and (2.29)) are amenable to empirical estimation if we have samples $(x_{0},x_{1})\sim\nu$, since in that case we can generate samples of $x_{t}=I(t,x_{0},x_{1})+\gamma(t)z$ at any time $t\in$. We will use this feature in the numerical experiments presented below.

Since $s$ is the score of $\rho$, an alternative objective to estimate it is

The derivation of (2.30) is standard: for the reader’s convenience we recall it at the end of Appendix B.1. The advantage of using (2.16) over (2.30) is that it does not require us to take the divergence of $\hat{s}$.

By definition, the score $s(t,x)=\nabla\log\rho(t,x)$ is a gradient field. As a result, if we model $\hat{s}(t,x)=-\nabla\hat{E}(t,x)$, we can turn (2.16) into an objective function for $\hat{E}(t,x)$

This objective is invariant to constant shifts in $\hat{E}$ and should therefore be minimized under some constraint, such as $\min_{x}\hat{E}(t,x)=0$ for all $t\in$. The minimizer of (2.31) provides us with an energy-based model (EBM) that can in principle be used to sample the PDF of the stochastic interpolant, $\rho(t,x)$, at any fixed $t\in$ using e.g. Langevin dynamics. We will not exploit this possibility here, and instead rely on generative models to sample $\rho(t,x)$, as discussed next in Sec. 2.3.

3 Generative models

Our next result is a direct consequence of Theorem 2.2, and it shows how to design generative models using the stochastic processes associated with the TE (2.9), the forward FPE (2.20), and the backward FPE (2.22):

[Generative models]corollarygenerative At any time $t\in$, the law of the stochastic interpolant $x_{t}$ coincides with the law of the three processes $X_{t}$, $X^{\mathsf{F}}_{t}$, and $X^{\mathsf{B}}_{t}$, respectively defined as:

The solutions of the probability flow associated with the transport equation (2.9)

solved either forward in time from the initial data $X_{t=0}\sim\rho_{0}$ or backward in time from the final data $X_{t=1}=x_{1}\sim\rho_{1}$.

The solutions of the forward SDE associated with the FPE (2.20)

solved forward in time from the initial data $X^{\mathsf{F}}_{t=0}\sim\rho_{0}$ independent of $W$.

The solutions of the backward SDE associated with the backward FPE (2.22)

solved backward in time from the final data $X^{\mathsf{B}}_{t=1}\sim\rho_{1}$ independent of $W^{\mathsf{B}}$; the solution of (2.34) is by definition $X^{\mathsf{B}}_{t}=Z^{\mathsf{F}}_{1-t}$ where $Z^{\mathsf{F}}_{t}$ satisfies

solved forward in time from the initial data $Z^{\mathsf{F}}_{t=0}\sim\rho_{1}$ independent of $W$.

To avoid repeated applications of the transformation $t\mapsto 1-t$, it is convenient to work with (2.34) directly using the reversed Itô calculus rules stated in the following lemma, which follows from the results in and is proven in Appendix B.2:

[Reverse Itô Calculus]lemmareversed If $X^{\mathsf{B}}_{t}$ solves the backward SDE (2.34):

We stress that the stochastic interpolant $x_{t}$, the solution $X_{t}$ to the ODE (2.32), and the solutions $X^{\mathsf{F}}_{t}$ and $X^{\mathsf{B}}_{t}$ of the forward and backward SDEs (2.33) and (2.34) are different stochastic processes, but their laws all coincide with $\rho(t)$ at any time $t\in$. This is all that matters when applying these processes as generative models. However, the fact that these processes are different has implications for the accuracy of the numerical integration used to sample from them at any $t$ as well as for the propagation of statistical errors (see also the next remark).

Generative models based on solutions $X_{t}$ to the ODE (2.32), solutions $X^{\mathsf{F}}_{t}$ to the forward SDE (2.33), and solutions $X^{\mathsf{B}}_{t}$ to the backward SDE (2.34) will typically involve drifts $b$, $b_{\mathsf{F}}$, and $b_{\mathsf{B}}$ that are, in practice, imperfectly estimated via minimization of (2.13) and (2.16) over finite datasets. It is important to estimate how this statistical estimation error propagates to errors in sample quality, and how the propagation of error depends on the generative model used, which is the object of our next section.

4 Likelihood control

In this section, we demonstrate that jointly minimizing the objective functions (2.29) and (2.16) (or the losses (2.13) and (2.16)) controls the $\mathsf{KL}$-divergence from the target density $\rho_{1}$ to the model density $\hat{\rho}_{1}$. We focus on bounds involving the score, but we note that analogous results hold for learning the denoiser $\eta_{z}(t,x)$ defined in (2.17) by the relation $\eta_{z}(t,x)=-s(t,x)/\gamma(t)$. The derivation is based on a simple and exact characterization of the $\mathsf{KL}$-divergence between two transport equations or two Fokker-Planck equations with different drifts. Remarkably, we find that the presence of a diffusive term determines whether or not it is sufficient to learn the drift to control $\mathsf{KL}$. This can be seen as a generalization of the result for score-based diffusion models described in to arbitrary generative models described by ODEs or SDEs. The proofs of the statements in this section are provided in Appendix B.3.

We first characterize the $\mathsf{KL}$ divergence between two densities transported by two different continuity equations but initialized from the same initial condition:

Then, the Kullback-Leibler divergence of $\rho(1)$ from $\hat{\rho}(1)$ is given by

where $\epsilon>0$. Then, the Kullback-Leibler divergence from $\rho(1)$ to $\hat{\rho}(1)$ is given by

Lemma 2.4 shows that, unlike for transport equations, the $\mathsf{KL}$-divergence between the solutions of two Fokker-Planck equations is controlled by the error in their drifts. The diffusive term in each Fokker-Planck equation provides an additional negative term in the $\mathsf{KL}$-divergence, which eliminates the need for explicit control on the Fisher divergence.

Putting the above results together, we can state the following result, which demonstrates that the losses (2.13) and (2.16) control the likelihood for learned approximations to the FPE (2.20).

where the function $\gamma$ satisfies the properties listed in Definition 2.1. Let $\hat{\rho}$ denote the solution to the Fokker-Planck equation

where $\mathcal{L}_{b}[\hat{b}]$ and $\mathcal{L}_{s}[\hat{s}]$ are the objective functions defined in (2.13) and (2.16), and

where $\mathcal{L}_{v}[\hat{v}]$ is the objective function defined in (2.29).

The above results have practical ramifications for generative modeling. In particular, they show that minimizing either the losses (2.13) and (2.16) or (2.29) and (2.16) maximize the likelihood of the stochastic generative model

but that minimizing the objective (2.13) is insufficient in general to maximize the likelihood of the deterministic generative model

Moreover, they show that, when learning $\hat{b}$ and $\hat{s}$, the choice of $\epsilon$ that minimizes the upper bound is given by

so that $\epsilon^{*}>1$ if the score is learned to higher accuracy than $\hat{b}$ and $\epsilon^{*}<1$ in the opposite situation. Note that (2.49) suggests to take $\epsilon=0$ if $\hat{b}$ is learned perfectly but $\hat{s}$ is not, and send $\epsilon\to\infty$ in the opposite situation. While taking $\epsilon=0$ is achievable in practice and leads to the ODE (2.32), taking $\epsilon\to\infty$ is not, as increasing $\epsilon$ increases the expense of the numerical integration in (2.33) and (2.34).

5 Density estimation and cross-entropy calculation

Then, given the PDFs $\rho_{0}$ and $\rho_{1}$:

The solution to (2.50) for the initial condition $\hat{\rho}(0)=\rho_{0}$ is given at any time $t\in$ by

The solution to (2.50) for the final condition $\hat{\rho}(1)=\rho_{1}$ is given at any time $t\in$ by

The proof of Lemma 2.5 can be found in Appendix B.4. Interestingly, we can obtain a similar result for the solution of the forward and backward FPEs in (2.20) and (2.22). These results make use of auxiliary forward and backward SDEs in which the roles of the forward and backward drifts are switched:

and let $Y^{\mathsf{F}}_{t}$ and $Y^{\mathsf{B}}_{t}$ denote solutions of the following forward and backward SDEs:

to be solved forward in time from the initial condition $Y^{\mathsf{F}}_{t=0}=x$ independent of $W$; and

to be solved backwards in time from the final condition $Y^{\mathsf{B}}_{t=1}=x$ independent of $W^{\mathsf{B}}$. Then, given the densities $\rho_{0}$ and $\rho_{1}$:

The proof of Theorem 2.5 can be found in Appendix B.4. Note that to generate data from either $\hat{\rho}_{\mathsf{F}}(1)$ or $\hat{\rho}_{\mathsf{B}}(0)$ assuming that we can sample exactly the PDF at the other end, i.e. $\rho_{0}$ and $\rho_{1}$ respectively, we would still rely on the equivalent of the forward and backward SDE in (2.33) and (2.34), now used with the approximate drifts in (2.54), i.e.

If we solve (2.61) forward in time from initial data $\hat{X}^{\mathsf{F}}_{t=0}\sim\rho_{0}$, we then have $\hat{X}^{\mathsf{F}}_{t=1}\sim\hat{\rho}_{\mathsf{F}}(1)$ where $\hat{\rho}_{\mathsf{F}}$ is the solution to the forward FPE (2.57). Similarly If we solve (2.62) backward in time from final data $\hat{X}^{\mathsf{B}}_{t=1}\sim\rho_{1}$, we then have $\hat{X}^{\mathsf{B}}_{t=0}\sim\hat{\rho}_{\mathsf{B}}(0)$ where $\hat{\rho}_{\mathsf{B}}$ is the solution to the backward FPE (2.59).

The results of Lemma 2.5 and Theorem 2.5 can be used to test the quality of samples generated by either the ODE (2.32) or the forward and backward SDEs (2.33) and (2.34). In particular, the following two results are direct consequences of Lemma 2.5 and Theorem 2.5, respectively:

corollarycrossentode Under the same conditions as Lemma 2.5, if $\hat{\rho}(0)=\rho_{0}$, the cross-entropy of $\hat{\rho}(1)$ relative to $\rho_{1}$ is given by

corollarycrossentsde Under the same conditions as Theorem 2.5, the cross-entropy of $\hat{\rho}_{\mathsf{F}}(1)$ relative to $\rho_{1}$ is given by

However, these bounds are not sharp in general – in fact, using calculations similar to the one presented in the proof of Theorem 2.5, we can derive exact expressions that capture precisely what is lost when applying Jensen’s inequality:

Unfortunately, since $\nabla\log\hat{\rho}_{\mathsf{F}}\not=\hat{s}$ and $\nabla\log\hat{\rho}_{\mathsf{B}}\not=\hat{s}$ in general due to approximation errors, we do not know how to estimate the extra terms on the right-hand side of (2.69) and (2.70). One possibility is to use $\hat{s}$ as a proxy for $\nabla\log\hat{\rho}_{\mathsf{F}}$ and $\nabla\log\hat{\rho}_{\mathsf{B}}$, which may be useful in practice, but this approximation is uncontrolled in general.

Instantiations and extensions

In this section, we instantiate the stochastic interpolant framework discussed in Section 2.

$I(t,x_{0},x_{1})$ is as in Definition 2.1;

$(x_{0},x_{1})\sim\nu$ with $\nu$ satisfying (2.3) in Definition 2.1;

$a(t)\in C^{2}()$ with $a(0)>0$ and $a(t)\geq 0$ for all $t\in(0,1]$, and;

$B_{t}$ is a standard Brownian bridge process, independent of $x_{0}$ and $x_{1}$.

As a result, (3.1) and (3.2) lead to the same generative models. Technically, it is easier to work with (3.2) than with (3.1), because it avoids the use of Itô calculus, and enables direct sampling of $x_{t}$ using samples from $\rho_{0}$, $\rho_{1}$, and $\mathsf{N}(0,\text{\it Id})$. However, (3.1) sheds light on some interesting properties of the generative models based on (3.2), i.e. stochastic interpolants with $\gamma(t)=\sqrt{2a(t)t(1-t)}$. To see why, we now re-derive the transport equation for the density $\rho(t,x)$ shared by (3.1) and (3.2) using the relation (3.1). For simplicity, we focus on the case where $a(t)$ is constant in time, i.e. we set $a(t)=a>0$ in (3.1).

To begin, recall that the Brownian Bridge $B_{t}$ can be expressed in terms of the Wiener process $W_{t}$ as $B_{t}=W_{t}-tW_{t=1}$. Moreover, it satisfies the SDE (obtained, for example, by conditioning on $B_{t=1}=0$ via Doob’s $h$-transform ):

A direct application of Itô’s formula implies that

Taking the expectation of this expression and using the independence between $(x_{0},x_{1})$ and $B_{t}$, we deduce that

Since for all fixed $t\in$ we have $B_{t}\stackrel{{\scriptstyle d}}{{=}}\sqrt{t(1-t)}z$ and $x^{\mathsf{d}}_{t}\stackrel{{\scriptstyle d}}{{=}}x_{t}$ with $x_{t}$ defined in (3.2), the time derivative (3.5) can also be written as

For the interpolant $x_{t}$ in (3.2), we have from the definitions of $b$ and $s$ in (2.10) and (2.14) that

As a result, $u-s=b$ and (3.7) can also be written as the TE (2.9) using $\Delta\rho=\nabla\cdot(s\rho)$.

Remarkably, the drift $u$ defined in (3.8) remains non-singular for all $t\in$ (including $t=0$) even if $\rho_{0}$ is replaced by a point mass at $x_{0}$; by contrast, both $b$ and $s$ are singular at $t=0$ in this case. Hence, the SDE associated with the FPE (3.7) provides us with a generative model that samples $\rho_{1}$ from a base measure concentrated at a single $x_{0}$ (i.e. such that the density $\rho_{0}$ is replaced by a point mass measure at $x=x_{0}$). We formalize this result in the following theorem:

theoremdiffgen Assume that $I(t,x_{0},x_{1})=x_{0}$ for $t\in[0,\delta]$ with some $\delta\in(0,1]$. Given any $a>0$, let

are such that $X^{\mathsf{d}}_{t=1}\sim\rho_{1}$.

Note that the additional assumption we make on $I(t,x_{0},x_{1})$ is consistent with the requirements in Definition 2.1 and Assumption 2.5: this additional assumption is made for simplicity and can probably be relaxed to $\partial_{t}I(t=0,x_{0},x_{1})=0$.

The proof of Theorem 3.1 is given in Appendix B.5. It relies on the calculations that led to (3.8), along with the observation that at $t=0$ and $x=x_{0}$,

In principle, the approach above can be generalized to any stochastic bridge $B_{t}^{x_{0},x_{1}}$, which can be obtained from any SDE by conditioning its solution to satisfy $B_{t=0}^{x_{0},x_{1}}=x_{0}$ and $B_{t=1}^{x_{0},x_{1}}=x_{1}$ with the help of Doob’s $h$-transform. In general, however, this construction cannot be made explicit, because the $h$-transform is typically not available analytically. One approach would be to learn it, as proposed e.g. in , but this adds an additional layer of difficulty that is avoided by the approach above.

A common choice of base density for generative modeling in the absence of prior information is to choose $\rho_{0}=\mathsf{N}(0,\text{\it Id})$. In this setting, we can group the effect of the latent variable $z$ with $x_{0}$. This leads to a simpler type of stochastic interpolant that, in particular, will enables us to instantiate score-based diffusion within our general framework (see Section 3.4).

$x_{1}$ and $z$ are independent random variables drawn from $\rho_{1}$ and ${\sf N}(0,\text{\it Id})$, respectively.

By construction, $x^{\mathsf{os}}_{t=0}=z\sim{\sf N}(0,\text{\it Id})$ and $x^{\mathsf{os}}_{t=1}=x_{1}\sim\rho_{1}$, so that the distribution of the stochastic process $x^{\mathsf{os}}_{t}$ bridges ${\sf N}(0,\text{\it Id})$ and $\rho_{1}$. It is easy to see that the one-sided stochastic interpolant defined in (3.15) will have the same density as the stochastic interpolant defined in (2.1) if we set $I(t,x_{0},x_{1})=J_{t}(x_{1})+\delta(t)x_{0}$ and take $\delta^{2}(t)+\gamma^{2}(t)=\alpha^{2}(t)$. Restricting to this case, our earlier theoretical results apply where the velocity field $b$ defined in (2.10) becomes

and the quadratic objective in (2.13) becomes

where $\eta_{z}(t,x)$ is the equivalent of the denoiser defined in (2.17). These functions are the unique minimizers of the objectives

Moreover, we can weaken Assumption 2.5 to the following requirement:

where the expectation is taken over $x_{1}\sim\rho_{1}$.

The construction above can easily be generalized to the case where $\rho_{0}=\mathsf{N}(0,C_{0})$ with $C_{0}$ a positive-definite matrix. Without loss of generality, we can then assume that $C_{0}$ can be represented as $C_{0}=\sigma_{0}\sigma_{0}^{\mathsf{T}}$ where $\sigma_{0}$ is a lower-triangular matrix and replace (3.15)

with $J$ and $\alpha$ satisfying the conditions listed in Definition 3.3 and where $z\sim{\sf N}(0,\text{\it Id})$.

3 Mirror interpolants

Another practically relevant setting is when the base and the target are the same density $\rho_{1}$. In this setting we can define a stochastic interpolant as:

$x_{1}$ and $z$ are random variables drawn independently from $\rho_{1}$ and ${\sf N}(0,\text{\it Id})$, respectively.

By construction, $x^{\mathsf{mir}}_{t=0}=x^{\mathsf{mir}}_{t=1}=x_{1}\sim\rho_{1}$, so that the distribution of the stochastic process $x^{\mathsf{mir}}_{t}$ bridges $\rho_{1}$ to itself. Note that a valid choice is $K(t,x_{1})=\alpha(t)x_{1}$ with $\alpha(0)=\alpha(1)=1$ (e.g. $\alpha(t)=1$): in this case, mirror interpolants are related to denoisers, as will be discussed in Section 5.2.

It is easy to see that our earlier theoretical results apply where the velocity field $b$ defined in (2.10) becomes

and the quadratic objective in (2.13) becomes

which are the unique minimizers of the objective functions

Moreover, we can weaken Assumption 2.5 to the following requirement:

where the expectation is taken over $x_{1}\sim\rho_{1}$.

Interestingly, if we take $K(t,x_{1})=x_{1}$, then $\partial_{t}K(t,x_{1})=0$, and the velocity field defined in (3.26) is completely defined by the denoiser $\eta_{z}$

Since the score $s$ also depends on $\eta_{z}$, this denoiser is the only quantity that needs to be learned.

If $\rho_{1}$ is only accessible via empirical samples, mirror interpolants do not enable calculation of the functional form of $\rho_{1}$. A notable exception is if we set $K(t,x_{1})=0$ for $t\in[t_{1},t_{2}]$ with $0<t_{1}\leq t_{2}<1$: in that case, $x^{\mathsf{mir}}_{t}=\gamma(t)z\sim\gamma(t){\sf N}(0,\text{\it Id})$ for $t\in[t_{1},t_{2}]$, which gives us a reference density for comparison. In this setup, mirror interpolants essentially reduce to two one-sided interpolants glued together (with the second one time-reversed), or in fact a regular stochastic interpolant when $\rho_{0}=\rho_{1}$ and we set $I(t,x_{0},x_{1})=0$ for $t\in[t_{1},t_{2}]$.

4 Stochastic interpolants and Schrödinger bridges

The stochastic interpolant framework can also be used to solve the Schrödinger bridge problem. For background material on this problem, we refer the reader to and the references therein. Consistent with the overall viewpoint of this paper, we consider the hydrodynamic formulation of the Schrödinger bridge problem, in which the goal is to obtain a pair $(\rho,u)$, that solves the following optimization problem for a fixed $\epsilon>0$

Under our assumptions on $\rho_{0}$ and $\rho_{1}$ listed in Assumption 2.5, it is known (see e.g. Proposition 4.1 in ) that (3.35) has a unique minimizer $(\rho,u=\nabla\lambda)$, with $(\rho,\lambda)$ classical solutions of the Euler-Lagrange equations:

To proceed we will make the additional assumption that the solution $\rho$ to (3.36) can be reversibly mapped to a standard Gaussian:

We stress that the actual form of the map $T$ is not important for the arguments below. Assumption 3.10 can be used to show the existence of a stochastic interpolant whose density solves (3.36): {restatable}lemmainterppdf If Assumption (3.10) holds, then the solution $\rho(t)$ to (3.36) is the density of the stochastic interpolant

as long as $\alpha^{2}(t)+\beta^{2}(t)+\gamma^{2}(t)=1$.

The proof is given in Appendix B.6: (3.38) corresponds to choosing $I(t,x_{0},x_{1})=T(t,\alpha(t)T^{-1}(t,x_{0})+\beta(t)T^{-1}(t,x_{1}))$ in (2.1). With the help of Lemma 3.10, we can establish the following result, which shows how to optimize over the function $I$ to solve the problem (3.35)

If Assumption 3.10 holds, then all the optimizers $(I,u)$ of (3.39) are such that the density of the associated $x_{t}=I(t,x_{0},x_{1})+\gamma(t)z$ is the solution $\rho$ to (3.36). Moreover, $u=\nabla\lambda$, with $\lambda$ the solution to (3.36).

The proof is also given in Appendix B.6. Note that if we fix $\hat{I}$, the velocity $u$ minimizing this objective is the forward drift $b_{\mathsf{F}}$ defined in (2.21). Note also that if we set $\epsilon\to 0$, the minimizing velocity field is $b$ as defined in (2.10), and the max-min problem formally reduces to solving the optimal transport problem. In this case, Assumption 3.10 becomes more stringent, as we need to assume that that system (3.36) with $\epsilon=0$ (i.e. in the absence of the diffusive terms) has a classical solution. Theorem (3.10) gives a practical route towards solving the Schrödinger bridge problem with stochastic interpolants, and we leave the numerical investigation of this formulation to future work.

Spatially linear interpolants

In this section, we study the stochastic interpolants that are obtained when we specialize the function $I$ used in (2.1) to be linear in both $x_{0}$ and $x_{1}$, i.e. we consider

where $(x_{0},x_{1})\sim\nu$ and $z\sim{\sf N}(0,\text{\it Id})$ with $(x_{0},x_{1})\perp z$, and $\alpha,\beta,\gamma^{2}\in C^{2}()$ satisfy the conditions

Despite its simplicity, this setup offers significant design flexibility. The discussion highlights how the presence of the latent variable $\gamma(t)z$ can simplify the structure of the intermediate density $\rho(t)$. Since our ultimate aim is to investigate the properties of practical generative models built upon ODEs or SDEs, we will also study the effect of time-dependent diffusion coefficient $\epsilon(t)$, which controls the amplitude of the noise in a generative SDE. Throughout, to build intuition, we choose $\rho_{0}$ and $\rho_{1}$ to be Gaussian mixture densities, for which the drift coefficients can be computed analytically (see Appendix A). This enables us to visualize the effect of each choice on the resulting generative models.

When the stochastic interpolant is of the form (4.1), the velocity $b$ and the score $s$ defined in (2.10) and (2.14) can both be expressed in terms of the following three conditional expectations (the third is the denoiser already defined in (2.17)):

This enables us to reduce computational expense: given two of the $\eta$’s, the third can always be calculated via (4.5). Finally, it is easy to see that the functions $\eta_{0}$, $\eta_{1}$, and $\eta_{z}$ are the unique minimizers of the objectives

where $x^{\mathsf{lin}}_{t}$ is defined in (4.1) and the expectation is taken independently over $(x_{0},x_{1})\sim\nu$ and $z\sim{\sf N}(0,\text{\it Id}).$

2 Some specific design choices

It is useful to assume that both $\rho_{0}$ and $\rho_{1}$ have been scaled to have zero mean and identity covariance (which can be achieved in practice, for example, by an affine transformation of the data). In this case, the time-dependent mean and covariance of (4.1) are given by

Preserving the identity covariance at all times therefore leads to the constraint

This choice is also sensible if $\rho_{0}$ and $\rho_{1}$ have covariances that are not the identity but are on a similar scale. In this case we no longer need to enforce (4.8) exactly, and could, for example, take three functions whose sum of squares is of order one. For definiteness, in the sequel we discuss possible choices that satisfy (4.8) exactly, with the understanding that the corresponding functions $\alpha$, $\beta$, and $\gamma$ could all be slightly modified without significantly affecting the conclusions.

One way to ensure that (4.8) holds while maintaining the influence of $\rho_{0}$ and $\rho_{1}$ everywhere on $$ except at the endpoints is to choose

This choice was advocated in , without the inclusion of the latent variable ($\gamma=0$). Another possibility that gives more leeway is to pick any $\gamma:\to$ and set

With $\gamma=0$, this was the choice preferred in . The PDF $\rho(t)$ obtained with the choices (4.9) and (4.10) when $\rho_{0}$ and $\rho_{1}$ are both Gaussian mixture densities are shown in Figure 6. As this example shows, when $\rho_{0}$ and $\rho_{1}$ have distinct complex features, these would be duplicated in $\rho(t)$ at intermediate times if not for the smoothing effect of the latent variable; this behavior is seen in Figure 6, where it is most prominent in the first row with $\gamma(t)=0$. From a statistical learning perspective, eliminating the formation of spurious features will simplify the estimation of the velocity field $b$, which becomes smoother as the formation of such features is suppressed.

Gaussian encoding-decoding.

A useful limiting case is to devolve the data from $\rho_{0}$ completely into noise by the halfway point $t=\tfrac{1}{2}$ and to reconstruct $\rho_{1}$ completely from noise starting from $t=\frac{1}{2}$. One choice that allows us to do so while satisfying (4.8) is

where $1_{A}(t)$ is the indicator function of $A$, i.e. $1_{A}(t)=1$ if $t\in A$ and $1_{A}(t)=0$ otherwise. With this choice, it is easy to see that $x_{t=\frac{1}{2}}=\gamma(\tfrac{1}{2})z\sim{\sf N}(0,\gamma^{2}(\tfrac{1}{2}))$, which seamlessly glues together two interpolants: one between $\rho_{0}$ and a standard Gaussian, and one between a standard Gaussian and $\rho_{1}$.

Even though the choice (4.11) encodes $\rho_{0}$ into pure noise on the interval $[0,\tfrac{1}{2}]$, which is then decoded into $\rho_{1}$ on the interval $[\tfrac{1}{2},1]$ (and vice-versa when proceeding backwards in time), the resulting velocity $b$ still defines a single continuity equation that maps $\rho_{0}$ to $\rho_{1}$ on $$. This is most clearly seen at the level of the probability flow (2.32), since its solution$X_{t}$is a bijection between the initial and final conditions$X_{t=0}$and$X_{t=1}$, but a similar pairing can also be observed in the solutions to the forward and backward SDEs (2.33) and (2.34), whose solutions at time$t=1$or$t=0$remain correlated with the initial or final condition used. This allows for a more direct means of image-to-image translation with diffusions when compared to the recent approach described in . The choice (4.11) is depicted in the final row of Figure 6, where no spurious modes form at all; individual sample trajectories of the deterministic and stochastic generative models based on ODEs and SDEs whose solutions have this$\rho(t)$as density can be seen in the panels forming the third column in Figure 7. We note that the elimination of spurious intermediate modes can also be implemented by use of a data-adapted coupling$\nu(dx_{0},dx_{1})$, as considered in .

Unsurprisingly, it is necessary to have $\gamma(t)>0$ for the choice (4.11): for $\gamma(t)=0$, the density $\rho(t)$ collapses to a Dirac measure at $t=\frac{1}{2}$. This consideration highlights that the inclusion of the latent variable $\gamma(t)z$ matters even for the deterministic dynamics (2.32), and its presence is distinct from the stochasticity inherent to the SDEs (2.33) and (2.34).

3 Impact of the latent variable γ​(t)​z𝛾𝑡𝑧\gamma(t)z and the diffusion coefficient ϵ​(t)italic-ϵ𝑡\epsilon(t)

The stochastic interpolant framework enables us to discern the independent roles of the latent variable $\gamma(t)z$ and the diffusion coefficient $\epsilon(t)$ we use in a generative model. As shown in Theorem 2.2, the presence of the latent variable $\gamma(t)z$ for $\gamma\neq 0$ smooths both the density $\rho(t)$ and the velocity $b$ defined in (2.10) spatially. This provides a computational advantage at sample generation time because it simplifies the required numerical integration of (2.32), (2.33), and (2.34). Intuitively, this is because the density $\rho(t)$ of $x_{t}$ can be represented exactly as the density that would be obtained with $\gamma(t)=0$ convolved with $\mathsf{N}(0,\gamma^{2}(t)Id)$ at each $t\in(0,1)$. A comparison between the density $\rho(t)$ obtained with trigonometric interpolants with $\gamma(t)=0$ and $\gamma(t)=\sqrt{2t(1-t)}$ can be seen in the first and second row of Figure 6.

By contrast, the diffusion coefficient $\epsilon(t)$ leaves the density $\rho(t)$ unchanged, and only affects the way we sample it. In particular, the probability flow ODE (2.32) results in a map that pushes every $X_{t=0}=x_{0}$ onto a single $X_{t=1}=x_{1}$ and vice-versa. The forward SDE (2.33) maps each $X^{\mathsf{F}}_{t=0}=x_{0}$ onto an ensemble $X^{\mathsf{F}}_{t=1}$ whose spread is controlled by the amplitude of $\epsilon(t)$ (and similarly for the reversed SDE (2.34) that maps each $X^{\mathsf{B}}_{t=1}=x_{1}$ onto an ensemble $X^{\mathsf{B}}_{t=0}$). This ensemble is not distributed according to $\rho_{1}$ for finite $\epsilon(t)$ – like with the ODE, we need to sample initial conditions from $\rho_{0}$ to get solutions at time $t=1$ that sample $\rho_{1}$ – but its density converges towards $\rho_{1}$ as $\epsilon(t)\to\infty$. These features are illustrated in Figure 7.

Another potential advantage of including the latent variable $\gamma(t)z$ is its impact on the velocity $b$ at the end points. Since $x_{t=0}=x_{0}$ and $x_{t=1}=x_{1}$, it is easy to see that the velocity $b$ of the linear interpolant $x_{t}$ defined in (4.1) satisfies

where $s_{0}=\nabla\log\rho_{0}$ and $s_{1}=\nabla\log\rho_{1}$. If $\gamma\in C^{2}()$, because $\gamma(0)=\gamma(1)=0$, the terms involving the scores $s_{0}$ and $s_{1}$ in these expressions vanish. Choosing $\gamma^{2}\in C^{1}()$ but $\gamma$ not differentiable at $t=0$ or $t=1$ leaves open the possibility that the limits remain nonzero. For example, if we take one of the choices discussed in Section 3.1, i.e.

As a result, the choice (4.13) ensures that the velocity $b$ encodes information about the score of the densities $\rho_{0}$ and $\rho_{1}$ at the end points. We stress however that, while the choice of $\gamma(t)$ given in (4.13) is appealing because of its nontrivial influence on the velocity $b$ at the endpoints, the user is free to explore a variety of alternatives. We present some examples in Table 8, specifying the differentiability of $\gamma$ at $t=0$ and $t=1$. The function $\gamma(t)$ specified in (4.13) is the only featured case for which the contribution from the score is non-vanishing in the velocity $b$ at the endpoints. In Section 7, we illustrate on numerical examples that there are tradeoffs between different choices of $\gamma$, which might be directly related to this fact. When using the ODE as a generative model, the score is only felt through $b$, whereas it is explicit when using the SDE as a generative model.

4 Spatially linear one-sided interpolants

Much of the discussion above generalizes to one-sided interpolants if we take the function $J(t,x_{1})$ in (3.15) to be linear in $x_{1}$ and define

where $\alpha^{2},\beta\in C^{2}()$ and $\alpha(0)=\beta(1)=1$, $\alpha(1)=\beta(0)=0$, and $\alpha(t)>0$ for all $t\in[0,1)$. The velocity $b$ and the score $s$ defined in (2.10) and (2.14) can now be expressed as

where the second expression holds for all $t\in[0,1)$ and we defined:

Note that, by definition of the conditional expectation, $\eta^{\mathsf{os}}_{z}$ and $\eta^{\mathsf{os}}_{1}$ satisfy

As a result, only one of them needs to be estimated. For example, we can express $\eta^{\mathsf{os}}_{1}$ as a function of $\eta^{\mathsf{os}}_{z}$ for all $t$ such that $\beta(t)\not=0$, and use the result to express the velocity (4.16) as

Assuming that $\beta(t)\not=0$ for all $t\in(0,1]$, this formula only needs to be supplemented at $t=0$ with

which follows from (4.16) since $x^{\mathsf{os,lin}}_{t=0}=z$. Later in Section 5.2 we will show that using the velocity $b$ in (4.19) to solve the probability flow ODE (2.32) can be seen as using a denoiser to construct a generative model.

Finally note that $\eta_{z}$ and/or $\eta_{1}$ can be estimated using the following two objective functions, respectively:

Connections with other methods

In this section, we discuss connections between the stochastic interpolant framework and the score-based diffusion method , the stochastic localization framework ), denoising methods , and the rectified flow method introduced in .

Score-based diffusion models (SBDM) are based on variants of the Ornstein-Uhlenbeck process

which has the property that the marginal density of its solution at time $\tau$ converges to a standard normal as $\tau$ tends towards infinity. By learning the score of the density of $Z_{\tau}$, we can write the associated backward SDE for (5.1), which can then be used as a generative model – this backwards SDE is also the one that is used in the stochastic localization process, see .

To see the connection with stochastic interpolants, notice that the solution of (5.1) from the initial condition $Z_{\tau=0}=x_{1}\sim\rho_{1}$ can be written exactly as

As a result, the law of $Z_{\tau}$ conditioned on $Z_{\tau=0}=x_{1}$ is given by

for any time $\tau\in[0,\infty)$. This is also the law of the process

If we let $x_{1}\sim\rho_{1}$ with $x_{1}\perp z$, the process $y_{\tau}$ is similar to a one-sided stochastic interpolant, except the density of $y_{\tau}$ only converges to ${\sf N}(0,\text{\it Id})$ as $\tau\to\infty$; by contrast, the one-sided interpolants we introduced in Section 3.2 converge on the finite interval $$. In SBDM, this is handled by capping the evolution of$Z_{\tau}$to a finite time interval$[0,T]$with$T<\infty$, and then by using the backward SDE associated with (5.1) restricted to$[0,T]$. However, this introduces a bias that is not present with one-sided stochastic interpolants, because the final condition used for the backwards SDE in SBDM is drawn from${\sf N}(0,\text{\it Id})$even though the density of the process (5.1) is not Gaussian at time$T$.

We can, however, turn (5.4) into a one-sided linear stochastic interpolant by defining $t=e^{-\tau}$ and by choosing $\alpha(t)$ and $\beta(t)$ in (4.15) to have a specific form. More precisely, evaluating (5.4) at $\tau=-\log t$,

With this choice of $\alpha(t)$ and $\beta(t)$, from (4.16) we get the velocity field

where $\eta_{z}^{\mathsf{os}}$ and $\eta_{1}^{\mathsf{os}}$ are defined in (4.17). This expression shows that the velocity $b$ used in the probability flow ODE (2.32) is well-behaved at all times, including at $t=1$ where $\dot{\alpha}(t)$ is singular. The same is true for the drift $b_{\mathsf{F}}(t,x)=b(t,x)+\epsilon(t)s(t,x)$ used in the forward SDE (2.33), regardless of the choice of $\epsilon\in C^{0}()$ with $\epsilon(t)\geq 0$. This shows that casting SBDM into a one-sided linear stochastic interpolant (4.15) allows the construction of unbiased generative models that operate on $t\in$. This comes at no extra computational cost, since only one of the two functions defined in (4.17) needs to be estimated, which is akin to estimating the score in SBDM.