COT-FM: Cluster-wise Optimal Transport Flow Matching

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

Source-linked AI summary

COT-FM: Cluster-wise Optimal Transport Flow Matching

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

arXiv:2603.13395v1Mar 11, 2026cs.CVcs.LGcs.RO

TL;DR

FM models can learn curved trajectories from random or batchwise couplings, increasing discretization error and limiting generation quality at low sampling budgets. COT-FM reshapes the probability path by clustering targets, assigning reversed source distributions, and solving local transport without changing the FM architecture. Across synthetic data, image generation, and robotic manipulation, it consistently accelerates sampling while improving quality.

Problem
Random and batchwise couplings can produce curved FM trajectories, increasing discretization error and reducing generation quality under low sampling budgets.
Method
COT-FM clusters target samples, obtains cluster-specific source distributions by reversing pretrained FM paths, and approximately solves optimal couplings within each cluster.
Results
COT-FM consistently improves generation across 2D datasets, image benchmarks, and robotic manipulation while accelerating one-step and few-step sampling.
Takeaways & Limitations
COT-FM is a plug-and-play approach that improves transport straightness and generation quality without altering FM model architectures or training pipelines.
Takeaways & Limitations
COT-FM faces scaling and generalization challenges because reverse ODE computation becomes inefficient with larger sample sizes and performance depends on clustering quality.

Abstract

from arXiv · show

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.

1. Introduction

Flow Matching can suffer from curved trajectories caused by random or batchwise couplings, increasing discretization error and reducing low-step generation quality. COT-FM clusters targets, assigns each cluster a reversed source distribution, and uses local transport to learn straighter fields without changing FM architectures.

Motivation: Flow Matching maps a simple source distribution to a complex data distribution by learning a deterministic vector field and integrating it during inference.The learned field induces the probability path between source and target distributions.
Motivation: Random couplings create conflicting velocity targets and curved fields, while batchwise approximations struggle with locality.Curved paths increase time-discretization error and can distort transport when samples reach unseen locations.
Method: COT-FM clusters target samples and assigns each cluster a source distribution obtained by reversing pretrained FM paths.The strategy reduces coupling problems to smaller cluster-level matches and restricts source distributions for more efficient field learning.
Method: COT-FM modulates only the target probability path, preserving FM architectures and input–output mechanisms while supporting conditional and unconditional clustering.Clustering can use class labels, textual descriptions, or unsupervised methods.
Results: COT-FM improves generation across 2D benchmarks, CIFAR-10, and robotic manipulation, including lower Wasserstein distance, lower FID, and strong one-NFE success rates.On Mixture of 5-Gaussians, Wasserstein distance decreases from 0.5421 to 0.1995; on CIFAR-10, Rectified Flow FID decreases from 12.6 to 8.23 at 10 steps.
Results: Optimal transport produces fewer intersections and straighter vector fields than random coupling, supporting more reliable low-step generation.The reported improvements span one-step and few-step generation and are accompanied by the lowest curvature on 2D benchmarks.

2. Preliminary of Flow Matching

Flow Matching learns a time-dependent vector field that transports samples from a simple source distribution to a complex target distribution through ODE integration. Its practical construction uses conditional paths, but random or batchwise couplings can produce curved marginal trajectories and discretization error.

Flow Matching regresses a neural vector field that transports samples from source distribution p0 to target distribution p1 along an ODE-defined probability path.The learned field supports generation by integrating the ODE during testing.
Conditional probability paths make otherwise intractable marginal vector fields constructible by mixing simpler paths over a variable z.Marginalizing the conditional paths yields the marginal probability path and vector field used by FM.
Random coupling averages conflicting pairwise directions, so its marginal vector field becomes curved even though each individual sample path is straight.This curvature can lead to distorted transport and lower-quality generation.
Batchwise optimal transport approximations reduce global computation but struggle with locality when mini-batches are small, resulting in curved paths.Global OT is computationally infeasible for large datasets because its map has cubic time complexity and continuous sources contain infinitely many samples.
COT-FM’s overview clusters the dataset, estimates cluster-specific source statistics by reversing a pretrained flow, then applies cluster-wise optimal transport.Inference samples a cluster, draws noise from its Gaussian source distribution, and generates an image with the finetuned flow.

3. Method

COT-FM divides target samples into clusters, recovers a source distribution for each cluster, and approximately solves local optimal-transport problems to construct straighter target vector fields. It alternates vector-field refinement with standard FM updates while preserving the model architecture and inference procedure apart from cluster-conditioned initialization.

Core framework: COT-FM clusters target samples and assigns each cluster a source distribution, reducing optimal-coupling calculations to smaller local matching problems.The framework is plug-and-play because it modulates the probability path without changing the FM architecture or input-output mechanisms.
Clustering choices: COT-FM supports fixed or non-fixed clustering, including class labels, textual descriptions, unsupervised clusters, and learned source distributions for growing robot-policy clusters.A learning-based module predicts source distributions when the cluster set changes during rollout.
Cluster-wise source distributions: Pretrained FM flows are reversed to trace source samples for each cluster, whose mean and covariance define an approximate Gaussian source distribution p0,k.The method relies on flow reversibility to estimate sources corresponding to the cluster’s data samples.
Cluster-wise optimal transport: Cluster-wise OT maps are computed between each cluster and its assigned source samples, making batchwise approximation more feasible by reducing sample counts.The resulting maps jointly define a target vector field that the FM model regresses.
Optimization: Training alternates between refining the cluster-wise target vector field and updating the FM model with conditional flow-matching regression.Minibatches sample clusters proportionally to cluster size and draw source-target pairs from the corresponding OT map.
Sampling: During inference, COT-FM samples a cluster index and then initializes generation from that cluster’s source distribution while retaining the standard FM integration procedure.The sole inference change is replacing one global source draw with a cluster-conditioned source draw.

4. Experiments

COT-FM is evaluated on 2D distributions, CIFAR-10 and ImageNet image generation, and LIBERO robotic manipulation. Across these settings, it improves distributional quality, flow straightness, or single-step success relative to baselines, with particularly strong gains at low sampling budgets.

2D Point Cloud Generation: COT-FM achieves the lowest Wasserstein distance and curvature among compared methods on the 2D synthetic benchmarks.The reported Wasserstein distances are 0.1995, 0.0266, and 0.2550, while curvature values are 0.0084, 0.0016, and 0.1505.
CIFAR-10 Image Generation: 82 FID points are gained over unclustered Rectified Flow, followed by an additional 91-point FID reduction from within-cluster batch optimal coupling at one sampling step.COT-FM also reduces FID from 2.92 to 2.60 and from 2.88 to 2.53 with the same backbone as MeanFlow.
ImageNet Generation: COT-FM consistently improves FID across all NFE settings for both ImageNet 256×256 backbone architectures, with larger gains in the low-NFE regime.The method scales to high-resolution class-conditional generation without modifying the core algorithm.
Robotic Manipulation: At one sampling step, COT-FM exceeds all single-step LIBERO baselines, gaining 0.4% on LIBERO-Spatial and 3% on LIBERO-Long over the second-best method.Its one-step performance surpasses FLOWER with four steps.
Ablation Analysis: Two alternating optimization iterations produce the best reported CIFAR-10 FID, improving the initial 4.45 to 3.97 before diminishing returns at iteration three.The held-out test results remain consistent with training results: Rectified Flow reports 8.55 versus COT-FM’s 8.19.

5. Conclusions

COT-FM reshapes Flow Matching paths by clustering targets and assigning each cluster a reversed source distribution, producing straighter and more accurate transport without changing model architectures or training pipelines. Experiments and ablations support faster, higher-quality generation, while scalability, clustering dependence, and the behavior of local transport remain open challenges.

Conclusions: COT-FM divides optimal-coupling computation across target clusters, estimates cluster-specific source distributions, and updates the FM model using straighter flows.The framework alternates between optimizing the target vector field and the FM model.
Conclusions: The method accelerates sampling and improves generation quality across 2D synthetic data, image generation, and robotic manipulation without altering model architectures or training pipelines.The conclusion characterizes COT-FM as plug-and-play across these domains.
Limitations and Future Work: Reverse ODE estimation becomes inefficient as sample size grows, and performance depends on clustering quality.The authors also identify incomplete understanding of how local OT captures global structure and why refinement converges within one or two cycles.

Supplementary Material

For conditional generation with unseen conditions, COT-FM learns a model that predicts a condition-specific noise distribution and trains it through a one-step reinforcement-learning formulation. The resulting samples are evaluated by rolling them through the pretrained FM model and rewarding terminal-point agreement with targets.

Conditional Source Distributions: Conditional COT-FM predicts the mean and standard deviation of a noise distribution dynamically for each condition, rather than precomputing one distribution per enumerated cluster.The covariance is approximated using only the predicted standard deviation because full high-dimensional covariance learning is difficult and unstable.
Conditional Model Training: The conditional model is trained as a one-step MDP in which the condition is the state, the sampled initial point is the action, and PPO maximizes expected return.The reward is based on mean squared error between generated and target terminal points.
Algorithmic Integration: During training and sampling, conditional outputs replace the predefined cluster mean and standard deviation in the original COT-FM procedure.The FM rollout produces a terminal point used to compute the transition reward and advantage.

B. Computational Cost Analysis

COT-FM adds cluster-wise OT computation and reverse-ODE preprocessing, trading extra preprocessing for improved generation quality at comparable per-epoch OT cost.

Computational complexity: O(n(n/K)^2) per epoch is COT-FM's exact cluster-wise OT cost, compared with O(nb^2) for batch OT.Exact global OT costs O(n^3), while COT-FM solves OT separately within K clusters.
Speed–quality tradeoff: At comparable per-epoch OT wall-clock time, COT-FM achieves significantly better FID than OT-CFM.The comparison uses OT-CFM with batch size b = 512 and COT-FM with cluster size n/K = 500.
Runtime overhead: 445 seconds of one-time CIFAR-10 preprocessing is dominated by reverse ODE computation, while per-epoch OT takes approximately 15 seconds.The reverse ODE accounts for 420 seconds, equivalent to roughly five training epochs.
Runtime overhead: 420 seconds of reverse-ODE preprocessing yields COT-FM FID 3.97, whereas training Rectified Flow for equivalent extra epochs raises FID from 4.45 to 4.49.OT-CFM shows no improvement under the same comparison, remaining at FID 4.78.

C. Comparison with 2-Rectified Flow

COT-FM complements 2-Rectified Flow by addressing transport through cluster-wise optimal transport after 2-RF straightens trajectories via synthetic couplings.

2-RF pipeline: 2-RF trains 1-RF, generates synthetic couplings with forward ODE, and retrains the model.This provides the two-stage rectification pipeline used for comparison.
Results: 10.91 FID at 1-step generation is achieved by 2-RF + COT-FM, versus 12.21 FID for 2-RF alone.The result demonstrates an additional gain from applying COT-FM on top of 2-RF.
Complementarity: 2-RF straightens trajectories through synthetic-coupling retraining, while COT-FM reduces path crossings through cluster-wise optimal transport.Because they address different aspects of transport, the methods can be combined for cumulative improvement.

D. Additional Metrics and Ablation on K

COT-FM improves CIFAR-10 distributional metrics while remaining robust to cluster count, although overly fine clustering reduces intra-cluster OT effectiveness.

Additional metrics: At 50 sampling steps, COT-FM improves FID, IS, and Recall over Rectified Flow while maintaining the same Precision.The pattern indicates broader target-distribution coverage without sacrificing sample fidelity.
Cluster-count ablation: FID 4.56 with K = 5 outperforms OT-CFM's FID 4.78, showing robustness to relatively few clusters.Performance peaks around K = 100 and slightly degrades at K = 120.
Cluster-count ablation: Overly fine-grained clustering can reduce samples per cluster and diminish intra-cluster OT effectiveness.The authors recommend the elbow method for selecting K in practice.
Clustering quality: For robot manipulation, COT-FM does not require explicit clustering because a conditional model learns source distributions from task embeddings.For unconditional generation, self-supervised DINO features provide semantically meaningful clusters.

F. Convergence of Alternating Optimization

Alternating optimization improves COT-FM over successive iterations, with monotonically decreasing training loss and diminishing FID gains by the third iteration.

Convergence results: FID decreases from 4.45 at 0 iterations to 4.23 after 1 iteration and 3.97 after 2 iterations, with diminishing returns at iteration 3.The training loss also decreases monotonically across alternating optimization iterations.
Convergence mechanism: The stable convergence behavior is attributed to non-intersecting reverse-ODE paths followed by lower-cost local transport within clusters.Each iteration therefore starts from a better-aligned coupling than the previous one, according to the paper's intuition.
Optimization procedure: COT-FM's alternating procedure refines the cluster-wise target vector field and FM model across successive rounds.The framework's broader acceleration strategy concerns learning straighter velocity fields rather than only skipping inference steps.

H.1. Compute Environment

The experiments span 2D benchmarks, CIFAR-10, ImageNet, and LIBERO robotics, using task-specific hardware and evaluation settings. Comparisons generally retain shared architectures or follow baseline configurations.

Compute environment: CIFAR-10 runs on one AMD MI300X GPU, ImageNet uses eight MI300X GPUs, and LIBERO uses one NVIDIA RTX 4090 GPU.The stated compute environment includes AMD EPYC 9684X processors and ROCm 7.2.0 for the AMD experiments.
Evaluation scope: Experiments cover 2D point clouds, unconditional CIFAR-10, conditional ImageNet, and LIBERO robotic manipulation tasks.The 2D benchmarks use Mixture of 5 Gaussians, Two Moons, and Checkerboard Grid; LIBERO evaluates Spatial and Long tasks.
2D training and metrics: All 2D methods use the same network architecture and are trained with Adam, learning rate 1e-3, batch size 512, and 500 epochs.Evaluation measures exact 2-Wasserstein distance and trajectory curvature, with lower curvature corresponding to straighter paths.
Image-generation settings: CIFAR-10 comparisons use Rectified Flow settings with 600 epochs and an NCSN++-style UNet, or MeanFlow settings with 240 epochs and a corresponding UNet configuration.Both configurations maintain exponential moving average weights, with decay rates of 0.999999 and 0.99999 respectively.
Fine-tuning: COT-FM is fine-tuned from a pretrained Rectified Flow model after 160 epochs, adding 30 training epochs within a 190-epoch total schedule.The implementation uses batch size 256 and Adam with β1 = 0.9 and β2 = 0.95.

I. Additional Generated Results

Additional results provide CIFAR-10 samples under matched-noise conditions and document the architectures and hyperparameters used in the experiments. Figures compare flow-based methods at different generation-step budgets, including MeanFlow at one step.

Additional generated images: Additional CIFAR-10 images use identical noise for baseline and COT-FM samples to support visual consistency.The stated comparison uses noise sampled from N(0, I).
Implementation references: Table 11 documents the UNet architecture used for CIFAR-10 experiments.Tables 12 and 13 document FLOWER model components and LIBERO hyperparameters, respectively.
CIFAR-10 comparisons: Figure 9 visualizes different flow-based methods on CIFAR-10 at 10-step and 50-step generation settings.The figure is intended for comparison across two sampling budgets.
One-step generation: Figure 10 visualizes MeanFlow on CIFAR-10 under a 1-step generation setting.This isolates the very-low-step sampling regime for MeanFlow.