Faster Inference of Flow-Based Generative Models via Improved Data-Noise Coupling
Authors: Aram Davtyan
Leello Tadesse Dadi
Volkan Cevher
Paolo Favaro
Venue: ICLR 2025 Poster (2025)
Keywords: generative models
flow matching
Primary area: generative models
TL;DR: the paper introduces LOOM-CFM, a method that straightens sampling trajectories in the flow matching framework via storing and exchanging the locally optimal data-noise couplings across minibatches

Abstract: Conditional Flow Matching (CFM), a simulation-free method for training continuous normalizing flows, provides an efficient alternative to diffusion models for key tasks like image and video generation. The performance of CFM in solving these tasks depends on the way data is coupled with noise. A recent approach uses minibatch optimal transport (OT) to reassign noise-data pairs in each training step to streamline sampling trajectories and thus accelerate inference. However, its optimization is restricted to individual minibatches, limiting its effectiveness on large datasets. To address this shortcoming, we introduce LOOM-CFM (Looking Out Of Minibatch-CFM), a novel method to extend the scope of minibatch OT by preserving and optimizing these assignments across minibatches over training time. Our approach demonstrates consistent improvements in the sampling speed-quality trade-off across multiple datasets. LOOM-CFM also enhances distillation initialization and supports high-resolution synthesis in latent space training.

Published as a conference paper at ICLR 2025
FASTER INFERENCE OF FLOW-BASED GENERATIVE MODELS VIA IMPROVED DATA-NOISE COUPLING
Aram Davtyan University of Bern Bern, Switzerland a aram. saannddaatyyannuubee..ct
Leello Tadesse Dadi EPFL.
Lausanne, Switzerland leello. dadi @epfl.ch ch
Volkan Cevher EPFL Lausanne, Switzerland vo volkan.ccvhet www.kkn..werrrrlll....o ch
Paolo Favaro
University of Bem
Bern, Switzerland paolo. wwwarouunibee..cooommmm
ABSTRACT
Conditional Flow Matching (CFM), a simulation method for training contin- uous normalizing flows, provides efficient alternative to diffusion models for key tasks like image and video generation. The performance of CFM in solving these tasks depends the way data coupled with noise. A recent approach uses minibatch optimal transport (OT) to reassign noise-data pairs in each training step streamline sampling trajectories and thus accelerate inference. However, its optimization is restricted to individual minibatches, limiting its effectiveness on large datasets. To address this shortcoming. we introduce LOOM-CEM, (Looking Out Of Minibatch-CFMM by preserving and optimizing these assignments across minibatches over training time. Our approach demonstrates consistent improvements in the sampling speed- quality trade-off across multiple datasets. LOOM-CFM also enhances distillation initialization and supports high-resolution synthesis in latent space training.
novel method to extend the scope of minibatch OT
INTRODUCTION
The bigh-qualie outputs and training stability of modern diffusion (Ho et al. 2020 Dhariwal & Nichol 2021) and flow-based models (Lipman al. 2022 Albergo & Vanden-Hijnden (023). or more generally iterative denoising methods, have led to their rapid adoption in nearly every area
of content creation. from image (Esser et al., 2024) and video generation (Davtyan al. 2023) to
motion (Hu et al. 2023) and audio synthesis (Guan et al. 2024).
Despite the superior performance of iterative methods, they require multiple evaluations of the derlying model during inference to generate content. This requirement stems from the gradual trans- formation of the initial sample into the desired data point, resulting in slower operation compared to single-pass methods like GANs (Goodfellow et al., 2020).
To mitigate this drawback, recent research has focused expediting the generation process through various strategies, including enhanced training techniques (Lee et al.) 2023 Bartosh et al. 2024).
model distillation (Luhman & Luhman 2021 Song et al., [023]] Liu et al. 2023] Salimans Ho
2022). and sampling modifications (Dhariwal & Nichol 2021 [u et al., 2022 Shaul et al.] 2024).
Specifically, iterative models employing probability flow ordinary differential equation (ODE)
framework promising, minimizing the curvature of their generative trajectories can signifi- cantly reduce the number of required network evaluations, thereby accelerating sampling.
This paper explores the Conditional Flow Matching (CFM) framework and its derivatives (Lipman) et al. 2022, Albergo Vanden- Vanden-Eijjdenn 2023 Liu et al. 2023 Tong et al. 2023). CFM provides straightforward, efficient, and versatile training method that not dependent distribution. Although data and noise are typically sampled independently. the joint distribution of data-noise pairs, data-noise coupling, significantly impacts the curvature of sampling trajectories
the initial noise

---

Published as a conference paper at ICLR 2025
in CFM-trained models (Lee et al.
2023 Sampling these pairs most effectively follows the Optimal
Transport (OT) plan (Villani et al.
2009) between the noise and data distributions (Liu et al. 2023
However, computing the OT plan for large datasets is not feasible, and while minibatch OT methods offer an approximation during training (Tong et al. 2023 Pooladian et al. 2023), their effectiveness decreases with increasing dataset size.
To overcome these limitations, we introduce a method that enhances the effectiveness of minibatch OT-based CFM. Central to our approach is the preservation and iterative refinement of noise-data pairings within each minibatch, facilitating implicit communication of local OT assignments across different minibatches, thus achieving a more accurate approximation of the global OT plan. To avoid model overfitting to static noise-data assignments, we further propose assigning multiple noise in- stances to each data point, selecting one randomly during each training step. We refer to our ap- proach as LOOM-CFM, standing for Looking Out Of Minibatch-CFN, We provide convergence analysis and demonstrate through experiments that LOOM-CFM outperforms existing methods on standard benchmarks. Additionally, LOOM-CFM serves distillation, further enhancing inference speed, and compatible with latent flow matching for gen- erating nigher-resoutte outputs. Our contributions can be summarized as follows:
effective initialization for model
We introduce LOOM-CFN - a novel iterative algorithm to boost the generation speed and accuracy of CFMs by optimizing the global data-noise assignments of minibatch OT;
We prevent overfitting of fixed data-noise assignments at no computational cost by allocat- ing multiple noise samples per data point (which artificially increases the dataset size);
We present convergence analysis of LOOM-CFM:
We evaluate LOOM-CFM extensively and demonstrate its superior performance over prior work. Specifically, LOOM-CFM reduces the FID with 12 NFE by 41 % on CIFAR10, 10, 46% on ImageNet-32, and 54% on ImageNet-64 compared to minibatch OT methods.
2 BACKGROUND
In this section, for completeness, we first introduce the required background and the prior work, which include CFM (in Section [. the existing approaches to speeding up its inference by straight- ening the sampling paths (in Section 2.2] and optimal transport (in Section [2.3]. We discuss the issues and the drawbacks of the existing methods and then explain our approach in Section 3.1
2.1 CONDITIONAL FLOW MATCHING
Generative modelling requires estimating the unknown target data distribution p(x) with some para- metric model (x)). A conventional choice for this estimator is pa (x) = (2x (2)) dz, where p(z) is given source noise distribution (typically the standard normal distribution) and ps (a|2) 22) is a learned conditional noise-to-dad distribution, ered to be deterministic, and in that case it can be written as po z) = 8(x -00())) where (=(=) is deterministic mapping from noise data, often referred the genenator, and 8(.) is the Dirac delta distribution.
generative distribution. Often the latter is consid-
Recent ODE-based methods (such denoising diffusion models (Ho et or conditional flow matching (Lipman et al., 2022 Albergo & Handen-Eijnded 2023 Liu et al.
2020 Song et al. 2020
2023
Tong et 2023) implicitly define the generator go go(2) through the following ODE
dø(z, t) vo(6(2,t),t)
((2, z,
where up is vector field, parameterized with neural network. to t) is trained in a way that go(2) = 0(2 1) induces a valid approximation of the data distribution. One way to achieve this is by minimizing the following objective
(C) = Ex, -- (... (2x + (1 - t)2 + +e,t) - (x - 2)1?
(3)
where p(z, 2) is some handerafted coupling distribution with marginals equal to p(x) and p(z). A
common choice of p(x;, z) is the so-called independent coupling p(x))(=))
2

---

Published as a conference paper at ICLR 2025
Sources
Sempling peth
rE (0,)), 6
Targets
Semoling paths
Targets
z a + (1+ -0
(a) Independent coupling
(b) Vector field averaging
(c) OT coupling
2D Gaussian- -oa--ssiann generation with CFM. This example illustrates the intuition behind the source of the curvature of the marginal sampling trajectories in the conditional flow matching framework and how optimal transport coupling could potentially solve this issue. (a) With independent coupling the sampling trajectories are highly curved. (b) This is due to the fact that in CFM the learned marginal vector field is average the Gaussian case the distribution is symmetric around the origin, the learned vector field for small values of L points towards the origin, since the target points are likely to be sampled from the other side. (c) OT coupling resolves the curvature of the trajectories and results the identity mapping via averaging only over the sources and largets that are close each other.
Figure 1:
all nource-to-targe directions. As in
To sample from model trained with the CFM objective, needs to first draw noise sample z from the source distribution and then numerically solve the ODE to obtain g=(=). This requires discretizing the ODE and hence calling the neural network (y () multiple times, which can be costly and can slow down the generation. the trajectories of the ODE had low curvature, fewer discretization steps would be needed in order achieve the same accuracy.
In fact, despite the convex interpolation between a single noise and data sample pair in Equation 3 being straight line, the sampling trajectory of the ODE actually far from straight. This can be seen even in toy example, where the source and the target distributions are identical and are both standard normal Gaussians. Intuitively, the optimal vector field in such case should act as an identity. However, the sampling trajectories of model trained in this setting first move towards the origin and then turn around and head back to the starting point (see Figure Tap. One might blame the convergence and numerical errors for such behavior. However, in this particular case there is closed form analytical solution to the CFM optimization problem that is given by
21
w(y, 1) y o² 12 (1 (1-
This expression clearly explains the above observation (see Appendix f for the derivation). The reason why we get such a solution is the averaging in the loss function over all directions between the source and target samples (see Figure 16.. This problem has already been noticed and pointed out in (Lee et al. 2023 Esser et al. 2024 Lee et al.] 2024). Despite being just toy example, the Gaussian to Gaussian case is actually important, the conventional preprocessing in image generation involves normalizing images to ensure zero-mean and unit standard deviation.
2.2 FASTER SAMPLING AND DATA COUPLINGS
Prior work has explored different approaches to speed up the integration of the ODE trajectories. One way to achieve faster generation is reduce the number of integration steps and hence the number of function evaluations. In order to retain accuracy with fewer integration steps straighter sampling trajectories are desired. It has been shown that the straightness of the sampling paths is highly dependent the number of criss-crosses of the data-to-noise interpolations that is induced by the coupling distribution p(z, 22 (Tong et al. 2023 Lee et al. 2023). Hence, by only changing the coupling distribution, one can obtain straighter sampling trajectories of the learned marginal flow- field. In Lee et al. (2023) the authors propose to model the coupling distribution as p(=))((2)))) and to also learn the encoder or the forward process p( (2))))). In Bartosh et al. (2024) the forward process
3

---

Published as a conference paper at ICLR 2025
is also learned and since the coupling distribution is flexible, they can additionally impose explicit constraints on the straightness of the sampling paths. However, both approaches require training of auxiliary models and incorporate additional loss functions, which complicates the training overall. In this paper we focus instead on approaches that involve only changing the coupling distribution. It has been shown that the flow obtained after optimizing Equation [3 induces data coupling with transport cost not larger than the cost of the initial coupling (Liu et al., 2023). More formally:
E Ex-p(s) Be-p(s)le(go(2)) (c(go()) S pm,ppp(( (tt.pp/......._,._____1_
(5)
where c(x; z) is an arbitrary convex function (e.g., | - = =113). Thus, using the induced coupling for the second training (the reflow algorithm Liu et al. (2023) tends to straighten the paths and lead to faster sampling. However, it comes with the burden of at least twice the training time. In Tong et al. (2023) and Pooladian et al. (2023) the authors noticed that one could instead use a coupling that is already optimal with respect to Equation [ (see Figure [c]. Such a coupling is given by the so-called optimal transport plan (we introduce optimal transport in Section 2.3
Unfortunately, obtaining an exact optimal transport plan at the modern data scale is computationally infeasible. Therefore, in Tong et al. (2023) and Pooladian et al. (2023) the authors propose to approximate it via minibatch optimal transport. Tong et al.] (2023) recover the soft permutation matrix with the Sinkhorn- Knopp alogirthm (Knight 2008) and sample from it as from the joint data-noise distribution, while PPolaaian et al. (2023) calculate the hard permutation matrix with the Hungarian algorithm (Kuhn) [955) and reassign the data-noise pairs accordingly. However, the effectiveness of minibatch OT decreases with the growing sizes of the datasets. In contrast, our method, LOOM-CFM, is designed to improve minibatch OT by finding a better approximation to the global optimal transport plan via exchanging information across different minibatches. We describe our approach in Section 3.1
2.3 OPTIMAL TRANSPORT
The optimal transport or Monge-Kantorovich problem refers the search for an optimal coupling between two distributions p(r) and p(z), over X' and 2, with respect to a cost function c(z, =) (Villani al., 2009). Formally, a coupling between p(x) and p(z) is a distribution p(a, 2) over L x Z whose x and marginals are exactly p(z) and p(z). In other words S' p(x, z)dz = p(x)
and Sx p(z, z)dz= = p(z). Among all such couplings. denoted II, the solution to the following optimization problem
inf E Fu, [c(2, =)]:
(6)
p(2,2)E11
referred to as the optimal coupling. Remarkably, the optimal coupling can be shown to be de- terministic for distributions over Rdd under minimal assumptions on Deterministic couplings p(x, z) are those that can be expressed as the joint distribution of (I, T(x) where x ~ p(x) and X + Z. The map T transports thus boils down to finding the optimal transport map T. In fact, further simplifications can be made since, in this paper, p(x)) will correspond to uniform distribution over n data samples, and p(2) will be a uniform distribution over n noise samples. Consequently, the problem considered in this paper is the optimal transport problem between two uniform distributions with finite support of equal size.
p(x) to T(x) p(z). In our setting, solving Equation 6
Several algorithms exist for tackling this optimization task, see Schrieber et al. (2016) for an exten- sive list. The problem can be cast as either linear program or graph-matching problem, and the most notable scheme for solving it is the Hungarian algorithm (Kuhn 1955) which can be under- stood either primal-dual method for solving the linear program an augmenting path approach to finding a minimum cost matching (see 3.6 in Peyré et al. (2019)). The complexity of solving Equation 6 between uniform distributions over n points is On l)g( This complexity quickly becomes prohibitive for large-scale problems, which motivated the search for tractable approxima- tions. The work of Cuturi (2013), for instance, showed that with additional regularization of (6) the resulting surrogate optimization problem can be solved with reduced (n²) complexity but remains out of reach for modern datasets.
Alternative approximations, more relevant to our work, are randomized block-coordinate (Xie et al.) 2024) and mini-batch approaches (mOT) (Fatras et al. (2021) which operate on subproblems of size m The mini-batch approach, which is used in Pooladian et al. (2023), consists of iteratively sampling a subset of m data points and m noise points independently at each iteration and deter- mining the optimal coupling with complexity of O(m²). This scheme does not converge to the
4

---

Published as a conference paper at ICLR 2025
optimal coupling but to a sub- sub-optimal, non-deterministic, averaged coupling whose bias has been analyzed in Sommerfeld et al. (2019) and greatly refined in Fatras et al. (202I). On the other hand, the work of Xie et (2024) proposes randomized coordinate selection scheme to solve the linear programming formulation of using sequence of linear programs over a reduced number of vari- ables. Convergence to optimality can be established as long as the subproblems are solved over 3n variables, which again remains prohibitively expensive.
Our work can be seen as an intermediate between those two approaches. At each iteration, we sample m data points and their corresponding noise samples determined at the previous iteration, instead of the independently sampled noise points of Fatras et al. (202T). The complexity of our scheme does not exceed ((m)) and is a block coordinate update scheme like Xie et al. (2024). Our method converges to stationary transport plan, instead of an averaged one mOT, albeit to sub-optimal one. We describe our method in more detail in the next section.
3 METHOD
In this section, we first introduce our approach to finding better data-to-noise couplings (Section 3.T Then, we propose simple modification to the method to prevent overfitting to the fixed source (Section 3.2). Finally, we provide some implementation details (Section 3.3
3.1 LOOKING OUT OF THE MINIBATCH
Suppose that are given set of training data points {x, heN, and noise samples (4 HeN. where Na = (1.....7) n and z,, 23 E Rd. We aim to find a bijection I" N, + N, that represents an optimal coupling between the data and the noise sets. That is, we seek permutation of N, that satisfies
arg min TES
E [((x)
SEN.
where c(r, 2) is the transport cost between x and 2, and Sn is the symmetric group (of degree n) of N,. In this paper, we consider the quadratic cost c(x, 2) = ||x=~ Let P/ be the permutation matrix corresponding to the mapping T and let X and 2 E Raxd be the data and noise matrices formed by stacking the corresponding samples row-wise. Then we can rewrite Equation 7 as:
Pr. arg min |X- - P
OOT)
TES..
Due to the finiteness of the feasible set, it is clear that solution to the above problem always exists. Ideally, we would like to use 7" as data coupling to train CFM. That is, we would define the coupling distribution
P+ (r 2)
ni Σ 63..2.
1
T2
(8)
and use it to optimize Equation However, algorithms to find 7" usually require O(n²) running
time (e.g., Hungarian algorithm (Kuhn (955)). which makes their use infeasible at the deep learning data scale. Existing work, such as (Tong et al. 2023] [Pooladian et al. 2023). approximates r" with minibatch version of Equation 7 At each iteration, new minibatch of noise is sampled and a local optimal transport is solved to assign noise samples to data within that minibatch. These assignments are used to perform the training step. Despite the simplicity of this method, it suffers from a limited scalability to larger datasets and data dimensionalities, as local assignments might be globally suboptimal. To address this issue, we propose an iterative procedure to better approximate the globally optimal assignment. First, {2 heN, are sampled and assigned to the corresponding (ox JEN, At the beginning, To set to the trivial identity permutation, ie., To(i) = i, Vi E N,, At the k-th training iteration, a minibatch of [mn,) and the correspoding (-1x-1(m,) ) 1m is sampled
from Pr. , (x, =). Tk-1 is then locally updated ensure the optimality of Equation 7 restricted to
the minibatch. That is, we find another permutation Wh that acts locally on the current minibatch
(i.e., (t) i, if i {Tk-1 (Ta-1(n,)) such that the following objective is minimized
Wik arg min WES...
(Ings Ru(Te- (u(()-
(9)
5

---

Published as a conference paper at ICLR 2025
Algorithm 1: LOOM-CFM
1: Input:
Set of data points {z, JeN, Set of noise samples hEN,* Initial assignment To = ld;
k-th heration
2: for in range(1, T) do
3:
Sample minibatch
{Inz ( 1
Pru_, (x,
4: Calculate we in Equation 9
At each iteration random minibatch is sampled according to the current data-to-noise assign- ments. Prior to taking the gradient step the current as- signment is locally updated to ensure optimality w.r.t. the current minibatch. This process resembles weav- ing on a loom, hence the name of the method LOOM- CFM, or Looking Out Of Minibatch CFM.
Figure 2:
5:
Update 7h < Wk o 1
6:
Take gradient descent update
w.r.t. Equation 7 on
{In,, =ra(n;)))==11
7: end for
8: Return: TT
Hence, the update for 7k-1 takes the form:
Tk = W/ o 7k-1-
(10)
Finally, the training step is performed using the updated assignments from Equation which are then saved for later iterations. Formally the method described in Algorithm and an illustration is provided in Figure 2
Notice that in contrast to the prior work, the solution to the local assignment problem not thrown away after each step, but is instead used to update the global assignment and affects the future mini- batches. This makes our approach strictly better than the prior work in terms of global optimality and approximation error.
The proposed Algorithm T guaranteed to converge to stable solution. The following theo- rem, whose proof can be found in Appendix B characterizes precisely the convergence behavior of C0OM-CFM. The analysis is based on casting our scheme as a randomized cycle climination method which converges stationary solution.
Theorem (Finite convergence). LOOM-CFM generates a sequence of assignments тыe with non- increasing costs. With probability 1 over the random batch selection, the iterates converge in finite number of steps to a final assignment Trinal whose associated matching Minal has no negative alternating cycles of length less than m.
Like all tractable approximation schemes, our algorithm is not guaranteed to recover the optimal coupling. However, unlike the minibatch OT schemes, our method outputs deterministic coupling. and more importantly, we show in Section 4.2 that it induces a significantly better straightening of sampling paths.
3.2 MULTIPLE NOISE CACHES
When the size of the dataset is not large enough. using fixed set of noise samples as the source distribution may lead to overfitting. As result, using new noise instances as starting points for the numerical integration of the ODE in Equation may lead to samples of poor quality. To overcome this issue, we propose to store more than one assigned noise sample per data point in the dataset. At each training iteration a minibatch of data-noise pairs is obtained by first sampling data points and then randomly picking one of the assigned noises. This corresponds to artificially enlarging the dataset by duplicating the data points and does not change the underlying data distribution. Later in the paper, refer to the number of the assigned noise samples as the number of noise caches and show empirically that this technique helps to prevent overfitting and allows for using new noise instances source points inference.

---

Published as a conference paper at ICLR 2025
3.3 SPEED AND MEMORY ANALYSIS
We use the Hungarian algorithm Kuhn 1955) to find locally optimal assignments at each train- ing step, making LOOM- CFM comparable in time complexity to OT-CFM (Tong et al., 2023) and BatchOT (Pooladian et al., 2023). J as it solves the same minibatch matching problem. While saving and loading assignments adds minor I/O overhead, is offset by LOOM-CFM's faster convergence. For instance, on ImageNet-32 and -64, LOOM-CFM required only 200 and 100 epochs, compared to BatchOT's 350 and 575. To minimize disk usage, we store the random numbers generator's seed values instead of the assigned noise instances. This incurs a negligible disk usage overhead com- pared to that required for storing and loading modern datasets sizes. For example, a IM-image dataset requires only under 4MB for the noise cache. For implementation details, see Appendix C
EXPERIMENTS
In this section, we present quantitative experiments to demonstrate the effectiveness of our method on real-world data. First, in Sec[4.T] we perform ablation studies to analyze the contribution of different components of LOOM-CFM. In 4.2 we compare our approach to prior work. showing that LOOM-CFM, as designed, generates higher-quality samples with fewer integration steps. Then we illustrate how LOOM- -FM enhances the initialization of the Reflow algorithm et al. 2023), removing the need for multiple Reflow iterations. Lastly, show that our method is compatible with training in the latent space of a pre-trained autoencoder, enabling higher-resolution synthesis.
In all experiments, the main reported metric is FID (Heusel al. 2017) which measures the distance between the distributions of some pre-trained deep features of real and generated data. It has been proven correlate well with human perception and established as a conventional metric for image quality. As commonly done in the literature, we report the FID on the validation set with respect to 50K generated samples. The other measure that is reported is NFE (Number of Function Evalu- ations) that refers to the number of integration steps used to produce the corresponding result. The combination of the FID and the NFE acts as a proxy measure for the curvature of the trajectories. Low FID values along with low NFE speak for the straightness of the sampling trajectories, as the model able to generate high-quality images using small number of steps. For qualitative results, see Appendix E
4.1 ABLATIONS
In this section, we test different versions of our algorithm to understand the impact of each com- ponent on the final performance. All ablations are conducted for unconditional generation on CI- FAR10 (Krizhevsky et al. 2009), dataset of 32 32 resolution images from 10 classes containing 50k training and 10k validation images. The metrics are reported on the validation set.
Training with fixed source: We start by training naive version of CFM, by fixing the initial random assignment between the collection of noise samples and the images. Interestingly, while the results worse than those of LOOM-CFM (see Figure and Appendix [ "Fixed source, w/o reassignments"), the FID at 4 and 8 NFE are better than those of OT-CFM (Tong et al. 2023).
Training after convergence: A logical variant of our approach would be to wait till the algorithm converges to the stable matching and only then start training the vector field. However, this ver- sion results in slightly worse results (see Figure and Appendix D "Train After Convergence"). This might be related to overfitting to the fixed collection of noise samples and the corresponding matching. Indeed, since the algorithm starts with a random matching, training the network from the very beginning, while the matching keeps improving, can be viewed as a soft transition between the independent coupling and the coupling induced by our method.
Training without saving reassignments: Next, we test the impact of updating the global matching with the local reassignments. Instead of using a fixed set of noises and exchanging the information between different minibatches, for a given minibatch of images, we always sample a new minibatch of noises from the standard Gaussian distribution and assign them to the images by calculating the optimal matching. This variant replicates the method of MFM // BatchOT (Pooladian et al. 2023). While the quality of generated images with this method is comparable to the quality of LOOM- CFM large NFE, with few sampling steps LOOM-CFM is significantly better (see Figure 3 and
7

---

Published as a conference paper at ICLR 2025
W/o saving local reassignments Fixed source, w/o reassigments Train after convergence LOOM-CFM
10
caches
FID
caches
NFES
Figure 3: Ablations different design choices on CIFAR10.
Figure 4: Ablation on the number of caches used on CIFAR10.
LOOM CFM
batch size 32
Se-4
saving local reassignments
batch size 64
batch size 128
Curvature
3e-4
10
Integration step
NFEs
Figure 6: Curvature of the sampling trajectories for the models trained on CIFARIO. Average of 1000 trajectories is reported.
Figure 5: Ablation the size of the minibatch used on CIFARIO.
Appendix D "W/o saving local reassignments"). Moreover, we also explicitly demonstrate that the curvature of the trajectories, measured as 1 (0(2 t),t) 6o(0(2, f At), + At). is smaller with LOOM-CFI (see Figure ( Here Ug is the normalized vector field and and + At are two consecutive integration timestamps.
Batch size: Further, we test LOOM-CFM with different batch sizes. As shown Figure 5 with large NFE different batch sizes lead to similar performance. However, with small NFE larger batch size significantly overperforms the smaller ones. This indicates that the method trained with larger batch size tends to yield straighter sampling trajectories. Nevertheless, we would like to point out that LOOM-CFN) trained with batch size 32 produces on par results for and 12 NFE with OT- CFM (Tong et al.) 2023) that was trained with batch size 128 and outperforms it with larger batch sizes. This demonstrates the effectiveness of our caching scheme.
Number of noise caches: Finally, ablate the number of noise caches used for training. Figure shows that larger number of noise caches improves the performance. We found empirically, that caches enough for dataset of size comparable ment diminishes with more caches, as the method becomes slower to converge and discovers worse couplings as the number of noise caches grows (see Appendix D). However, if the dataset is large enough (e.g. ImageNet (Russakovsky et al. 2015) ), even a single cache may be sufficient.
CIFAR10. We observed that the improve-
4.2 RESULTS
Unconditional Image Generation. We train LOOM-CFM for unconditional generation on CI- FAR10 (Krizhevsky al. 2009) and ImageNet32/64 (Russakovsky al. 2015). with the results presented in Tables and 2 respectively. LOOM-CE -FM consistently achieves better few-step FID scores compared to minibatch OT methods, such OT-CFM (Tong et al.] 2023) and MFM "/
8

---

Published as a conference paper at ICLR 2025
Table 1: Unconditional generation results on CIFARIO. (*) indicates that the reported numbers are
calculated by us using the official code, as they were not reported in the corresponding papers.
Method
Solver
CIFARIO
NFE
FIDL
ODE/SDE-based methods
DDPM (Ho
2020
1000
3.17
DDPM (Ho et al
2020
dopri5
274
7.48
St. Interpolatns (Albergo & Vanden-Bijnden,
2023
10.27
FM
OT (Lipman et al
2022
dopri5
142
6.35
I-CFM
Tong al.
2023
culer
100
4,46
L-CFM
Tong cl al. 2025
euler
1000
3.64
I-CFM
Tong et al
2023
dopri5
146
3.66
Improved sampling
DDIM (Song
2021
10
13.36
DPM Solver-2 [Lu ct al.
2022
12
5.28
DPM Solver-3 [Lu ct al 2022
24
2.75
RK2-BES (Shail et al 2024]
10
2.73
RK2-BES (Shall et al.
2024
20
2.59
Rectified Flows
Rectified Flow
(Liu et al.
2023
culer
378
+ distill
euler
6.18
2- Rectified Flow (Liu
2023
euler
12.21
+ distill
euler
4.85
Rectified Flow (Liu et al.
2023
culer
8.15
+ distill
culer
5.21
Improved truining for struighter trajectories
OT-CFM
Tong
2023
midpoint
4
15.95
OT-CFM
Tong ct al
2023
midpoint
8
9.73
OT-CFM
Tong et al.
2023
midpoint
12
7.77
OT-CFM
Tong
2025
euler
100
4,44
OT-CFM Töng
2023
dopri5
134
3.57
Minimizing Trajectory Curvature
(Lee et al
2023
heun
5
18.74
NFDM-OT Bartosh et al.
2024
euler
2
12.44
NFDM-OT Bartosh et al.
2024
euler
4
7.76
NFDM-OT Bartosh et al.
2024
euler
12
5.20
Our results
LOOM-CFM (ours) 4 caches
midpoint
4
11.60
LOOM-CFM (ours) 4 caches
midpoint
8
5.38
LOOM-CFM (ours) 4 caches
midpoint
12
4.60
LOOM-CFM (ours) caches
dopri5
134
4.41
LOOM-CFM (ours) caches + reflow
culer
1
7.63
LOOM-CFM (ours) 4 caches + reflow
midpoint
2
4.49
BatchOT (Pooladian al. 2023). Additionally, LOOM-CFM outperforms Lee al. (2023) and de- livers on par quality with NFDM-OT (Bartosh et al. 2024), surpassing it with 12 NFE while slightly underperforming lower NFE. Unlike these methods, which require training additional compo- nents, LOOM-CFM optimizes the original CFM objective with a modified coupling distribution.
Rectified Flows. As mentioned earlier, LOOM-CFM offers improved initialization for the reflow algorithm (i.e.. retraining the CFM using the coupling induced by the first-stage model). To validate this, we generated 1M noise-data pairs by sampling from our model trained with caches on CI- FAR10 and used those samples to train new CFM model. The results, shown in Table reflow), indicate that LOOM-CFM combined with reflow outperforms not only 2-Rectified Flow (L.iu al.) 2023) but also 3- 3-Rectified Flow. This demonstrates that performing additional reflows is unneces- sary as long as the first reffow well-initialized.
High Righ-Resolution Image Synthesis. Finally, we show that our method can be directly applied to training in the latent space a pre- traitrd autoencoder, similar to the approaches in Rombach et al.
2022: Dao et al. (2023). We trained LOOM- CFM on FFHQ 256 x 256 (Karras al. (019) using a pre-trained autoencoder from Rombach et al. (2022). As shown in Table LOOM achieves lower FID score with order of magnitude fewer NFE compared previous methods using

---

Published as a conference paper at ICLR 2025
Table 2: Unconditional generation results on ImageNet-32 and ImageNet-64
Method
Solver
ImageNet-32
ImageNet-64
NFE
FID
NFE
FID
ODE/SDE-busee methods
St. Interpolants (Albergo & Landen-Eijnder
2023
8.49
FM " OT (Lipman al. 2022)
dopri5
122
5.02
138
14.45
Distilled models Dedicated solvery
Improved training for straighter trajectories
MFM
BatchOT (Pooladian al.
2023
midpoint
4
17.28
4
38.45
MFM
BatchOT (Pooladian
202.
midpoint
8.73
8
20.85
MFM
BatchOT (Pooladian
al. 202.5
midpoint
7.18
12
18.27
MFM
Stable [Pooladian ef al.
2023
midpoint
21.82
4
46.08
MFM
Stable (Pooladian et al.
2023
midpoint
8
9,99
8
21.36
MFM
Stable (Pooladian
2023
midpoint
12
7.84
12
17.60
NEDM-OT
(Bartosh
euler
2
9.83
2
27.70
NFDM-OT (Bartosh et al.
2024
culer
6.13
4
17.28
NEDM-OT Bartosh ct al.
2024)
culer
12
4.11
12
11,58
Our results
LOOM-CFM (ours)
cache
midpoint
13.47
4
37.24
LOOM-CFM (ours)
cache
midpoint
8
5.08
8
11.76
LOOM-CFM (ours)
cache
midpoint
12
3.89
12
8.49
LOOM-CFM (ours)
cache
dopri5
137
2.75
133
6.63
Table 3: Unconditional generation results on FFHQ 256x256
Method
Solver
FFHQ-256
NFE
FID|
Prior Work
LDM (Rombach et al.
2022
50
4.98
ImageBAKI (Esser
2021
9.57
Latent Flow Matching (Do el al.
2023
(ADM)
dopri5
84
8.07
Latent Flow Matching (Do ct al.
2023
LL/2
dopri5
4.55
Our results
LOOM-CFM (ours) caches (ADM)
midpoint
2
14.89
LOOM-CFN (ours) caches (ADM)
midpoint
6.76
LO0M-CFM (ours) caches (ADM)
midpoint
5.50
LOOM-CFM (ours) caches (ADM)
dopri5
77
4.77
same model architecture. This highlights compatibility our method with conventional techniques for high-resoluties synthesis and paves the way for exploring its application to high- resolution and large-scale datasets.
5 CONCLUSIONS
In this paper, introduced LOOM-CFM, method improve data-noise coupling in training gen- erative models with the CFM framework. Our method is based on finding locally optimal match- ings between data and noise each minibatch. In contrast to previous work, these matchings are stored and recycled in future iterations to obtain a better data-noise assignment at the global level. Through an extensive experimental section, we established that LOOM-CFM achieves better sam- pling speed-quality trade-off than prior work with no additional computational cost and negligible disk usage overhead during training. LOOM-CFM effective and simple enough to be composed with other techniques, such as rectified flows, to further enhance the sampling speed.
ACKNOWLEDGMENTS
This work has been supported by Swiss National Science Foundation Projects 188690 and 10001278.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Michael Albergo and Eric Vanden-Eijnden, Building normalizing flows with stochastic interpolants. In ICLR 2023 Conference, 2023.
Grigory Bartosh, Dmitry Vetrov, and Christian Naesseth. Neural flow diffusion models: Learnable forward process for improved diffusion modelling. arXiv preprint arXiv:2404.12%4 2024.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
Quan Dao, Hao Phung, Binh Nguyen, and Anh Tran, Flow matching in latent space. arXiv preprint arXiv:2307.08698, 2023.
Aram Davtyan, Sepehr Sameni, and Paolo Favaro. Efficient video prediction via sparsely condi- tioned flow matching. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 23263-23274, 2023.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780-8794, 2021.
Patrick Esser, Robin Rombach, Andreas Blattmann, and Bjorn Ommer. Imagebart: Bidirectional context with multinomial diffusion for autoregressive image synthesis. Advances in neural infor- mation processing systems, 34:3518-3533 2021.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first International Conference Machine Learning, 2024.
Kilian Fatras, Younes Zine, Szymon Majewski, Rémi Flamary, Rémi Gribonval, and Nicolas Courty. Minibatch optimal transport distances; analysis and applications. arXiv preprint arXiv:2101.01792. 2021.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 64441)4139-144 2020.
Wenhao Guan, Kaidi Wang, Wangjin Zhou, Yang Wang, Feng Deng, Hui Wang, Lin Li, Qingyang Hong, and Yong Qin. Lafma: latent flow matching model for fext-to-audio generation. arXiv preprint arXiv:2406.064 08203, 2024.
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. Advances in neural information processing systems, 30, 2017.
Jonathan Ho and Tim Salimans. arXiv:2207.12598 2022.
Classifier-free
diffusion guidance.
arXiv preprint
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
Vincent Tao Hu, Wenzhe Yin, Pingehuan Ma, Yunlu Chen, Basura Fernand,, Yuki M Asano, Ef- stratios Gavves. Pascal Mettes, Bjorn Ommer, and Cees GM Snock. Motion flow matching for human motion synthesis and editing. arXiv preprint arXiv:2312.08895 2023.
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 4401-4410, 2019.
Philip A. Knight. The sinkhorn knopp algorithm: Convergence and applications. SIAM Journal on Matrix Analysis and Applications, 30(1):261-275, 2008. doi: 10/06666666 URL
https://dof www.a/o...gg....00 10..//0.....11376669692

---

Published as a conference paper at ICLR 2025
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Harold W Kuhn. The hungarian method for the assignment problem. Naval research logistics quarterly, 2(1-2):33-97, 1955.
Sangyun Lee, Beomsu Kim, and Jong Chul Ye. Minimizing trajectory curvature of ode-based gen- erative models. In International Conference on Machine Learning, pp. 18957-18973. PMLR, 2023.
Sangyun Lee, Zinan Lin, and Giulia Fanti. Improving the training of rectified flows. arXiv preprint 2xXiv:2405.20320. 2024.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint 2xXiv:2200002447, 2022.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Kleventh International Conference Learning Repre- sentations (ICLR), 2023.
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. Advances in Neural Information Processing Systems, 35:5775-5787, 2022.
Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv:2101.02388, 2021.
Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications data science. Foundations and Trends@ in Machine Learning. 1165-6):35-607, 2019.
Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lip- man, and Ricky TQ Chen. Multisample flow matching: Straightening flows with minibatch cou- plings. arXiv preprint arXiv:2304. 14772, 2023.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and puttern recognition, pp. 10684-10695, 2022.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomed- ical image segmentation. In Medical image computing and computer-assisted intervention- MICCAI 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceed- ings, part III 18, pp. 234-241. Springer, 2015.
Tim Roughgarden. Cs261: A second course in algorithms lecture# 5: Minimum-cost bipartite matching, 2016.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang. Andrej Karpathy. Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJCV), 115(3)2211222 15222,,, 2015. doi: 10. 100/// (000//11263/-120 1263-015-0816-y;
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. In International Conference on Learning Representations, 2022. URL htttpss://ppnreeieew net Sttt/EoouumiiddTTLLXXx.
Jörn Schrieber, Dominic Schuhmacher, and Carsten Gottschlich. Dotmark-a benchmark for discrete optimal transport. IEEE Access, 5:271-282, 2016.
Neta Shaul, Juan Perez, Ricky T. Q. Chen, Ali Thabet, Albert Pumarola, and Yaron Lipman. Be- spoke solvers for generative flow models. The Twelfth International Conference on Learning Representations, 2024. URL httpsttt:: : / /ooenrevv
n feretfforung
Max Sommerfeld, Jöm Schrieber, Yoav Zemel, and Axel Munk. Optimal transport: Fast proba- bilistic approximation with exact solvers. Journal of Machine Learning Research, 20(105):1-23,, 2019.
12

---

Published as a conference paper at ICLR 2025
Jiaming Song. Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502, 2020.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising difflusion implicit models. In Interna-
tional Conference Learning Representations, 2021. URL httpss//pppeneeeieww.... net/
Forum?id=St 1g1arCHL1 arCHLP
Yang Song. Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint DrXiv:2303.01469, 2023.
Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Kilian Fatras, Guy Wolf, and Yoshua Bengio. Conditional flow matching: Simulation- dynamic optimal transport. arXiv preprint arXiv:2302.00482, 2(3), 2023.
Cédric Villani et al. Optimal transport: old and new, volume 338. Springer, 2009.
Yue Xie, Zhongjian Wang, and Zhiwen Zhang. Randomized methods for computing optimal trans- port without regularization and their convergence analysis. Journal of Scientific Computing. 100 (2):37, 2024.
13

---

Published as a conference paper at ICLR 2025
APPENDIX
Here we provide some additional details and results that could not be included in the main paper due the page number limit. Those include: the derivation of the optimal vector field for the Gaussian to Gaussian case (Appendix A); the proof of the finite convergence of our algorithm (Appendix B); implementation details (Appendix C): more quantitative results (Appendix D); more qualitative re- sults (Appendix E.
A GAUSSIAN TO GAUSSIAN CASE
In this section, we are looking for the closed form solution for the optimal flow in the case of both source and target distributions being standard normal distributions. That is, we are aiming to finding f(y, i) such that:
oy, 1) = arg min (01, ||o + (1 +),+ (x - =>]l (11)
u(g.t)
If write down the expectation explicitly through integrals, obtain
1 100 Rd. 3d Jo 1 [10(,,)) () (or ==)ll Nyyllt + (1 =) 2 (= ()N((, 1) dydidadz
(12)
Let us define
f(u,t,v) [ / Io(y, t) (x =)]l Nyylt + (1 - 0)2, drdz.
(13)
The above functional in the objective takes the form
S. Cly, (ly,,,v(y,1) ()) dydl.
(14)
In order to find the optimal v(v, t). needs to write down the Euler-Lagrange equation (notice that C does not depend on the derivatives of v, and hence the Euler-Lagrange equation only includes the derivative w.r.t v)
aCly,t,v6 t)) 0,
(15)
dv
which in this special case takes the form (here and later on we omit the integration domains for saving space and assume IRd everywhere)
[ʃ 2* ( - (x
drdz = 0.
(16)
And hence, solving for u(y. t), gives
( - -/ dxdz
(17)
I) dædz
First, let calculate the denominator:
14

---

Published as a conference paper at ICLR 2025
/f te (1-1))) (0, dz
(18)
Co
(19)
(2T) =h4 0d
= =Co /|
3 [.)
x x + 2 2T2
0² ((V tx) (1 +)=) ((y - tz) (1 1)2)) de dz
(20)
// exp
(-2
11 (1-1)*
2 212 2
(1- t)
(y tx) Tx
2² 20 (v - tx) ( - tr)
1
dx
0²
(21)
1
Ci
=
(22)
(2 + +(1 - 1)2)
13 =
1 x
1
(1 1)2 V + ++xxxx= x- 22y-2yy x) dr
1
(23)
2
/
2
20
+
(1 f)2
x
0² + (1
t)zu
2 d² + (1 1/2 v 0) dr
1
(24)
(2
1 1/2 2 02 (1 t)2
1
1
=
exp
y
2 (1 t)2 2/ y
(25)
(1
7111-13
e
(26)
Here we used the well-known Gaussian integrals of the form
Ax fexp (exp(-3aTAn) Ax +b x dr = dee((2m--1)exx 1i A TT A b-1++c)) c).
(27)
Finally, the last expression be simplified
(
y y
(2m)- (0² + + + -)²) exp
(28)
20 + t2 (1 t)2
which is the probability density function of
1, (+++++(1))+)))))))
(29)
Now let us move to the nominator, which can be rewritten as Fly, t) Fly, t), where
Fly.0) 1)
1) drdz.
(30)
Thus, we only need to calculate Fly, (). Similarly to the denominator, using the well-known Gaus- sian integrals of the form
/ x. cxpp
2TT (-1x² + bT aee dx
b
bTb
exp
+ +c)
(31)
a
a
2a
we can obtain
yl
Fly,t) 02 0²+ 12 + (1- (1- Nor10 + +(1 -1)3).
(32)
15

---

Published as a conference paper at ICLR 2025
1:00
0
0.30
50
0.00
-0.50
-0.75
-1.00
Figure 7: The form of the multiplier s(t) by y in the optimal vector field vyy.t) = y s(t) for different values of a.
Hence, by combining all the results, we get the optimal vector field
2t 1 0² + 12 (1 t) 1)2
0(y, y. s()), where s(l)
(33)
The vector field of such form is parallel the line connecting y and the origin and changes its direction halfway through time [0 (see Figure 7). This analytical result coincides with the empirical observation fitting the vector field as a neural network via gradient descent the two-
dimensional setting (see Section 2.
B ANALYSIS
In this section we provide proof for the theorem established in section 3.] and include a counter- example, where only using batch size equal to the size of the dataset could solve for the globally optimal matching (see Fig. 88
As noted section 2.3 OT can be understood as minimum cost perfect matching problem in
complete bipartite graph. The assignment task is captured by the graph G = (V X UUZ 7, E), where X is the set of nodes representing the data points and 2 is the set of nodes representing the noise samples. The set of edges E X E, where the cost of an edge (z,,
E is given by the Gy = ||= ž, 112, defines the complete bipartite graph (see Figure 2 for a visual representation).
A perfect matching M is set of edges in G that cover all the vertices. The cost of matching the sum of the costs of the edges. An optimal matching is one achieving the minimal cost. The optimality of matching M is entirely captured by the absence of negative M alltternating cycles.
Definition B. (M-alternating cycles). Given perfect matching M, any cycle in G given by (Itz. Zie Iz,, where for any p [k)]
sequence of edges C (zi, 24, (2,, 2ty) we have (Ii,, 2,,)) E M is called an alternating M Meyylee
M-alternating cycles as their name suggests, pass from X to back and forth while passing through edges in M to go from X to Z. An M-alternating cycle is said to be negative if the sum of the costs of the edges in M larger than the cost of the edges in absent from M:
Σ C. < Σ C..
(34)
CEC\M
eECпm
Consequently, for any matching M, if negative M-alternating cycle C' exists, then matching M' with lower cost can be constructed by taking M' = (C M) U (M \ C.). In other words, by swapping the going edges in C(from X to Z) with the return edges (from Z to X), can diminish the cost thanks to Equation 34 A necessary condition for optimality is thus the absence negative M- M-alternating cycles. The following theorem shows that it is in fact sufficient condition.
16

---

Published as a conference paper at ICLR 2025
Figure 8: Counter-example showing the existence of stationary sub-optimal matchings. The data points (red) are placed at every other vertex of a regular polygon, the noise points (green) are slightly offset in counterclockwise fashion from the vertices of the polygon. The optimal as- signment matches each data points to the following noise point going in the clockwise direction. The assignment represented blue is sub-optimal but stationary for any batch size m < n.
Proposition B.1 (Thm 2.2 (Roughgarden 2016). A matching is optimal if and only if there are no negative M M-alternating cycles.
Given the exposition above, LOOM-CFM can be understood as an iterative negative-cycle climina- tion scheme. Indeed the algorithms starts from matching Mo given by the edges (I,, (7)()), then at each iteration k; a subgraph Ga. = (V:- { ym U (n,)) } Ex) consisting of only the 2m vertices sampled at step
Algorithm [ is considered. An optimal matching computed within that subgraph and the matching M is obtained by updating the edges involving vertices in Vg. As computing optimal matchings are equivalent to climination of all negative cycles, the update of MM corresponds to the climination of negative Mæ-alternating cycles that are contained in Gk- Such cycles are of length at most m. The process described by algorithm ] is thus a stepwise elimination of negative Ma-alternating cycles of length at most m. This allows us characterize the convergence of our method in the following theorem.
Theorem (Finite convergence). LOOM-CFM generates sequence of assignments тıe with increasing costs. With probability 1 over the random batch selection, the iterates converge in a finite number of steps to a final assignment Tinal whose associated matching Minal has no negative alternating cycles of length less than m.
Proof. Let c(T) ;= EiEr, c(E;, Z+))). Each update of LOOM-CFM (step 4 and 5 of Algorithm
improves the matching over sub-graph selected at step 3. The sequence c(Tx)ι is thus non- increasing. Moreover, for all k the sequence of costs takes its values in the finite set {c(T) T e S. It then follows that (c(Tk))s becomes constant after finite number of iterations Kinat since a monotone sequence taking values in a finite set.
Suppose now that there existed a negative alternating cycle of length less than in the matching induced by TKme. At each iteration, a subgraph containing a negative cycle has a non-zer zero proba- bility of being selected step 3. Since (c(Tx))& constant for all l: K this implies that such a subgraph is never sampled at step for all k: >> Kínat- The probability of such a sub-graph never being sampled for all k Kinat is O by simple application of Borell-Cantelliss lemma. Consequently, the probability of TKaa containing negative alternating cycle is 0.
Although this theorem shows that LOOM-CFM converges to stationary solution, it does not guar- antee global optimality, However, as already mentioned in Section 3.1 the main paper, this is the case for all minibatch-based methods. We exhibit the counter-example Xie et al. (002)) (see approach would fail to recover the optimal coupling
Figure () showing that any sub- caab-roobemm--ee a full cyclic permutation of all the points is necessary to solve the problem. Despite this, the coupling induced by LOOM-CFM is more optimal than that of minibatch OT ba-be methods and leads to straighter sampling trajectories, as evidenced by the experiments real data.
17

---

Published as a conference paper at ICLR 2025
Table ADM network architecture and training parameters of LOOM-CFM for each model.
CIFARIO
ImageNet-32
ImageNet-64
FFHQ256
Input shape
[3, 32, 32]
[3, 32, 32] 128
[3, 64, 64]
[2, 32]
Channels
128
192
256
Number of Res blocks
2
3
2
2
Channels multipliers
[1. 2, 2, 2] 4
[1, 2, 2, 21
[1, 2, 3, 4]
[1, 2, 3, 4] 4
Heads
4
4
Heads channels
64
64
64
64
Attention resolution
[16]
[16. 8]
[16]
[16, 8, 4]
Dropout
0.1
0.1
0.1
0.1
Effective batch size
128
512
96
128
GPUs
4
4
4
4
Epochs
1000
200
100
500
Iterations
391k
500k
1334k
273k
Learning rate
0.0002
0.0001
0.0001
0,00002
Learning rate scheduler
Constant
Constant
Constant
Constant
Warmup steps
5k
20k
20k
3.5k
EMA decay
0.9999
0.9999
0.9999
0.9999
Training time (hours)
17.3
73.5
190.6
66.8
CFM
1e-7
1e-7
1e-7
1e-7
Number of noise caches
4
4
C IMPLEMENTATION DETAILS
This section includes the implementation details of our algorithm, including model architectures, training parameters, etc.
In practice, to be comparable with the prior work, we parametrize the learned vector field with neural network that has an improved UNet (Ronneberger al. 2015) architecture (ADM) from Dhariwal & Nichol (2021). To make fair comparisons, we used the same model architectures and training parameters as the prior work, when possible. See Table | for the network architecture configurations and training parameters. All models (except for the ablations) were trained on Nvidia RTX 3090 GPUs. For the latent space models we used the pretrained autoencoder from Stable Diffusion (Rombach al. 2022) provided at the Hugging Face model hut The training code can be found at /etttts://githiiee com/ araachie/ loom- cfm
D QUANTITATIVE RESULTS
In this section, we provide more quantitative evaluations of our method.
Ablations. For a better visual perception, the results of ablations were presented as plots in the main paper. Here we include the numbers used to build those plots (see T)
Convergence. Throughout the training log the number of reassignments minibatch. We report those For different batch sizes and in Figure different number of noise caches (mimicking different dataset sizes). Additionally, we also report the minibatch OT cost around the training time where the reassignments are quite rare (Figures 10 and 12). It can be seen that LOOM- CFM with different batch sizes despite starting at different number of reassignments (around the batch size, since the intial assignments are random) converges roughly in the same amount time (see Figure J.. However, with larger batch size the method converges be inferred from the minibatch OT' cost (see Figure TO). At the same time, increasing the dataset size leads to slower convergence, since the algorithm has visit more minibatches sort the assignments (see Figure II And the minibatch OT cost behaves accordingly (see Figure 12
better matching,
https: https::/huuger https:://hugginggfacs co.st..o///a//oil tittiii////i /sddii//d-ee vae- ftaaa--f--me--riiinie /bl o0/oman/n ft-mse- 840000- ema- moa--pruned.cog
18

---

Published as a conference paper at ICLR 2025
Training time. As for the time complexity, LOOM-CFM is comparable OT-CFM (Tong al. 2023) or BatchOT (Pooladian et al. 2023) as it solves the same minibatch matching problem as the prior methods. However, saving and loading the current assignments from the disk introduces a small i/o overhead. Hence, the training time of LOOM-CFM is slightly longer than that of the baselines. For instance, training for 1000 epochs CIFAR 10 with batch size 128 scattered across 4 Nvidia RTX 3090 GPUs takes 17.3 hours with LOOM-CFM compared to 13.9 hours when the assignments are not stored.
Table 5: Detailed ablation results on CIFARIO.
Solver Fixed source, w/o reassignuments
NFE
FID
midpoint
22.57
Number of rearzigrmeent betch
120
batch
midpoint
midpoint
12
6.17
dopri5
142
5.98
W/o saving local reasxignuments
midpoint
21.47
midpoint
8
8.69
midpoint
12
6.56
10000
70200
30000
M000
dopriS
131
4.52
Train after convergence
midpoint
4
12.62
midpoint
6.45
25000
50000
75000
100000
175000 200000
midpoint
5.53
dopri5
139
5.07
Figure 9: Number reassignments within a mini- batch along the training progress depending the batch size.
caches, bs 128
midpoint
12.28
midpoint
midpoint 12
9.81
dopri5
126
9.75
61.2
caches, bs 128
betch size 32 balcch size 128
midpoint
11.30
61.1
midpoint
8
6.40
midpoint
dopri5
134
12
5.82
5.92
Min batchh OT cost
81.0
Why
Am
Mym nhahh
caches, bs 128
60.9
midpoint midpoint
4
11.60
5.38
60.8
Wy
midpoint
dopriS
134
4.41
60.7
4 caches, bs 64
midpoint
19.39
80.6
midpoint
8.85
Training step
midpoint
12
6.68
Figure 10: OT cost per minibatch along the train- ing progress depending the batch size. The lines corresponding to the exponential moving av- erage of the cost shown.
dopriS
134
4.51
caches, bs 32
midpoint
4
20.83
midpoint
8
9.56
midpoint
12
7.40
dopri5
122
5.02
19

---

Published as a conference paper at ICLR 2025
Number within Battch
2.20
caches caches cachee
00.52
caches caches
caches
69
Mivibatch
00.74
M
40500
00.72
Wit
60.70
60.08
150000 175000 200000
200000
225000
275000
250000
stop
Figure 11: Number of reassignments within a Figure 12: OT cost per minibatch along the train- minibatch along the training progress depending ing progress depending the number of noise on the number noise caches (the size of the dataset). The lines corresponding to the expo- nential moving average of the cost are shown.
caches (the size of the dataset). The lines corre- sponding to the exponential moving average of the cost are shown.
E QUALITATIVE RESULTS
In this section, we provide some qualitative results that could not be included in the main paper due to the page number limit.
Sampling Paths. We provide examples of sampling trajectories from models. Starting from a randomly sampled noise instance at t = O, our models iteratively denoise those via numerically integrating the ODE in Equation to obtain clean images. In Figures [3 [4] 15 and I6 we show intermediate points visited by the solver at f S 1. For the model trained on CIFARIO, we also compare the sampling paths to the version of the model that does not store reassignments. For FFHQ, in Figure 17 we additionally provide samples using different number of steps in the ODE solver, resulting in different NFEs.
Unconditional Generation. Lastly, we provide uncurated samples from our models in Figures 18
19 20
F ALTERNATIVE APPROACHES
In the early stages of developing LOOM-CFM, we explored various approaches to prevent overfit- ting to fixed set of noise samples. Although none of these approaches yielded satisfactory out- comes, we summarize them here for clarity and completeness.
Completely refreshing the noise. We experimented with replacing all the noise samples every N epochs. While intuitive, this approach caused training instability since the newly assigned noise samples did not necessarily align with the previous ones. Consequently. the network's targets could shift significantly when the noise refreshed, impairing conver- gence.
2. Gradual noise injection. In this approach, we introduced a hyperparameter, ó. to control the noise refreshing, For each data point in minibatch, the assigned noise was replaced with new sample with probability é. Although this method allowed for a smoother re- fresh, choosing optimal proved to be challenging. For example, setting é 0.1 caused approximately one third of each minibatch to be reshuffled, which weakened the effect of LOOM-CFM. In contrast, lower value of d 0.01 insufficient to prevent overfitting.
3. Interpolation between LOOM-CFM and independent coupling. For each data point, we coupled it with its assigned noise with probability & and with freshly sampled noise with probability - ó. Unlike approach 2, the new noise did not replace the cached noise. We found this technique to be quite effective and the results indeed interpolated the results of I. 100M-CFM and the independent coupling. We also explored making a function of
20

---

Published as a conference paper at ICLR 2025
the curvature of sampling paths depends (as shown in Figure 6 Unfortunately, this hasn't 't lead to substantial improvements. Nevertheless, we found this approach interesting for its generality it can interpolate any pair of coupling techniques. We leave further exploration of this method for future work.
In contrast to these more complex methods, the approach with multiple noise caches in LOOM-CFM straightforward, artificially increases the dataset size and equalizes the settings for problems with small and large dataset sizes.
G LIMITATIONS
One limitation of LOOM-CFLI and other OT-based methods is that they are not directly compatible with conditional generation, especially when the conditioning signal is complex. In any conditional setup, the marginals of all label-conditional probability paths at / = 0 must match the source dis- tribution. However, naive implementation model without adapting the couplings, may introduce sampling bias, as certain labels could dispro- portionately align with specific regions the noise space. Although techniques like classifier (-reee) guidance (Ho & Salimans, 2022) could be adapted, we leave the exploration of those extensions to future work.
any coupling-based method, by conditioning the
21

---

Published as a conference paper at ICLR 2025
Figure 13: Sampling trajectories of LOOM-CFM (top) and the version of the algorithm that does not save the reassignments (bottom) trained CIFAR 10. The i-th row in both grids starts with the same initial noise. LOOM-CFM tends to produce sharper genereations and converges carlier.
22

---

Published as a conference paper at ICLR 2025
Figure 14: Sampling trajectories of the model trained on ImageNet-32.
Figure 15: Sampling trajectories of the model trained on ImageNet-64.
23

---

Published as a conference paper at ICLR 2025
Figure 16: Sampling trajectories of the model trained on FFHQ-256.
Figure 17: Samples from the FFHQ-256 model with different NFE (from left to right: 2, 4, 8 and 12 function evaluations).
Figure 18: Uncurated samples from the model trained on CIFARIO.
24

---

Published as a conference paper at ICLR 2025
Figure 19: Uncurated samples from the model trained on ImageNet-32.
Figure 20: Uncurated samples from the model trained on ImageNet-64.
25

---

Published as a conference paper at ICLR 2025
Figure 21: Uncurated samples from the model trained FFHQ-256.
26


Reviews:
[{"review_id": "ckBeNEBhPN", "reviewer": "Reviewer_sou1", "summary": "This work investigates coupling methods for drawing paired samples of data and noise to during the training of Conditional Flow Matching (CFM) models to improve model performance. In particular, well-chosen coupling during training can make sampling paths straighter at inference time so that generation requires fewer model evaluations. Methods for achieving such a coupling are known but prohibitively expensive for large training datasets. Prior work performed the coupling only within a minibatch for Gaussian noise drawn at each step. This work improves upon this idea by using persistent noise samples whose pairings with data are gradually refined. Initially, one noise sample is drawn for each data sample at random. In each minibatch, noise and data are drawn together, and data-noise pairs are reassigned according to an efficient optimal pairing algorithm acting on the small minibatch. These new pairs are used to update the model, after which the pairs are stored until the data sample is revisited and pairings recomputed within the new batch. Multiple noise caches are used for small-scale datasets to prevent overfitting. Experimental results show significant improvement over the baseline independent coupling method and superior performance to other coupling methods, especially in the low-step regime.", "questions": "* How scalable is this approach? Could it be used datasets much larger than ImageNet, such as LAION-5B?\n* Is the use of fixed noise an issue? Could there be a benefit from gradually refreshing the noise?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* This paper has a clear focus and motivation. Coupling methods with little computational cost that can improve the performance of CFMs, especially in the low-step regime, have the potential to become widely adopted. The proposed method is scalable and the extra I/O cost seems reasonable.\n* The method is clearly presented and natural. While the method remains unable to guarantee an optimal pairing (a fact which holds for any method which does not scan the entire dataset at once), it provides a richer pairing mechanism than previous minibatch approaches and remains low cost.\n* The experimental results show significant improvement over the original coupling, along with strong performance relative to other few-step generation methods, and in particular among other coupling approaches.", "weaknesses": "* It seems somewhat limited to use only a fixed set of pre-sampled noise to train the model, in contrast to the typical approach of drawing fresh noise at every iteration. I would guess that the need for using caching for small datasets is related to the fact that the noise is not refreshed, leaving the method with a bit of a loose end. The work claims that in practice this is not an issue for large enough datasets. Nonetheless, I am curious if it is possible to define a variant of the method where noise is refreshed, eliminating the need for caches for smaller datasets.\n* While the work maintains a strong focus on an important problem, actual method is fairly straightforward and heavily influence by prior work. This is not a major drawback in my view as long as the benefits that are presented in the experimental section are accurate and reproducible, because the proposed method is very easy to adopt.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GtBS3ZNeP8", "reviewer": "Reviewer_5N5h", "summary": "This paper introduced LOOM-CFM to extend the scope of minibatch OT by proposing a novel iterative algorithm to optimize the global data-noise assignments of minibatch OT. This method has consistent improvements in the sampling speed-quality trade-off and also helps enhance distillation initialization of rectified flow and supports high-resolution synthesis in latent space training.", "questions": "See weakness", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "This paper introduces an iterative approach to more accurately approximate the global optimal assignment between noise and data samples, resulting in a more precise estimation of the global OT plan compared to minibatch-CFM. It also present a convergence analysis of LOOM-CFM.", "weaknesses": "1. The paper lacks an analysis of the training time overhead and convergence rate of LOOM-CFM compared to minibatch-CFM. Additionally, while finite convergence is intuitive, providing a detailed analysis of the convergence rate is crucial—especially to understand the scalability of this method for large-scale text-to-image (T2I) and text-to-video (T2V) diffusion model training.\n\n2. While the use of multiple noise caches empirically helps reduce overfitting, the one-to-many correspondence may affect the properties of probability flow. It would be valuable if the authors could provide further explanation or analysis on how multiple noise caches might or might not influence the probability flow.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oe3wYzz04t", "reviewer": "Reviewer_VAZG", "summary": "This paper proposes an improved method (LOOM) for computing the minibatch optimal transport coupling used to train conditional flow matching models. Unlike previous OT-CFM approaches which throw away the minibatch OT coupling after each neural net update, LOOM uses minibatch OT coupling to update the global coupling between training datasets.", "questions": "**Q1** : to what degree does LOOM accelerate convergence (in terms of number of model updates) compared to, e.g., independent coupling, Hungarian OT coupling, or Sinkhorn OT coupling?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**Originality** : this paper introduces an alternative (LOOM) to the Hungarian algorithm or Sinkhorn for approximating the OT coupling between two empirical measures. LOOM is suitable for conditional flow matching / minibatch flow matching in the sense that it uses local OT coupling to improve the global OT coupling.\n\n**Clarity** : method and experiment results are clearly presented. I had no problem following the writing.\n\n**Quality** : proposed method is supported with both theory and experiments on CIFAR10, ImageNet-32&64, and FFHQ 256x256. The experiment results support the authors' claim that conditional flow matching with LOOM yields faster flow models.\n\n**Significance** : the paper takes a step towards scaling up OT minibatch flow matching.", "weaknesses": "**Weak Performance** : the only weakness of the paper stopping me from giving Accept is the weak empirical performance of the proposed method compared to relevant baselines such as [1] and [2] (which also happen to be missing from Section 4). For instance, [1] achieves 1.97 FID on CIFAR10 with 35 NFEs, whereas LOOM achieves 4.41 FID with 134 NFEs. On ImageNet-64, [2] achieves around 1.4 FID with 63 NFEs while LOOM yields 2.75 FID with 133 NFEs.\n\n[1] Elucidating the Design Space of Diffusion-Based Generative Models\n\n[2] Analyzing and Improving the Training Dynamics of Diffusion Models", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "U3p8OjZYzw", "reviewer": "Reviewer_r9FF", "summary": "The paper proposes a novel method for reassigning noise-data, which overcomes the limitations of minibatch OT reassigning. This method allows noise-data pairing beyond the confines of minibatch, reducing the curvature of the sampling trajectory and thereby accelerating the adoption process.", "questions": "1. From the comparison of results between LOOM-CFM with 4 caches + reflow and LOOM-CFM with 4 caches, it can be observed that there are still many intersections in the noise-data coupling generated by LOOM-CFM. This indirectly reflects the low efficiency of the algorithm mentioned in the weaknesses. Therefore, does the issue of low efficiency in continuously exploring the global optimal pairing using minibatches impose demanding requirements on the number of training steps and training time? Could you provide the time required for training the model?\n2. The meaning of the subscript $j$ in Equation 9 does not seem very clear. Could you please provide a clearer explanation?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper presents a novel approach that overcomes the limitations of minibatch optimal transport (OT) by retaining coupling information within minibatches, thus facilitating the transfer of coupling information across minibatches. This advancement enables the construction of noise-data pairings that are more closely aligned with the global optimal transport framework. The proposed LOOM-CFM offers a new solution for developing superior noise-data couplings in the training of continuous flow models (CFMs). This work contributes to the application of computationally complex OT calculations in high-dimensional spaces within generative models, providing insights and references for future research in this area. Besides, the paper is well-written.", "weaknesses": "While LOOM-CFM demonstrates superior performance in unconditional image generation tasks compared to methods like minibatch OT, its approach of updating global pairings through optimal pairings within minibatches may not be very efficient. There exists a vast solution space for potential noise-data pairings, which poses a significant challenge. For instance, in the CIFAR-10 dataset, even when using just one cache, there are around $50000!$ possible noise-data pairing configurations. However, with a batch size of 128, each minibatch can explore a maximum of only $128!$ coupling configurations at a time. Consequently, the number of training steps required to identify the optimal pairing becomes substantially large.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)