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.
Figures & tables
Figure 1: 2D Gaussian-to-Gaussian 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 an average over all source-to-target directions. As in the Gaussian case the distribution is symmetric around the origin, the learned vector field for small values of t points towards the origin, since the target points are more likely to be sampled from the other side. (c) OT coupling resolves the curvature of the trajectories and results in the identity mapping via averaging only over the sources and targets that are close to each other.
Figure 2Figure 3
Method
Solver
CIFAR10
NFE
FID ↓
ODE/SDE-based methods
DDPM ( Ho et al., 2020 )
-
1000
3.17
DDPM ( Ho et al., 2020 )
dopri5
274
7.48
St. Interpolatns ( Albergo and Vanden-Eijnden, 2023 )
-
-
10.27
FM w/ OT ( Lipman et al., 2022 )
dopri5
142
6.35
Table 1: Unconditional generation results on CIFAR10. ( ∗ ) indicates that the reported numbers are calculated by us using the official code, as they were not reported in the corresponding papers.
Method
Solver
ImageNet-32
ImageNet-64
NFE
FID ↓
NFE
FID ↓
ODE/SDE-based methods
St. Interpolants ( Albergo and Vanden-Eijnden, 2023 )
-
-
8.49
-
-
FM w/ OT ( Lipman et al., 2022 )
dopri5
122
5.02
138
14.45
Distilled models / Dedicated solvers
Improved training for straighter trajectories
Table 2: Unconditional generation results on ImageNet-32 and ImageNet-64
Method
Solver
FFHQ-256
NFE
FID ↓
Prior Work
LDM ( Rombach et al., 2022 )
-
50
4.98
ImageBART ( Esser et al., 2021 )
-
3
9.57
Latent Flow Matching ( Dao et al., 2023 ) (ADM)
dopri5
84
8.07
Latent Flow Matching ( Dao et al., 2023 ) (DiT L/2)
dopri5
88
4.55
Table 3: Unconditional generation results on FFHQ 256 × 256
Appendix figures & tables12 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 7: The form of the multiplier s(t) by y in the optimal vector field v(y,t)=y⋅s(t) for different values of σ .
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 a counterclockwise fashion from the vertices of the polygon. The optimal assignment matches each data points to the following noise point going in the clockwise direction. The assignment represented in blue is sub-optimal but it is stationary for any batch size m<n .
CIFAR10
ImageNet-32
ImageNet-64
FFHQ256
Input shape
[3, 32, 32]
[3, 32, 32]
[3, 64, 64]
[4, 32, 32]
Channels
128
128
192
256
Number of Res blocks
2
3
2
2
Channels multipliers
[1, 2, 2, 2]
[1, 2, 2, 2]
[1, 2, 3, 4]
[1, 2, 3, 4]
Heads
4
4
4
4
Heads channels
64
64
64
64
Appendix
Table 4: ADM network architecture and training parameters of LOOM-CFM for each model.
Figure 13: Sampling trajectories of LOOM-CFM (top) and the version of the algorithm that does not save the reassignments (bottom) trained on CIFAR10. The i -th row in both grids starts with the same initial noise. LOOM-CFM tends to produce sharper genereations and converges earlier.
Figure 14: Sampling trajectories of the model trained on ImageNet-32.
Figure 15: Sampling trajectories of the model trained on ImageNet-64.
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 CIFAR10.
Figure 19: Uncurated samples from the model trained on ImageNet-32.
Figure 20: Uncurated samples from the model trained on ImageNet-64.
Figure 21: Uncurated samples from the model trained on FFHQ-256.
Flow-based generative modeling learns probability transports from simple noise to data for fast sampling, density estimation, and controllable generation.
Flow Matching (FM) is a simulation-free method for learning a continuous, invertible flow that interpolates between two distributions, and in particular generates data from noise. Inspired by the variational nature of the diffusion process as a gradient flow, we introduce a stepwise FM model, Local Flow Matching (LFM), which sequentially learns a sequence of FM submodels, each matching a diffusion process up to the time-step size in the data-to-noise direction. In each step, the two distributions to be interpolated by the sub-flow model are closer than those in the full-flow matching model, which interpolates data to noise distributions, enabling smaller models with more efficient training. This variational perspective also allows us to prove a theoretical generation guarantee for the proposed flow model in terms of the χ2-divergence between the generated and true data distributions, leveraging the contraction property of the diffusion process. In practice, the stepwise structure of LFM is naturally amenable to model distillation, and various distillation techniques can be applied to accelerate generation. We empirically demonstrate that LFM achieves competitive generative performance compared to FM on unconditional generation of tabular and image datasets, and on conditional generation of robotic manipulation policies.
Chen Xu, Xiuyuan Cheng, Yao Xie
H. Milton Stewart School of Industrial and Systems Engineering, Georgia Institute of Technology. · Department of Mathematics, Duke University
Diffusion and flow matching models generate samples by learning time-dependent vector fields whose integration transports noise to data, requiring tens to hundreds of network evaluations at inference. We instead learn the transport map directly. We propose Optimal Transport Neural Flow Matching (OT-NFM), an ODE-free generative framework that parameterizes the flow map with neural flows, enabling true one-step generation with a single forward pass. We show that naive flow-map training suffers from mean collapse, where inconsistent noise-data pairings drive all outputs toward the data mean. We prove that consistent coupling is necessary for non-degenerate learning and address this using optimal transport pairings with scalable minibatch and online coupling strategies. Experiments on synthetic benchmarks and image generation tasks (MNIST and CIFAR-10) demonstrate competitive sample quality while reducing inference to a single network evaluation.
We introduce COT-FM, a general framework that reshapes the probability path in Flow Matching (FM) to achieve faster and more reliable generation. FM models often produce curved trajectories due to random or batchwise couplings, which increase discretization error and reduce sample quality. COT-FM fixes this by clustering target samples and assigning each cluster a dedicated source distribution obtained by reversing pretrained FM models. This divide-and-conquer strategy yields more accurate local transport and significantly straighter vector fields, all without changing the model architecture. As a plug-and-play approach, COT-FM consistently accelerates sampling and improves generation quality across 2D datasets, image generation benchmarks, and robotic manipulation tasks.