1 Introduction↩︎

Generative modeling considers the problem of approximating and sampling from a probability distribution. Normalizing flows, which have emerged as a competitive generative modeling method, construct an invertible and efficiently differentiable mapping between a fixed (e.g., standard normal) distribution and the data distribution [1]. While original normalizing flow work specified this mapping as a static composition of invertible modules, continuous normalizing flows (CNFs) express the mapping by a neural ordinary differential equation (ODE) [2]. Unfortunately, CNFs have been held back by difficulties in training and scaling to large datasets [2]–[4].

Meanwhile, diffusion models, which are the current state of the art on many generative modeling tasks [5]–[8], approximate a stochastic differential equation (SDE) that transforms a simple density to the data distribution. Diffusion models owe their success in part to their simple regression training objective, which does not require simulating the SDE during training. Recently, [9]–[11] showed that CNFs could also be trained using a regression of the ODE’s drift similar to training of diffusion models, an objective called flow matching (FM). FM was shown to produce high-quality samples and stabilize CNF training, but made the assumption of a Gaussian source distribution, which was later relaxed in generalizations of FM to more general manifolds [12], arbitrary sources [13], and couplings between source and target samples that are either part of the input data or are inferred using optimal transport. The first main contribution of the present paper is to propose a unifying conditional flow matching (CFM) framework for FM models with arbitrary transport maps, generalizing existing FM and diffusion modeling approaches (1).

A major drawback of both CNF (ODE) and diffusion (SDE) models compared to other generative models (e.g., variational autoencoders [14], (discrete-time) normalizing flows, and generative adversarial networks [15]), is that integration of the ODE or SDE requires many passes through the network to generate a high-quality sample, resulting in a long inference time. This drawback has motivated work on enforcing an optimal transport (OT) property in neural ODEs [4], [11], [16]–[18], yielding straighter flows that can be integrated accurately in fewer neural network evaluations. Such regularizations have not yet been studied for the full generality of models trained with FM-like objectives, and their properties with regard to solving the dynamic optimal transport problem were not empirically evaluated. Our second main contribution is a variant of CFM called optimal transport conditional flow matching (OT-CFM) that approximates dynamic OT via CNFs. We show that OT-CFM not only improves the efficiency of training and inference, but also leads to more accurate OT flows than existing neural OT models based on ODEs [16], [17], SDEs [19], [20], or input-convex neural networks [21]. Furthermore, an entropic variant of OT-CFM can be used to efficiently train a CNF to match the probability flow of a Schrödinger bridge. We show that in the case where the true transport plan is sampleable, our methods approximate the dynamic OT maps and Schrödinger bridge probability flows for arbitrary source and target distributions with simulation-free training.

In summary, our contributions are:

1. We introduce a generalized formulation of the recent conditional flow matching framework (3.1), and prove its correctness encompassing many existing flow matching methods [9], [10], [22] (See 1.)

2. We consider a special case of CFM that draws source and target samples according to an optimal transport plan, allowing us to solve the dynamic OT and Schrödinger bridge problems in a simulation-free way, using only static OT maps between marginal distributions. We show that efficient minibatch approximations to the OT map still yield correct solutions to the generative modeling problem while incurring a low detriment to the dynamic OT solution (3.2).

3. We evaluate CFM and OT-CFM in experiments on single-cell dynamics, image generation, unsupervised image translation, and energy-based models. We show that the OT-CFM objective leads to more efficient training and decreases inference time while finding better approximate solutions to the dynamic OT and Schrödinger bridge problems. For high-dimensional image generation, we also propose improved and reproducible training practices for flow-based models that significantly improve the performance of algorithms from past work (5).

4. We release a Python package, torchcfm, that unifies new and existing algorithms for training flow-based generative models under a shared interface and provides implementations of our main experiments. The Python code is available at https://github.com/atong01/conditional-flow-matching.

2 Background: Optimal transport and neural ODEs↩︎

Throughout the paper, we consider the setting of a pair of data distributions over \(\mathbb{R}^d\) with (possibly unknown) densities \(q(x_0)\) and \(q(x_1)\) (also denoted \(q_0\), \(q_1\)). Generative modeling considers the task of fitting a mapping \(f\) from \(\mathbb{R}^d\) to \(\mathbb{R}^d\) that transforms \(q_0\) to \(q_1\), that is, if \(x_0\) is distributed with density \(q_0\) then \(f(x_0)\) is distributed with density \(q_1\). This includes both the typical case when \(q_0\) is an easily sampled density, such as a Gaussian, and the case when \(q_0\) and \(q_1\) are empirical data distributions available as finite sets of samples.

2.1 ODEs and probability flows↩︎

A smooth¹ time-varying vector field \(u : [0, 1] \times \mathbb{R}^d \to \mathbb{R}^d\) defines an ordinary differential equation: \[\label{eq:ode} dx = u_t(x)\,dt,\tag{1}\] where we use the notation \(u_t(x)\) interchangeably with \(u(t,x)\). Denote by \(\phi_t(x)\) the solution of the ODE (1 ) with initial condition \(\phi_0(x)=x\); that is, \(\phi_t(x)\) is the point \(x\) transported along the vector field \(u\) from time \(0\) up to time \(t\).

Given a density \(p_0\) over \(\mathbb{R}^d\), the integration map \(\phi_t\) induces a pushforward \(p_t:=[\phi_t]_\#(p_0)\), which is the density of points \(x\sim p_0\) transported along \(u\) from time \(0\) to time \(t\). The time-varying density \(p_t\), viewed as a function \(p:[0,1]\times\mathbb{R}^d\to\mathbb{R}\), is characterized by the well-known continuity equation: \[\label{eq:continuity} \frac{\partial p}{\partial t}=-\nabla\cdot(p_tu_t)\tag{2}\] and the initial conditions \(p_0\). Under these conditions, \(u\) is said to be a probability flow ODE for \(p\), and \(p\) is the (marginal) probability path² generated by \(u\).

2.1.0.1 Approximating ODEs with neural networks.

Suppose the probability path \(p_t(x)\) and the vector field \(u_t(x)\) generating it are known and \(p_t(x)\) can be tractably sampled. If \(v_\theta(\cdot,\cdot):[0,1]\times\mathbb{R}^d\to\mathbb{R}^d\) is a time-dependent vector field parametrized as a neural network with weights \(\theta\), \(v_\theta\) can be regressed to \(u\) via the flow matching (FM) objective: \[\label{eq:fm} {\mathcal{L}}_{\rm FM}(\theta) := \mathbb{E}_{t\sim {{\mathcal{U}}(0,1)},x\sim p_t(x)} \| v_\theta(t, x) - u_t(x) \|^2.\tag{3}\] [9] used a version of this objective with a stochastic regression target to fit ODEs that map a Gaussian density \(q_0\) to a target \(q_1\). However, this objective becomes intractable for general source and target distributions. In 3, we develop generalizations that allow more flexible and efficient generative modeling.

2.1.0.2 The case of Gaussian marginals.

Consider the special case of an ODE whose marginal densities are Gaussian: \(p_t(x) = \mathcal{N}(x \mid \mu_t, \sigma_t^2)\). While the ODE that generates these marginal densities is not unique, one of the simplest is the one that satisfies \[\label{eq:flow} \phi_t(x_0) = \mu_t + \sigma_t \left(\frac{x_0-\mu_0}{\sigma_0}\right),\tag{4}\] which is unique by the following theorem.

2.2 Static and dynamic optimal transport↩︎

The (static) optimal transport problem seeks a mapping from one measure to another that minimizes a displacement cost. The case of greatest interest is the \(2\)-Wasserstein distance between distributions \(q_0\) and \(q_1\) on \(\mathbb{R}^d\) with respect to the Euclidean distance cost \(c(x,y)=\|x-y\|\). The corresponding optimization problem is \[\label{eq:ot} W(q_0, q_1)_2^2 = \inf_{\pi \in \Pi} \int_{\mathbb{R}^d \times \mathbb{R}^d} c(x, y)^2\,d\pi(x, y),\tag{6}\] where \(\Pi\) denotes the set of all joint probability measures on \(\mathbb{R}^d\times\mathbb{R}^d\) whose marginals are \(q_0\) and \(q_1\). For compactly supported distributions and for the ground cost \(c(x,y)=\|x-y\|\), the set of solutions of (6 ) is not empty [23], and \(W_2\) is a metric on the space of probability distributions on \(\mathbb{R}^d\) with finite second moment.

The dynamic form of the 2-Wasserstein distance is defined by an optimization problem over vector fields \(u_t\) that transform one measure to the other: \[\label{eq:dot} W(q_0, q_1)^2_2 = \inf_{p_t, u_t} \int_{\mathbb{R}^d} \int_0^1 p_t(x) \|u_t(x) \|^2\,dt\,dx,\tag{7}\] with \(p_t \ge 0\) and subject to the boundary conditions \(p_0 = q_0\), \(p_1 = q_1\), and the continuity equation (2 ). The equivalence between the dynamic and static optimal transport formulations was first proven in [24] under the assumptions that \(q_0\) and \(q_1\) are compactly supported distributions with bounded density. We refer to [25] for a recent overview on optimal transport the relation between the two formulations.

[16], [17] showed that CNFs with \(L^2\) regularization approximate dynamic optimal transport. For general marginals, however, these models required integrating over and backpropagating through tens to hundreds of function evaluations, resulting in both numerical and efficiency issues. We aim to avoid these issues by directly regressing to the vector field in a simulation-free way.

Optimal transport is also related to the Schrödinger bridge (SB) problem [26]. We show in 3.2.4 that a variant of the algorithm we propose recovers the probability flow of the solution to a SB problem with a Brownian motion reference process.

3 Conditional flow matching: ODEs from static couplings↩︎

3.1 Vector fields generating mixtures of probability paths↩︎

Suppose that the marginal probability path \(p_t(x)\) is a mixture of probability paths \(p_t(x|z)\) that vary with some conditioning variable \(z\), that is, \[\label{eq:ppath} p_t(x) = \int p_t(x | z) q(z)\, dz,\tag{8}\] where \(q(z)\) is some distribution over the conditioning variable. If the probability path \(p_t(x|z)\) is generated by the vector field \(u_t(x|z)\) from initial conditions \(p_0(x|z)\) (see 2.1), then the vector field \[\label{eq:mvec} u_t(x) := \mathbb{E}_{q(z)} \frac{u_t(x | z) p_t(x | z)}{p_t(x)}\tag{9}\] generates the probability path \(p_t(x)\), under some mild conditions:

All proofs appear in 7. This result extends [9] to general conditioning variables and delineates some minor conditions on \(q(z)\).

3.1.0.1 A regression objective for mixtures.

We are interested in the case where conditional probability paths \(p_t(x|z)\) and vector fields \(u_t(x|z)\) are known and have a simple form, and we wish to recover the vector field \(u_t(x)\), defined by (9 ), that generates the probability path \(p_t(x)\). Exact computation via (9 ) is generally intractable, as the denominator \(p_t(x)\) is defined by an integral (8 ) that may be difficult to evaluate. Instead, we develop an unbiased stochastic objective for regression of a learned vector field to \(u_t(x)\), which generalizes the unconditional flow matching objective (3 ).

Let \(v_\theta(\cdot, \cdot):[0,1]\times\mathbb{R}^d\to\mathbb{R}^d\) be a time-dependent vector field parametrized as a neural network with weights \(\theta\). Define the conditional flow matching (CFM) objective: \[\label{eq:CFM} {\mathcal{L}}_{\rm CFM}(\theta) := \mathbb{E}_{t, q(z), p_t(x | z)} \|v_\theta(t, x) - u_t(x | z)\|^2.\tag{10}\] The CFM objective describes how to regress against the marginal vector field \(u_t(x)\) given by (9 ) with access only to samples from the conditional probability path \(p_t(x|z)\) and conditional vector fields \(u_t(x|z)\). This is formalized in the following theorem.

The CFM objective is useful when the marginal vector field \(u_t(x)\) is intractable but the conditional vector field \(u_t(x | z)\) is tractable. As long as we can efficiently sample from \(q(z)\) and \(p_t(x|z)\) and calculate \(u_t(x|z)\), we can use this stochastic objective to regress \(v_\theta\) to the marginal vector field \(u_t(x)\).

We discuss the variance arising from the stochastic regression target, and ways to reduce it, in 9.1, 3, 4.

3.2 Sources of conditional probability paths↩︎

In this section, we introduce several forms of CFM depending on the choices of \(q(z)\), \(p_t(\cdot | z)\), and \(u_t(\cdot | z)\). All of the CFM variants and related objectives from prior work are summarized in 1.

• 3.2.1: We interpret the algorithm of [9] (FM from a Gaussian) as a special case of CFM.

• 3.2.2: We relax the Gaussian source requirement by letting the condition \(z\) be a pair \((x_0, x_1)\) of an initial and a terminal point. In the basic form of CFM (I-CFM), we take the distribution \(q(z)\) to equal \(q(x_0)q(x_1)\), allowing generative modeling with an arbitrary source distribution.

• 3.2.3: We consider joint distributions \(q(z)=q(x_0,x_1)\) that are given by minibatch optimal transport maps, causing the learned flow to be an (approximate) OT flow.

• 3.2.4: we consider \(q(z)\) given by an entropy-regularized OT map and show that the CFM objective with this \(q(z)\) solves the Schrödinger bridge problem.

3.2.1 FM from the Gaussian↩︎

[9] considered the problem of unconditional generative modeling given a training dataset. Identifying the condition \(z\) with a single datapoint \(z \mathrel{\vcenter{:}}= x_1\), and choosing a smoothing constant \(\sigma>0\), one sets \[\begin{align} \label{eq:32fm95density95vectorfield} p_t(x | z) &= \mathcal{N}(x\mid t x_1, (t \sigma - t + 1)^2), \\ u_t(x | z) &= \frac{x_1 - (1 - \sigma) x}{1 - (1 - \sigma) t}, \end{align}\tag{11}\] which is a probability path from the standard normal distribution (\(p_0(x|z)={\mathcal{N}}(x;0,{\boldsymbol{I}})\)) to a Gaussian distribution centered at \(x_1\) with standard deviation \(\sigma\) (\(p_1(x|z)={\mathcal{N}}(x;x_1,\sigma^2)\)). If one sets \(q(z)=q(x_1)\) to be the uniform distribution over the training dataset, the objective introduced by [9] is equivalent to the CFM objective (10 ) for this conditional probability path.

We emphasize that although the conditional probability path \(p_t(x|z)\) is an optimal transport path from \(p_0(x|z)\) to \(p_1(x|z)\), the marginal path \(p_t(x)\) is not in general an OT path from the standard normal \(p_0(x)\) to the data distribution \(p_1(x)\).

3.2.2 Basic form of CFM: Independent coupling↩︎

In the basic form of CFM (I-CFM), we identify \(z\) with a pair of random variables, a source point \(x_0\) and a target point \(x_1\), and set \(q(z)=q(x_0)q(x_1)\) to be the independent coupling. We let the conditionals be Gaussian flows between \(x_0\) and \(x_1\) with standard deviation \(\sigma\), defined by \[\begin{align} p_t(x | z) &= \mathcal{N}(x\mid t x_1 + (1 - t) x_0, \sigma^2),\tag{12}\\ u_t(x | z) &= x_1 - x_0. \tag{13} \end{align}\] We note that the formulation of \(u_t(x | z)\) follows from an application of 1 to the conditional probability path with \(\mu_t = t x_1 + (1 - t) x_0\) and \(\sigma_t = \sigma\). Furthermore, we note that \(p_t(x | z)\) is efficiently samplable and \(u_t\) is efficiently computable, thus gradient descent on \({\mathcal{L}}_{\rm CFM}\) is also efficient. For this choice of \(z\), \(p_t(\cdot | z)\), and \(u_t(\cdot | z)\), we know the marginal boundary probabilities approach \(q_0\) and \(q_1\) respectively as \(\sigma \to 0\). This is made explicit in the following Proposition:

In particular, as \(\sigma\to0\), the marginal vector field \(u_t(x)\) approaches one that transports the distribution \(q(x_0)\) to \(q(x_1)\) and can thus be seen as a generative model of \(x_1\). Note that there is no requirement for \(q(x_0)\) to be Gaussian. Conditioning on \(x_0\) allows us to generalize flow matching to arbitrary source distributions with intractable densities. In the case of FM from a Gaussian, while each conditional flow is the dynamic optimal transport flow from \({\mathcal{N}}(x_0,\sigma^2)\) to \({\mathcal{N}}(x_1,\sigma^2)\), the marginal vector field \(u_t(x)\) is not necessarily an OT flow.

3.2.2.1 Connection with related Rectified Flow and Stochastic interpolants methods.

We note that I-CFM is closely related to the algorithms proposed by [10], [11]. In the case where the conditional probability path \(p_t\) is a Dirac (i.e., \(\sigma = 0\)), I-CFM is equivalent to [11]. Furthermore, if we consider the Gaussian mean \(\mu_t = \cos(\frac{1}{2} \pi t) x_0 + \sin (\frac{1}{2} \pi t) x_1\) instead of the linear interpolation, I-CFM would be equivalent to the variance preserving stochastic interpolant in [10], which has also been further generalized.

3.2.2.2 Connection to FM from the Gaussian.

There exists a set of conditional probability paths conditioned on \(x_1\) and \(x_0 \sim \mathcal{N}(0, 1)\) that have an equivalent probability flow to the marginal \(p_t\) of flow matching from the Gaussian (3.2.1), which is only conditioned on \(x_1\). These paths are defined by \[p_t(x | z) = \mathcal{N}(x\mid t x_1 + (1 - t) x_0, (\sigma t)^2 + 2 \sigma t (1-t)).\label{eq:scfm:pt}\tag{14}\] 2 states an equivalence between I-CFM with these paths and the objective from 3.2.1.

3.2.3 Optimal transport CFM↩︎

In this section, we present our second main contribution. The formulation in the previous section can readily be generalized to distributions \(q(z)=q(x_0,x_1)\) in which \(x_0\) and \(x_1\) are not independent, as long as \(q(z)\) has marginals \(q(x_0)\) and \(q(x_1)\). Therefore, we propose to set \(q(z)\) to be the 2-Wasserstein optimal transport map \(\pi\) achieving the infimum in (6 ), namely, \[q(z) \mathrel{\vcenter{:}}= \pi(x_0, x_1).\label{eq:otcfm:qz}\tag{15}\] In this case, \(z\) is still a tuple of points, but instead of \(x_0,x_1\) being sampled independently from their marginal distributions, they are sampled jointly according to the optimal transport map \(\pi\). We call this method optimal transport CFM (OT-CFM). If one uses the \(p_t(x | z)\) defined by (12 ) and \(u_t(x | z)\) in (13 ), OT-CFM is equivalent to dynamic optimal transport in the following sense.

We consider two cases: (1) when the data set is small enough and we know the static optimal transport plan (e.g.single cell data). (2) when the data is too large (or continuous) (e.g.image data) and the static OT plan is computationally infeasible to determine exactly. In the first case we are able to extend the transport map to unseen data similar to the task presented in [31]. In the second case we show an approximation with minibatch OT improves over a random plan in terms of generative modelling performance and training time.

3.2.3.1 Minibatch OT approximation.

For large datasets, the transport plan \(\pi\) can be difficult to compute and store due to OT’s cubic time and quadratic memory complexity in the number of samples [16], [32]. Therefore, we rely on a minibatch OT approximation similar to [33]. Although minibatch OT incurs an error relative to the exact OT solution, it has been successfully used in many applications like domain adaptation or generative modeling [34], [35]. Specifically, for each batch of data \((\{x_0^{(i)}\}_{i=1}^B,\{x_1^{(i)}\}_{i=1}^B)\) seen during training, we sample pairs of points from the joint distribution \(\pi_{\rm batch}\) given by the OT plan between the source and target points in the batch. (The OT batch size need not match the optimization batch size, but we keep them equal for simplicity.) Thus, we solve a minibatch approximation of dynamic optimal transport. However, when the OT batch size equals the support size of \((q_0, q_1)\), we recover exact OT and therefore, by [lem:otcfm], learn the exact dynamic optimal transport. We show empirically that the batch size can be much smaller than the full dataset size and still give good performance, which aligns with prior studies [36], [37]. Concurrently, a similar framework and theoretical results appeared in [13].

3.2.4 Schrödinger bridge CFM↩︎

Recently, there has been significant effort in learning diffusion models with general source distributions, formulated as a Schrödinger bridge problem [19], [20], [38] or bridge matching [39]–[41]. Here we show that SB-CFM, an entropic variant of OT-CFM, can be used to train an ODE to match the probability flow of a Schrödinger bridge with a Brownian motion reference process.

Let \(p_{\rm ref}\) be the standard Wiener process scaled by \(\sigma\) with initial-time marginal \(p_{\rm ref}(x_0)=q(x_0)\). The Schrödinger bridge problem [42] seeks the process \(\pi\) that is closest to \(p_{\rm ref}\) while having initial and terminal marginal distributions specified by the data distribution \(q(x_0)\) and \(q(x_1)\): \[\label{eq:schrodinger95bridge} \pi^*:=\mathop{\mathrm{arg\,min}}_{\pi(x_0)=q(x_0),\pi(x_1)=q(x_1) }{\rm KL}(\pi\,\|\,p_{\rm ref}).\tag{19}\] We define the joint distribution \[\label{eq:sbcfm:qz} q(z) \mathrel{\vcenter{:}}= \pi_{2\sigma^2}(x_0, x_1)\tag{20}\] where \(\pi_{2\sigma^2}\) is the solution of the entropy-regularized optimal transport problem [32] with cost \(\|x_0-x_1\|\) and entropy regularization \(\lambda=2 \sigma^2\) (see (27 ) for the background on entropic OT). We set the conditional path distribution to be a Brownian bridge with diffusion scale \(\sigma\) between \(x_0\) and \(x_1\), with probability path and generating vector field \[\begin{align} p_t(x \mid z) &= \mathcal{N}( x \mid t x_1 + (1 - t) x_0, t(1-t)\sigma^2) \tag{21} \\ u_t(x \mid z) &= \frac{1-2t}{2t(1-t)}(x - ( t x_1 + (1-t)x_0) ) + (x_1 - x_0) \tag{22}, \end{align}\] where \(u_t\) is computed by (5 ) as the vector field generating the probability path \(p_t(x|z)\). The marginal coupling \(\pi_{2\sigma^2}\) and \(u_t(x|z)\) define \(u_t(x)\), which is approximated by the regression objective in 7. The solution of the SB is known to be the map which is the solution of the entropically-regularized OT problem, motivating the next proposition.

While we define SB-CFM with an entropic regularization coefficient of \(\varepsilon=2\sigma^2\), the flow still matches the marginals for any choice of \(\varepsilon\). Interestingly, we recover OT-CFM when \(\varepsilon \to 0\) and I-CFM when \(\varepsilon \to \infty\). A similar result was proven in a concurrent work [13].

4 Related work↩︎

4.0.0.1 Simulation-free continuous-time modeling.

Simulation-free training is common in stochastic flow models where backpropagating through the simulation is numerically challenging and has high variance [44]. While these diffusion models have recently achieved exceptional generative performance on many tasks [5], [8], [27], [28], [45]–[47], their simulation requires an inherently costly SDE simulation with many follow-up works to improve inference efficiency [48]–[52]. These methods generally consider a simple Gaussian diffusion process, and do not consider generalizing the source distribution. Other works consider general source distributions but this makes optimization and inference more challenging, needing multiple iterations or other tricks to perform well [19], [20], [53].

Prior work considering simulation-free training of CNFs considers algorithms that are equivalent to CFM with Gaussian source distribution [9], [54], [55] or independent samples from \(q_0\), \(q_1\) [10], [56], [57]. Recent work also studies Schrödinger bridges from unpaired samples [58] and regularization of flows using dynamic OT [18]. We also note the work [13], concurrent with the preprint version of this paper. Other concurrent works explore various solutions to approximate Schrödinger bridges [59]–[61].

4.0.0.2 Dynamic optimal transport.

There are a variety of methods that consider dynamic OT between continuous distributions with neural networks; however, these require constrained architectures [21], [62], [63] or use a regularized CNF, which is challenging to optimize [4], [16], [17], [64]. With our work it is possible to achieve optimal transport flows without either of these constraints.

5 Experiments↩︎

In this section we empirically evaluate the I-CFM, OT-CFM, and SB-CFM objectives, as well as algorithms from prior work, with respect to both optimal transport and generative modeling criteria. All experiment details can be found in 11.

5.1 Low-dimensional data: Optimal transport and faster convergence↩︎

We evaluate how well various models perform dynamic optimal transport and generative modeling in low dimensions. We train ODEs mapping between four pairs of two-dimensional datasets: between a standard Gaussian and , , and and between and .

5.1.0.1 OT-CFM approximates dynamic OT.

To measure how well a model solves the OT problem we use normalized path energy (NPE), defined via the 2-Wasserstein distance as \(\mathrm{NPE}(v_\theta) = |{\rm PE}(v_\theta) - W^2_2(q_0, q_1)| / W^2_2(q_0, q_1)\), where the path energy (PE) is \(\mathrm{PE}(v_\theta) = \mathbb{E}_{x(0) \sim q(x_0)} \int_0^1 \|v_\theta(t, x(t))\|^2 dt\). 2 summarizes our results showing that OT-CFM flows generalize better to the test set and are very close to the dynamic OT paths as measured by normalized path energy. We find transforming \(\leftrightarrow\) to be particularly challenging to learn for I-CFM as compared to OT-CFM; the learned paths are depicted in 1 (bottom). Although OT-CFM uses a minibatch OT map, we find that OT-CFM requires surprisingly small batches to approximate the OT map well, suggesting some generalization advantages of the network optimization (9).

5.1.0.2 OT-CFM yields faster training.

By conditioning on minibatch optimal transport flows, OT-CFM is substantially easier to train, which we posit is due to the variance reduction of the conditional flow. In 3 (left), we evaluate the performance over time of OT-CFM against CFM and FM objectives. For the same number of steps OT-CFM has better performance on the validation set. In 3, we compare the training times for various Neural OT methods whose performance can be seen in 2. Simulation-free optimization is significantly faster to train with equal or superior performance.

5.1.0.3 OT-CFM yields faster inference.

We next evaluate the quality of samples during inference time. In 3 (right), we compare the quality of samples for different number of function evaluations (NFEs) across different flow matching objectives. In this experiment we sample from the source distribution test set and simulate the ODE over time for different solvers. We find that OT-CFM consistently requires fewer evaluations to achieve the same quality and achieves better quality with the same NFEs. This is consistent with previous work, which found OT paths lead to faster, higher quality inference in regularized CNFs [4], [17] and flow matching vs.standard variance-preserving and variance-exploding probability paths [9].

5.1.0.4 SB-CFM reproduces Schrödinger bridge flows.

There are a number of methods which theoretically converge to a Schrödinger bridge between two datasets. In ¿tbl:tab:sb95comparison? we compare SB-CFM and the diffusion Schrödinger bridge (DSB) method introduced in [19] on the quality of the learnt Schrödinger bridges based on the average 2-Wasserstein distance to ground truth Schrödinger bridge samples over 18 time steps. Furthermore, SB-CFM is also significantly faster than DSB (3).

5.2 Application to single-cell interpolation↩︎

As a specific application, we consider the task of single-cell trajectory interpolation. In this task we use leave-one-out validation over the timepoints. From times data at times \([0,t-1], [t+1,T]\) we try to interpolate its distribution at time \(t\) following the setup of [16], [64], [65]. Low error means we model individual cells well, which is useful in a number of downstream tasks such as gene regulatory network inference [66], [67]. Following [68], we repurpose the CITE-seq and Multiome datasets from a recent NeurIPS competition for this task [69]. We also include the Embryoid body data from [16], [70]. ¿tbl:tab:eb95comparison? shows the average earth mover’s distance (1-Wasserstein) on left–out timepoints for three datasets. On all three datasets OT-CFM outperforms other methods and baselines on average.

5.3 High-dimensional data: Lower-cost training and inference↩︎

We perform an experiment on unconditional CIFAR-10 generation from a Gaussian source to examine how OT-CFM performs in the high-dimensional image setting. We use a similar setup to that of [9], including the time-dependent U-Net architecture from [71] that is commonly used in diffusion models. We were not able to reproduce the results reported from [9] with the parameters specified in the paper. ⁴ Therefore, we selected different training hyperparameters.

The main differences with [9] are that we use a constant learning rate, set to \(2\times10^{-4}\), instead of a linearly decreasing one (from \(5\times10^{-4}\) to \(10^{-8}\)). To prevent training instabilities and variance, we clip the gradient norm to 1 and rely on exponential moving average with a decay of 0.9999. Regarding the architecture, we used the same as [9], but with a smaller number of channels (128 instead of 256), leading to much faster training, as well as 10% dropout. Furthermore, our batch size was 128 instead of 256, which leads to a reduced memory cost.

We train our OT-CFM, as well as I-CFM and the original FM, with this new training procedure and report the Fréchet inception distance (FID) in ¿tbl:tab:cifar?. In 4 (left), we show the FID over training time with the Dormand-Prince fifth-order adaptive solver [72] using a relative and absolute error threshold of \(10^{-5}\) similarly to [9], and in 4 (right), we present the FID as a function of the numbers of function evaluations (NFE) using Euler integration.

We find that:

• With improved hyperparameters, we achieve a significantly better FID with the FM training objective than the one reported by [9] at a lower cost.

• For a short computation budget, OT-CFM outperforms FM and (non-OT) I-CFM (¿tbl:tab:cifar?, left).

• After a long training time, all methods achieve similar performance at a high number of function evaluations using fixed-step ODE integration, but OT-CFM performs significantly better with a small number of function evaluations (i.e., allows more efficient inference), indicating straighter, easily integrable flows (¿tbl:tab:cifar?, right).

• FM and I-CFM are equivalently computationally efficient per iteration and OT-CFM comes with a low (\(<\)​1%) computational overhead during training.

Our new training procedures, available at https://github.com/atong01/conditional-flow-matching, allow us to outperform the previous reported results from [9], while the results with our OT-CFM are state-of-the-art for simulation-free neural ODE training algorithms.

5.4 OT-CFM for unsupervised translation↩︎

We show how CFM can be used to learn a mapping between two unpaired datasets in high-dimensional space using the CelebA dataset [73], [74], which consists of \(\sim200\)k images of faces together with 40 binary attribute annotations. For each attribute, we wish to learn an invertible mapping between images with and without the attribute (e.g., ‘not smiling’\(\leftrightarrow\)’smiling’).

To reduce dimensionality, we first train a VAE on the images and encode them as 128-dimensional latent vectors. For each attribute, we learn a flow to map between the embeddings of images without the attribute and those of images with the attribute. After the CNF is learned, we push forward a held-out set of negative vectors by the CNF and compare them to the held-out positive vectors and vice versa. As a metric of divergence, we use maximum mean discrepancy (MMD) with a broad Gaussian kernel (\(\exp(-\|x-y\|^2/(2\cdot128))\)). The results aggregated over all attributes are shown in ¿tbl:tab:vae95results?, showing that OT-CFM discovers a better mapping than other methods. Although MMD is lower for larger \(\sigma\), we found that the alignment is less natural when \(\sigma\) is large, and performance begins to degrade when \(\sigma>1\). 17 shows several visualizations of the learned trajectories.

Finally, while here we work in a latent space, future work should consider learning flows directly in image space, where GAN-based approaches [75] continue to dominate.

5.5 Additional experiments and extensions↩︎

We present numerous other extensions, applications, and evaluations of CFM in 10, notably:

5.5.0.1 OT-CFM reduces variance in the regression target.

To accompany the theoretical results in 9.1, in 10.1 we empirically study the variance of the stochastic regression objective in (OT-)CFM. The results suggest an explanation for the faster convergence of models trained with OT-CFM.

5.5.0.2 Energy-based CFM.

In 9.2 and 10.2 we show how CFM can be used to fit samplers for unnormalized density functions, where exact samples from \(q(x_0)\) or \(q(x_1)\) are not available.

5.5.0.3 Extension to stochastic dynamics.

[76] extends CFM to allow learning stochastic dynamics from unpaired source and target data.

6 Conclusion↩︎

We have introduced a novel class of simulation-free objectives for learning continuous-time flows with a general source distribution. Our approach to training continuous normalizing flows and conditional flow models does not require integration over time during training. We have shown that lifting the static optimal transport problem to the dynamic setting leads to simulation-free solutions to the dynamic OT and SB problems, while also allowing more efficient training and inference of flow models by lowering variance of the objective and straightening flows. One limitation of CFM is that it requires closed-form conditional flows, which hinders its application to situations where we want to regularize the marginal vector field \(u_t(x)\) based on prior information [16]. In addition, the minibatch approximation to OT can incur error in high dimensions; subsequent work can consider the use of neural-network approximations to OT maps [77], [78] in conjunction with CFM. We expect future work to overcome these limitations and hope that ideas from conditional flow matching will improve high-dimensional generative models.

Contribution statement↩︎

A.T.initially conceived the idea. Y.Z., G.H., and N.M.led the development of the theory. High-dimensional experiments and open-source code were led by A.T.and K.F. Additional experiments were contributed by Y.Z., J.R., G.H., N.M., and A.T. All authors contributed to designing the experiments. N.M.and A.T.drove the writing of the paper, with contributions from all other authors. G.W.and Y.B.guided the project.

Acknowledgments↩︎

We would like to thank Stefano Massaroli for productive conversations as well as thank Xinyu Yuan, Marco Jiralerspong, Tara Akhound-Sadegh and Joey Bose for their helpful comments and feedback on the manuscript. We are also grateful to the anonymous reviewers for suggesting numerous improvements. This research was enabled in part by compute resources provided by Mila (mila.quebec) and NVIDIA Corporation. The authors acknowledge funding from CIFAR, Genentech, Samsung, and IBM. In addition, K.F.acknowledges funding from NSERC (RGPIN-2019-06512) and G.W.acknowledges funding from NSERC Discovery grant 03267 and NIH grant R01GM135929.

7 Proofs of theorems↩︎

8 Additional theoretical results↩︎

9 Algorithm extensions↩︎

In 2 we presented the general algorithm for conditional flow matching given \(q(z)\), \(p_t(x|z)\), \(u_t(x|z)\). In 1 we presented a number of settings of these leading to interesting probability paths. In practice, we may wish to compute \(q(z)\) on the fly. Therefore in 5, 6, and 7, we give algorithms for the simplified conditional flow matching, and minibatch versions of OT conditional flow matching and Schrödinger bridge conditional flow matching. In general these consist of first sampling a batch of data from both the source and the target empirical distributions, then resampling pairs of data either randomly (CFM) or according to some OT plan (OT-CFM and SB-CFM).

9.1 Variance reduction by averaging across batches↩︎

An interesting consequence of introducing optimal transport to conditional flow matching is that it greatly reduces variance of the regression target. Informally, as \(\sigma \to 0\), \(\mathbb{E}_{x,t,z} \|u_t(x | z) - u_t(x)\|^2 \to 0\) for OT-CFM and SB-CFM, which is not true of previous probability paths in 1 (See 3 for a precise statement). As flow models get larger, more powerful, and more costly, reducing objective variance, and thereby faster training may lead to significant cost savings [8]. To this end we also explore reducing the variance of the objective by averaging over a batch. This is not feasible in score matching where the flow conditioned on multiple datapoints is complex. Our CFM framework naturally extends from a pair of datapoints to a batch of pairs. Instead of conditioning on a single pair of datapoints we can condition on a batch of pairs. As the batch increases in size, we trade higher cost in computing the target for lower variance in the target as the batch size increases, the variance in the target goes to zero (see 4 for a precise statement).

As formalized in 4, we can reduce variance in the target by averaging over multiple datapoints. Specifically, in this case we let \(\bar{{\boldsymbol{z}}} \mathrel{\vcenter{:}}= \{z^i \mathrel{\vcenter{:}}= (x_0^i, x_1^i)\}_{i=1}^m\), where \(z^i\) are i.i.d.from \(q(z)\) and \[\begin{align} p_t(x | \bar{{\boldsymbol{z}}}) &= \frac{\sum_i^m p_t(x | z^i)q(z^i)}{\sum_i^m q(z^i)} \tag{23}\\ u_t(x | \bar{{\boldsymbol{z}}}) &= \frac{\sum_i^m u_t(x | z^i) p_t(x | z^i) q(z^i)}{p_t(x | \bar{{\boldsymbol{z}})}}\tag{24} \end{align}\] It takes roughly \(m\) times as long to compute the conditional target \(u_t(x | \bar{{\boldsymbol{z}}})\) but reduces the variance. As the evaluation and backpropagation through \(v_\theta\) gets more difficult this tradeoff can be beneficial.

9.2 Modeling energy functions↩︎

If we have access to an energy function two (unnormalized) energy functions \(\mathcal{R}_{\{0,1\}} : \mathbb{R}^d \to \mathbb{R}^+\) at the endpoints instead of i.i.d.samples \({\boldsymbol{X}}_t \sim q_t(x_t)\), then the objective must be slightly modified. We formulate the Energy Conditional Flow Matching Objective as \[\begin{align} \label{eq:ecfm} {\mathcal{L}}_{\rm ECFM}= &\mathbb{E}_{t, \hat{q}_0(x_0), \hat{q}_1(x_1), p_t(x | x_0, x_1)} \bigg [ \frac{R_0(x_0) R_1(x_1)}{\hat{q}_0(x_0) \hat{q}_1(x_1)} \|v_\theta(t, x) - u_t(x | {\boldsymbol{x}}_0, {\boldsymbol{x}}_1)\|^2_2 \bigg ] \end{align}\tag{25}\] We can use this object to train a flow which matches the energies without access to samples. This is formalized in the following theorem.

Of course \({\mathcal{L}}_{\rm ECFM}\) leaves the question of sampling open for high-dimensional spaces. Sampling uniformly does not scale well to high dimensions, so for practical reasons we may want a different sampling strategy.

We use this objective in 16 with a uniform proposal distribution as a toy example of this type of training.

10 Additional results↩︎

We start this section by the definition of the entropy regularized OT problem: \[\label{eq:sot} W(q_0, q_1)^2_{2, \lambda} = \inf_{\pi_\lambda \in \Pi} \int_{\mathcal{X}^2} c(x, y)^2 \pi_\lambda(dx, dy) - \lambda H(\pi),\tag{27}\] where \(\lambda \in \mathbb{R}^+\) and \(H(\pi) = \int \ln \pi(x,y) d \pi(dx,dy)\).

10.0.0.1 Regularized CNF tuning

Continuous normalizing flows with a path length penalty optimize a relaxed form of a dynamic optimal transport problem [4], [16], [17]. Where dynamic optimal transport solves for the optimal vector field in terms of average path length where the marginals at time \(t=0\) and \(t=1\) are constrained to equal two input marginals \(q_0\) and \(q_1\). Instead of this pair of hard constraints, regularized CNFs instead set \(q_0 \mathrel{\vcenter{:}}= \mathcal{N}( x \mid 0, 1)\) and optimize a loss of the form \[L(x(t)) = -\log p(x(t)) + \lambda_e \int_0^1 \|v_\theta(t, x(t))\|^2 dt\] where \(\frac{dx}{dt} = v_\theta(t, x(t))\) and \(\log p(x(T))\) is defined as \[\log p(x(T)) = p(x(0)) + \int_0^T \frac{\partial \log p(x(t))}{\partial t} dt = p(x(0)) + \int_0^T -\mathrm{tr}\left ( \frac{ d v_\theta}{d x(t)} \right ) dt\] where the second equality follows from the instantaneous change of variables theorem [2]. In practice it is difficult to pick a \(\lambda_e\) which both produces flows with short paths and allows the model to fit the data well. We analyze the effect of this parameter over three datasets in 8. In this figure we analyze the Normalized 2-Wasserstein to the target distribution (which approaches 1 with good fit), and the Normalized Path Energy (NPE). We find a tradeoff between short paths (Low NPE) and good fit (Low 2-Wasserstein). We choose \(\lambda_e = 0.1\) as a good tradeoff across datasets, which has paths that are not too much longer than optimal but also fits the data well.

10.0.0.2 Ablation results on batch size.

Since we use Minibatch-OT for OT-CFM, when the minibatch size is equal to one, then OT-CFM is equivalent to CFM. This effect can be seen in 9, where over four datasets, OT-CFM starts with equal path length and approximately equal 2-Wasserstein. Then the normalized path energy decreases surprisingly quickly plateauing after batchsize reaches \(\sim\)​64. While the minibatch size needed to approximate the true dynamic optimal transport paths will vary with dataset (for example in the moon-8gaussian case we need a larger batch size) it is still somewhat surprising that such small batches are needed as this is less than 0.5% of the entire 10k point dataset per batch.

10.0.0.3 The effect of \(\sigma\) on fit and path length.

Next we consider \(\sigma\), the bandwidth parameter of the Gaussian conditional probability path. In 10 we study the effect of \(\sigma\) on the fit (top) and the path energy (bottom). With \(\sigma > 1\) methods start to underfit with high 2-Wasserstein error and either very long or very short paths. As for specific models, SB-CFM becomes unstable with \(\sigma\) too small due to the lack of convergence for the static Sinkhorn optimization with small regularization. FM and CFM follow similar trends where they fit fairly well with \(\sigma \le 1\) but have paths that are significantly longer than optimal by 2-3x. OT-CFM maintains near optimal path energies and near optimal fit until \(\sigma > 1\).

10.0.0.4 Schrödinger bridge fit over simulation time.

In 14 we compare the fit of Diffusion Schrödinger Bridge model with SB-CFM conditioned on time. The Diffusion Schrödinger Bridge seems to outperform SB-CFM early in the trajectory, however fails to fit the bridge after many integration steps.

10.1 Objective variance.↩︎

We consider the variance of the objective \(u_t(x | z)\) with respect to \(z\). While for any \(x\) we have \(\mathbb{E}_q(z) u_t(x | z) = u_t(x)\), we find a lower second moment speeds up training. Specifically, we seek to understand the effect of the second moment which we call the objective variance defined as \[\label{eq:objective95variance} OV = \mathbb{E}_{t \sim U(0,{\boldsymbol{I}}), x \sim p_t(x), z \sim q(z)} \|u_t(x | z) - u_t(x)\|^2\tag{28}\] on training speed for different objectives in 1. We estimate the variance on a small data with a known \(u_t(x)\). We examine this estimated objective variance and its effect on training convergence in 15, showing that either OT-CFM or variance reduced CFM with averaging over the batch results in lower variance of the objective. This in turn leads to faster training times as shown on the right. Averaging over a batch of data leads to faster training particularly for methods with high objective variance (CFM) and less so for those with low (OT-CFM), which already trains quickly.

Variance in the conditional objective target \(u_t(x | z)\) varies across models. In 11 we study the objective variance across CFM objective functions. Here we estimate the objective variance in (28 ) as \[\mathbb{E}_{x, t, z} \| u_t(x | z) - v_\theta(t, x) \|^2\] after training has converged. After training has converged \(v_\theta\) should be very close to \(u_t(x)\) so we use it as an empirical estimator of \(u_t(x)\) to compute the variance. We find that across all datasets OT-CFM and SB-CFM have at least an order of magnitude lower variance than CFM and FM objectives. This correlates with faster training as measured by lower validation error in fewer steps for lower variance models as seen in 3 (left).

We examine the objective variance OV by conditioning \(u_t(x | {\boldsymbol{z}})\) on a batch of pairs of data points, \(\bar{{\boldsymbol{z}}} \mathrel{\vcenter{:}}= \{z^i \mathrel{\vcenter{:}}= (x_0^i, x_1^i)\}_{i=1}^m\), we can reduce the variance of the OV objective to 0 for all models as batchsize goes to population size. For the batchsize \(m\) range from 1 to the number of the population, we uniformly sample \(m\) pairs of points \(z^i\) and compute the probability \(p_t(x| \bar{{\boldsymbol{z}}})\) and the objective \(u_t(x|\bar{{\boldsymbol{z}}})\) from (23 ) and (24 ).

We also find that averaging over batches makes the network acheive a lower validation error in fewer steps and in less walltime (12).

10.2 Energy-based CFM↩︎

We show how CFM and OT-CFM can be adapted to the case where we do not have access to samples from the target distribution, but only an unnormalized density (equivalently, energy function) of the target, \(R(x_1)\) (16). We consider the 10-dimensional funnel dataset from [80]. We aim to learn a flow from the 10-dimensional standard Gaussian to the energy function of the funnel. We consider two algorithms, each of which has certain advantages:

1. Reweighted importance sampling (RWIS): We construct a weighted batch of target points \(x_1\) by sampling \(x_1\sim{\mathcal{N}}(\boldsymbol{0},{\rm\mathbf{I}})\) and assigning it a weight of \(R(x_1)/{\mathcal{N}}(x_1;\boldsymbol{0},{\rm\mathbf{I}})\) normalized to sum to 1 over the batch. The FM and CFM objectives handle weighted samples in a trivial way (by simply using the weights as \(q(x_1)\) in 1), while OT-CFM treats the weights as target marginals in constructing the OT plan between \(x_0\) and \(x_1\). We expect RWIS to perform well when batches are large and the proposal and target distributions are sufficiently similar; otherwise, numerical explosion of the importance weights can hinder learning.

2. MCMC: We use samples from a long-run Metropolis-adjusted Langevin MCMC chain on the target density as approximate target samples. We expect this method to perform well when the MCMC mixes well; otherwise, modes of the target density may be missed.

As an evaluation metric, we use the estimation bias of the log-partition function using a reweighted variational bound, following prior work that studied the problem using SDE modeling [81]. The computation of this metric for CNFs is given in 11.6.

The results are shown in 4. When an adaptive ODE integrator is used, all algorithms achieve similar results (no pair of mean log-partition function estimates is statistically distinguishable with \(p<0.1\) under a Welch’s \(t\)-test) but OT-CFM is about twice as efficient as CFM and FM. However, with a fixed computation budget for ODE integration, OT-CFM performs significantly better.

11 Experiment and implementation details↩︎

11.1 Physical experimental setup↩︎

All experiments were performed on a shared heterogenous high-performance-computing cluster. This cluster is primarily composed of GPU nodes with RTX8000, A100, and V100 Nvidia GPUs. Since the network and nodes are shared, other users may cause high variance in the training times of models. However, we believe that the striking difference between the convergence times in 3 and combined with the CFM training setup with a single CPU and the baseline models trained with two CPUS and a GPU, paints a clear picture as to how efficient CFM training is. Qualitatively, we feel that most CFMs converge quite a bit more rapidly than these metrics would suggest, often converging to a near optimal validation performance in minutes.

11.2 2D, single-cell, and Schrödinger bridge experimental setup↩︎

For all experiments we use the same architecture implemented in PyTorch [82]. We concatenate the flattened input \(x \in \mathbb{R}^d\) and the time \(t\) as the \(d+1\) inputs to a network with three hidden layers of width \(64\) interspersed with SELU activations [83] followed by a linear output layer of width \(d\). This forms our \(v_\theta\) for all experiments. For all 2D and single-cell experiments we train for 1000 epochs and implement early stopping on the validation loss which checks the loss on a validation set every 10 epochs and stops training if there is no improvement for 30 epochs. We also set a time limit of 100 minutes for each CFM model. This is hit almost exclusively for SB-CFM models with small \(\sigma\) which are unstable to train due to instabilities and non-convergence of the Sinkhorn [32] transport plan optimization. We use the AdamW [84] optimizer with weight decay \(10^{-5}\) with batchsize 512 by default in 2D experiments and 128 in the single cell datasets. For OT-CFM and SB-CFM we use exact linear programming EMD and Sinkhorn algorithms from the python optimal transport package [85] For evaluation of trajectories unless otherwise noted we use the Runge-Kutta45 (rk4) ODE solver with 101 timesteps from 0 to 1.

11.3 Variance reduction by averaging↩︎

We tackle the exploration of the effects of reducing variance of the target \(u_t(x|z)\) from two directions. The first is for small example where we can compute the ground truth \(u_t(x)\) quickly, and the second is in the setting of trained models where we can estimate \(u_t(x)\) with \(v_\theta(t,x)\) after \(v_\theta\) has converged.

We first consider the convergence of each flow matching objective (OT-CFM, CFM, FM, SB-CFM) to zero as a function of the batch size relative to the dataset size. This is done by first sampling \(t,x,z\) then computing the true objective variance across many samples. This appears in 15.

We next consider the effect of averaging over a batch to reduce the variance of the objective in 15 (right). Here Batchsize refers to the size of the batch we are averaging over. We aggregate this into a single target so that the model sees a single \(d\) dimensional target vector for one sampled \(x,t\). This means that we can compare different aggregation sizes fairly.

11.4 Schrödinger bridge evaluation setup↩︎

To evaluate how well Schrödinger Bridge models actually model a Schrödinger Bridge, we constrain ourselves to a small example with 1000 points. We note that the closed-form Schrödinger marginals are known for discrete densities, for Gaussians [86], and can be constructed for two approximate datasets [87], which present other ways of evaluating Schrödinger bridge performance. For any time \(t\) we can sample from the ground truth Schrödinger bridge density \(p_t(x)\) as \[\begin{align} (x_0, x_1) &\sim \pi_{2 \sigma^2}\\ X_t &\sim \mathcal{N}(x \mid t x_1 + (1-t) x_0, \sigma t (1-t)) \end{align}\] We sample trajectories of length 20 from \(t=0\) to \(t=1\) by integrating over time from \(t=0\) to \(t=1\). At each of the 18 intermediate timepoints we compute the 2-Wasserstein distance between a sample of size 1000 from the trajectories at that time and the ground truth \(X_t\) as above at that time. We reported the average across the 18 intermediate timepoints in ¿tbl:tab:sb95comparison? and plot the 2-Wasserstein distance over time in 14.

11.4.0.1 SB-CFM Model

We train SB-CFM with \(\sigma=1\) and batchsize=512 for each of the datasets. We save 1000 trajectories from a test set integrated with the tsit5 solver with atol=rtol=1e-4.

11.4.0.2 Diffusion Schrödinger bridge model implementation details

We use the implementation from [19]. Only the networks were changed for a fair comparison with CFM. The forward and backward networks are composed of an MLP with three hidden layers of size 64, with SELU activations in between layers. We used a time and a positional encoders composed of two layers of size 16 and 32 with LeakyReLU activations has inputs to the score network. The architectures are the same for the 2D examples and the single-cell examples (except for the input dimension). During training, we set the variance (\(\gamma\) in the author’s code) to 0.001 and did 20 steps to discretize the Langevin dynamic. We trained for 10k iterations with 10k particles and batch size of 512, for 20 iterative proportional fitting steps, and a learning rate set to 0.0001. For the interpolation task we used the tenth timepoint from the Langevin dynamic with the backward network trained to go from the distribution at time \(t-1\) to \(t+1\). All trajectories are evaluated from the backward dynamic. We use \(\sigma=1\) and batchsize=512.

11.5 Single-cell experimental setup↩︎

We strove to be consistent with the experimental setup of [16]. For the Embryoid body (EB) data, we use the same processed artifact which contains the first 100 principal components of the data. For our tests we truncate to the first five dimensions, then whiten (subtract mean and divide by standard deviation) each dimension. For the Embryoid body dataset which consists of 5 timepoints collected over 30 days we train separate models leaving out times \(1,2,3\) in turn. We train a CFM over the full time scale (0-4). During testing we push forward all points \(X_{t-1}\) to time \(t\) as a distribution to test against.

For the Cite and Multi datasets these are sourced from the Multimodal Single-cell Integration challenge at NeurIPS 2022, a NeurIPS challenge hosted on Kaggle where the task was multi-modal prediction [69]. In this competition they used this data to investigate the predictability of RNA from chromatin accessibility and protein expression from RNA. Here, we repurpose this data for the task of time series interpolation. Both of these datasets consist of four timepoints from CD34+ hematopoietic stem and progenitor cells (HSPCs) collected on days 2, 3, 4, and 7. For more information and the raw data see the competition site.⁵ We preprocess this data slightly to remove patient specific effects by focusing on a single donor (donor 13176), then we again compute the first five principal components and again whiten each dimension to further normalize the data.

11.6 Energy-based CFM↩︎

The 10-dimensional funnel dataset is defined by \(\mathbf{x}_0\sim{\mathcal{N}}(0,{\boldsymbol{I}})\), \(\mathbf{x}_{1,\ldots,9}\sim{\mathcal{N}}(\boldsymbol{0},\exp(\mathbf{x}_0){\rm \mathbf{I}})\). We attempted to mimic the SDE model architecture from [81] for the flow model \(v_\theta(t,x)\). The time step \(t\) is encoded with 128-dimensional Fourier features, then both \(x\) and \(t\) are independently processed with two-layer MLPs. The two representations are concatenated and processed through another three-layer MLP to make the prediction. All MLPs use GELU activation and have 128 units per hidden layer. We trained all models with \(\sigma=0.05\) and learning rate \(10^{-2}\), the highest at which they were table, for 1500 batches of size 300, to be consistent with the settings from [81].

The importance-weighted estimate of the log-partition function is defined \[\log\hat{Z}=\log\frac{1}{K}\sum_{i=1}^K\frac{R(x_1^{(i)})}{{\mathcal{N}}(x_0^{(i)};0,\mathbf{I})}\left|\frac{\partial x_1}{\partial x_0}\right|_{x_0=x_0^{(i)}},\] where \(x_0^{(i)}\) are independent samples from the source distribution and \(x_1^{(i)}\) is \(x_0^{(i)}\) pushed forward by the flow (note that the Jacobian can be computed by differentiating the ODE integrator). We used \(K=6000\) samples.

For MCMC, to be consistent with [81], we generated 15000 samples, each of which was seen 30 times in training. We used 1000 steps of Metropolis-adjusted Langevin sampling with \(\epsilon\) linearly decaying from 0.1 to 0.

The flow network used to generate 16 followed similar settings to those used in 5.1.

11.7 Unsupervised translation↩︎

We trained a vanilla convolutional VAE, with about 7 million parameters in the encoder, on CelebA faces scaled to \(128\times128\) resolution.

For the flow network \(v_\theta(t,x)\), we used a MLP with four hidden layers of 512 units and leaky ReLU activations taking the 129-dimensional concatenation of \(x\) and \(t\) as input. All models CFM and OT-CFM were trained for 5000 batches of size 256 and the Adam optimizer with learning rate \(10^{-3}\). Integration was performed using the Dormand-Prince integrator with tolerance \(10^{-3}\). For each attribute, 1000 positive and negative images each were used as a held-out test set.

17 shows some examples of the learned trajectories.

11.8 Unconditional CIFAR-10 experiments↩︎

For the CIFAR-10 experiments we followed the setup as described in [9]. All methods were trained with the same setup, only differeing in the choice of probability path. Since code has not been released for this work, there are a few parameters which may differ. We summarize the setup here, where the exact parameter choices can be seen in the source code.

We used the Adam optimizer with \(\beta_1 = 0.9\), \(\beta_2 = 0.999\), \(\epsilon=10^{-8}\), and no weight decay. To reproduce [9], we used the UNet architecture from [5] with channels \(=256\), depth \(=2\), channels multiple \(=[1,2,2,2]\), heads \(=4\), heads channels \(=64\), attention resolution \(=16\), dropout \(=0.0\), batch size per gpu \(= 128\), gpus \(=2\), epochs \(=2000\), maximum learning rate \(= 5 \times 10^{-4}\), minimum learning rate \(= 0\), with a learning schedule that increases linearly from the minimum to the maximum learning rate over the first \(200\) epochs, and decays linearly from back to the minimum after that. We use\(\sigma = 10^{-6}\). For sampling, we use Euler integration using the torchdyn package and dopri5 from the torchdiffeq package. Since the dopri5 solver is an adaptive step size solver, it uses a different number of steps for each integration. We use a batch size of 500 for a 100 total batches and average the number of function evaluations (NFE) over batches.

For our improved models we use the same parameters as above except we use channels \(=128\), dropout \(=0.1\), a single A100 GPU, steps=\(400000\), \(\sigma = 0\), constant learning rate of \(2 \times 10^{-4}\), gradient clipping with norm \(=1.0\), and exponential moving average weights with decay \(=0.9999\).

References↩︎