hkchengrex/c2ot-imagenet32-ot

C$^2$OT: Conditional Optimal Transport for Flow-Based Generation

This model implements C$^2$OT, an algorithm for computing prior-to-data couplings for flow-matching-based generative models, as presented in the paper The Curse of Conditions: Analyzing and Improving Optimal Transport for Conditional Flow-Based Generation. It is specifically trained on ImageNet-32x32 for conditional image generation, leveraging minibatch optimal transport data coupling.

Abstract

Minibatch optimal transport coupling straightens paths in unconditional flow matching. This leads to computationally less demanding inference as fewer integration steps and less complex numerical solvers can be employed when numerically solving an ordinary differential equation at test time. However, in the conditional setting, minibatch optimal transport falls short. This is because the default optimal transport mapping disregards conditions, resulting in a conditionally skewed prior distribution during training. In contrast, at test time, we have no access to the skewed prior, and instead sample from the full, unbiased prior distribution. This gap between training and testing leads to a subpar performance. To bridge this gap, we propose conditional optimal transport C$^2$OT that adds a conditional weighting term in the cost matrix when computing the optimal transport assignment. Experiments demonstrate that this simple fix works with both discrete and continuous conditions in 8gaussians-to-moons, CIFAR-10, ImageNet-32x32, and ImageNet-256x256. Our method performs better overall compared to the existing baselines across different function evaluation budgets.

Links

• Paper: https://arxiv.org/abs/2503.10636
• Project Page: https://hkchengrex.github.io/C2OT
• Code Repository: https://github.com/hkchengrex/C2OT

Sample Usage

For a practical demonstration and to get started with C$^2$OT, we highly recommend trying out the official Colab notebook:

https://colab.research.google.com/drive/1uhYPqnGlPoMTEqEgzpPvFQEcnr0faSBA?usp=sharing

This notebook provides a comprehensive guide to using the model and reproducing key results.

Citation

If you use C$^2$OT in your research, please cite the original paper:

@inproceedings{cheng2025curse,
  title={The Curse of Conditions: Analyzing and Improving Optimal Transport for Conditional Flow-Based Generation},
  author={Cheng, Ho Kei and Schwing, Alexander},
  booktitle={ICCV},
  year={2025}
}