Faster Inference of Flow-Based Generative Models via Improved Data-Noise Coupling

Aram Davtyan; Leello Tadesse Dadi; Volkan Cevher; Paolo Favaro

cs.LGMar 16, 2026

Faster Inference of Flow-Based Generative Models via Improved Data-Noise Coupling

Authors: Aram Davtyan, Leello Tadesse Dadi, Volkan Cevher, Paolo Favaro

Organizations: University of Bern, Bern, Switzerland · EPFL, Lausanne, Switzerland

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.

Figures & tables

Panel of Figure 1 — 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.

Method	Solver	CIFAR10
		NFE	FID $\downarrow$
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 ${}^{\text{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 $\downarrow$	NFE	FID $\downarrow$
ODE/SDE-based methods
St. Interpolants ( Albergo and Vanden-Eijnden, 2023 )	-	-	8.49	-	-
FM ${}^{\text{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 $\downarrow$
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

\times

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\cdot s(t)$ for different values of $\sigma$ .

Panel of Figure 8 — 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.

Panel of Figure 13 — 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

Flow-based generative modeling learns probability transports from simple noise to data for fast sampling, density estimation, and controllable generation.

See all 538 papers

Local Flow Matching Generative Models
Oct 3, 2024Chen Xu, Xiuyuan Cheng, Yao XieFlow Matching Generative Modeling
H. Milton Stewart School of Industrial and Systems Engineering, Georgia Institute of Technology. · Department of Mathematics, Duke University
ODE-free Neural Flow Matching for One-Step Generative Modeling
Apr 7, 2026Xiao ShouOne-Step Generative Modeling Flow Matching
Baylor University
COT-FM: Cluster-wise Optimal Transport Flow Matching
Mar 11, 2026Chiensheng Chiang, Kuan-Hsun Tu, Jia-Wei Liao +2Flow Matching Generative Modeling
National Taiwan University

Oct 3, 2024J/K move · Enter open · S save

Local Flow Matching Generative Models

Chen Xu, Xiuyuan Cheng, Yao Xie

H. Milton Stewart School of Industrial and Systems Engineering, Georgia Institute of Technology. · Department of Mathematics, Duke University

Flow Matching Generative Modeling Continuous Normalizing Flows Flow-Based Generative Modeling

Open paper AskPDF arXiv

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.

Abstract

Flow-Based Generative Modeling

Explore similar work

Local Flow Matching Generative Models

ODE-free Neural Flow Matching for One-Step Generative Modeling

COT-FM: Cluster-wise Optimal Transport Flow Matching