COT-FM: Cluster-wise Optimal Transport Flow Matching

Chiensheng Chiang; Kuan-Hsun Tu; Jia-Wei Liao; Cheng-Fu Chou; Tsung-Wei Ke

cs.CVMar 11, 2026

COT-FM: Cluster-wise Optimal Transport Flow Matching

Authors: Chiensheng Chiang, Kuan-Hsun Tu, Jia-Wei Liao, Cheng-Fu Chou, Tsung-Wei Ke

Organizations: National Taiwan University

Abstract

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 $\mathcal{C}_{k}$ , we reverse the ODE back to the noise space using a pretrained flow model to compute the new mean $\bm{\mu}_{0,k}$ and covariance $\bm{\Sigma}_{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 $\bm{x}$ from $\mathcal{N}(\bm{\mu}_{0,k},\bm{\Sigma}_{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
Dataset: 2D Examples	NFE	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
SiT-B/2	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
SiT-B/4	COT-FM (Ours)	7.65	7.81	9.87	114.10	241.18

Table 3: Conditional Image Generation on ImageNet 256

\times

256. FID (

\downarrow

) at different NFE steps.

(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 $\bm{\mu}_{0,k}$ and $\bm{\Sigma}_{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 ( $\downarrow$ )
2-Rectified Flow [ 32 ]	12.21
2-Rectified Flow [ 32 ] + COT-FM (Ours)	10.91

Table 8: 1-step FID (

\downarrow

) comparison on CIFAR-10 between 2-Rectified Flow and 2-Rectified Flow enhanced with COT-FM.

Method	FID $\downarrow$	IS $\uparrow$	Precision $\uparrow$	Recall $\uparrow$
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

\downarrow

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.

Settings	LIBERO
Action Space Encoders	2-layered MLP
Action Space Decoders	Linear Attention
Number of Flow-T Layers	18
Latent Dimension	1024
Number of Heads	16
Position Embedding	1D Rope

Table 13: LIBERO hyperparameters of FLOWER.

Flow Matching

Flow matching learns continuous-time vector fields that transport simple noise distributions into data for efficient generative modeling.

See all 639 papers

Optimal Transport Flow Matching by Design
Jun 2, 2026Shimon Malnick, Matan Rusanovsky, Ohad Fried +1Flow Matching One-Step Generative Modeling
1Tel Aviv University · 2Reichman University
Isokinetic Flow Matching for Pathwise Straightening of Generative Flows
Apr 6, 2026Tauhid KhanFlow Matching Flow-Based Generative Modeling
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

Jun 2, 2026J/K move · Enter open · S save

Optimal Transport Flow Matching by Design

Shimon Malnick, Matan Rusanovsky, Ohad Fried, Shai Avidan

1Tel Aviv University · 2Reichman University

Flow Matching One-Step Generative Modeling Flow-Based Generative Modeling Optimal Transport

Open paper AskPDF arXiv

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\times

compared to existing flow matching methods, yielding better generation quality in the few-step regime.

Abstract

Flow Matching

Explore similar work

Optimal Transport Flow Matching by Design

Isokinetic Flow Matching for Pathwise Straightening of Generative Flows

Local Flow Matching Generative Models