Weighted Conditional Flow Matching

## Abstract

Abstract Conditional flow matching (CFM) has emerged as a powerful framework for training continuous normalizing flows due to its computational efficiency and effectiveness. However, standard CFM often produces paths that deviate significantly from straight-line interpolations between prior and target distributions, making generation slower and less accurate due to the need for fine discretization at inference. Recent methods enhance CFM performance by inducing shorter and straighter trajectories but typically rely on computationally expensive mini-batch optimal transport (OT). Drawing insights from entropic optimal transport (EOT), we propose Weighted Conditional Flow Matching (W-CFM), a novel approach that modifies the classical CFM loss by weighting each training pair ( x , y ) (x,y) ( italic_x , italic_y ) with a Gibbs kernel. We show that this weighting recovers the entropic OT coupling up to some bias in the marginals, and we provide the conditions under which the marginals remain nearly unchanged. Moreover, we establish an equivalence between W-CFM and the minibatch OT method in the large-batch limit, showing how our method overcomes computational and performance bottlenecks linked to batch size. Empirically, we test our method on unconditional generation on various synthetic and real datasets, confirming that W-CFM achieves comparable or superior sample quality, fidelity, and diversity to other alternative baselines while maintaining the computational efficiency of vanilla CFM.

## 1 Introduction

Generative modeling seeks to learn a parameterized transformation that maps a simple prior (e.g., a Gaussian) to a complex data distribution. Continuous normalizing flows (CNFs) instead train a time-dependent vector field to solve an ordinary differential equation (ODE) that transports base samples to data samples with exact likelihood computation and invertibility. However, for some time, CNFs suffered from training instability and failed to scale efficiently to large or high-dimensional datasets (Chen et al., 2018; Grathwohl et al., 2018; Onken et al., 2021). Flow Matching (FM) (Lipman et al., 2023) reframes CNF training as a simple regression problem: a vector field is learned to match the endpoint displacement between a prior sample and its paired data point, yielding near-optimal transport trajectories when the prior is Gaussian. However, the independent pairing of FM cannot ensure that the marginal flow follows an optimal transport geodesic, leading to suboptimal paths in practice. Conditional Flow Matching (CFM) (Lipman et al., 2023; Tong et al., 2024) generalizes FM by learning a vector field that transports samples from an arbitrary transport map, conditioned on paired source and target samples. This method allows for simulation-free training of continuous normalizing flows and can learn conditional generative models from any sampleable source distribution, extending beyond the Gaussian source.

Thanks to its flexibility, CFM has been applied in many areas of science, such as molecule generation (Irwin et al., 2024; Geffner et al., 2025), sequence and time-series modeling (Stark et al., 2024; Zhang et al., 2024; Rohbeck et al., 2025), and text-to-speech translation (Guo et al., 2024).
A refinement of CFM is minibatch Optimal Transport CFM (OT-CFM) (Lipman et al., 2023; Tong et al., 2024), which uses an entropic or exact plan as the coupling so that each training pair is drawn according to the optimal transport solution between the minibatch source and target samples. This yields substantially straighter, lower-cost trajectories in the learned flow and improves sample quality with fewer integration steps. However, computing these OT plans—even approximately via Sinkhorn—for every minibatch incurs substantial per-iteration overhead, scaling cubically with batch size (or quadratically under entropic regularization). Moreover, requiring well-balanced class representation in each batch to approximate the global OT plan makes this approach impractical for large, multi-class datasets.

As an alternative that addresses these limitations, we introduce Weighted Conditional Flow Matching (W-CFM), which replaces costly batch-level transport computations by simply weighting each independently sampled pair $(x,y)$ with the entropic OT (EOT) Gibbs kernel, $w(x,y)=\exp(-c(x,y)/\varepsilon)$ (Cuturi, 2013). This importance weighting provably recovers the entropic OT (EOT) plan up to a controllable bias in the marginals. As a result, the learned flow follows straight paths without ever explicitly solving an OT problem during training. Moreover, we show that W-CFM matches OT-CFM in the large-batch limit, thereby not incurring any of the batch size-related limitations or any extra costs. In practice, W-CFM delivers straight flows and high‐quality samples consistently outperforming CFM and achieving comparable performance to OT-CFM, but with no extra overhead. In short, our contributions are the following:

- •
We introduce a novel CFM variant that, inspired by EOT, incorporates a Gibbs kernel weight on each sample pair and show that our method serves as a new way to approximate the EOT coupling without any additional cost during training.
- •
We discuss practical design choices to alleviate the change of the marginals with the new loss and derive sufficient conditions under which this change becomes trivial, so the true marginals are approximately preserved.
- •
We show that as the batch size grows and assuming that the bias in the marginals remains negligible, W-CFM converges to an entropy-regularized version of OT-CFM, retaining some trajectory straightness without the computational scaling issues associated with the OT plan computations.
- •
We demonstrate on toy and image-generation benchmarks that W-CFM matches or outperforms existing CFM and OT-CFM methods in sample quality, fidelity, and diversity under a sensible choice of the hyperparameter ${\varepsilon}$ in the Gibbs kernel.

## 2 Background

### 2.1 Entropic Optimal Transport

We refer the reader to Nutz (2021) for a comprehensive introduction to the topic and only mention the relevant results for our work.
For $\mu,\nu\in\mathcal{P}(\mathbb{R}^{d})$, we recall the definition of the Kullback–Leibler divergence $D_{\mathrm{KL}}(\mu\|\nu)\coloneqq\int_{\mathbb{R}^{d}}\log\frac{d\mu}{d\nu}(x)d\mu(x)$ if $\mu\ll\nu$ and $+\infty$ otherwise.
Assume that $\mu,\nu$ have finite first moment.
We denote the set of couplings between $\mu$ and $\nu$ by $\Pi(\mu,\nu)\coloneqq\{\pi\in\mathcal{P}(\mathbb{R}^{d}\times\mathbb{R}^{d}):\pi(\mathbb{R}^{d},dy)=\nu(dy),\pi(dx,\mathbb{R}^{d})=\mu(dx)\}$.
We consider the following entropic optimal transport (EOT) problem with parameter ${\varepsilon}>0$:

$$ $\min_{\pi\in\Pi(\mu,\nu)}\int_{\mathbb{R}^{d}\times\mathbb{R}^{d}}c(x,y)d\pi(x,y)+{\varepsilon}D_{\mathrm{KL}}(\pi\|\mu\otimes\nu),$ (1) $$

where $c:\mathbb{R}^{d}\times\mathbb{R}^{d}\to\mathbb{R}_{+}$ denotes a cost function, typically $c(x,y)=\|x-y\|$, which is taken such that equation [1](https://arxiv.org/html/2507.22270v1#S2.E1) is finite. When ${\varepsilon}=0$, we recover the classical Monge–Kantorovich transportation problem of moving a distribution $\mu$ to a distribution $\nu$ by minimizing the transport cost as measured by $c$, whose solution is not necessarily unique; we denote by $\pi^{\star}\in\Pi(\mu,\nu)$ such a solution.
EOT is of significant importance in machine learning and scientific computing (Genevay et al., 2018; Peyré & Cuturi, 2019), as it approximates the original Monge–Kantorovich transport problem and can be solved tractably with Sinkhorn’s algorithm (Cuturi, 2013; Altschuler et al., 2017).
It is a classical result (see, e.g., Chapter 4 in Peyré & Cuturi (2019)) that solving equation [1](https://arxiv.org/html/2507.22270v1#S2.E1) is equivalent to solving the following projection problem
$\min_{\pi\in\Pi(\mu,\nu)}D_{\mathrm{KL}}(\pi\|\mathcal{K}_{{\varepsilon}})$,
with the Gibbs kernel $\mathcal{K}_{{\varepsilon}}(dx,dy)\coloneqq e^{-c(x,y)/{\varepsilon}}\mu(dx)\nu(dy)$. A convexity argument can be made to prove that there exists a unique minimizer $\pi_{\varepsilon}\in\Pi(\mu,\nu)$ of equation [1](https://arxiv.org/html/2507.22270v1#S2.E1).
More specifically, this projection formulation allows us to write the optimal coupling in terms of $\mathcal{K}_{\varepsilon}$.

###### Theorem 1 (Theorem 4.2 in Nutz ( 2021 ) ) .

If $c(x,y)<\infty$ $\mu\otimes\nu$-almost-surely, there exist measurable functions $\phi,\psi:\mathbb{R}^{d}\rightarrow\mathbb{R}$, referred to as Schrödinger potentials, such that the EOT plan for any ${\varepsilon}>0$ is given by

$$ $\pi_{{\varepsilon}}(dx,dy)=\exp\left(\phi_{{\varepsilon}}(x)+\psi_{{\varepsilon}}(y)-\frac{c(x,y)}{{\varepsilon}}\right)\mu(dx)\nu(dy).$ (2) $$

In other words, we have $\pi_{{\varepsilon}}(dx,dy)=f_{\varepsilon}(x)g_{\varepsilon}(y)\mathcal{K}_{{\varepsilon}}(dx,dy)$ for some positive functions $f_{\varepsilon}(x)=\exp(\phi_{{\varepsilon}}(x)),g_{\varepsilon}(y)=\exp(\psi_{\varepsilon}(y))$. In particular, the component capturing the dependence in the EOT plan $\pi_{\varepsilon}$ is exactly given by the Gibbs kernel, which is easy to compute, and one only needs to adjust the marginals independently to obtain $\pi_{\varepsilon}$.

###### Theorem 1 (Theorem 4.2 in Nutz ( 2021 ) ) .

### 2.2 Conditional Flow Matching

The flow matching methodology Lipman et al. (2023) is a simulation-free method of training a continuous normalizing flow (Chen et al., 2018), i.e., a smooth vector field $v_{\theta}(t,x)$, for generating samples from a target distribution $\nu\in\mathcal{P}(\mathbb{R}^{d})$ given a source distribution $\mu\in\mathcal{P}(\mathbb{R}^{d})$. Assume that there exists a vector field $v_{t}(x)$ such that the flow given by $dx_{t}=v_{t}(x_{t})dt$ with initial condition $x_{0}\sim\mu$ satisfies $x_{1}\sim\nu$. The flow matching loss is given by

$$ $\mathcal{L}_{\mathrm{FM}}(\theta)\coloneqq\mathbb{E}_{t\sim\mathcal{U}(0,1),X_{t}\sim p_{t}}\left[\|v_{\theta}(t,X_{t})-v_{t}(X_{t})\|^{2}\right],$ (3) $$

where $p_{t}$ denotes the distribution of $x_{t}$, i.e. a probability path between $\mu$ and $\nu$. Under technical assumptions, $v_{t}$ generates $p_{t}$ if and only if they satisfy the continuity equation

$$ $\frac{\partial p_{t}}{\partial t}+\nabla\cdot(p_{t}v_{t})=0,\quad p_{0}=\mu,\quad p_{1}=\nu.$ (4) $$

Upon finding a minimizer of equation [3](https://arxiv.org/html/2507.22270v1#S2.E3), one integrates the ODE $d\tilde{x}_{t}=v_{\theta}(t,\tilde{x}_{t})dt$ from a source sample $\tilde{x}_{0}\sim\mu$, so that $\tilde{x}_{1}$ is approximately distributed according to $\nu$. In its most general form, conditional flow matching (Lipman et al., 2023) replaces the intractable FM loss equation [3](https://arxiv.org/html/2507.22270v1#S2.E3) by an equivalent loss involving a conditional vector field $v_{t}(x\mid z)$ and conditional probability path $p_{t}(x\mid z)$ satisfying equation [4](https://arxiv.org/html/2507.22270v1#S2.E4) between $\mu(dx\mid z)$ and $\nu(dx\mid z)$

$$ $\mathcal{L}_{\mathrm{CFM}}(\theta;q)\coloneqq\mathbb{E}_{t\sim\mathcal{U}(0,1),Z\sim q}\mathbb{E}_{X_{t}\sim p_{t}(\cdot\mid Z)}\left[\|v_{\theta}(t,X_{t})-v_{t}(X_{t}\mid Z)\|^{2}\right],$ (5) $$

where $z$ denotes a latent variable and $q$ the prior distribution, typically choosing $z=x_{1}$. Tong et al. (2024) have generalized the approach of Lipman et al. (2023) by considering arbitrary latent variables $z$, showing that minimizing the CFM loss equation [5](https://arxiv.org/html/2507.22270v1#S2.E5) is equivalent to minimizing equation [3](https://arxiv.org/html/2507.22270v1#S2.E3). Hence, the conditional flow matching method calls for two important design choices:

- •
The latent variable $z$ and prior $q$: in this paper we focus on the popular choice of $z=(x_{0},x_{1})$, and we therefore require $q$ to be a coupling, that is $q\in\Pi(\mu,\nu)$.
- •
The conditional vector field $v_{t}(x\mid z)$ and probability path $p_{t}(x\mid z)$: in this paper we consider the linear interpolation path $X_{t}=(1-t)X_{0}+tX_{1}$, i.e. $p_{t}(x\mid x_{0},x_{1})=\delta_{(1-t)x_{0}+tx_{1}}(x)$ together with $v_{t}(x\mid x_{0},x_{1})=x_{1}-x_{0}$, this is the implicit choice made for Rectified Flow (Liu et al., 2023).

Choosing $q=\mu\otimes\nu$ in equation [5](https://arxiv.org/html/2507.22270v1#S2.E5) leads to the independent conditional flow matching algorithm (I-CFM), which is straightforward to compute but yields irregular trajectories, requiring fine discretizations at inference. On the contrary, choosing $q=\pi^{\star}$ (i.e., an optimal transport plan between the source and target distributions), leads to the optimal transport conditional flow matching algorithm (OT-CFM), which induces straighter trajectories for the learned model $v_{\theta}(t,x)$ (Tong et al., 2024). As $\pi^{\star}$ is a priori difficult to sample from, many works have explored approximations by computing short distance couplings at the batch level during training (Tong et al., 2024; Pooladian et al., 2023).

## 3 Weighted Conditional Flow Matching

By writing $L_{\theta}(t,X,Y)=\|v_{\theta}(t,X_{t})-v_{t}(X_{t}\mid X,Y)\|^{2}$ where $X_{t}=(1-t)X+tY$, the I-CFM loss can be written as $\mathcal{L}_{\mathrm{I-CFM}}(\theta)=\mathbb{E}_{t\sim\mathcal{U}(0,1),(X,Y)\sim\mu\otimes\nu}\left[L_{\theta}(t,X,Y)\right]$. In order to enforce straightness of the learned velocity field, one would like to bias this loss towards training sample pairs $(x,y)$ that are close to each other. OT-CFM achieves this by computing the OT plan between discrete sets of points within a batch. We propose to introduce a bias directly inside the expectation by considering modifications of the I-CFM loss function of the form

$$ $\mathcal{L}_{w}(\theta)\coloneqq\mathbb{E}_{t\sim\mathcal{U}(0,1),(X,Y)\sim\mu\otimes\nu}\left[w(X,Y)L_{\theta}(t,X,Y)\right],$ (6) $$

where $w:\mathbb{R}^{d}\times\mathbb{R}^{d}\to\mathbb{R}_{+}^{*}$ denotes a positive weighting function, which should be large whenever $X,Y$ are close and small when $X,Y$ are far apart. This weighting can be understood in terms of a change of measure. In particular, by defining

$$ $\pi_{w}(dx,dy)\propto w(x,y)\mu(dx)\nu(dy),$ (7) $$

we have that $\mathcal{L}_{w}(\theta)=\mathbb{E}_{t\sim\mathcal{U}(0,t),(X,Y)\sim\pi_{w}}\left[L_{\theta}(t,X,Y)\right]$. In other words, weighting the I-CFM loss yields a CFM loss with a new prior distribution $q=\pi_{w}$.
This technique can be thought of as a form of importance sampling, where the baseline distribution $\mu\otimes\nu$ is easy to sample from, and $w(x,y)$ corresponds to the importance sampling weight.

### 3.1 Approximating EOT with Weighted Conditional Flow Matching

Let ${\varepsilon}>0$. We propose a new loss function that can be used as a drop-in replacement within conditional flow matching and call it Weighted Conditional Flow Matching (W-CFM). Given the form of the EOT plan equation [2](https://arxiv.org/html/2507.22270v1#S2.E2),
and a cost function $c:\mathbb{R}^{d}\times\mathbb{R}^{d}\to\mathbb{R}$, we consider $\mathcal{L}_{w}$ equation [6](https://arxiv.org/html/2507.22270v1#S3.E6)
with the weighting function $w_{\varepsilon}(x,y)\coloneqq\exp(-c(x,y)/{\varepsilon})$ and introduce the weighted
conditional flow matching loss

$$ $\boxed{\mathcal{L}_{\mathrm{W-CFM}}(\theta;{\varepsilon})\coloneqq\mathbb{E}_{t\sim\mathcal{U}(0,1)}\mathbb{E}_{(X,Y)\sim\mu\otimes\nu}\left[w_{\varepsilon}(X,Y)||v_{\theta}(t,X)-(Y-X)||^{2}\right].}$ (8) $$

In particular, Theorem [1](https://arxiv.org/html/2507.22270v1#Thmtheorem1) implies that $\mathcal{L}_{\mathrm{W-CFM}}(\theta;{\varepsilon})=Z_{\varepsilon}\mathcal{L}_{\mathrm{CFM}}(\theta;q_{\varepsilon})$, where $q_{\varepsilon}$ is the following prior

$$ $q_{\varepsilon}(dx,dy)\coloneqq\frac{\mathcal{K}_{{\varepsilon}}(dx,dy)}{Z_{\varepsilon}}=\pi_{\varepsilon}(dx,dy)\frac{\exp(-\phi_{\varepsilon}(x)-\psi_{\varepsilon}(y))}{Z_{\varepsilon}},$ (9) $$

and $Z_{\varepsilon}=\int\exp(-c(x,y)/{\varepsilon})\mu(dx)\nu(dy)$ is the normalizing constant.
Thus, training a CNF model using the W-CFM loss given by equation [8](https://arxiv.org/html/2507.22270v1#S3.E8) is equivalent to training a CNF using the EOT plan as the prior distribution, up to a change (a.k.a. tilt) in the marginals given by the Schrödinger potentials $\phi_{\varepsilon},\psi_{\varepsilon}$. Hence, $\mathcal{L}_{\mathrm{W-CFM}}$ can be thought of as an approximation of the following loss function

$$ $\mathcal{L}_{\mathrm{EOT-CFM}}(\theta;{\varepsilon})=\mathbb{E}_{t\sim\mathcal{U}(0,1)}\mathbb{E}_{(X,Y)\sim\pi_{\varepsilon}}\left[||v_{\theta}(t,X_{t})-(Y-X)||^{2}\right],$ (10) $$

with the approximation quality depending on the Schrödinger potentials $\phi_{\varepsilon},\psi_{\varepsilon}$.

### 3.2 Marginal Tilting under W-CFM

Using $\mathcal{K}_{{\varepsilon}}$ for the prior leads to the following tilted marginals, which are obtained by integrating equation [9](https://arxiv.org/html/2507.22270v1#S3.E9) with respect to $y$ and $x$ respectively:

$$ $\tilde{\mu}_{\varepsilon}(dx)=\frac{\exp(-\phi_{\varepsilon}(x))}{Z_{{\varepsilon}}^{1}}\mu(dx),\quad\tilde{\nu}_{\varepsilon}(dy)=\frac{\exp(-\psi_{\varepsilon}(y))}{Z_{{\varepsilon}}^{2}}\nu(dy),$ (11) $$

where $Z_{\varepsilon}^{1},Z^{2}_{\varepsilon}$ are normalizing constants.
Consequently, training a CNF using the W-CFM loss induces a vector field mapping $\tilde{\mu}_{\varepsilon}$ to $\tilde{\nu}_{\varepsilon}$.
We formalize this result in the following proposition.

###### Proposition 1 (Marginal tilting and continuity equation) .

Assume $\mu,\nu\in\mathcal{P}(\mathbb{R}^{d})$ have finite second moment. Consider the variational problem

$$ $\min_{v}\;\mathbb{E}_{t\sim\mathcal{U}(0,1)}\,\mathbb{E}_{(X,Y)\sim\mu\otimes\nu}\bigl{[}w_{\varepsilon}(X,Y)\,\|v(t,X_{t})-(Y-X)\|^{2}\bigr{]},\quad X_{t}=(1-t)\,X+t\,Y.$ (12) $$

Let $\rho_{t}$ denote the law of $X_{t}$ under $(X,Y)\sim q_{{\varepsilon}}$. Then, equation [12](https://arxiv.org/html/2507.22270v1#S3.E12) admits a minimizer $v_{\varepsilon}\in L^{2}([0,1]\times\mathbb{R}^{d};\rho_{t}(dx)dt)$, which is unique in that space.
Moreover $(\rho,v_{\varepsilon})$ solve the continuity equation in the weak sense

$$ $\partial_{t}\rho_{t}+\nabla\!\cdot(\rho_{t}\,v_{\varepsilon})=0,\qquad\rho_{0}=\tilde{\mu}_{\varepsilon},\;\rho_{1}=\tilde{\nu}_{\varepsilon}.$ $$

In other words, under mild regularity conditions, the flow generated by $v_{\varepsilon}$ pushes $\tilde{\mu}_{\varepsilon}$ forward onto $\tilde{\nu}_{\varepsilon}$.
We now present a way to evaluate the marginal tilting. Using equation [9](https://arxiv.org/html/2507.22270v1#S3.E9), the density ratios between tilted and original marginals are given by

$$ $f_{\varepsilon}(x)=\frac{d\tilde{\mu}_{\epsilon}}{d\mu}(x)\propto\int_{\mathbb{R}^{d}}\exp\left(-\frac{c(x,y)}{{\varepsilon}}\right)\nu(dy),\>\>g_{\varepsilon}(y)=\frac{d\tilde{\nu}_{\epsilon}}{d\nu}(y)\propto\int_{\mathbb{R}^{d}}\exp\left(-\frac{c(x,y)}{{\varepsilon}}\right)\mu(dx).$ (13) $$

These integrals can be estimated by Monte Carlo sampling. If $f_{\varepsilon}(x)$ is constant $\mu$ almost-everywhere, then one is guaranteed that the source marginal is preserved, i.e., that $\tilde{\mu}_{\varepsilon}=\mu$. Similarly, if $g_{\varepsilon}$ is constant $\nu$ almost-everywhere, then $\tilde{\nu}_{\varepsilon}=\nu$.
We now give an example which features preservation of the $\nu$ marginal.

###### Proposition 2 .

Let $\mathbb{S}^{d-1}_{R}=\{z\in\mathbb{R}^{d}:\|z\|=R\}$. Assume $c(x,y)=\|x-y\|$, and that $\mu$ is a rotation-invariant measure, for instance $\mu=\mathcal{N}(0,\sigma^{2}I)$. Then, $g_{\varepsilon}(y)\equiv C(R,{\varepsilon})$ for all $y\in\mathbb{S}^{d-1}_{R}$, where $C(R,{\varepsilon})$ is a nonnegative constant. In particular, if $\nu$ is supported on $\mathbb{S}^{d-1}_{R}$, then $\tilde{\nu}_{\varepsilon}=\nu$.

Proposition [2](https://arxiv.org/html/2507.22270v1#Thmproposition2) implies that using W-CFM with an isotropic distribution as the source, and a target distribution which is supported on a $d-1$-dimensional sphere of fixed radius will induce a coupling that does not tilt the target marginal. In particular, when using a smooth cost and an appropriate ${\varepsilon}$, we expect that the target distribution will not be tilted significantly provided that its mass is concentrated on a thin annulus, which typically happens when normalizing high-dimensional data. The proofs can be found in Appendix [A](https://arxiv.org/html/2507.22270v1#A1).

###### Proposition 1 (Marginal tilting and continuity equation) .

###### Proposition 2 .

#### 3.2.1 On the Choice of ε {\varepsilon} italic_ε

The entropy regularization constant ${\varepsilon}$ controls the trade-off between geometric bias (shorter, straighter flows) and marginal distortion (changing $\mu,\nu$ into $\tilde{\mu},\tilde{\nu}$). As suggested in the previous section, equation [13](https://arxiv.org/html/2507.22270v1#S3.E13) can be used to build a proxy to identify sensible values for ${\varepsilon}$. In particular, if the functions $f_{\varepsilon}$ and $g_{\varepsilon}$ are approximately constant over the supports of $\mu$ and $\nu$, the reweighted marginals $\tilde{\mu}$ and $\tilde{\nu}$ remain close to the original ones. In this case, the W-CFM loss in equation [8](https://arxiv.org/html/2507.22270v1#S3.E8) closely approximates the EOT-CFM loss in equation [10](https://arxiv.org/html/2507.22270v1#S3.E10). To quantify this, we estimate $f_{\varepsilon}$ and $g_{\varepsilon}$ by Monte Carlo sampling over a small number of batches and compute their relative variance, defined as $\frac{\mathrm{Var}(f_{\varepsilon}(X))}{\mathbb{E}[f_{\varepsilon}(X)]^{2}}$—and analogously for $g_{\varepsilon}$. This metric measures how close the functions are to being constant and is invariant to scaling by a constant factor. This invariance ensures the metric is comparable across datasets, allowing consistent evaluation of how much the importance weights distort the marginals. Low relative variance indicates that the induced marginals are close to $\mu$ and $\nu$, suggesting that the selected ${\varepsilon}$ yields a good approximation. Large variance indicates significant distortion and a potential mismatch between training and target distributions.

As an initial heuristic for choosing ${\varepsilon}$ in high-dimensional settings (e.g., images or language embeddings) with Euclidean cost $c(x,y)=\|x-y\|$, we rely on the observation that normalized high-dimensional data typically concentrates near a thin spherical shell of radius $\sqrt{d}$, where $d$ is the data dimension (see concentration of measures in Vershynin (2018)). Consequently, the typical inter-sample distance (and hence the typical value of the $L^{2}$ cost function) is $\mathcal{O}(\sqrt{d})$. Selecting ${\varepsilon}$ on this same scale ensures that the ratio inside the exponential of the weighting function in equation [8](https://arxiv.org/html/2507.22270v1#S3.E8) is $\mathcal{O}(1)$(^1^11The same logic applies for different cost functions, e.g., when the cost function is the squared Euclidean norm, then our epsilon should be $O(d)$.). In this way, the kernel varies slowly with respect to typical data variations since most pairwise distances become comparable to ${\varepsilon}$. This reasoning is similar to the heuristic commonly used in kernel methods (e.g., SVMs), where the Gaussian kernel width is set proportional to the median pairwise distance between points (Christianini et al., 2000).

Following this rationale, in all our experiments we set ${\varepsilon}=\kappa\sqrt{d}$, with the scalar $\kappa$ tuned efficiently (within seconds and a few lines of code) using the relative variance proxy described previously. Specifically, we search over a grid of $\kappa$ values spaced uniformly in log scale and select the smallest value for which the relative variance starts flattening, following an ”elbow rule” heuristic akin to the selection of the number of principal components in PCA (Jolliffe, 2002). Additionally, we experimented with various schedulers for ${\varepsilon}$ (including cosine, exponential, and linear), none of which showed a significant improvement over a fixed constant ${\varepsilon}$. The exact values of ${\varepsilon}$ we used are reported along with the results in Section [5](https://arxiv.org/html/2507.22270v1#S5).

### 3.3 Equivalence to OT-CFM in the Large Batch Limit

OT–CFM relies on solving a mini‐batch optimal transport problem for each batch, which (i) requires batch sizes large enough to represent every mode or class—otherwise the empirical OT plan might be a poor approximation of the true OT plan—and (ii) incurs at least cubic (or quadratic for the entropic approximation) cost in the batch size. By contrast, W-CFM uses a simple and virtually free per-pair Gibbs weight $w_{\varepsilon}(x,y)=\exp(-c(x,y)/{\varepsilon})$, avoiding any coupling step. Under mild regularity (no marginal tilt, bounded support), one can show that as the batch size goes to infinity, the batch-level EOT-CFM loss converges to a limit which is proportional to the W-CFM loss. We formalize this in Proposition [3](https://arxiv.org/html/2507.22270v1#Thmproposition3) below—proof is given in Appendix [A](https://arxiv.org/html/2507.22270v1#A1). A more detailed discussion can be found in Appendix [B](https://arxiv.org/html/2507.22270v1#A2).

###### Proposition 3 .

Let ${\varepsilon}>0$. Suppose that $\mu,\nu,c$ are such that equation [1](https://arxiv.org/html/2507.22270v1#S2.E1) is finite and $\mu,\nu$ have bounded support. Let $(t_{n},x_{n},y_{n})_{n\geq 1}$ be iid samples of $\mathcal{U}(0,1)\otimes\mu\otimes\nu$. Assume that $\tilde{\mu}_{\varepsilon}=\mu$ and $\tilde{\nu}_{{\varepsilon}}=\nu$. Let $\pi_{\varepsilon}$ be the optimal EOT plan between $\mu$ and $\nu$. Let $\pi_{\varepsilon}^{n}$ be the optimal EOT plan between the empirical distributions $\mathbf{x}_{n}=\frac{1}{n}\sum_{i=1}^{n}\delta_{x_{i}}$ and $\mathbf{y}_{n}=\frac{1}{n}\sum_{i=1}^{n}\delta_{y_{i}}$. Then, $\pi^{n}_{\varepsilon}\to\pi_{\varepsilon}$ almost surely as $n\to\infty$ in the weak sense. In particular, if $\mathcal{B}_{n}=\{(t_{i},x_{i},y_{i}):1\leq i\leq n\}$ and $v_{\theta}(t,z)$ is uniformly integrable in $t\in[0,1]$, continuous in $z\in\mathbb{R}^{d}$, we have, for any $\theta$

$$ $\mathbb{E}\left[L_{\mathrm{EOT-CFM}}(\mathcal{B}_{n},\theta;{\varepsilon})\right]\to l(\theta;{\varepsilon})\propto\mathcal{L}_{\mathrm{W-CFM}}(\theta;{\varepsilon}),\>\text{as $n\to\infty$},$ $$

where the expectation is taken over the random batch $\mathcal{B}_{n}$ and

$$ $L_{\mathrm{EOT-CFM}}(\mathcal{B}_{n},\theta;{\varepsilon})=\frac{1}{n}\sum_{i,j=1}^{n}\pi_{\varepsilon}^{n}(x_{i},y_{j})\|v_{\theta}(t_{i},(1-t_{i})x_{i}+t_{i}y_{j})-(y_{j}-x_{i})\|^{2}.$ $$

###### Proposition 3 .

## 4 Related Work

##### Straightening Sample Paths with Optimal Transport.

We review the existing literature on techniques to learn straighter sample paths for flow-based models. Including OT priors within the maximum likelihood training of CNFs has been considered by Onken et al. (2021).
Zhang et al. (2025) propose to learn an acceleration field to capture a random vector field, thereby allowing sample trajectories to cross and leading to straighter trajectories.
As highlighted above, our work is closely related to the minibatch optimal transport method proposed by Lipman et al. (2023) and Pooladian et al. (2023) concurrently, which leads to OT-CFM.
Within OT-CFM, Klein et al. (2025) have advocated for using the Sinkhorn algorithm (hence computing EOT) to scale to larger batch sizes.
Other works have explored an adaptation of OT-CFM to conditional generation through conditional optimal transport (Kerrigan et al., 2024; Cheng & Schwing, 2025).
Finally, recent efforts have been made to learn straighter marginal probability paths directly using Wasserstein gradient flows and the JKO scheme (Choi et al., 2024). Compared to OT-CFM, our method has no issue scaling with the batch size, see the discussion in Section [3.3](https://arxiv.org/html/2507.22270v1#S3.SS3).

##### Other Approaches for Faster Inference.

One important line of research has been to distill existing models (Luhman & Luhman, 2021; Salimans & Ho, 2022), either diffusion or flow-based, and learn the corresponding flow maps so that inference can be done in very few steps. Liu et al. (2023) propose to learn straighter trajectories via ReFlow steps, where one trains a student model using couplings generated by a base teacher model trained by flow matching.
The ReFlow paradigm is still of interest and continues to be improved Kim et al. (2025). Consistency models Song et al. (2023) were developed to overcome the cost of discretizing the probability flow ODE in diffusion models, by using the self-consistency property of the underlying flow map. This approach has been translated to flow matching models by Yang et al. (2024).
Finally, some great effort has been put into using efficient integrators tailored for these models (Lu et al., 2022; Liu et al., 2022; Zhang & Chen, 2023; Sabour et al., 2024; Williams et al., 2024). Compared to ReFlow, our method does not rely on synthetic data for training and does not require integrating a base neural ODE to generate the training pairs.

Finally, we acknowledge that a potential drawback of our approach is that it effectively learns a mapping between tilted versions of the source and target distributions. Although we conclude that the tilting is negligible in high-dimensional settings with a careful choice of ${\varepsilon}$ in Section [3.2.1](https://arxiv.org/html/2507.22270v1#S3.SS2.SSS1), one could benefit from being able to adjust for the tilting at inference, in order to get a genuine mapping between $\mu$ and $\nu$. We leave this challenge for future research.

## 5 Experiments

We evaluate our flow matching framework across three complementary domains: 2D toy transports, unconditional image generation, and a fidelity and diversity analysis of the generated samples. In these experiments, we demonstrate our framework’s performance and competitiveness in comparison to well-established baselines.

### 5.1 Experimental Setup

To visually probe the benefits of our weighted loss, we design similar low‐dimensional transport benchmarks as in Tong et al. (2024). First, we focus on mapping a distribution concentrated on an annulus to a configurable Mixture of Gaussians (MoG). The second setup consists of recovering the moons 2D dataset from a MoG source.
We compare W-CFM for different choices of ${\varepsilon}$ with the cost $c(x,y)=\|x-y\|$ against both OT-CFM and I-CFM, training a two-layer ELU‐MLP with 64 hidden units per layer via Adam with a learning rate of $10^{-3}$ for 60,000 iterations with a default batch size of 48. We evaluate sample quality, path straightness, and marginal density estimates—using KDE contours—to highlight how the added bias in the marginals W-CFM is negligible.

To validate our approach in higher-dimensional settings, we evaluate on CIFAR-10, CelebA64, and ImageNet64-10—a 64×64 version of 10 ImageNet classes (Deng et al., 2009). We use a UNet backbone (Ronneberger et al., 2015) adapted to each dataset: for CIFAR-10, a smaller model with two residual blocks, 64 base channels, and 16×16 attention; for the rest of the datasets, a deeper UNet with three residual blocks, 128 base channels, a [1, 2, 2, 4] channel multiplier, and additional attention at 32×32 for ImageNet64-10, Food20, and Intel. All models are trained with Adam, a learning rate of $2\times 10^{-4}$, cosine learning rate scheduling with 5,000 warmup steps, and EMA (decay 0.9999), for 400,000 steps using batch sizes of 128 for CIFAR-10, 64 for CelebA and ImageNet64-10, and 48 for Food20 and Intel. Our goal is not to reach state-of-the-art performance, but to compare flow matching variants under matched computational budgets and architectures.

### 5.2 Illustrative Comparison on Toy Datasets

Figure: Figure 1: Sample trajectories for MoG. From left to right, the models used are trained with: I-CFM, W-CFM (${\varepsilon}=0.4$), W-CFM (${\varepsilon}=0.2$), OT-CFM (batch size 16), and OT-CFM.
Refer to caption: plots/trajs_mog_otbs16.png

##### Sample Trajectories on Mixture of Gaussians.

We consider the task of mapping a mixture of many low-variance Gaussians, whose support is concentrated on an annulus, to a target distribution consisting of a mixture of five low-variance Gaussians. We train W-CFM with ${\varepsilon}=0.2$ (small ${\varepsilon}$) and ${\varepsilon}=0.4$ (large ${\varepsilon}$). A specific instance where OT-CFM faces challenges is when a typical batch is not representative of the true target distribution (Kornilov et al., 2024; Klein et al., 2025). Hence, on top of training OT-CFM with the default parameters, we train OT-CFM with a smaller batch size to emulate a scenario with a high number of clusters to batch size ratio (we compensate for the smaller batch size by increasing the number of iterations for this version). We plot sample trajectories of the trained models in Figure [1](https://arxiv.org/html/2507.22270v1#S5.F1), and report the performance of the models in Table [1](https://arxiv.org/html/2507.22270v1#S5.T1). We use the $W_{2}^{2}$ distance between generated samples and the true target for overall sample quality, and compute the normalized path energy $\mathrm{NPE(\theta)}=W_{2}^{-2}(\mu,\nu)\left|\mathbb{E}\left[\int_{0}^{1}||v_{\theta}(t,x_{t})||^{2}dt\right]-W_{2}^{2}(\mu,\nu)\right|$ for straightness of the trajectories (Tong et al., 2024). Overall, W-CFM leads to better sample quality than OT-CFM with trajectories of similar straightness. We also observe a reduction in straightness for some of the paths when training OT-CFM with a smaller batch size.

**Table 1: Comparison of CFM training algorithms’ performance on 2D datasets generation on 5 random seeds. $W_{2}^{2}$ measures the overall quality of sample generation (lower is better), NPE measures the straightness of trajectories (lower is better), using the true optimal transport cost as a reference.**
| Dataset $\to$ | $\text{Circular MoG}\to 5\text{ Gaussians}$ | $8\text{ Gaussians}\to\text{moons}$ |  |  |
| --- | --- | --- | --- | --- |
| Algorithm $\downarrow$  Metric $\to$ | $W_{2}^{2}$ | NPE | $W_{2}^{2}$ | NPE |
| I-CFM | 0.056 $\pm$ 0.030 | 1.690 $\pm$ 0.184 | 0.703 $\pm$ 0.316 | 1.049 $\pm$ 0.098 |
| OT-CFM | 0.038 $\pm$ 0.010 | 0.077 $\pm$ 0.010 | 0.294 $\pm$ 0.131 | 0.147 $\pm$ 0.022 |
| OT-CFM (B = 16) | 0.040 $\pm$ 0.016 | 0.176 $\pm$ 0.032 | 0.633 $\pm$ 0.062 | 0.063 $\pm$ 0.043 |
| W-CFM (small ${\varepsilon}$) | 0.021 $\pm$ 0.003 | 0.082 $\pm$ 0.011 | 1.703 $\pm$ 0.230 | 0.302 $\pm$ 0.041 |
| W-CFM (large ${\varepsilon}$) | 0.031 $\pm$ 0.005 | 0.135 $\pm$ 0.026 | 0.671 $\pm$ 0.308 | 0.549 $\pm$ 0.084 |

The W-CFM method outperforms OT-CFM and I-CFM in terms of sample quality, and it exhibits straightness on par with OT-CFM. In particular, we observe that the trajectories of W-CFM are straighter than those in OT-CFM with a small batch size. Overall, this suggests that our method can be well-suited for unconditional generation involving a target distribution with many clusters.

##### Marginal Tilting on Moons Generation.

As highlighted above, our method might induce a tilting of the marginal distributions. We investigate this phenomenon when generating moons from a mixture of 8 Gaussians, and training W-CFM with ${\varepsilon}=2$ (small ${\varepsilon}$) and ${\varepsilon}=10$ (large ${\varepsilon}$). As seen in Figure [2](https://arxiv.org/html/2507.22270v1#S5.F2), using W-CFM leads to straighter paths compared to I-CFM as measured by NPE, even for a large value of ${\varepsilon}$; however, it also causes some bias in the marginals as ${\varepsilon}$ decreases, as shown in Figure [3](https://arxiv.org/html/2507.22270v1#S5.F3), which leads to poor sample generation quality, see Table [1](https://arxiv.org/html/2507.22270v1#S5.T1). This can be mitigated by choosing a larger value of ${\varepsilon}$, while maintaining some straightening effect. We present results for ${\varepsilon}\in\{2,4,6,8,10\}$ for further validation of this tradeoff in Appendix [C](https://arxiv.org/html/2507.22270v1#A3). Finally, although OT-CFM with a small batch size seems to outperform the other models in NPE, we want to highlight that this metric measures straightness conditioned on the fact that the model does generate good samples of the target distribution, which is arguable for this instance of the model.

Figure: Figure 2: Sample trajectories for moons generation. Blue dots are source samples, red dots are generated samples. From left to right, the models used are trained with: I-CFM, W-CFM (${\varepsilon}=10$), W-CFM (${\varepsilon}=2$), and OT-CFM.
Refer to caption: plots/moons_eps10_2.png

Figure: Figure 3: Contour plots of learned density for moons (using 50,000 generated samples). The leftmost plot corresponds to the true target distribution. Then, from left to right, the models used are trained with: I-CFM, W-CFM (${\varepsilon}=10$), W-CFM (${\varepsilon}=2$), and OT-CFM. We observe that choosing a small value of ${\varepsilon}$ for W-CFM leads to a ”disentanglement” of the two generated moons.
Refer to caption: plots/kde_moons_eps10_2.png

### 5.3 Comparison on Image Datasets

To evaluate the generation capabilities of W-CFM beyond low-dimensional toy examples, we conducted unconditional image generation experiments across five different benchmarks: CIFAR-10, CelebA64, ImageNet64-10, Intel Image Classification, and Food20 (a subset of Food101 (Bossard et al., 2014)). Table [2](https://arxiv.org/html/2507.22270v1#S5.T2) summarizes the Fréchet Inception Distance (FID, Heusel et al. (2017)) scores obtained by our proposed W-CFM method compared to I-CFM and OT-CFM with an adaptive solver. W-CFM consistently matches or outperforms the baselines, achieving the best FID scores on CIFAR-10, ImageNet64-10, Intel, and Food20, while remaining competitive on CelebA64. The slightly better performance of OT-CFM on CelebA64 is likely explained by the relatively unimodal nature of the dataset (Zhang, 2023), which alleviates the mode representation issues intrinsic to minibatch OT methods (as discussed in Section [3.3](https://arxiv.org/html/2507.22270v1#S3.SS3)). Conversely, the multimodal structure of datasets like CIFAR-10, ImageNet64-10, Intel, and Food20 highlights the advantage of W-CFM. Batch sizes were 128 for CIFAR-10, 64 for ImageNet64-10 and CelebA64, and 48 for the remaining datasets. Applying the proxy described in Section [3.2.1](https://arxiv.org/html/2507.22270v1#S3.SS2.SSS1) for selecting ${\varepsilon}$, we used ${\varepsilon}=5$ for CIFAR-10, ${\varepsilon}=13$ for CelebA64, ${\varepsilon}=14$ for ImageNet64-10, ${\varepsilon}=14$ for Intel, and ${\varepsilon}=15$ for Food20.

**Table 2: FID $\downarrow$ (lower is better) across datasets for different flow matching models with the Dopri5 solver.**
| Model | CIFAR-10 | CelebA64 | ImageNet64-10 | Intel | Food20 |
| --- | --- | --- | --- | --- | --- |
| I-CFM | 7.44 | 21.99 | 13.86 | 27.54 | 8.15 |
| OT-CFM | 7.60 | 20.93 | 14.39 | 25.63 | 8.23 |
| W-CFM | 7.33 | 21.96 | 13.56 | 25.22 | 7.93 |

We further assessed model efficiency by comparing FID scores at various numbers of neural function evaluations (NFEs) during Euler integration. Table [3](https://arxiv.org/html/2507.22270v1#S5.T3) shows that W-CFM consistently achieves lower or comparable FID scores at fewer NFEs compared to both I-CFM and OT-CFM. On CelebA64, OT-CFM outperforms both models, which again reflects the dataset’s unimodal structure that mitigates OT-CFM’s batch limitations. Illustrative examples of generated samples for each dataset are included in Appendix [C](https://arxiv.org/html/2507.22270v1#A3).

**Table 3: FID $\downarrow$ at varying number of neural function evaluations (NFE) using Euler discretization. Each column reports FID at 50, 100, and 120 NFE.**
| FID @ NFE |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Dataset | I-CFM | OT-CFM | W-CFM |  |  |  |  |  |  |
|  | 50 | 100 | 120 | 50 | 100 | 120 | 50 | 100 | 120 |
| CIFAR-10 | 10.87 | 9.76 | 8.68 | 11.03 | 9.89 | 8.53 | 10.53 | 9.28 | 8.08 |
| CelebA64 | 29.49 | 25.26 | 24.50 | 27.76 | 23.86 | 22.93 | 29.32 | 25.22 | 24.37 |
| ImageNet64-10 | 14.94 | 13.91 | 13.86 | 15.67 | 14.78 | 14.68 | 15.82 | 14.17 | 13.71 |
| Intel | 26.72 | 26.40 | 26.20 | 25.45 | 25.98 | 24.26 | 25.01 | 24.47 | 24.08 |
| Food20 | 10.10 | 8.98 | 8.85 | 10.16 | 9.17 | 8.95 | 10.01 | 8.97 | 8.57 |

### 5.4 Fidelity and Diversity of the Generated Samples

A natural concern with the importance weighting in W-CFM is that it could skew the learned flow toward easier-to-transport pairs, potentially under-representing low-probability modes or degrading sample quality in certain regions of the data manifold. To test this, we evaluate sample fidelity and diversity using the precision-recall-density-coverage (PRDC) suite of metrics (Naeem et al., 2020) on $10{,}000$ generated samples and $5{,}000$ held-out real images per dataset. Precision measures the fraction of generated samples near the real data manifold, recall assesses coverage of the real distribution, and density and coverage estimate support concentration and breadth, respectively. F1 summarizes the trade-off via the harmonic mean of precision and recall.

**Table 4: Sample quality and diversity metrics on ImageNet64-10.**
| Model | Precision ($\uparrow$) | Recall ($\uparrow$) | Density ($\uparrow$) | Coverage ($\uparrow$) | F1 ($\uparrow$) |
| --- | --- | --- | --- | --- | --- |
| I-CFM | 0.75 | 0.69 | 0.91 | 0.94 | 0.72 |
| OT-CFM | 0.74 | 0.67 | 0.91 | 0.96 | 0.71 |
| W-CFM | 0.75 | 0.68 | 0.94 | 0.92 | 0.72 |

Table [4](https://arxiv.org/html/2507.22270v1#S5.T4) reports the results on ImageNet64-10, showing that W-CFM maintains comparable or better recall, F1, and density than both I-CFM and OT-CFM, without sacrificing coverage or precision. We observe the same qualitative trends across the remaining datasets, with all three methods exhibiting nearly identical PRDC metrics. For completeness, we also include the corresponding tables for CIFAR-10 and CelebA64 in Appendix [C](https://arxiv.org/html/2507.22270v1#A3). These results confirm that W-CFM does not impair diversity or mode coverage, even when trained with independent pairwise weighting.

## 6 Conclusion

In this work, we introduced Weighted Conditional Flow Matching (W-CFM), a novel method that leverages insights from entropic optimal transport (EOT) to efficiently improve path straightness and sample quality in continuous normalizing flows. By weighting each training pair using a Gibbs kernel, our approach approximates the entropic OT plan without incurring the computational overhead associated with explicit OT computations, and crucially, without being limited by the batch size when solving the optimal transport problem. We derived theoretical conditions under which our approximation preserves the original marginals and demonstrated equivalence with OT-CFM in the large-batch limit, given that the marginals remain unchanged. Empirical evaluations across low-dimensional toy problems, unconditional image generation tasks, and fidelity-diversity analyses confirm that W-CFM achieves performance competitive with OT-CFM while maintaining the computational efficiency of standard CFM.

## References

- Altschuler et al. (2017)
Jason Altschuler, Jonathan Niles-Weed, and Philippe Rigollet.
Near-linear time approximation algorithms for optimal transport via sinkhorn iteration.
In *Advances in Neural Information Processing Systems*, 2017.
- Bossard et al. (2014)
Lukas Bossard, Matthieu Guillaumin, and Luc Van Gool.
Food-101 – mining discriminative components with random forests.
In *European Conference on Computer Vision*, 2014.
- Chen et al. (2018)
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud.
Neural ordinary differential equations.
In *Advances in Neural Information Processing Systems*, 2018.
- Cheng & Schwing (2025)
Ho Kei Cheng and Alexander Schwing.
The curse of conditions: Analyzing and improving optimal transport for conditional flow-based generation.
*arXiv preprint arXiv:2503.10636*, 2025.
- Choi et al. (2024)
Jaemoo Choi, Jaewoong Choi, and Myungjoo Kang.
Scalable wasserstein gradient flow for generative modeling through unbalanced optimal transport.
*arXiv preprint arXiv:2402.05443*, 2024.
- Christianini et al. (2000)
Nello Christianini, John Shawe-Taylor, et al.
*An introduction to support vector machines and other kernel-based learning methods*.
Cambridge university press Cambridge, 2000.
- Cuturi (2013)
Marco Cuturi.
Sinkhorn distances: Lightspeed computation of optimal transport.
In *Advances in Neural Information Processing Systems*, 2013.
- Deng et al. (2009)
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.
Imagenet: A large-scale hierarchical image database.
In *2009 IEEE Conference on Computer Vision and Pattern Recognition*, 2009.
- Geffner et al. (2025)
Tomas Geffner, Kieran Didi, Zuobai Zhang, Danny Reidenbach, Zhonglin Cao, Jason Yim, Mario Geiger, Christian Dallago, Emine Kucukbenli, Arash Vahdat, et al.
Proteina: Scaling flow-based protein structure generative models.
*arXiv preprint arXiv:2503.00710*, 2025.
- Genevay et al. (2018)
Aude Genevay, Gabriel Peyre, and Marco Cuturi.
Learning generative models with sinkhorn divergences.
In *International Conference on Artificial Intelligence and Statistics*, 2018.
- Ghosal et al. (2022)
Promit Ghosal, Marcel Nutz, and Espen Bernton.
Stability of entropic optimal transport and schrödinger bridges.
*Journal of Functional Analysis*, 2022.
- Grathwohl et al. (2018)
Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud.
Ffjord: Free-form continuous dynamics for scalable reversible generative models.
*arXiv preprint arXiv:1810.01367*, 2018.
- Guo et al. (2024)
Yiwei Guo, Chenpeng Du, Ziyang Ma, Xie Chen, and Kai Yu.
Voiceflow: Efficient text-to-speech with rectified flow matching.
In *ICASSP 2024-2024 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)*, 2024.
- Heusel et al. (2017)
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter.
Gans trained by a two time-scale update rule converge to a local nash equilibrium.
In *Advances in Neural Information Processing Systems*, 2017.
- Irwin et al. (2024)
Ross Irwin, Alessandro Tibo, Jon Paul Janet, and Simon Olsson.
Efficient 3d molecular generation with flow matching and scale optimal transport.
In *International Conference on Machine Learning*, 2024.
- Jolliffe (2002)
I.T. Jolliffe.
*Principal Component Analysis*.
Springer Series in Statistics. Springer, 2002.
- Kerrigan et al. (2024)
Gavin Kerrigan, Giosue Migliorini, and Padhraic Smyth.
Dynamic conditional optimal transport through simulation-free flows.
In *Advances in Neural Information Processing Systems*, 2024.
- Kim et al. (2025)
Beomsu Kim, Yu-Guan Hsieh, Michal Klein, Marco Cuturi, Jong Chul Ye, Bahjat Kawar, and James Thornton.
Simple reflow: Improved techniques for fast flow models.
In *International Conference on Learning Representations*, 2025.
- Klein et al. (2025)
Michal Klein, Alireza Mousavi-Hosseini, Stephen Zhang, and Marco Cuturi.
On fitting flow models with large sinkhorn couplings.
*arXiv preprint arXiv:2506.05526*, 2025.
- Kornilov et al. (2024)
Nikita Kornilov, Petr Mokrov, Alexander Gasnikov, and Aleksandr Korotin.
Optimal flow matching: Learning straight trajectories in just one step.
In *Advances in Neural Information Processing Systems*, 2024.
- Lipman et al. (2023)
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le.
Flow matching for generative modeling.
In *International Conference on Learning Representations*, 2023.
- Liu et al. (2022)
Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao.
Pseudo numerical methods for diffusion models on manifolds.
In *International Conference on Learning Representations*, 2022.
- Liu et al. (2023)
Xingchao Liu, Chengyue Gong, and Qiang Liu.
Flow straight and fast: Learning to generate and transfer data with rectified flow.
In *International Conference on Learning Representations*, 2023.
- Lu et al. (2022)
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu.
Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps.
In *Advances in Neural Information Processing Systems*, 2022.
- Luhman & Luhman (2021)
Eric Luhman and Troy Luhman.
Knowledge distillation in iterative generative models for improved sampling speed.
*arXiv preprint arXiv:2101.02388*, 2021.
- Naeem et al. (2020)
Muhammad Ferjad Naeem, Seong Joon Oh, Youngjung Uh, Yunjey Choi, and Jaejun Yoo.
Reliable fidelity and diversity metrics for generative models.
In *International Conference on Machine Learning*, 2020.
- Nutz (2021)
Marcel Nutz.
Introduction to entropic optimal transport.
*Lecture notes, Columbia University*, 2021.
- Onken et al. (2021)
Derek Onken, Samy Wu Fung, Xingjian Li, and Lars Ruthotto.
Ot-flow: Fast and accurate continuous normalizing flows via optimal transport.
In *Proceedings of the AAAI Conference on Artificial Intelligence*, 2021.
- Peyré & Cuturi (2019)
Gabriel Peyré and Marco Cuturi.
*Computational Optimal Transport: With Applications to Data Science*.
Now Publishers, Inc., 2019.
- Pooladian et al. (2023)
Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, and Ricky T. Q. Chen.
Multisample flow matching: Straightening flows with minibatch couplings.
In *International Conference on Machine Learning*, 2023.
- Rohbeck et al. (2025)
Martin Rohbeck, Charlotte Bunne, Edward De Brouwer, Jan-Christian Huetter, Anne Biton, Kelvin Y. Chen, Aviv Regev, and Romain Lopez.
Modeling complex system dynamics with flow matching across time and conditions.
In *International Conference on Learning Representations*, 2025.
- Ronneberger et al. (2015)
Olaf Ronneberger, Philipp Fischer, and Thomas Brox.
U-net: Convolutional networks for biomedical image segmentation.
In *Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18*. Springer, 2015.
- Sabour et al. (2024)
Amirmojtaba Sabour, Sanja Fidler, and Karsten Kreis.
Align your steps: Optimizing sampling schedules in diffusion models.
In *International Conference on Machine Learning*, 2024.
- Salimans & Ho (2022)
Tim Salimans and Jonathan Ho.
Progressive distillation for fast sampling of diffusion models.
*arXiv preprint arXiv:2202.00512*, 2022.
- Santambrogio (2015)
Filippo Santambrogio.
*Optimal Transport for Applied Mathematicians*.
Springer, 2015.
- Song et al. (2023)
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever.
Consistency models.
In *International Conference on Machine Learning*, 2023.
- Stark et al. (2024)
Hannes Stark, Bowen Jing, Chenyu Wang, Gabriele Corso, Bonnie Berger, Regina Barzilay, and Tommi Jaakkola.
Dirichlet flow matching with applications to dna sequence design.
*arXiv preprint arXiv:2402.05841*, 2024.
- Tong et al. (2024)
Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Guy Wolf, and Yoshua Bengio.
Improving and generalizing flow-based generative models with minibatch optimal transport.
*Transactions on Machine Learning Research*, 2024.
- Varadarajan (1958)
Veeravalli S Varadarajan.
On the convergence of sample probability distributions.
*Sankhyā: The Indian Journal of Statistics (1933-1960)*, 1958.
- Vershynin (2018)
Roman Vershynin.
*High-Dimensional Probability: An Introduction with Applications in Data Science*.
Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2018.
- Williams et al. (2024)
Christopher Williams, Andrew Campbell, Arnaud Doucet, and Saifuddin Syed.
Score-optimal diffusion schedules.
In *Advances in Neural Information Processing Systems*, 2024.
- Yang et al. (2024)
Ling Yang, Zixiang Zhang, Zhilong Zhang, Xingchao Liu, Minkai Xu, Wentao Zhang, Chenlin Meng, Stefano Ermon, and Bin Cui.
Consistency flow matching: Defining straight flows with velocity consistency.
*arXiv preprint arXiv:2407.02398*, 2024.
- Zhang (2023)
Na Zhang.
A study on the impact of face image quality on face recognition in the wild.
*arXiv preprint arXiv:2307.02679*, 2023.
- Zhang & Chen (2023)
Qinsheng Zhang and Yongxin Chen.
Fast sampling of diffusion models with exponential integrator.
In *International Conference on Learning Representations*, 2023.
- Zhang et al. (2024)
Xi Nicole Zhang, Yuan Pu, Yuki Kawamura, Andrew Loza, Yoshua Bengio, Dennis Shung, and Alexander Tong.
Trajectory flow matching with applications to clinical time series modelling.
In *Advances in Neural Information Processing Systems*, 2024.
- Zhang et al. (2025)
Yichi Zhang, Yici Yan, Alex Schwing, and Zhizhen Zhao.
Towards hierarchical rectified flow.
In *International Conference on Learning Representations*, 2025.

## Appendix A Proofs of Theoretical Results

### A.1 Proof of Proposition 1

Recall the prior $q_{\varepsilon}$ induced by the Gibbs kernel $\mathcal{K}_{\varepsilon}(dx,dy)=\exp(-c(x,y)/{\varepsilon})\mu(dx)\nu(dy)$ defined in equation [9](https://arxiv.org/html/2507.22270v1#S3.E9). Recall that $\rho_{t}$ denotes the distribution of $X_{t}=(1-t)X+tY$ under $(X,Y)\sim q_{\varepsilon}$. For any $v\in L^{2}([0,1]\times\mathbb{R}^{d};\rho_{t}(dx)dt)$, we have

$$ $\displaystyle\mathbb{E}_{t\sim\mathcal{U}(0,1)}\mathbb{E}_{(X,Y)\sim\mu\otimes\nu}\left[w_{\varepsilon}(X,Y)||v(t,X_{t})-(Y-X)||^{2}\right]$ $\displaystyle=\mathbb{E}_{t\sim\mathcal{U}(0,1)}\mathbb{E}_{(X,Y)\sim q_{{\varepsilon}}}\left[\frac{d(\mu\otimes\nu)}{dq_{{\varepsilon}}}(X,Y)\exp(-c(X,Y)/{\varepsilon})||v(t,X_{t})-(Y-X)||^{2}\right]$ $\displaystyle=Z_{\varepsilon}\mathbb{E}_{t\sim\mathcal{U}(0,1)}\mathbb{E}_{(X,Y)\sim q_{{\varepsilon}}}\left[||v(t,X_{t})-(Y-X)||^{2}\right],$ $$

where $Z_{\varepsilon}$ denotes the normalizing constant $Z_{\varepsilon}\coloneqq\mathbb{E}_{(X,Y)\sim\mu\otimes\nu}[w_{\varepsilon}(X,Y)]>0$.
Hence the variational problem given by equation [12](https://arxiv.org/html/2507.22270v1#S3.E12) is equivalent to

$$ $\min_{v}\mathbb{E}_{t\sim\mathcal{U}(0,1)}\mathbb{E}_{(X,Y)\sim q_{\varepsilon}}\left[||v(t,X_{t})-(Y-X)||^{2}\right].$ (14) $$

By the $L^{2}$-projection property of conditional expectations, the variational problem of equation [14](https://arxiv.org/html/2507.22270v1#A1.E14) is solved by the function $v_{\varepsilon}:[0,1]\times\mathbb{R}^{d}\to\mathbb{R}^{d}$ defined by

$$ $v_{\varepsilon}(t,z)=\mathbb{E}_{(X,Y)\sim q_{\varepsilon}}[Y-X\mid X_{t}=z].$ (15) $$

Note that this definition is unique in $L^{2}([0,1]\times\mathbb{R}^{d};\rho_{t}(dx)dt)$.
We now check that $v_{\varepsilon}$ generates a valid probability path between $\tilde{\mu}_{\varepsilon}$ and $\tilde{\nu}_{{\varepsilon}}$, i.e., that $(\rho,v_{\varepsilon})$ satisfy the continuity equation equation [4](https://arxiv.org/html/2507.22270v1#S2.E4) in the weak sense. Clearly, $v_{\varepsilon}(t,\cdot)\in L^{1}(\mathbb{R}^{d},\rho_{t})$ and $\int_{0}^{1}\int_{\mathbb{R}^{d}}|v_{\varepsilon}(t,x)|p(t,x)dxdt<\infty$. By Proposition 4.2 in Santambrogio (2015), it is enough to check that the continuity equation is satisfied in the sense of distributions. Let $\phi\in C_{c}^{1}((0,1)\times\mathbb{R}^{d})$, then

$$ $\displaystyle\int_{0}^{1}\int_{\mathbb{R}^{d}}\partial_{t}\phi(t,x)\rho_{t}(dx)dt+\int_{0}^{1}\int_{\mathbb{R}^{d}}\nabla\phi(t,x)\cdot v_{\varepsilon}(t,x)\rho_{t}(dx)dt$ $\displaystyle=\int_{0}^{t}\mathbb{E}[\partial_{t}\phi(t,X_{t})+\nabla\phi(t,X_{t})\cdot(Y-X)]dt=\mathbb{E}[\phi(1,Y)-\phi(0,X)]=0.$ $$

∎

### A.2 Proof of Proposition 2

Let $y_{2},y_{2}\in\mathbb{S}^{d-1}_{R}$. We can get hold of $\varphi\in O(d)$ (i.e. a distance-preserving transformation in $\mathbb{R}^{d}$) such that $\varphi(y_{1})=y_{2}$. Since $||\varphi(y_{1})-\varphi(x)||=||y_{1}-x||$ for all $x\in\mathbb{R}^{d}$, we have

$$ $\displaystyle\int_{\mathbb{R}^{d}}\exp\left(-\frac{\|y_{2}-x\|}{{\varepsilon}}\right))\mu(dx)$ $\displaystyle=\int_{\mathbb{R}^{d}}\exp\left(-\frac{\|\varphi(y_{1})-x\|}{{\varepsilon}}\right))\mu(dx)$ $\displaystyle=\int_{\varphi^{-1}(\mathbb{R}^{d})}\exp\left(-\frac{\|\varphi(y_{1})-\varphi(x)\|}{{\varepsilon}}\right)(\varphi^{-1}_{\#}\mu)(dx)$ $\displaystyle=\int_{\varphi^{-1}(\mathbb{R}^{d})}\exp\left(-\frac{\|y_{1}-x\|}{{\varepsilon}}\right)(\varphi^{-1}_{\#}\mu)(dx)$ $\displaystyle=\int_{\mathbb{R}^{d}}\exp\left(-\frac{\|y_{1}-x\|}{{\varepsilon}}\right)\mu(dx),$ $$

which directly implies that $g_{\varepsilon}(y_{2})=g_{{\varepsilon}}(y_{1})$
∎

### A.3 Proof of Proposition 3

Let ${\varepsilon}>0$. Recall that $(t_{n},x_{n},y_{n})_{n\geq 1}$ are iid samples of $\mathcal{U}(0,1)\otimes\mu\otimes\nu$, and that we assume $\tilde{\mu}_{\varepsilon}=\mu$ and $\tilde{\nu}_{{\varepsilon}}=\nu$. Let $\pi_{\varepsilon}$ be the optimal EOT plan between $\mu$ and $\nu$. Let $\pi_{\varepsilon}^{n}$ be the optimal EOT plan betwen $\mathbf{x}_{n}=\frac{1}{n}\sum_{i=1}^{n}\delta_{x_{i}}$ and $\mathbf{y}_{n}=\frac{1}{n}\sum_{i=1}^{n}\delta_{y_{i}}$. In this proof, the convergence of probability measures is understood in the weak sense.

First, the almost-sure convergences $\mathbf{x}_{n}\rightarrow\mu$ and $\mathbf{y}_{n}\rightarrow\nu$ come from a classical result in probability theory on the convergence of empirical distributions to the true distribution, see Varadarajan (1958).

Since the minimization problem of equation [1](https://arxiv.org/html/2507.22270v1#S2.E1) is non-trivial, an application of Theorem 1.4 in Ghosal et al. (2022) shows that the empirical EOT plan satisfies $\pi^{n}_{\varepsilon}\to\pi_{\varepsilon}$ almost surely. Now, for any $n\geq 1$, we have

$$ $\displaystyle\mathbb{E}\left[L_{\mathrm{EOT-CFM}}(\mathcal{B}_{n},\theta;{\varepsilon})\right]$ $\displaystyle=\mathbb{E}\left[\int_{\mathbb{R}^{d}\times\mathbb{R}^{d}}\int_{0}^{1}\left\|v_{\theta}(t,(1-t)x+ty)-(y-x)\right\|^{2}dt\pi_{\varepsilon}^{n}(dx,dy)\right]$ $\displaystyle=\mathbb{E}\left[\int_{\mathrm{s}(\mu)\times\mathrm{s}(\nu)}\int_{0}^{1}\left\|v_{\theta}(t,(1-t)x+ty)-(y-x)\right\|^{2}dt\pi_{\varepsilon}^{n}(dx,dy)\right],$ $$

where $\mathrm{s}(\mu),\mathrm{s}(\nu)$ denote the support of $\mu$ and $\nu$ respectively, which are assumed to be bounded. Since $(x,y)\mapsto\int_{0}^{1}\left\|v_{\theta}(t,(1-t)x+ty)-(y-x)\right\|^{2}dt$ is continuous and bounded on $\mathrm{s}(\mu)\times\mathrm{s}(\nu)$ by our assumption on $v_{\theta}$, we have

$$ $\displaystyle\int_{\mathrm{s}(\mu)\times\mathrm{s}(\nu)}\left(\int_{0}^{1}\left\|v_{\theta}(t,(1-t)x+ty)-(y-x)\right\|^{2}dt\right)\pi_{\varepsilon}^{n}(dx,dy)$ $\displaystyle\to\int_{\mathrm{s}(\mu)\times\mathrm{s}(\nu)}\left(\int_{0}^{1}\left\|v_{\theta}(t,(1-t)x+ty)-(y-x)\right\|^{2}dt\right)\pi_{\varepsilon}(dx,dy)$ $$

almost surely as $n\to\infty$. Now, the dominated convergence theorem ensures that this convergence also holds in expectation, i.e.

$$ $\mathbb{E}\left[L_{\mathrm{EOT-CFM}}(\mathcal{B}_{n},\theta;{\varepsilon})\right]\to\mathbb{\int}_{\mathrm{s}(\mu)\times\mathrm{s}(\nu)}\left(\int_{0}^{1}\left\|v_{\theta}(t,(1-t)x+ty)-(y-x)\right\|^{2}dt\right)\pi_{\varepsilon}(dx,dy).$ $$

Finally, we want to prove that this integral is proportional to $\mathcal{L}_{\mathrm{W-CFM}}(\theta;{\varepsilon})$. Since we assume no tilting of the marginals, i.e. $q_{\varepsilon}=\pi_{\varepsilon}$, we have

$$ $\displaystyle\mathcal{L}_{\mathrm{W-CFM}}(\theta;{\varepsilon})$ $\displaystyle=Z_{\varepsilon}\mathbb{E}_{t\sim\mathcal{U}(0,1),(X,Y)\sim\pi_{\varepsilon}}\left[||v_{\theta}(t,X_{t})-(Y-X)||^{2}\right]$ $\displaystyle=Z_{\varepsilon}\int_{\mathrm{s}(\mu)\times\mathrm{s}(\nu)}\left(\int_{0}^{1}\left\|v_{\theta}(t,(1-t)x+ty)-(y-x)\right\|^{2}dt\right)\pi_{\varepsilon}(dx,dy).$ $$

by using the same change of measure argument as in the proof of Proposition [1](https://arxiv.org/html/2507.22270v1#Thmproposition1). ∎

## Appendix B Details on the Equivalence to OT-CFM in the Large Batch Limit

We recall the mini-batch optimal transport technique that is central in the OT-CFM algorithm of Tong et al. (2024). Given a batch of i.i.d. samples $\mathcal{B}=\{(t_{i},x_{i},y_{i}):i=1,\ldots,B\}$, where $t_{i}$ are i.i.d. according to $\mathcal{U}(0,1)$, $x_{i}$ are i.i.d. according to $\mu$, $y_{i}$ are i.i.d. according to $\nu$, and $t_{i},x_{i},y_{i}$ are drawn independently, one can compute the optimal transport plan between the two corresponding discrete distribution, i.e. one computes

$$ $\pi_{\mathcal{B}}\in\arg\min_{\pi\in\Pi_{\mathcal{B}}}\sum_{i=1}^{B}\sum_{j=1}^{B}c(x_{i},y_{j})\pi(x_{i},y_{j}),$ (16) $$

where $\Pi_{\mathcal{B}}$ is the set of couplings between the empirical measures

$$ $\mathbf{x}_{\mathcal{B}}=\frac{1}{B}\sum_{i=1}^{B}\delta_{x_{i}},\quad\mathbf{y}_{\mathcal{B}}=\frac{1}{B}\sum_{i=1}^{B}\delta_{y_{i}}.$ $$

In particular, any $\pi\in\Pi_{\mathcal{B}}$ must satisfy $\sum_{j}\pi(x_{i},y_{j})=\sum_{i}\pi(x_{i},y_{j})=\frac{1}{B}$.
Then, given an optimal $\pi_{\mathcal{B}}$, one computes the following

$$ $L_{\mathrm{OT-CFM}}(\mathcal{B},\theta)=\frac{1}{B}\sum_{i=1}^{B}(v_{\theta}(t_{i},(1-t_{i})x_{i}+t_{i}y_{\sigma(i)})-(y_{\sigma(i)}-x_{i}))^{2},$ (17) $$

where $\sigma$ is a permutation corresponding to a Monge map for the problem equation [16](https://arxiv.org/html/2507.22270v1#A2.E16), i.e., for some $T:\{x_{i}:i=1,\ldots,B\}\to\{y_{i}:i=1,\ldots,B\}$ such that $T(x_{i})=y_{\sigma(i)}$ and $\pi_{\mathcal{B}}\coloneqq(\mathrm{Id},T)_{\#}\mathbf{x}_{\mathcal{B}}$ is a solution to equation [16](https://arxiv.org/html/2507.22270v1#A2.E16) (Peyré & Cuturi, 2019). This sample loss is used as an approximation of the following OT-CFM loss

$$ $\mathcal{L}_{\mathrm{OT-CFM}}(\theta)\coloneqq\mathbb{E}_{t\sim\mathcal{U}(0,1)}\mathbb{E}_{(X,Y)\sim\pi^{\star}}\left[||v_{\theta}(t,X_{t})-(Y-X)||^{2}\right],$ (18) $$

where $\pi^{\star}$ solves the unregularized optimal transport problem, that is equation [1](https://arxiv.org/html/2507.22270v1#S2.E1) with ${\varepsilon}=0$.

The sample OT-CFM loss in equation [17](https://arxiv.org/html/2507.22270v1#A2.E17) is a low bias approximation of equation [18](https://arxiv.org/html/2507.22270v1#A2.E18) only when the batch size is large enough.
The actual samples for which we compute equation [17](https://arxiv.org/html/2507.22270v1#A2.E17) are not exactly distributed according to a genuine OT plan between $\mu$ and $\nu$, since the OT plan $\pi^{\star}$ and the product measures $\mu\otimes\nu$ might be mutually singular.
Additionally, computing the exact batch OT plan becomes prohibitively expensive as the batch size grows.
A solution is to compute an approximate OT plan, by using the Sinkhorn algorithm (Cuturi, 2013), which is an efficient way of computing the entropic OT plan between two discrete sets of measures. In that case, as the batch size increases, the sample OT-CFM loss equation [17](https://arxiv.org/html/2507.22270v1#A2.E17) approximates $\mathcal{L}_{\mathrm{EOT-CFM}}$ given by equation [10](https://arxiv.org/html/2507.22270v1#S3.E10). Nevertheless, approximating the OT at the batch level is particularly challenging in datasets with multiple modes, as it becomes unrealistic to faithfully approximate the global OT if not all modes are adequately represented within each (or the average) batch.
Consequently, the batch size must scale with the number of models or classes present in the dataset.

Our method does not have these scaling issues with the batch size, since it only involves computing a simple weighting factor $w_{\varepsilon}(x_{i},y_{i})=\exp(-c(x_{i},y_{i})/{\varepsilon})$ for every training sample pair $(x_{i},y_{i})$ in a batch. In other words, if one assumes that the weight does not tilt the marginals, the weighted CFM method corresponds to a large batch limit of OT-CFM (where batch EOT is used).

## Appendix C Additional Results

**Table 5: Comparison of CFM training algorithms’ performance on $\text{8 Gaussians}\to\text{moons}$ on 5 random seeds. $W_{2}^{2}$ measures the overall quality of sample generation (lower is better), NPE measures the straightness of trajectories (lower is better), using the true optimal transport cost as a reference. We emphasize on the tradeoff incurred by the choice of ${\varepsilon}$.**
| Model | $W_{2}^{2}$ ($\downarrow$) | NPE ($\downarrow$) |
| --- | --- | --- |
| I-CFM | 0.703 $\pm$ 0.316 | 1.049 $\pm$ 0.098 |
| OT-CFM | 0.294 $\pm$ 0.131 | 0.147 $\pm$ 0.022 |
| OT-CFM ($B=16$) | 0.633 $\pm$ 0.062 | 0.063 $\pm$ 0.043 |
| W-CFM (${\varepsilon}=2$) | 1.703 $\pm$ 0.230 | 0.302 $\pm$ 0.041 |
| W-CFM (${\varepsilon}=4$) | 1.180 $\pm$ 0.206 | 0.048 $\pm$ 0.033 |
| W-CFM (${\varepsilon}=6$) | 0.886 $\pm$ 0.178 | 0.239 $\pm$ 0.037 |
| W-CFM (${\varepsilon}=8$) | 0.943 $\pm$ 0.253 | 0.311 $\pm$ 0.018 |
| W-CFM (${\varepsilon}=10$) | 0.671 $\pm$ 0.308 | 0.549 $\pm$ 0.084 |

Figure: Figure 4: Sample trajectories on $\text{8 Gaussians}\to\text{moons}$ with W-CFM. From left to right, the models used are trained with the following values of ${\varepsilon}$: 10,8,6,4,2.
Refer to caption: plots/moons_epsall.png

Figure: Figure 5: Contour plots of learned target density for $\text{8 Gaussians}\to\text{moons}$. The leftmost plot corresponds to the true target distribution. Then, from left to right, the models used are trained with the following values of ${\varepsilon}$: 2,4,6,8,10.
Refer to caption: plots/kde_moons_eps.png

Figure: Figure 6: Generated samples from W-CFM trained on CIFAR-10.
Refer to caption: x1.png

Figure: Figure 7: Generated samples from W-CFM trained on CelebA64.
Refer to caption: x2.png

Figure: Figure 8: Generated samples from W-CFM trained on ImageNet-10.
Refer to caption: x3.png

Figure: Figure 9: Generated samples from W-CFM trained on Food-101.
Refer to caption: x4.png

Figure: Figure 10: Generated samples from W-CFM trained on Intel Image Classification.
Refer to caption: x5.png

**Table 6: Sample quality and diversity metrics on CIFAR-10.**
| Model | Precision ($\uparrow$) | Recall ($\uparrow$) | Density ($\uparrow$) | Coverage ($\uparrow$) | F1 ($\uparrow$) |
| --- | --- | --- | --- | --- | --- |
| I-CFM | 0.83 | 0.75 | 0.98 | 0.91 | 0.78 |
| OT-CFM | 0.80 | 0.75 | 1.00 | 0.92 | 0.77 |
| W-CFM | 0.81 | 0.76 | 0.94 | 0.91 | 0.78 |

**Table 7: Sample quality and diversity metrics on CelebA64.**
| Model | Precision ($\uparrow$) | Recall ($\uparrow$) | Density ($\uparrow$) | Coverage ($\uparrow$) | F1 ($\uparrow$) |
| --- | --- | --- | --- | --- | --- |
| I-CFM | 0.86 | 0.66 | 1.26 | 0.98 | 0.74 |
| OT-CFM | 0.84 | 0.65 | 1.23 | 0.96 | 0.73 |
| W-CFM | 0.83 | 0.66 | 1.19 | 0.98 | 0.74 |

## Peer Reviews

confidence: 4
contribution: 1
initial_rating: 2
presentation: 3
questions: - Can authors explain the methodology for choosing $\epsilon$ more thoroughly and extrapolate on lines 253-255?

- Which mini batch pairing algorithm, i.e., mini batch Optimal Transport or mini batch Entropic Optimal Transport with some $\epsilon$, was used as a part of OT-CFM in the experimental part of the paper?
soundness: 2
strengths: - The authors propose a method for training a CFM model that lets one to approximate an entropic optimal transport plan for training data pairs, without extra compute as in OT-CFM/EOT-CFM.

- The authors test their approach on several data types and analyze different aspects of CFM models.

- The method delivers decent results on image data, outperforming other CFM approaches by both fidelity and diversity.
summary: The authors propose the method for learning the Conditional Flow Matching models with straighter and shorter trajectories. The authors propose to modify the optimization objective by incorporating the Gibbs kernel, i.e., weightening of training pairs (x, y), and call their method Weighted Conditional Flow Matching (W-CFM). The modified objective coincides with the objective of Conditional Flow Matching with Entropic Optimal Transport plan for training pairs (x, y), but with different “tilted” marginals. Authors address the issue with “tilted” marginals in practice by optimizing over the stochasticity of EOT parameter $\epsilon$ and show examples where marginals are not “tilted”. In addition, the authors show that under some assumptions their method W-CFM coincides with the EOT-CFM method in the limiting batch size case.

On practice authors, compare their method with other Conditional Flow Matching methods with different data couplings. The methods are compared on toy 2D datasets as well as on high dimensional image datasets.
weaknesses: The main weakness of the method is **“tilted” marginals**:

- There are no clear types of cases where marginals are not tilted (except for very limited case described in Proposition 2). In contrast the OT-CFM/EOT-CFM methods are asymptotically unbiased and when batch size grows to infinity they recover the ground truth OT or EOT plan. While in W-CFM there are no such parameters that guarantee the lack of tilting.

- In that light the claims of authors that distributions are not “tilted” significantly remain unsubstantiated.

- Proposition 3 is built on a “mild regularity” assumption (no marginal tilting and bounded support) and in my opinion these assumptions are not “mild” at all, since authors show very few cases where marginals are not “tilted” and bounded support is also a strong assumption.

The practical methodology for choice of $\epsilon$ **remains unclear.** Authors say “we search over a grid of k values spaced uniformly in log scale and select the smallest value for which the relative variance starts flattening, following an ”elbow rule” heuristic akin to the selection of the number of principal components in PCA”. But still I do not catch the overall algorithm. Can authors explain in more details?:
- As far as I understand the bigger the $\epsilon$ the closer ratios between “tilted” and original marginals, i.e, $f_\epsilon$ and $g_\epsilon$ in Eq 12, to constant. In that sense the bigger $\epsilon$ one takes the less “tilted” the marginals, but this undermines W-CFM purpose, because then the EOT plan is closer to an independent plan. Then the choice of $\epsilon$ remains even more controversial and unclear.

The experiment results are **ambiguous.**
- On the MoG -> 5 Gaussians problem W-CFM delivers not the straightest trajectories (Table 1, NPE), but shows the best distributions fit (Table 1, W2)
- While on 8 gaussians -> moons the W-CFM is worse in both fitting the target distribution and trajectories straightness.
- While on image dataset the W-CFM in majority cases indeed deliver better fidelity and diversity then I-CFM and OT-CFM, it holds not in all the cases and the gains w.r.t. other methods are rather marginal.
- The authors compare their method with OT-CFM, but I couldn't find information wherever they use mini batch Optimal Transport or mini batch Entropic Optimal Transport (Sinkhorn) to generate training data pairs. While the comparison with mini batch EOT with the same $\epsilon$  would be desirable for clean experimental comparison.

---

confidence: 3
contribution: 2
initial_rating: 2
presentation: 2
questions: - In Table 2, the performance of the proposed method is compared with I-CFM and OT-CFM. Could the authors also report the training time to clarify how large is the computational advantage of the proposed method?
- It would also be helpful if the authors included results with different cost functions, such as the $\ell_1$ distance, though this is minor.
- In addition, could the authors provide an analysis of the trade-off with respect to $\varepsilon$ on image data?
- Finally, how is this work related to [1] (Theorem 4.3) and [2] (Eq. 10)?

[1] Zhang, Shiyuan, Weitong Zhang, and Quanquan Gu. "Energy-Weighted Flow Matching for Offline Reinforcement Learning." ICLR 2025.
[2] Park, Seohong, Qiyang Li, and Sergey Levine. "Flow q-learning."  ICML 2025.
soundness: 2
strengths: The authors propose a method that approximates EOT-CFM by simply reweighting pairs sampled from an independent plan, thus avoiding the need for mini-batch OT computations. They provide a theoretical justification for its equivalence to OT-CFM in the large-batch limit. However, this reweighting results in tilted marginals, i.e., $\tilde{\mu}_\epsilon(dx)=\frac{\exp{(-\phi\_\epsilon(x))}}{Z^1\_\varepsilon}\mu(dx)$ (eq. 10), where $\mu(dx)$ is the original marginal, so the authors analyze the conditions under which this tilting can be neglected. They also introduce a procedure for selecting the regularization coefficient $\varepsilon$, which controls the trade-off between the straightness of trajectories and the degree of marginal tilting in the proposed framework. Extensive experiments show that the proposed method achieves performance comparable to previous approaches.
summary: The paper introduces Weighted Conditional Flow Matching (W-CFM). The key idea is to weight each training pair using a Gibbs kernel, allowing the method to approximate the Entropic Optimal Transport (EOT) plan without the high computational cost of explicitly solving the OT problem and without being constrained by batch size. The authors derive theoretical conditions under which this approximation preserves the original marginals and prove that W-CFM becomes equivalent to OT-CFM in the large-batch limit, provided the marginals remain unchanged. The method is evaluated on both toy and image datasets.
weaknesses: ### **Terminology**

- In the statement of Theorem 1, there are $\phi, \psi$, whereas in Equation (2) they appear as $\phi_\varepsilon, \psi_\varepsilon$. Moreover, Theorem 4.2 from [1] states that $\phi, \psi$ are called _EOT potentials_, not _Schrödinger potentials_ and Remark 4.3 therein notes that there is already some inconsistency in the naming.  

### **Theory**
- As I understand it, $\mathcal{L}\_{\text{W-CFM}}$ can be considered an accurate approximation of $\mathcal{L}\_{\text{EOT-CFM}}$ only if the EOT potentials corresponding to the entropic plan satisfy

$$-\phi\_\varepsilon(x)\psi_\varepsilon(y) \approx C$$

which seems a very restrictive assumption meaning that the potentials are constants.
  
- Furthermore, there is a notation overload in Equation (12), since $f_\varepsilon(x)$ was already defined as $f_\varepsilon = \exp(\phi_\varepsilon(x))$. The conditions below (12) for marginal preservation are rather broad and obvious. It is not clear how the constant between the two measures affects the tilting – it seems that this constant should be close to 1, but the sensitivity of this constant is not discussed.

- Proposition 3 appears rather trivial, as the assumption that the tilted marginals are not tilted is too strong. In this case, $q_\varepsilon = \pi_\varepsilon$, since $\exp(-\phi_\varepsilon(x)) = Z^1_\varepsilon$ and $\exp(-\psi_\varepsilon(y)) = Z^2_\varepsilon$, and therefore, from Equation (8), it follows that $Z_\varepsilon = Z^1_\varepsilon Z^2_\varepsilon$. For the case $q_\varepsilon = \pi_\varepsilon$, the mini-batch approximation of the true entropic plan is a standard result [2]. Moreover, the proof relies mainly on intuitive reasoning and does not provide references for rigorous derivations.

---

### **Practice**

- W-CFM has a notable limitation – a trade-off controlled by $\varepsilon$ between preserving the marginals and maintaining straight trajectories. In Table 5, larger values of $\varepsilon$ prevent marginal distortion but result in less straight trajectories. Since $\varepsilon$ appears under the exponent, the tilting of marginals seems to be very sensitive to its value. Moreover, I did not find a sensitivity analysis for $\varepsilon$ on the image domain.

- Finally, the performance improvement is modest, as shown in Table 2, where even I-CFM achieves comparable results.

[1] Nutz, Marcel. "Introduction to entropic optimal transport." Lecture notes, Columbia University (2021).

[2] Hundrieser, Shayan, Marcel Klatt, and Axel Munk. "Limit distributions and sensitivity analysis for empirical entropic optimal transport on countable spaces." The Annals of Applied Probability 34.1B (2024): 1403-1468.

---

confidence: 4
contribution: 2
initial_rating: 2
presentation: 2
questions: 1. Why cannot such a straight line accelerate the sampling process? The motivation for this question is that Rectified flow [1] mentions that a straighter path can reduce the NFE, which is the essence of Reflow. In this paper, a straighter path performs similarly to FM under the same NFE, according to Table 3, which is confusing.

2. Can the author report the performance of W-CFM in the hyperparameter=1 situation? The motivation for this question is that Fig. 2 shows that the smaller the hyperparameter, the straighter the path. A natural question is: what about using 1 directly? 

3. Can you offer the experiments on different solvers, such as the Euler solver? The motivation for this question is that the Euler solver is the mainstream for FM, so we need to ensure W-CFG works well with it.

[1] Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow. Liu et al. ICLR 2022.

To sum up, the topic of this paper deserves more attention. But because of the weaknesses, it is difficult for me to be convinced that W-CFM can work. Therefore, I rate it as reject temporarily. If the author can clarify these concerns, I am willing to increase my score.
soundness: 1
strengths: 1. The topic of this paper is essential. Making the FM path straighter is key to reducing NFE and accelerating the sampling process while maintaining high generation quality. Reflow is one way, but it suffers from two-stage training and requires traversing the entire reverse process. W-CFM is one-stage and does not necessarily traverse the entire reverse process, which is interesting.
summary: This paper proposes W-CFM, which leverages the Gibbs kernel as a weight function to reduce FM loss. The primary motivation for W-CFM is to straighten FM's path and thus improve FM's performance. The experimental results show that W-CFM performs well on unconditional generation in simulation and toy datasets such as CIFAR-10.
weaknesses: However, this paper contains some major concerns:

1. The motivation for introducing the Gibbs kernel to achieve a straighter path is unconvincing. Firstly, this paper does not offer any theoretical justification for this motivation. Meanwhile, the experimental results cannot support the motivation either. For example, Fig.2 shows that the path of OT-CFM is straighter than that of W-CFM. Meanwhile, image generation tasks also show that W-CFM cannot achieve this motivation. Under the same NFE, if the path is straight enough, the FID should show a "lerp". However, Fig. 3 indicates that the FID of W-CFM is similar to that of OT-CFM. 

2. The motivation for using the Gibbs kernel itself is unclear. The related content spans 155-161 lines. But the author just claims that I-CFM can be written as a weight-function formulation. Then, they use the Gibbs kernel as another invariant of the weight function. Is there any specific reason here? In my view, directly leveraging a $l_{2}$ norm between GT and $x$ as the weight function can still work.

3. W-CFM is sensitive to the hyperparameter. As shown in Fig. 2, the hyperparameter directly determines W-CFM's performance, hindering its implementation in i2v and t2v tasks.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: Some minor questions:
1. How senstive of the W-CFM to the choice of cost function?
2. How senstive to $\epsilon$? According to simulation, looks quite sensitive?
soundness: 3
strengths: 1. The paper is well written and easy to follow.
2. The training of W-CFM is as cheap as I-CFM, while recovering OT-like straightness.
3. The formal link to EOT and marginal tilt validate the weighting on loss.
summary: To overcome the computational burden of OT-CFM, this paper proposes to weight the FM loss by per-pair Gibbs weight $w_{\epsilon}(x,y)$ and it leads to the W-CFM. The W-CFM loss approximates the entropic OT up to marginal tilting. They also show large-batch equivalence to OT-CFM. On the choice of tuning parameter $\epsilon$, the paper discuss a huristic way. By simulation and experiements on image datasets, they show the W-CFM has competitive or better FID and path straightness.
weaknesses: The choice of $\epsilon$ can still be relatively cumbersome? In the paper, they set $\epsilon = \kappa \sqrt{d}$, while $\kappa$ is chosen by grid-search.


Decision: Reject