Stochastic Sampling from Deterministic Flow Models

## Abstract

Abstract Deterministic flow models, such as rectified flows, offer a general framework for learning a deterministic transport map between two distributions, realized as the vector field for an ordinary differential equation (ODE).
However, they are sensitive to model estimation and discretization errors and do not permit different samples conditioned on an intermediate state, limiting their application.
We present a general method to turn the underlying ODE of such flow models into a family of stochastic differential equations (SDEs) that have the same marginal distributions.
This method permits us to derive families of stochastic samplers , for fixed (e.g., previously trained) deterministic flow models, that continuously span the spectrum of deterministic and stochastic sampling, given access to the flow field and the score function.
Our method provides additional degrees of freedom that help alleviate the issues with the deterministic samplers and empirically outperforms them.
We empirically demonstrate advantages of our method on a toy Gaussian setup and on the large scale ImageNet generation task.
Further, our family of stochastic samplers provide an additional knob for controlling the diversity of generation, which we qualitatively demonstrate in our experiments.

## 1 Introduction

Deterministic flow models, including Rectified Flow (Liu et al., ), Flow Matching (Lipman et al., ; Tong et al., ), Stochastic Interpolants (Albergo & Vanden-Eijnden, ; Albergo et al., ), and probability flow ODE (Song et al., ) learn a reversible deterministic transport between two end distributions $p_{0}(x_{0})$ and $p_{1}(x_{1})$.
Diffusion models require one of the distributions to be a Gaussian distribution, though generalizations exist (Yoon et al., ).
In contrast, Rectified Flows, Stochastic Interpolants, and Flow Matching offer a general framework for learning deterministic transports, without this restriction.
While deterministic transport enables efficient deterministic sampling, e.g. by the rectification procedure suggested by Liu et al. (), stochastic sampling may be desirable for:
(1) robustness to estimation errors in the learned flow model,
(2) ability to produce random samples conditioned on an intermediate state $x_{t},t\in[0,1]$,
and (3) robustness to discretization error resulting from discrete step sampling from a continuous time model.
We present a new theorem (Section 3.2) that provides a recipe to create an infinite family of parameterized stochastic samplers, given access to the flow field and the score function for the marginal distributions.
Our result provides a general and unified view, while including a few existing proposals (e.g. in Huang et al. (); Berner et al. ()) as special cases.

The deterministic transport specifies a deterministic mapping between the samples from the two distributions and is realized as a learned vector field corresponding to an ordinary differential equation (ODE).
However, if one distribution is chosen to be a Gaussian, these Flow models can be viewed as reparameterizations of other deterministic models that also choose a Gaussian as one of the distributions e.g. probability flow ODEs arising from Gaussian diffusion models.
We refer to such models as *Gaussian flow* models.
Transport map learning algorithms such as Gaussian flows are practical to train and enable applications like generative modeling (Ramesh et al., ; Lu et al., ; Saharia et al., ; Esser et al., ), stylization (Isola et al., ; Meng et al., ), and image restoration (Delbracio & Milanfar, ; Rombach et al., ; Lugmayr et al., ; Kawar et al., ), to name a few.
However, corresponding deterministic sampler has limitations that we empirically demonstrate on a toy Gaussian task, where it exhibits a bias and consistently underestimates the variance of the target distribution, as seen in Figure 2.
To enable stochastic sampling from such deterministic models, we provide a special case of our general result to turn the underlying ODE of Gaussian flow models into a family of stochastic differential equations (SDEs) that have the same marginal distributions.
Our stochastic samplers allow trading the bias of deterministic sampler for increased variance in the estimated mean and variance parameters (Figure 4).
Since, our method requires access to the score function for the marginal distributions, we impute it directly from the given flow model, alleviating the need for learning it separately.
This method permits us to derive families of *stochastic samplers*, for fixed (e.g., previously trained) *deterministic* Gaussian flow models, that allow flexible and time dependent injection of stochasticity during sampling, enabling both deterministic and stochastic sampling.
This additional degree of freedom allows exploration of stochastic samplers that can help alleviate the issues with the deterministic samplers and outperform them.
We demonstrate this empirically on a toy Gaussian setup, as well as on the large scale ImageNet generation task.
The stochastic samplers also provide an additional knob for controlling the diversity of generation as we qualitatively demonstrate in our experiments, and are compatible with classifier-free guidance (Ho & Salimans, ), as can be seen in Figures 1 and 12.

Our key contributions are:
(1) Specification of a flexible family of SDEs (Section 3.2) that have the same marginal distributions as a given SDE or a flow model, enabling exploration of sampling schemes for a given fixed model,
(2) Derivation of new as well as existing special cases directly from Section 3.2 (Section 3.2 and Section 3.2) demonstrating generality of Section 3.2,
(3) Study of a set of SDE families corresponding to Gaussian flow models, derived using Section 3.2, on both a toy as well as a large scale ImageNet setup, demonstrating flexible stochastic sampling and controllable diversity in generation, *without requiring retraining* (Table 1, Figures 1 and 12).

## 2 Background

Figure: Figure 1: Stochastic sampling improves diversity at all classifier-free guidance levels. We visualize samples from a rectified flow model at four classifier-free guidance levels $\lambda$ (Section 3.3) and at four stochasticity scales $\alpha$ for NonSingular (Table 1). Three samples are shown for each configuration where the sampling starts at the same draw from $p_{1}(x_{1})$. When $\alpha=0$, the sampler is deterministic and samples are the same (therefore we show only one). When $\lambda=0$, there is no classifier-free guidance. Note the increased diversity as $\alpha$ increases. More examples in Figure 12.
Refer to caption: x1.png

#### Notation.

Throughout this work we use small Latin letters $t,x,y$ etc. to represent scalar and vector variables, $f,g$ etc. to represent functions, Greek letters $\alpha,\beta$ etc. to represent (hyper-)parameters, and capital letters $G$ to represent matrices.
With a slight abuse of notation we use lower case letters $x$ to represent both the random variable and a particular value $x\sim p(x)$.
Whenever unambiguous, we suppress the dependence of state $x_{t}$ on time $t$ as $x\equiv x_{t}$, and dependence of functions on state $x_{t}$ and time $t$ as $f\equiv f(x_{t},t)$ to simplify notation.

We briefly discuss rectified flow and continuous time diffusion models.
Refer to Liu et al. (); Song et al. () for details.

### 2.1 Rectified Flow

Let $x_{0}\sim p_{0}(x_{0})\in\mathbb{R}^{d}$ be the draws from the data distribution $p_{0}$ that we are interested in learning and sampling from.
Let $x_{1}\sim p_{1}(x_{1})\in\mathbb{R}^{d}$ be an easy to sample source distribution.
Loosely, the key idea behind the diffusion and flow family of models is to learn a mapping that either stochastically or deterministically transforms a sample from $p_{1}$, in an iterative manner, to produce a sample from $p_{0}$.
Let $\nu(x_{0},x_{1})$ be an arbitrary coupling distribution for the two random variables $x_{0},x_{1}$ such that $p_{0}(x_{0})=\int\nu(x_{0},x_{1})dx_{1},p_{1}(x_{1})=\int\nu(x_{0},x_{1})dx_{0}$.
A simple choice is the product of the two: $\nu(x_{0},x_{1})\equiv p_{0}(x_{0})p_{1}(x_{1})$.
To construct a rectified flow first an interpolation between the two variables is defined as $x_{t}\equiv h(x_{0},x_{1},t)$ that is differentiable w.r.t. time.
The default interpolation proposed and studied in Liu et al. () is:

$$ $x_{t}=(1-t)x_{0}+tx_{1},\quad t\in[0,1].$ (1) $$

With the above, rectified flow learns a vector field $v(x_{t},t)$ by minimizing the following objective:

$$ $v(x,t)=\operatorname*{arg\,min}_{v^{\prime}}\mathbb{E}_{(x_{0},x_{1})\sim\nu} \left[\int_{0}^{1}\left\lVert\frac{dx_{t}}{dt}-v^{\prime}(x_{t},t)\right\rVert ^{2}dt\right].$ (2) $$

The solution to the above optimization problem is $v(x,t)\equiv\mathbb{E}[x_{1}-x_{0}|x,t]$ and is referred to as 1-Rectified flow.
Since $v(x,t)$ is not available in closed-form in general, $v$ is typically parameterized with parameters $\theta$ and optimization in Equation 2 is performed w.r.t. $\theta$.
In the rest of the paper, we drop this dependence on the parameters in notation as we assume a model $v(x,t)$ to be given.
Note that a closed-form expression is available when $p_{0},p_{1}$ are Gaussian (see Appendix F).
We use this expression for the toy setup in our experiments. For example, the biased deterministic sampler in Figure 2 is using the ground truth flow field.
Once the flow $v(x_{t},t)$ is estimated, samples from $p_{0}(x_{0})$ can be produced by drawing a sample from $p_{1}(x_{1})$ and simulating the flow backward in time, using:

$$ $\displaystyle dx=v(x,t)dt$ (3) $$

### 2.2 Score based diffusion with stochastic differential equations

The general idea in this family of methods is to specify a forward stochastic process that slowly transforms the data density $p_{0}(x_{0})$ into an easy to sample source density $p_{1}(x_{1})$.
Song et al. () specified such a process using an Itô SDE of the following form:

$$ $\displaystyle dx=f(x,t)dt+G(x,t)dW_{t}$ (4) $$

where $f(x,t):\mathbb{R}^{d}\times[0,1]\rightarrow\mathbb{R}^{d}$ is the drift coefficient, $G(x,t):\mathbb{R}^{d}\times[0,1]\rightarrow\mathbb{R}^{d}\times\mathbb{R}^{d}$ is state and time dependent diffusion coefficient and $W_{t}$ is the Wiener process.
Choosing(^1^11Song et al. () provide general results for $G(x,t)$ which we omit here for brevity.) $G\equiv g(t):[0,1]\rightarrow\mathbb{R}$ and using results from Anderson (), a reverse time SDE can be specified that has the same marginals as Equation 4:

$$ $\displaystyle dx=[f(x,t)-g^{2}(t)\nabla_{x}\ln p_{t}(x)]dt+g(t)d\tilde{W}_{t}$ (5) $$

where $\tilde{W}_{t}$ is a standard Wiener process with time running backwards.
Note that the time reversal requires access to the score function $\nabla_{x}\ln p_{t}(x)$.
Score matching (Vincent, ) can be used to learn an estimate for the score for all $t$ (Song et al., ), which can then be used to simulate reverse time dynamics starting with a sample from $p_{1}(x_{1})$ to produce a sample from $p_{0}(x_{0})$ at $t=0$.
A forward deterministic process can also be derived from the above that has the same marginal densities $p_{t}(x)$:

$$ $\displaystyle dx=\left[f(x,t)-\frac{1}{2}g^{2}(t)\nabla_{x}\ln p_{t}(x)\right]dt$ (6) $$

The above ODE is also referred to as the probability flow ODE.
Samples can be generated using the above ODE in a similar fashion as rectified flow, by simulating the ODE backwards in time.

## 3 Deriving stochastic samplers

Figure: Figure 2: Discretization of deterministic flow leads to bias. Comparison of samplers from Table 1 on the two Gaussian toy problem (Appendix G). Deterministic underestimates the variance parameter, but the stochastic samplers avoid that issue, in exchange for variance in the parameter estimation. Singular’s variance diverges if we start from $t=1$, so instead we start the sampler at $t=1-10^{-3}$, which allows it to eventually converge by $t=0$.
Refer to caption: x2.png

#### Method intuition.

Probability flow ODEs (Song et al., ), proposed in the context of diffusion models, provide a deterministic sampling method for diffusion models.
These ODEs have the same marginal distribution $p_{t}(x)$ at all $t$ as the original SDE from which they are derived.
Here, we take the reverse path: we start from an ODE (corresponding to the Gaussian flow model) and deduce the family of SDEs that have the same marginal distributions at all $t$ as the original ODE.
Before we introduce the general result, we will show a naive approach that gives an SDE with a problematic singularity, motivating the need for the generalization.

### 3.1 A singular SDE corresponding to Gaussian flow

For Gaussian flow, $p_{1}(x_{1})\equiv N(x_{1};\mu_{1},\sigma_{1}^{2}I)$ is assumed to be Gaussian.
With interpolation $x_{t}=(1-t)x_{0}+tx_{1}$, the perturbation kernel $p(x_{t}|x_{0})=N(x_{t};(1-t)x_{0}+t\mu_{1},t^{2}\sigma_{1}^{2}I)$ is also Gaussian.

Note that since $x_{0},x_{1}$ are independent, we can directly infer the first and second moments $\mu_{t},\Sigma_{t}$ for the marginals $p_{t}(x)$ as $\mu_{t}=(1-t)\mu_{0}+t\mu_{1}$ and $\Sigma_{t}=(1-t)^{2}\Sigma_{0}+t^{2}\sigma_{1}^{2}I$.
With these constraints and choosing $\mu_{1}\equiv 0,\sigma_{1}\equiv 1$, we can solve for drift and diffusion coefficients that yield the same marginal distributions:

$$ $\displaystyle f(x,t)$ $\displaystyle=-\frac{x}{1-t}$ $\displaystyle g(t)$ $\displaystyle=\sqrt{\frac{2t}{1-t}}$ (7) $$

See Appendix A for the details and a more general expression for arbitrary $\mu_{1},\sigma_{1}$.
The coefficients $f(x,t),g(t)$ are singular at the boundary $t=1$ of the interval.
Consequently, simulation methods such as Euler-Maruyama, that need $f(x,t),g(t)$ to be Lipschitz are not guaranteed to work at the boundary (see Figure 2 and Section 4.1).
We refer to this SDE as the Singular SDE.
An empirical trick that is often used in such cases is to assume $p_{1-\epsilon}(x_{1-\epsilon})\approx p_{1}(x_{1}),\epsilon\ll 1$.
However, this can lead to unpredictable behavior and we show how to avoid it in the following section.

### 3.2 Set of SDEs that share the same marginal distribution p t ⁢ ( x ) subscript 𝑝 𝑡 𝑥 p_{t}(x) italic_p start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ( italic_x )

First we state our general result with the diffusion coefficient $G(x,t)$ a function of both the state $x$ and time $t$, and then state simpler forms more directly applicable to models used in practice.

theoremmainresult
Let $p_{t}(x)$ be the probability density of the solutions of the SDE in Equation 4
evolving as $\frac{\partial p_{t}}{\partial t}$.
Then, the probability density of solutions of the following set of SDEs, indexed by $\tilde{G},\gamma_{t}$, also evolves as $\frac{\partial p_{t}}{\partial t}$.

$$ $\displaystyle dx=\bar{f}(x,t)dt+\bar{G}(x,t)dW_{t}$ (8) $$

where

$$ $\displaystyle\bar{f}$ $\displaystyle=f-\frac{1}{2}\left(\nabla\cdot[(1-\gamma_{t})GG^{T}-\tilde{G} \tilde{G}^{T}]+[(1-\gamma_{t})GG^{T}-\tilde{G}\tilde{G}^{T}]\cdot\nabla\ln p_{ t}\right)$ (9) $\displaystyle\bar{G}$ $\displaystyle=[\gamma_{t}GG^{T}+\tilde{G}\tilde{G}^{T}]^{\frac{1}{2}}$ (10) $$

and $\tilde{G}\equiv\tilde{G}(x,t),\gamma_{t}\geq 0$ are arbitrary functions such that the solutions of Equation 8 exist and are unique.

Proof of Section 3.2 is given in Appendix C and follows from manipulations of Fokker-Planck-Kolmogorov (FPK) equations corresponding to Equation 8.

Section 3.2 implies that given the same initial density $p_{0}(x)$, evolution according to both Equation 4 and Equation 8 will have the same marginal densities $p_{t}(x)$ for all times $t$.
Further, Equation 8 can be simulated backward in time using the result from Anderson (), again with the same marginal densities $p_{t}(x)$.
Consequently, Equation 4 can be simulated forward or backward in time using any member of the family specified by Equation 8.
Note that setting $\gamma_{t}=1,\tilde{G}=0$ recovers the original SDE in Equation 4, while setting $\gamma_{t}=0,\tilde{G}=0$ recovers the general probability flow ODE from Song et al. (, eq. 37).
Additionally, $\tilde{G}$ is particularly useful for deterministic flow models, further discussed in Section 3.2.
Section 3.2 gives a recipe for developing particular samplers, such as those in the remainder of this section, some of which have appeared in the literature.
A priori, Section 3.2 cannot determine which concrete sampler will be best for a given application, but since the samplers do not require any training to use, it is possible to postpone the choice of sampler to an empirical analysis at test time.

The flexibility afforded by Equation 8 is particularly useful
(1) in the presence of singularities in the drift and diffusion coefficients $f$ and $G$ respectively of Equation 4,
(2) in the presence of errors resulting from finite discretization,
and (3) for flexible specification of the diffusion coefficient in generative applications.
Our experimental evaluations primarily focus on these aspects of Section 3.2.

A direct consequence of Section 3.2, by defining $\tilde{G}\equiv 0,G\equiv g(t)I$, is the following corollary applicable to commonly used generative diffusion models with additive noise:
{restatable}corollarycorollaryone
For the SDE in Equation 4 with $G\equiv g(t)I$, a subset of SDEs prescribed by Section 3.2 and indexed by $\gamma_{t}$ is:

$$ $\displaystyle dx$ $\displaystyle=\left[f(x,t)-\frac{(1-\gamma(t))g^{2}(t)}{2}\nabla_{x}\ln p_{t}( x)\right]dt+\sqrt{\gamma(t)}g(t)dW_{t}$ (11) $$

Proof in Appendix D.
Note that choosing $\gamma_{t}=0$ results in the probability flow ODE specified in Equation 6.
Intuitively, the members in the family differ in terms of the amount of noise injected as a function of time.
$\gamma_{t}=0$ yields a purely deterministic simulation; $\gamma_{t}>0$ yields a variety of stochastic simulations.
Further, similar special cases discussed in Huang et al. () and Berner et al. () also directly follow from Section 3.2 as well.

Some properties of Section 3.2:

- 1.
$\gamma(t)$ can be chosen at sampling time and doesn’t affect the training of the score function.
- 2.
With $\gamma(t)=\hat{\gamma}^{2}(t)g^{-2}(t)$, where $\hat{\gamma}(t)$ is an arbitrary function (satisfying constraints of Section 3.2), we can choose an arbitrary diffusion term at sampling time.
For example, choosing $\gamma_{t}=\gamma^{2}/g^{2}(t)$ leads to a constant diffusion coefficient.
- 3.
For the SDE specified by Equation 7, we can choose $\gamma(t)=(1-t)\hat{\gamma}^{2}(t)g^{-2}(t)$ to get rid of the singularity in the diffusion term.

Note that Section 3.2 can be used whenever we have access to the score function $\nabla_{x}\ln p_{t}$.
Next, we first construct a specialized solution based on Section 3.2 for deterministic flow models that enables flexible control of both drift and diffusion coefficients, and apply it to the special case of deterministic Gaussian flows where the score function can be imputed from the velocity field (Section 3.3).
Recall that deterministic flows specify a transport via the ODE $dx=v(x,t)dt$.
This ODE can be viewed as an SDE where the diffusion term has been set to zero.
Choosing $G\equiv 0,\tilde{G}\equiv\tilde{g}(t)I$ in Section 3.2 gives Section 3.2, which enables deriving stochastic samplers for Gaussian flow models:
{restatable}corollarycorollarytwo
For the ODE in Equation 3, a subset of SDEs prescribed by Section 3.2 and indexed by $\tilde{g}(t)$ is

$$ $\displaystyle dx$ $\displaystyle=\left[v(x,t)+\frac{\tilde{g}^{2}(t)}{2}\nabla_{x}\ln p_{t}(x) \right]dt+\tilde{g}(t)dW_{t}$ (12) $$

Proof in Appendix E.
Section 3.2 enables flexible specification of a time dependent diffusion coefficient $\tilde{g}(t)$, allowing the introduction of stochasticity in the simulation of otherwise deterministic models, *purely at sampling time*.
Note that Equation 12 requires access to the score function $\nabla_{x}\ln p_{t}(x)$ for the marginal distributions $p_{t}(x)$. In Section 3.3, we describe how the score function can be imputed from the learned flow model $v(x,t)$ for the special case of Gaussian flow models.
It can be verified that the particular choice of $f$ and $g$ in Equation 7 satisfy Equation 12 by using the expression for the score from Equation 13.

**Table 1: Examples of SDEs that have the same marginal distribution $p_{t}(x)$ as a given Gaussian flow specified by $v\equiv v(x,t)$. $\alpha\geq 0$ is a scale parameter that varies the magnitude of the diffusion coefficient $g$. Each of these behaves differently when discretized and simulated (Figure 2 and Section J.2). These and infinitely many more can be constructed using the scheme in Equation 12.**
| Name | $\tilde{g}(t)$ | $\tilde{f}(x,t)$ | Description |
| --- | --- | --- | --- |
| Deterministic | 0 0 | $v$ | Base flow model |
| Constant | $\alpha$ | $v+\frac{\alpha^{2}}{2}\nabla_{x}\ln p_{t}$ | Constant $g$, singular $f$ |
| Singular | $\alpha\sqrt{t/(1-t)}$ | $v+\frac{\alpha^{2}}{2}\frac{t}{1-t}\nabla_{x}\ln p_{t}$ | Singular $g,f$ |
| NonSingular | $\alpha\sqrt{t},$ | $v+\frac{\alpha^{2}}{2}t\nabla_{x}\ln p_{t}$ | Non-singular $g,f$ |
| ZeroEnds | $\alpha\sqrt{t(1-t)},$ | $v+\frac{\alpha^{2}}{2}t(1-t)\nabla_{x}\ln p_{t}$ | Non-singular $g,f,g(0)=g(1)=0$ |

While infinitely many choices are available for $\tilde{g}$, we consider a few interesting ones listed in the Table 1, constructed by choosing integer powers of $t$ and $1-t$ and introducing a scaling coefficient $\alpha$, for experimental evaluations.
Note that the only degree of freedom in Table 1 is the choice of $\tilde{g}(t)$, which determines $\tilde{f}(x,t)$, given the flow field $v(x,t)$ and the score $\nabla_{x}\ln p_{t}(x)$.
The $f(x,t)$ is singular in Constant because the score $\nabla_{x}\ln p_{t}(x_{t})$, as computed in Equation 13, has $t$ in the denominator, making $f(x,t)$ singular at $t=0$.
The choice in NonSingular precisely eliminates this singularity.
Figure 2 compares these choices in a toy setup; Section 4 has comparisons on ImageNet.

### 3.3 Score function and classifier free guidance for a Gaussian flow model

Recall that Section 3.2 requires access to the score function.
For Gaussian flows, the score function can be inferred from the velocity field itself, alleviating the need to learn it separately.
This result is known (see e.g. Zheng et al. () in the context of flow matching) and we present it here in our setting.
For Gaussian flows, with $p_{1}(x_{1})\equiv N(x_{1};\mu_{1},\sigma_{1}^{2}I)$ and interpolation specified in Equation 1, the score can be computed as:

$$ $\displaystyle\nabla_{x}\ln p_{t}(x)$ $\displaystyle=\frac{-(1-t)v(x,t)+\mu_{1}-x}{t\sigma_{1}^{2}}$ (13) $$

where $v(x,t)=\mathbb{E}[x_{1}-x_{0}|x,t]$ is the estimated flow.
Proof is provided in Appendix B.
Note that the score function can also be estimated given $\mathbb{E}[x_{0}|x,t]$ or $\mathbb{E}[x_{1}|x,t]$.
In summary, the expression follows directly from using results from Denoising Score Matching (Vincent, ) and the Gaussianity of $p_{1}(x_{1})$.
Similar expressions can be derived for other interpolations that are linear in $x_{0},x_{1}$.
With access to the score function and linearity of Equation 13 in $v$ we can define classifier free guided (Ho & Salimans, ) Gaussian flow as:

$$ $\displaystyle v_{\text{cfg}}(x,t,c)=(1+\lambda)v(x,t,c)-\lambda(v(x,t,c=\varnothing)$ (14) $$

where $c$ indicates extra conditioning information as in classifier free guidance, $\varnothing$ indicates no conditioning and $\lambda$ specifies the relative strength of the guidance.
$\lambda=0$ reduces to class conditional sampling, while $\lambda>0$ puts a larger weight on the conditioning, biasing the sample towards the modes of the conditional distribution.
Using classifier-free guidance with a stochastic sampler will, of course, give diversity that isn’t possible with a deterministic sampler, as can be seen in Figure 1.
Note that Xie et al. (); Dao et al. (); Zheng et al. () also consider related definitions in the context of flow matching.

## 4 Experiments

Our method allows us to identify a family of SDEs that correspond to a given deterministic Gaussian flow model, enabling construction of stochastic samplers with flexible diffusion coefficients.
In our experiments we compare various samplers derived from the corresponding SDEs in Table 1, using Euler-Maruyama, for a given Gaussian flow model *without any additional training*.

Figure: Figure 3: Stochasticity is most helpful at coarser discretizations. We visualize the effect of coarseness of discretization by sampling for 100 and 500 sampling steps. See Figure 2 for the same plots at 50 steps, which shows more extreme bias in variance for Deterministic and Singular.
Refer to caption: x3.png

Figure: Figure 4: Stochasticity helps mitigate bias. We plot the error in mean and error in variance for NonSingular for a set of diffusion coefficient scales $\alpha\in\{0.0,0.5,1.0,1.5,2.0,2.5\}$. Estimates for variance at $t=0$ improve as $\alpha$ increases, leading to a drop in KL divergence from the true distribution. However, with very high $\alpha$ values intermediate marginals develop a bias.
Refer to caption: x5.png

### 4.1 Comparison on estimating marginal statistics for a two Gaussian toy problem

We start by considering a toy problem where both $p_{0}$ and $p_{1}$ are Gaussian.
See Appendix G for details of the experimental setup and Appendix I for a JAX (Bradbury et al., ) implementation of NonSingular.

#### Discretization of deterministic flow leads to bias.

In Figure 2, with 50 sampling steps, we observe that the estimate for the mean is fairly accurate for all samplers for the entirety of the interval $t\in[0,1]$.
However, the samplers differ in their behavior for variance.
Deterministic exhibits a noticeable bias and underestimates the variance (with zero variance in its estimate), with the worst estimate at $t=0$.
Stochastic samplers provide noticeably better estimates at $t=0$, but with increased variance.

#### Stochasticity is most helpful at coarser discretizations.

In Figure 3 we study the effect of the number of discretization steps on the different samplers (also see Figure 2 for 50 steps).
While mean estimates are accurate for all methods, Deterministic gets increasingly biased for variance estimates as the number of sampling steps is decreased.
Stochastic samplers perform consistently well at various discretization levels for $t=0$, with significantly better estimates for fewer sampling steps.
Note that Singular has very large bias as well as variance closer to $t=1$; those improve with finer discretization.
Since Constant also has a singularity, but only in the drift term $f$, we conclude that the instability is primarily due to the singularity in Singular’s diffusion term.

#### Stochasticity helps mitigate bias.

In Figure 4 we study the effect of diffusion coefficient scale $\alpha$ on the NonSingular sampler at 100 sampling steps.
Finite discretization introduces a bias in the deterministic sampler (when $\alpha=0$), where the variance is consistently underestimated and is worst at $t=0$.
Increased stochasticity with increasing diffusion coefficient scale ($\alpha>0$) helps mitigate this bias at the cost of increased variance.
This can be seen in the figure with larger $\alpha$ values yielding better estimate of the variance, although with larger variance in the estimate.

### 4.2 Comparison of SDEs for rectified flows on ImageNet generation

We compare the behavior of various SDEs on the sampling quality for large scale image generation using the ImageNet (2012) dataset (Deng et al., ; Russakovsky et al., ).
We train rectified flow models at two different image resolutions ($64\times 64$ and $128\times 128$) and compare their sample quality using the Frechet Inception Distance (FID) metric (Heusel et al., ) for class conditional samples.
See Appendix H for setup details.
The results show that small but statistically significant differences exist between samplers even for metrics like FID, but the optimal sampler is likely to be application and model specific.

#### Stochasticity can improve FID.

In Table 2 we report the best FID using each SDE in Table 1 for two image resolutions using 300 sampling steps, along with the corresponding diffusion term scale $\alpha$ and one standard deviation confidence interval.
Two key observations stand out:
(1) stochastic samplers tend to produce better FIDs,
and (2) the two non-singular samplers have much better FIDs than Deterministic or Singular.
Note that observation (1) has also been made previously for probability flow ODEs (Song et al., ).
The addition of a parameter $\alpha$ to control the strength of the stochasticity while keeping the marginal distribution $p_{t}$ unchanged (Section 3.2), permits principled post-training optimization of the metrics like FID.

#### Non-singular samplers work well over a broad range of α 𝛼 \alpha italic_α .

In Figures 5 and 7 we show how the FID varies with $\alpha$ for each sampler for two different image resolution models.
NonSingular and ZeroEnds attain better FID in general and are better behaved over a much larger range of the diffusion coefficient scale $\alpha$ at both resolutions.
These samplers both have small diffusion coefficients $g(t)$ close to $t=0$; their performance indicates that noise near $t=0$ is particularly harmful.
The low variance of ZeroEnds in comparison to NonSingular indicates that a large diffusion coefficient near $t=1$ tends to introduce variance in the final FID.

**Table 2: Stochasticity can improve FID. Comparison of various samplers at their best $\alpha$ values with 300 sampling steps for ImageNet image generation task at two resolutions.**
|  | $64\times 64$ | $128\times 128$ |  |  |
| --- | --- | --- | --- | --- |
| Sampler | FID | $\alpha$ | FID | $\alpha$ |
| Deterministic | $3.07\pm 0.01$ | $0.0$ | $5.19\pm 0.02$ | $0.0$ |
| Singular | $3.07\pm 0.01$ | $0.08$ | $5.13\pm 0.04$ | $0.14$ |
| Constant $g$ | $2.97\pm 0.04$ | $0.08$ | $5.17\pm 0.05$ | $0.1$ |
| NonSingular | $\bf{2.95\pm 0.01}$ | $0.56$ | $\bf{4.93\pm 0.06}$ | $0.42$ |
| ZeroEnds | $\bf{2.95\pm 0.01}$ | $0.54$ | $5.03\pm 0.01$ | $0.52$ |

Figure: (a) 64 $\times$ 64
Refer to caption: x6.png

#### Stochasticity makes FID robust to discretization.

In Figure 6 we compare the effect of the number of sampling steps on FID for various samplers at two image resolutions.
We set $\alpha^{2}$ proportional to the number of sampling steps with the maximum value provided by Table 2.
Again the non-singular samplers perform better than Deterministic at all discretization levels.

#### Stochastic sampling improves diversity at all classifier-free guidance levels.

In Figures 1 and 12 we show samples from NonSingular using classifier-free guidance (Section 3.3), varying both $\alpha$ and $\lambda$, the guidance weight.
In all cases, we can see that diversity increases with $\alpha$, and class typicality increases with $\lambda$.

#### Qualitative comparisons.

For qualitative comparisons, we visualize a few samples at various diffusion coefficient scales using different SDEs in Figures 9, 10 and 11.
All samples in a column are generated by starting at the same draw $x_{1}\sim p_{1}(x_{1})$; different columns start from different draws.
Noise scale $\alpha$ gets progressively larger as we move down the rows.
For Constant, we observe that samples get increasingly noisy with increasing $\alpha$ indicating accumulating errors with increasing scale.
The samples from NonSingular look better, as expected from Figure 5.
Lastly, samples from Singular change much more rapidly in comparison to the other samplers, indicating that the singularities in the SDE coefficients increase the effect of noise.

Figure: (a) 64 $\times$ 64
Refer to caption: x8.png

## 5 Related work

Transport learning methods learn a mapping between two distributions, where the learned model can transform a sample from one distribution into a sample from the other one.
Typically, one of the distributions is easy to sample (such as a Gaussian) and the other one is the data distribution that one is interested in modeling.
The learned mapping can either be deterministic or stochastic.
A thorough overview of related areas can be found in Yang et al. ().

#### Deterministic transport.

Deterministic transport methods implement a change of variable, either explicitly or approximately, that can be used to uniquely map a sample from one distribution to the other.
The normalizing flow family (Rezende & Mohamed, ; Dinh et al., ; Kingma & Dhariwal, ) of methods construct an explicit invertible model that realizes this map either in one step or a finite number of discrete steps.
Neural ODEs (Chen et al., ; Grathwohl et al., ) generalize from discrete steps to a continuous time mapping by inferring and learning the gradient field for all times.
However, Neural ODEs are difficult to train due to the need for simulating the ODEs as part of the training.
Rectified flows, flow matching, and iterative denoising methods (Liu et al., ; Lipman et al., ; Tong et al., ; Heitz et al., ; Delbracio & Milanfar, ) either implicitly or explicitly specify a continuous mapping and learn a model to approximate the continuous time mapping.
Similarly, probability flow ODEs (Song et al., ) learned by diffusion models (Sohl-Dickstein et al., ) approximate an implicitly defined continuous mapping.
Our work is useful for flexible sampling from such pre-trained continuous time deterministic Gaussian flows, or more generally where the score function for all the marginal distributions is either provided or can be deduced from the learned flow model.

#### Stochastic transport.

Stochastic transport methods learn a stochastic mapping, where a sample from one distribution gets stochastically mapped to a sample from the other.
Gaussian diffusion models are a salient example of such discrete Sohl-Dickstein et al. (); Ho et al. () or continuous time Song et al. (); Kingma et al. () mappings where one of the distributions is constrained to be Gaussian.
Several generalizations have been proposed that extend from Gaussian to more general families of distributions Yoon et al. ().
The stochastic interpolants framework (Albergo & Vanden-Eijnden, ; Albergo et al., ; Ma et al., ) further generalizes to a larger family of distributions by introducing a random latent variable allowing efficient estimation of the score function at all times.
Our work is directly applicable to models learned with such methods where the score function is accessible and can be used to construct and explore a large family of samplers.
The convergence rates of diffusion models have been studied in Chen et al. (); Benton et al. () with respect to the number of data samples and dimensionality.
However, since our method does not require retraining, it does not affect these properties of the original training algorithms.

#### Schrödinger bridge and optimal transport.

These methods consider a more general problem of learning transport maps with additional constraints.
k-Rectified flows Liu et al. () provide an iterative procedure for tackling deterministic optimal transports for a family of costs, while the more general Schrödinger bridge problem, viewed as an entropy regularized optimal transport, is an active area of research Shi et al. (); Liu et al. ().
Our work is complementary to these methods as we focus on flexible sampling, given the access to the score function for the marginal distributions.

## 6 Conclusion

We introduced a general method to identify a family of SDEs that have the same marginal distribution as a particular SDE, including the special case where the diffusion coefficient of the given SDE is zero.
This special case corresponds to flow models which naively only support deterministic sampling.
Our method enables flexible construction of stochastic samplers for such deterministic models where the diffusion coefficient can be chosen at sampling time from an infinitely large set of possibilities in an application and evaluation metric dependent way.
Our method requires explicit access to the score function, in absence of which it is limited to a subset of flow models where the score function can be derived from the given flow model.
However, this set includes currently popular rectified flow and diffusion models where one of the distributions is Gaussian.

## Appendix A Derivation of singular SDE

We consider the following SDE with additive noise; i.e., the diffusion coefficient $g$ is only a function of time.

$$ $\displaystyle dx=f(x,t)dt+g(t)dW_{t}$ (15) $$

The perturbation kernel $p(x_{t}|x_{0})$ corresponding to rectified flow is Gaussian, with $p(x_{t}|x_{0})=N(x_{t};(1-t)x_{0}+t\mu_{1},t^{2}\sigma_{1}^{2}I)$.
Since the perturbation kernel is Gaussian, following Song et al. (), we assume that the drift term is affine; i.e. $f(x,t)\equiv f(t)x$.
Further since $X_{0},X_{1}$ are independent, we can directly infer the first and second moments $\mu_{t},\Sigma_{t}$ for the marginals $p_{t}(x)$ as $\mu_{t}=(1-t)\mu_{0}+t\mu_{1}$ and $\Sigma_{t}=(1-t)^{2}\Sigma_{0}+t^{2}\sigma_{1}^{2}I$.

From Eq (5.50) of Särkkä & Solin () we have

$$ $\displaystyle\frac{d\mu_{t}}{dt}$ $\displaystyle=\mathbb{E}_{p_{t}(x)}[f(t)x]$ (16) $\displaystyle=f(t)\mu_{t}$ (17) $$

where $\mu_{t}$ is the mean at time $t$. Rearranging and integrating both sides:

$$ $\displaystyle\ln\frac{\mu_{t}}{\mu_{0}}$ $\displaystyle=\int_{0}^{t}f(s)ds$ (18) $\displaystyle\ln\frac{(1-t)\mu_{0}+t\mu_{1}}{\mu_{0}}$ $\displaystyle=\int_{0}^{t}f(s)ds$ $\displaystyle\text{Substituting}~{}\mu_{t}=(1-t)\mu_{0}+t\mu_{1}$ (19) $\displaystyle\frac{\mu_{1}-\mu_{0}}{(1-t)\mu_{0}+t\mu_{1}}$ $\displaystyle=f(t)$ Differentiating both sides w.r.t. $t$ (20) $$

Substituting $\mu_{1}=0$, we get as in Equation 7:

$$ $\displaystyle f(x,t)$ $\displaystyle=-\frac{x}{1-t}$ (22) $$

Similarly, from Eq. (5.51) of Särkkä & Solin ():

$$ $\displaystyle\frac{d\Sigma_{t}}{dt}$ $\displaystyle=\mathbb{E}_{p_{t}(x)}\left[f(x,t)(x-\mu_{t})^{T}+(x-\mu_{t})f(x, t)^{T}+G(x,t)QG(x,t)^{T}\right]$ (23) $$

Substituting $Q\equiv I$ (we are assuming isotropic dispersion), $G(x,t)\equiv g(t)I$ (symmetric, time-dependent diffusion coefficient), and $f(x,t)$ from Equation 22:

$$ $\displaystyle\frac{d\Sigma_{t}}{dt}$ $\displaystyle=\mathbb{E}_{p_{t}(x)}\left[-\frac{x}{1-t}(x-\mu_{t})^{T}-(x-\mu_ {t})\frac{x^{T}}{1-t}+g^{2}(t)I\right]$ (24) $\displaystyle=\frac{2}{1-t}\mathbb{E}_{p_{t}(x)}\left[-xx^{T}+\mu_{t}\mu_{t}^{ T}\right]+g^{2}(t)I$ (25) $\displaystyle=-\frac{2\Sigma_{t}}{1-t}+g^{2}(t)I$ (26) $\displaystyle\implies\frac{d\Sigma_{t}}{dt}+\frac{2\Sigma_{t}}{1-t}$ $\displaystyle=g^{2}(t)I$ (27) $$

Above is an inhomogenous differential equation.
The integrating factor $I(t)$ can be calculated as:

$$ $\displaystyle I(t)$ $\displaystyle=\exp\left(\int_{0}^{t}\frac{2}{1-s}ds\right)=\frac{1}{(1-t)^{2}}$ (28) $$

Multiplying both sides of Equation 27, we can write:

$$ $\displaystyle\frac{d}{dt}\left[\frac{\Sigma_{t}}{(1-t)^{2}}\right]$ $\displaystyle=\frac{g^{2}(t)I}{(1-t)^{2}}$ (29) $$

Integrating both sides:

$$ $\displaystyle\left[\frac{\Sigma_{s}}{(1-s)^{2}}\right]_{0}^{t}$ $\displaystyle=\int_{0}^{t}\frac{g^{2}(s)I}{(1-s)^{2}}ds$ (30) $\displaystyle\frac{\Sigma_{t}}{(1-t)^{2}}-\Sigma_{0}$ $\displaystyle=\int_{0}^{t}\frac{g^{2}(s)I}{(1-s)^{2}}ds$ (31) $$

Substituting $\Sigma_{t}=(1-t)^{2}\Sigma_{0}+t^{2}\sigma_{1}^{2}I$:

$$ $\displaystyle\frac{(1-t)^{2}\Sigma_{0}+t^{2}\sigma_{1}^{2}I}{(1-t)^{2}}-\Sigma _{0}^{2}$ $\displaystyle=\int_{0}^{t}\frac{g^{2}(s)I}{(1-s)^{2}}ds$ (32) $$

Differentiating both sides w.r.t. $t$ and simplifying yields:

$$ $\displaystyle g^{2}(t)=\frac{2t\sigma_{1}^{2}}{1-t}$ (33) $$

Substituting $\sigma_{1}=1$ to the result in Equation 7:

$$ $\displaystyle g(t)=\sqrt{\frac{2t}{1-t}}$ (34) $$

## Appendix B Score function from rectified flow

Given a base data distribution $p(x)$ and a conditional noising distribution $p_{\sigma}(\tilde{x}|x)$, Denoising score matching Vincent () learns the score for the marginal $p_{\sigma}(\tilde{x})$ by optimizing:

$$ $\nabla_{\tilde{x}}\ln p_{\sigma}(\tilde{x})=\operatorname*{arg\,min}_{\psi} \mathbb{E}_{p_{\sigma}(x_{0},\tilde{x})}\left[\frac{1}{2}\left\lVert\psi( \tilde{x})-\frac{\partial\ln p_{\sigma}(\tilde{x}|x_{0})}{\partial\tilde{x}} \right\lVert^{2}\right]$ (35) $$

where $p_{\sigma}(x_{0},\tilde{x})\equiv p(x)p_{\sigma}(\tilde{x}|x_{0})$.
The solution to the above optimization problem can be written as:

$$ $\nabla_{\tilde{x}}\ln p_{\sigma}(\tilde{x})=\mathbb{E}_{p_{\sigma}(x_{0}| \tilde{x})}\frac{\partial\ln p_{\sigma}(\tilde{x}|x_{0})}{\partial\tilde{x}}$ (36) $$

Mapping the above to rectified flow with $\sigma\equiv t,\tilde{x}\equiv x_{t}$ we get:

$$ $\nabla_{x_{t}}\ln p_{t}(x_{t})=\mathbb{E}_{p_{t}(x_{0}|x_{t})}\frac{\partial \ln p_{t}(x_{t}|x_{0})}{\partial x_{t}}$ (37) $$

Next if:

$$ $\displaystyle p_{t}(x_{t}|x_{0})$ $\displaystyle=N(x_{t};\mu(x_{0},t),\sigma^{2}(x_{0},t)I)$ (38) $\displaystyle\frac{\partial\ln p_{t}(x_{t}|x_{0})}{\partial x_{t}}$ $\displaystyle=\frac{\partial}{\partial x_{t}}\frac{-||x_{t}-\mu(x_{0},t)||^{2} }{2\sigma(x_{0},t)^{2}}$ (39) $\displaystyle=\frac{-(x_{t}-\mu(x_{0},t))}{\sigma(x_{0},t)^{2}}$ (40) $$

Now:

$$ $\nabla_{x_{t}}\ln p_{t}(x_{t})=\mathbb{E}_{p_{t}(x_{0}|x_{t})}\frac{-(x_{t}- \mu(x_{0},t))}{\sigma(x_{0},t)^{2}}$ (41) $$

Next, assume the covariance $\sigma(x_{0},t)$ doesn’t depend on $x_{0}$ – i.e., $\sigma(x_{0},t)\equiv\sigma(t)$ – and the mean $\mu(x_{0},t)$ is linear in $x_{0}$.
Then:

$$ $\displaystyle\mathbb{E}_{p_{t}(x_{0}|x_{t})}\frac{-(x_{t}-\mu(x_{0},t))}{ \sigma(x_{0},t)^{2}}$ $\displaystyle=\mathbb{E}_{p_{t}(x_{0}|x_{t})}\frac{-(x_{t}-\mu(x_{0},t))}{ \sigma(t)^{2}}$ (42) $\displaystyle=\frac{-(x_{t}-\mu(\mathbb{E}[x_{0}|x_{t}],t))}{\sigma(t)^{2}}$ (43) $$

### B.1 Gaussian Rectified Flow

Consider the special case where $x_{t}=(1-t)x_{0}+tx_{1},x_{1}\sim N(x_{1};\mu_{1},\sigma_{1}^{2}I)$.
We have $p_{t}(x_{t}|x_{0})=N((1-t)x_{0}+t\mu_{1},t^{2}\sigma_{1}^{2})$.
Using the result from Equation 43 we get:

$$ $\displaystyle\nabla_{x_{t}}\ln p_{t}(x_{t})$ $\displaystyle=\frac{-(x_{t}-\mu(\mathbb{E}[x_{0}|x_{t}],t))}{\sigma(t)^{2}}$ (44) $$

From this, we can write:

$$ $\displaystyle x_{1}$ $\displaystyle=\frac{x_{t}-(1-t)x_{0}}{t}$ (45) $\displaystyle\mathbb{E}[x_{1}-x_{0}|x_{t}]$ $\displaystyle=\mathbb{E}\left[\frac{x_{t}-(1-t)x_{0}}{t}-x_{0}\bigg{\lvert}x_{ t}\right]$ (46) $\displaystyle=\mathbb{E}\left[\frac{x_{t}-(1-t)x_{0}-tx_{0}}{t}|x_{t}\right]$ (47) $\displaystyle=\mathbb{E}\left[\frac{x_{t}-x_{0}}{t}|x_{t}\right]$ (48) $\displaystyle=\frac{x_{t}-\mathbb{E}[x_{0}|x_{t}]}{t}$ (49) $\displaystyle\mathbb{E}[x_{0}|x_{t}]$ $\displaystyle=x_{t}-t\mathbb{E}[x_{1}-x_{0}|x_{t}]$ (50) $\displaystyle\nabla_{x_{t}}\ln p_{t}(x_{t})$ $\displaystyle=\frac{\mu(\mathbb{E}[x_{0}|x_{t}],t)-x_{t}}{\sigma(t)^{2}}$ (51) $\displaystyle=\frac{(1-t)\mathbb{E}[x_{0}|x_{t}]+t\mu_{1}-x_{t}}{t^{2}\sigma_{ 1}^{2}}$ (52) $\displaystyle=\frac{(1-t)(x_{t}-t\mathbb{E}[x_{1}-x_{0}|x_{t}])+t\mu_{1}-x_{t} }{t^{2}\sigma_{1}^{2}}$ (53) $\displaystyle=\frac{-(1-t)t\mathbb{E}[x_{1}-x_{0}|x_{t}]+t\mu_{1}-tx_{t}}{t^{2 }\sigma_{1}^{2}}$ (54) $\displaystyle=\frac{-(1-t)\mathbb{E}[x_{1}-x_{0}|x_{t}]+\mu_{1}-x_{t}}{t\sigma _{1}^{2}}$ (55) $$

### B.2 General Rectified Flow

First recall the change of variables formula for a density $p(x)$ with $y=g(x)$ where $g$ is invertible and $g^{-1}$ is differentiable:

$$ $\displaystyle p(y)$ $\displaystyle=p(g^{-1}(y))\left\lvert\det\left[\frac{\partial g^{-1}(z)}{ \partial z}\right]_{z=y}\right\rvert$ (56) $$

Now, with $x_{0}\sim p_{1}(x_{0})$ and $x_{1}\sim p_{1}(x_{1})$ and $x_{0},x_{1}\in\mathbb{R}^{d}$, let $x_{t}=g(x_{1};x_{0})$ be a function that is invertible in first argument and whose inverse $g^{-1}(x_{t};x_{0})$ is differentiable w.r.t. the first argument.
Note that for simple rectified flows, $x_{t}=(1-t)x_{0}+tx_{1}$ satisfies these conditions.

We can now express the conditional density $p(x_{t}|x_{0})$ as:

$$ $\displaystyle p(x_{t}|x_{0})=p_{1}(g^{-1}(x_{t};x_{0}))\left\lvert\det\left[ \nabla_{z}g^{-1}(z;x_{0})\right]_{z=x_{t}}\right\rvert$ (57) $$

The score for the conditional density can then be calculated as

$$ $\displaystyle\frac{\partial\ln p_{t}(x_{t}|x_{0})}{\partial x_{t}}$ $\displaystyle=\nabla_{z}\ln p_{1}(z)|_{z=g^{-1}(x_{t};x_{0})}+\nabla_{x_{t}} \ln\left\lvert\det\left[\nabla_{z}g^{-1}(z;x_{0})\right]_{z=x_{t}}\right\rvert$ (58) $$

and the score for the marginal density as:

$$ $\displaystyle\nabla_{x_{t}}\ln p_{t}(x_{t})$ $\displaystyle=\mathbb{E}_{p_{t}(x_{0}|x_{t})}\frac{\partial\ln p_{t}(x_{t}|x_{ 0})}{\partial x_{t}}$ (59) $\displaystyle=\mathbb{E}_{p_{t}(x_{0}|x_{t})}\left[\nabla_{z}\ln p_{1}(z)|_{z= g^{-1}(x_{t};x_{0})}+\nabla_{x_{t}}\ln\left\lvert\det\left[\nabla_{z}g^{-1}(z; x_{0})\right]_{z=x_{t}}\right\rvert\right]$ (60) $$

For the specific case of Rectified flows, define $g(x_{1};x)=(1-t)x+tx_{1}$.
Then:

$$ $\displaystyle x_{t}$ $\displaystyle=g(x_{1};x_{0})$ (61) $\displaystyle g^{-1}(x_{t};x_{0})$ $\displaystyle=\frac{x_{t}-(1-t)x_{0}}{t}$ Inverse is w.r.t. first argument (62) $\displaystyle\frac{\partial g^{-1}(x_{t};x_{0})}{\partial x_{t}}$ $\displaystyle=\frac{1}{t}I$ (63) $\displaystyle\det\frac{1}{t}I$ $\displaystyle=\frac{1}{t^{d}}$ $I$ is $d\times d$ (64) $\displaystyle p_{t}(x_{t}|x_{0})$ $\displaystyle=\frac{1}{t^{d}}p_{1}\left(\frac{x_{t}-(1-t)x_{0}}{t}\right)$ (65) $$

Substituting into Equation 60:

$$ $\displaystyle\nabla_{x_{t}}\ln p_{t}(x_{t})$ $\displaystyle=\mathbb{E}_{p_{t}(x_{0}|x_{t})}\left[\frac{1}{t}\nabla_{z}\ln p_ {1}(z)|_{z=g^{-1}(x_{t};x_{0})}\right]$ (66) $$

It can be verified that with the choice of $p_{1}(x_{1})\equiv N(\mu_{1},\sigma_{1}^{2}I)$, we recover Equation 55.

## Appendix C Proof of Section 3.2

*

###### Proof.

The evolution of the marginal probability density $p_{t}(x)$ is then described by the Fokker-Planck-Kolmogorov (FPK) equation (Särkkä & Solin, ) as:

$$ $\displaystyle\frac{\partial p_{t}}{\partial t}=-\sum_{i=1}^{d}\frac{\partial}{ \partial x_{i}}[\bar{f}p_{t}]+\frac{1}{2}\sum_{i=1}^{d}\sum_{j=1}^{d}\frac{ \partial^{2}}{\partial x_{i}x_{j}}\left[\sum_{k=1}^{d}\bar{G}_{ik}\bar{G}_{jk} p_{t}\right]$ (67) $$

We write the above more succinctly as:

$$ $\displaystyle\frac{\partial p_{t}}{\partial t}=-\nabla\cdot[\bar{f}p_{t}]+ \frac{1}{2}\nabla\cdot\left[\bar{G}\bar{G}^{T}p_{t}\right]\cdot\nabla^{T}$ (68) $$

Where $\nabla\cdot$ is the divergence operator.
Next for an arbitrary $R\equiv R(x,t)$ consider the following identity:

$$ $\displaystyle\left[RR^{T}p_{t}\right]\cdot\nabla^{T}$ $\displaystyle=\nabla\cdot\left[RR^{T}p_{t}\right]$ $\displaystyle\quad RR^{T}~{}\text{ is symmetric}$ (69) $\displaystyle=\left[\nabla\cdot RR^{T}\right]p_{t}+RR^{T}\cdot\nabla p_{t}$ (70) $\displaystyle=\left[\nabla\cdot RR^{T}\right]p_{t}+RR^{T}\cdot p_{t}\nabla\ln p _{t}$ (71) $\displaystyle=\left(\nabla\cdot RR^{T}+RR^{T}\cdot\nabla\ln p_{t}\right)p_{t}$ (72) $$

Expanding out $\bar{f}p_{t}$ by substituting for $\bar{f}$:

$$ $\displaystyle\bar{f}p_{t}$ $\displaystyle=\left[f-\frac{1}{2}\left(\nabla\cdot[(1-\gamma_{t})GG^{T}-\tilde {G}\tilde{G}^{T}]+[(1-\gamma_{t})GG^{T}-\tilde{G}\tilde{G}^{T}]\cdot\nabla\ln p _{t}\right)\right]p_{t}$ (73) $\displaystyle\begin{split}&=fp_{t}-\frac{1-\gamma_{t}}{2}\left(\nabla\cdot GG^ {T}+GG^{T}\cdot\nabla\ln p_{t}\right)p_{t}\\ &+\frac{1}{2}\left(\nabla\cdot\tilde{G}\tilde{G}^{T}+\tilde{G}\tilde{G}^{T} \cdot\nabla\ln p_{t}\right)p_{t}\end{split}$ (74) $$

Using Equation 72 and rewriting:

$$ $\displaystyle\begin{split}\bar{f}p_{t}&=fp_{t}-\frac{1-\gamma_{t}}{2}[GG^{T}p_ {t}]\cdot\nabla^{T}+\frac{1}{2}[\tilde{G}\tilde{G}^{T}p_{t}]\cdot\nabla^{T} \end{split}$ (75) $$

Next we revisit Equation 68, and substitute for $\bar{f}p_{t}$ and $\bar{G}$ with $\bar{G}\bar{G}^{T}=\gamma_{t}GG^{T}+\tilde{G}\tilde{G}^{T}$:

$$ $\displaystyle\begin{split}\frac{\partial p_{t}}{\partial t}&=-\nabla\cdot[fp_{ t}-\frac{1-\gamma_{t}}{2}[GG^{T}p_{t}]\cdot\nabla^{T}+\frac{1}{2}[\tilde{G} \tilde{G}^{T}p_{t}]\cdot\nabla^{T}]\\ &+\frac{1}{2}\nabla\cdot\left[(\gamma_{t}GG^{T}+\tilde{G}\tilde{G}^{T})p_{t} \right]\cdot\nabla^{T}\end{split}$ (76) $\displaystyle\begin{split}&=-\nabla\cdot[fp_{t}]+\frac{1-\gamma_{t}}{2}\nabla \cdot[GG^{T}p_{t}]\cdot\nabla^{T}-\frac{1}{2}\nabla\cdot[\tilde{G}\tilde{G}^{T }p_{t}]\cdot\nabla^{T}\\ &+\frac{\gamma_{t}}{2}\nabla\cdot\left[GG^{T}p_{t}\right]\cdot\nabla^{T}+\frac {1}{2}\nabla\cdot\left[\tilde{G}\tilde{G}^{T}p_{t}\right]\cdot\nabla^{T}\end{split}$ (77) $$

With cancellations, we arrive at:

$$ $\displaystyle\frac{\partial p_{t}}{\partial t}$ $\displaystyle=-\nabla\cdot[fp_{t}]+\frac{1}{2}\nabla\cdot[GG^{T}p_{t}]\cdot \nabla^{T}$ (78) $$

which is the FPK equation describing the time evolution of the marginal probability density $p_{t}(x)$ of the solutions of the SDE in Equation 4.

∎

###### Proof.

## Appendix D Proof of Section 3.2

*

###### Proof.

Starting with Theorem 1, let’s define $\tilde{G}\equiv 0$ and $G\equiv g(t)I$, where $g(t)$ is a scalar valued function. These choices lead to following

$$ $\displaystyle\bar{G}$ $\displaystyle=[\gamma_{t}(g(t)I)^{2}]^{\frac{1}{2}}=\sqrt{\gamma_{t}}g(t)I$ (79) $\displaystyle\bar{f}$ $\displaystyle=f-\frac{1}{2}\left(\nabla\cdot[(1-\gamma_{t})g^{2}(t)I]+[(1- \gamma_{t})g^{2}(t)I]\cdot\nabla\ln p_{t}\right)$ (80) $\displaystyle=f-\frac{1}{2}\left([(1-\gamma_{t})g^{2}(t)I]\cdot\nabla\ln p_{t}\right)$ (81) $\displaystyle=f-\frac{(1-\gamma_{t})g^{2}(t)}{2}\nabla\ln p_{t}$ (82) $$

Note that Equation 81 follows from $\nabla\cdot[(1-\gamma_{t})g^{2}(t)I]=0$ since neither $\gamma_{t}$ nor $g(t)$ are functions of $x$.
∎

###### Proof.

## Appendix E Proof of Section 3.2

*

###### Proof.

First note that $f\equiv v(x,t)$ by definition from equation (3) in the paper. Now, again starting with Theorem 1, let’s define $\tilde{G}\equiv\tilde{g}(t)I$ and $G\equiv 0$, where $\tilde{g}(t)$ is a scalar valued function. These choices lead to following

$$ $\displaystyle\bar{G}$ $\displaystyle=[(\tilde{g}(t)I)^{2}]^{\frac{1}{2}}=\tilde{g}(t)I$ (83) $\displaystyle\bar{f}$ $\displaystyle=v(x,t)-\frac{1}{2}\left(\nabla\cdot[-\tilde{g}^{2}(t)I]+[-\tilde {g}^{2}(t)I]\cdot\nabla\ln p_{t}\right)$ (84) $\displaystyle=v(x,t)+\frac{1}{2}\left([\tilde{g}^{2}(t)I]\cdot\nabla\ln p_{t}\right)$ (85) $\displaystyle=v(x,t)+\frac{\tilde{g}^{2}(t)}{2}\nabla\ln p_{t}$ (86) $$

Note that Equation 85 follows from $\nabla\cdot[-\tilde{g}^{2}(t)I]=0$ since $\tilde{g}(t)$ is not a function of $x$.
∎

###### Proof.

## Appendix F Closed form rectified flow expression for the two Gaussian case

Our empirical studies use a two Gaussian toy problem setup.
We state the closed form expression for the rectified flow for this case.
Consider $x_{0}\sim N(\mu_{0},\sigma^{2}_{0}I),x_{1}\sim N(\mu_{1},\sigma^{2}_{1}I)$:

$$ $x_{t}=\alpha_{t}x_{0}+\beta_{t}x_{1},\quad\alpha_{t}>0,\alpha_{0}=1,\alpha_{1} =0,\beta_{t}>0,\beta_{0}=0,\beta_{1}=1$ (87) $$

The marginal density $p_{t}(x_{t})$ is also Gaussian:

$$ $p_{t}(x_{t})=N(x_{t};\alpha_{t}\mu_{0}+\beta_{t}\mu_{1},\alpha_{t}^{2}\sigma_{ 0}^{2}+\beta_{t}^{2}\sigma_{1}^{2})$ (88) $$

We have:

$$ $\displaystyle v(x,t)=\mathbb{E}[x_{1}-x_{0}|x_{t}]\equiv\mathbb{E}_{p(x_{0},x_ {1}|x_{t})}[x_{1}-x_{0}]$ (89) $$

Using the following:

$$ $\displaystyle p(x_{0},x_{1}|x_{t})$ $\displaystyle=\frac{p(x_{t}|x_{0},x_{1})p(x_{0},x_{1})}{p(x_{t})}=\frac{p(x_{t }|x_{0},x_{1})p_{0}(x_{0})p_{1}(x_{1})}{p(x_{t})}$ (90) $\displaystyle p(x_{t}|x_{0},x_{1})$ $\displaystyle=\delta(x_{t}-(1-t)x_{0}-tx_{1})$ (91) $$

and elementary properties of Gaussian and Dirac delta distributions, it can be verified that:

$$ $v(x,t)=\frac{(k_{t}\mu_{1}-x_{t})\alpha_{t}\sigma_{0}^{2}+(x_{t}-k_{t}\mu_{0}) \beta_{t}\sigma_{1}^{2}}{\alpha_{t}^{2}\sigma_{0}^{2}+\beta_{t}^{2}\sigma_{1}^ {2}}$ (92) $$

where $k_{t}=\alpha_{t}+\beta_{t}$.

## Appendix G Toy Gaussian Experiment Details

In the experiments in Section 4.1 we study how various SDEs in Table 1 behave on a toy problem where both $p_{0}\equiv N(-1,0.3)$ and $p_{1}\equiv N(0,1.0)$ are Gaussian.
In this case the marginal distributions $p_{t}$ for Gaussian flow are Gaussian with $p_{t}=N(\mu_{t},\sigma_{t}^{2})$ and the true statistics $\mu_{t},\sigma_{t}^{2}$ can be easily computed.
In addition, the rectified flow $v(x,t)$ is available in closed form (see Appendix F).
The SDEs are simulated backwards in time from $t=1$ with draws from $p_{1}$ using Equation 5.
The drift $\tilde{f}$ and diffusion $\tilde{g}$ terms are calculated using Table 1 by setting $\alpha=1$ and using the closed form $v(x,t)$ from Equation 92.
We simulate 10 trials of 10000 trajectories using Euler-Maruyama discretization with varying number of steps. Estimates for mean $\mu_{t}$ and variance $\sigma_{t}^{2}$ at each timestep for various SDEs are calculated, along with their standard deviation across trials.
Error, calculated as estimate - truth, is then plotted in Figures 2, 3 and 4 for both the mean and the variance estimates along with the KL-Divergence from the true marginal distribution.

## Appendix H ImageNet experiment training/evaluation details

We train two base Rectified flow models to yield $v(x,t)$ at two resolutions of 64 $\times$ 64 and 128 $\times$ 128, on the entire ImageNet training dataset containing roughly 1.2 million images.
Our model is based on the architecture described in Hoogeboom et al. ().
The model is structured such that the lower feature map resolution is 16 $\times$ 16.
Therefore, for 64 $\times$ 64 resolution two downsamplings are performed, while for 128 $\times$ 128 three downsamplings are performed.
The model is trained with SGD using adamw (Kingma & Ba, ; Loshchilov & Hutter, ) with $\beta_{1}=0.9,\beta_{2}=0.99,\epsilon=10^{-12}$ for $1000$ epochs.
We use center crop and left-right flips as the only augmentations.
An exponential moving average, with a decay of $0.9999$, of parameters is used for evaluation.
FIDs are reported over the training dataset with reference statistics computed with center crop but without any augmentation, but with class conditioning.
Samplers were evaluated for all $\alpha\in\{0.0,0.02,0.04,\dots,1.0\}$.

The $64\times 64$ model trained for 500 epochs in 4 days, 8 hours on $8\times 8$ Google Cloud TPUs v3.
The $128\times 128$ model trained for 500 epochs in 4 days, 20 hours on $8\times 8$ Google Cloud TPUs v3.

## Appendix I Example Implementation

See Figure 8 for an example implementation of the NonSingular sampler.

## Appendix J Additional Experimental Results

Figure: (a) 64 $\times$ 64
Refer to caption: x10.png

### J.1 FID vs α 𝛼 \alpha italic_α

In Figure 7 we show a larger range FID for various samplers compared in Figure 5.
We observe that the Singular sampler tends to perform well only at low scales with an intriguing behavior for higher scales where the FID starts to improve again after worsening significantly.

### J.2 Effect of diffusion coefficient magnitude on samples

We qualitatively visualize the effect of diffusion coefficient magnitude for the three SDEs discussed in the main paper.
Figure 9 visualizes samples for the constant diffusion term SDE as a function increasing coefficient magnitude.
Each column is a different sample starting with the same random draw from $p_{1}(x_{1})$.
Each row corresponds to a different magnitude for the diffusion coefficient $g(t)$.
Figures 10 and 11 visualize samples with a similar scheme for the non-singular and singular SDE.

### J.3 Classifier-free guidance samples

Figure 12 shows additional samples using classifier-free guidance with NonSingular at different values of $\alpha$ and $\lambda$, as in Figure 1.

Figure: Figure 8: NonSingular Sampler written in JAX.

Figure: Figure 9: Constant samples with increasing scaling $\alpha$. Each row displays samples at a particular $g$-scale, from 0 0 increasing to $1$ in the increments of $0.1$ from top to bottom. Sampling for each columns starts off with the same initial noise image and conditioning class.
Refer to caption: x12.png

Figure: Figure 10: NonSingular samples with increasing scaling $\alpha$. Each row displays samples at a particular $g$-scale, from 0 0 increasing to $1$ in the increments of $0.1$ from top to bottom. Sampling for each columns starts off with the same initial noise image and conditioning class.
Refer to caption: x13.png

Figure: Figure 11: Singular samples with increasing scaling $\alpha$. Each row displays samples at a particular $g$-scale, from 0 0 increasing to $1$ in the increments of $0.1$ from top to bottom. Sampling for each columns starts off with the same initial noise image and conditioning class.
Refer to caption: x14.png

Figure: Figure 12: Stochastic sampling improves diversity at all classifier-free guidance levels. Additional results as in Figure 1, described in Section 4.2.
Refer to caption: x15.png

## Peer Reviews

confidence: 4
contribution: 1
initial_rating: 3
presentation: 3
questions: Can the stochastic sampling be generalized to 2-Rectified Flows (i.e. after rectification / training with a non-independent coupling)?
soundness: 3
strengths: Clearly written.
summary: This paper discusses the case of using linear interpolants for training deterministic flow models. The paper motivates using stochastic sampling instead of deterministic sampling.
weaknesses: Lack of novelty. The conversion between SDEs with different state-independent diffusion coefficients (including zero for ODE) through the Fokker-Planck equation is very well-understood now. While Theorem 1 does cover a general case (state-dependent), this general case is not applied. An example of a paper that discussed stochastic sampling and classifier-free guidance is SiT [1].


[1] https://arxiv.org/abs/2401.08740

---

confidence: 5
contribution: 1
initial_rating: 3
presentation: 3
questions: I think it's important for the author's to clarify what the novel contribution is. Could the authors address that?
soundness: 3
strengths: The paper provides a useful empirical exploration of the trade-offs/biases that emerge depending on how the sampling SDE is discretized and which diffusion coefficient is used, e.g. in Table 2. 

Section 4.2 itemizes useful takeaways for the tradeoffs in SDE-based sampling, and provides nice experiments to demonstrate how diffusivity alleviates e.g. sample degeneracy.
summary: This paper describes how deterministic dynamical transport models can be changed into stochastic samplers without retraining when the base distribution is Gaussian. They then provide a series of experiments exploring the benefits and trade-offs of using this velocity model under different settings of the SDE: extent of diffusivity and time-dependent schedule of its magnitude.
weaknesses: First off, thank you for your submission! It would be great if the authors could address the following *primary concern* about the work

Corollary 1 is already well known -- it is one of the main points of the interpolant procedure, and is the essential knob explored in 2303.08797, Corollary 2.18. This was then extensively tested (the variety of potential diffusion coefficients from the same velocity model) in 2401.08740. It's unclear to the reviewer how what this paper is proposing is different from these essential elements. In addition Corollary 1.2 in the given paper is already noticed by Song et al 2011.13456. The authors make remarks about these works in the "related work" section but do not seem to recognize this connection.

Because this is the main point of the paper given in this work, it is unclear to the reviewer what this paper sees as its main contribution to the literature.

---

confidence: 5
contribution: 1
initial_rating: 3
presentation: 3
questions: I do not have any further questions. And also, feel free to let me know if I misunderstood anything!
soundness: 3
strengths: 1. This paper is well-written, with thorough ablation studies that carefully compare deterministic and stochastic samplers, which I greatly appreciate.

2. The paper is well-written and easy to follow.

3. The motivation is clear.
summary: The paper introduces a method to enhance deterministic flow models, such as rectified flows, by converting their underlying ordinary differential equations (ODEs) into a family of stochastic differential equations (SDEs) with the same marginal distributions. This approach allows for a continuous transition between deterministic and stochastic sampling, using both the flow field and the score function. The method addresses limitations of deterministic samplers, improving their performance on tasks like Gaussian modeling and large-scale ImageNet generation. Additionally, it offers a mechanism to control the diversity of generated samples, providing more flexibility in the sampling process.
weaknesses: 1.To the best of my knowledge, the paper lacks novelty. The core result of Theorem 1 appears straightforward to derive by combining existing results (Equation 4 and Appendix D in [1], Doob's h-transform, and Equation 37 in [2]).

Let me explain more. The Theorem 1 can be understood as
 
$\bar{f}=f-\frac{1}{2}[\nabla(GG^T-(\gamma_t GG^T+\tilde{G}\tilde{G}^T)]-\frac{1}{2}[GG^T-(\gamma_t GG^T+\tilde{G}\tilde{G}^T)]\nabla \log p$, 

$=f-\frac{1}{2}[\nabla(GG^T-\bar{G}\bar{G}^T)]-\frac{1}{2}[GG^T-\bar{G}\bar{G}^T)]\nabla \log p$, 
and

$\bar{G}=[\gamma_t GG^T+\tilde{G}\tilde{G}^T]^{1/2}$

According to [1], The original SDE eq.4 can be written as,
$dx_t= [f-\frac{1}{2}\nabla[GG^T]-\frac{1}{2}GG^T\nabla\log p]dt$

Ok then, all the tricks happen in the diffusion term of FPK. You can easily extend eq.28 in [1] with $x_t$ dependent $G(x,t)$ (it may give even more flexibility?), and absorbing the additional term into the drift term as shown in eq.37 in [2]. 

My understanding is that, once an equivalence is established between flow/bridge models and diffusion models, the remaining work can be readily extended. 

However, I want to acknowledge the authors' efforts; despite the challenges of the problem, the community still benefits from having researchers dedicated to addressing it.

[1]Zhang, Qinsheng, and Yongxin Chen. "Fast sampling of diffusion models with exponential integrator." arXiv preprint arXiv:2204.13902 (2022).

[2]Song, Yang, et al. "Score-based generative modeling through stochastic differential equations." arXiv preprint arXiv:2011.13456 (2020).

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
questions: None in addition to the points raise above.
soundness: 4
strengths: - The paper's central observation, given in Theorem 1, that a learned deterministic flow can be converted into a family of stochastic flows targeting the same distribution is a neat theoretical result. It also has obvious practical relevance since, as the authors claim, stochastic sampling methods tend to perform better in practice.
- The presentation of the paper is extremely clear throughout, and I had no trouble understanding the authors arguments and contributions. Thank you for making it easy to engage with your work!
- The authors provide a compelling empirical demonstration of the practical success of their method, on both toy and real-world data distributions.

Overall, I believe this to be a very technically solid paper worthy of acceptance to ICLR. While the core idea of the contribution is not totally novel - the idea that one can convert between SDE and ODE sampling methods via rearranging the Fokker-Plank PDEs governing the marginals is implicit in previous work - this is the first work I'm aware of that demonstrates this practically in the ODE to SDE direction. This is a useful idea and I find the authors' empirical demonstrations of its relevance compelling.
summary: This paper introduces a way to take a trained deterministic flow model and find a family of stochastic flow models that induce the same marginals and which can be simulated using the original learned deterministic flow. This allows one to sample from a much broader range of flow models targeting the same distribution without retraining the network.

The authors claim and then demonstrate that this can be desirable, since stochastic flows have been demonstrated to have better empirical performance in many settings. In particular, the authors show that their stochastic flows have perform better at estimating the variance of the target distribution in a toy setting, and have better FID score when targeting the ImageNet distribution.
weaknesses: There are no significant weaknesses with this paper that I would want to raise. I have just a couple of minor editorial points:
- $\bar G$ and $\tilde G$ are very hard to distinguish in print, and you might want to consider alterantive notation here.
- It took me a while to understand the plots in Figures 1, 2 and 3; for example what exactly was meant by "error in mean" and "error in variance", and which distributions you were measuring the KL between. If there's an easy way to clarify this, that might be worth adding.

Meta-review: This paper presents a stochastic sampling algorithm for deterministic flow (e.g., Rectified flow) models. The main contribution is Theorem 1 which identifies a class of SDEs that share the same marginal distributions. It includes the widely used marginal-equivalent SDE as a special case (Corollary 1.1). The main criticism is that only the algorithm associated with the special case is tested in experiments and used in practice. It is not clearly what the benefits Theorem 1 brings over Corollary 1.1. Reviewer 7s8u also points out that Theorem 1 is a simple exercise in the subject of SDE.

Decision: Reject