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.
Figures & tables
Figure 2: Vector field results from different coupling strategies. Random coupling forces the model to regress inconsistent (gray) velocity targets, creating ambiguous intersections and pushing the model toward an averaged (purple) direction, which produces curved velocity fields. In contrast, optimal transport provides consistent couplings with fewer intersections, enabling the model to learn much straighter velocity fields.
Figure 3: Overview of our proposed method. In Stage 1, we cluster the dataset, and for each image within a cluster Ck , we reverse the ODE back to the noise space using a pretrained flow model to compute the new mean μ0,k and covariance Σ0,k for that cluster. In Stage 2, we apply optimal transport to finetune the pretrained flow model across all clusters, encouraging it to align with the distributional structure captured in Stage 1. During inference, we first sample a cluster index k , then draw a noise x from N(μ0,k,Σ0,k) , and use this noise to generate an image through the finetuned flow model.
Dataset: 2D Examples
NFE
Mixture of 5-Gaussian
Two Moon
Checker Board
Wasserstein 2
Curvature
Wasserstein 2
Curvature
Wasserstein 2
Curvature
Rectified Flow [ 32 ]
100
0.5421
0.0316
0.1006
0.0111
0.3900
9.1946
OT-CFM [ 56 ]
100
0.6582
0.0104
0.1074
0.0020
0.3188
0.1741
MeanFlow [ 17 ]
1
0.7612
0.9170
0.1233
-
0.9170
-
COT-FM (Ours)
100
0.1995
0.0084
0.0266
0.0016
0.2550
0.1505
Table 1: 2D synthetic data distributions. COT-FM outperforms all baselines in both Wasserstein distance and curvature.
Figure 4: Visualization of 2D checkerboard data under different methods. Blue points denote the source distribution, gray points the target distribution, and orange points the generated samples.
Model
Method
100
50
10
2
1
SiT-B/2
Rectified Flow [ 32 ]
5.82
5.86
8.25
119.57
264.36
COT-FM (Ours)
5.11
5.28
7.52
101.66
231.99
SiT-B/4
Rectified Flow [ 32 ]
8.30
8.39
11.16
134.99
276.13
COT-FM (Ours)
7.65
7.81
9.87
114.10
241.18
Table 3: Conditional Image Generation on ImageNet 256 × 256. FID ( ↓ ) at different NFE steps.
Figure 6
(c) Different cluster sampling probabilities
Method
FID
Uniform
4.26
Proportional to cluster size
3.97
Table 5: Ablation studies on CIFAR-10 at 50-step generation. (a) Additional alternating iterations progressively reduce FID. (b) COT-FM maintains consistent performance on held-out test data. (c) Size-proportional cluster sampling outperforms uniform sampling.
Figure 6: CIFAR-10 visualization for different flow-based methods under 50-step generation settings.
Figure 7: Visualization of a CIFAR-10 cluster. The first row shows representative images from the cluster. Using these images, we apply reverse ODE to estimate the μ0,k and Σ0,k (Fig. 3 , Stage 1). The second row shows images generated by sampling noise from this estimated distribution. The last row presents the results after finetuning the pretrained flow model.
OT-CFM (Batch OT)
COT-FM (Ours)
b
Time (sec.)
FID
n/K
Time (sec.)
FID
512
6.2
4.78
500 (50000/100)
6.3
3.97
Table 6: Speed–quality tradeoff comparison on CIFAR-10. OT-CFM with batch size b=512 and COT-FM with cluster size n/K=500 yield comparable per-epoch OT computation time, but COT-FM achieves significantly better FID.
One-time Preprocessing
Time Cost (sec.)
Reverse ODE
420
DINO feature extraction
15
K-means clustering
10
Table 7: Wall-clock time breakdown for one-time preprocessing on CIFAR-10.
Method
FID ( ↓ )
2-Rectified Flow [ 32 ]
12.21
2-Rectified Flow [ 32 ] + COT-FM (Ours)
10.91
Table 8: 1-step FID ( ↓ ) comparison on CIFAR-10 between 2-Rectified Flow and 2-Rectified Flow enhanced with COT-FM.
Method
FID ↓
IS ↑
Precision ↑
Recall ↑
Rectified Flow [ 32 ]
4.45
9.11
0.78
0.65
COT-FM (Ours)
3.97
9.36
0.78
0.67
Table 9: Additional metrics on CIFAR-10 (50-step) with K=100 .
Method
K=5
K=50
K=100
K=120
COT-FM
4.56
4.10
3.97
4.06
Table 10: Ablation on number of clusters K (CIFAR-10, 50-step FID ↓ ).
Figure 8: Training loss across alternating optimization iterations, showing monotonic decrease.
Figure 9: Visualization of CIFAR-10 for different flow-based methods under 10-step and 50-step generation settings.
Figure 10: Visualization of CIFAR-10 for MeanFlow under 1-step generation settings.
Component
Rectified Flow
MeanFlow
Channels
128
128
Channel multiple
(1, 2, 2, 2)
(2, 2, 2)
Residual blocks per resolution
4
4
Normalization
GroupNorm
GroupNorm
Nonlinearity
Swish
Swish
Attention resolution
16
16
Table 11: UNet architecture for CIFAR-10 experiments.
Components
Number of Parameters
ViT
360M
VLM
205M
Action Encoders
3.2M
Action Heads
31.8K
Global-AdaLN
28.3M
Cond Linear Proj.
1.0M
Table 12: The parameters of all model components in FLOWER.
Flow matching models learn to transport samples from a simple prior distribution to a complex data distribution. When prior-data pairs are coupled via optimal transport (OT), the learned trajectories are straight and non-crossing, enabling fast, even single-step, generation. However, computing the OT coupling in high dimensions is intractable, and existing methods attempt to solve the OT problem, at the cost of persistent bias or significant overhead. Rather than solving for the OT coupling, we reformulate the problem. Once the prior is treated as a design choice rather than a fixed input, the OT coupling between prior and data is no longer unique. Many priors admit an OT-optimal identity coupling to the data, leaving us free to choose one that is also tractable to sample. We identify low-frequency projection of natural images as such a choice. The identity coupling between data and its low-frequency representation is empirically OT-optimal, the prior is structured enough to be sampled by a lightweight model at inference, and the remaining flow-matching task reduces to synthesizing high-frequency detail. Interpolating the prior with Gaussian noise further improves generation quality while preserving the OT coupling. The approach requires no modifications to the flow model itself, and integrates naturally with latent-space models, classifier-free guidance, and one-step generation frameworks. Across all benchmarks, our method reduces trajectory curvature by more than 2× compared to existing flow matching methods, yielding better generation quality in the few-step regime.
Flow Matching (FM) constructs linear conditional probability paths, but the learned marginal velocity field inevitably exhibits strong curvature due to trajectory superposition. This curvature severely inflates numerical truncation errors, bottlenecking few-step sampling. To overcome this, we introduce Isokinetic Flow Matching (Iso-FM), a lightweight, Jacobian-free dynamical regularizer that directly penalizes pathwise acceleration. By using a self-guided finite-difference approximation of the material derivative Dv/Dt, Iso-FM enforces local velocity consistency without requiring auxiliary encoders or expensive second-order autodifferentiation. Operating as a pure plug-and-play addition to single-stage FM training, Iso-FM dramatically improves few-step generation. On CIFAR-10 (DiT-S/2), Iso-FM slashes conditional non-OT FID at 2 steps from 78.82 to 27.13 - a 2.9x relative efficiency gain - and reaches a best-observed FID at 4 steps of 10.23. These results firmly establish acceleration regularization as a principled, compute-efficient mechanism for fast generative sampling.
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