https://github.com/hkchengrex/c2ot
The Curse of Conditions: Analyzing and Improving Optimal Transport for Conditional Flow-Based Generation
https://github.com/hkchengrex/c2ot
flow-matching generative-modeling generative-models
Last synced: about 1 year ago
JSON representation
The Curse of Conditions: Analyzing and Improving Optimal Transport for Conditional Flow-Based Generation
- Host: GitHub
- URL: https://github.com/hkchengrex/c2ot
- Owner: hkchengrex
- License: mit
- Created: 2025-03-13T04:55:29.000Z (over 1 year ago)
- Default Branch: main
- Last Pushed: 2025-03-14T03:32:54.000Z (over 1 year ago)
- Last Synced: 2025-05-08T02:53:24.021Z (about 1 year ago)
- Topics: flow-matching, generative-modeling, generative-models
- Language: Jupyter Notebook
- Homepage: http://hkchengrex.com/C2OT/
- Size: 5.79 MB
- Stars: 7
- Watchers: 1
- Forks: 1
- Open Issues: 1
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
The Curse of Conditions: Analyzing and Improving Optimal Transport for Conditional Flow-Based Generation
Paper | Webpage | Colab
Ho Kei Cheng and
Alexander Schwing
University of Illinois Urbana-Champaign

## High-Level Summary
C2OT is an algorithm for computing prior-to-data couplings for flow-matching-based generative models during training.
Our goal is to achieve straighter flows, enabled by optimal transport (OT) couplings, while mitigating the test-time degradation that OT encounters in the conditional setting (see figure above).
The key idea is that OT samples from a condition-skewed prior distribution at test time, whereas C2OT unskews the prior by incorporating a condition-dependent term into the OT cost.
## Installation
We have only tested this on Ubuntu.
### Prerequisites
We recommend using a [miniforge](https://github.com/conda-forge/miniforge) environment.
- Python 3.9+
- PyTorch 2.5.1+ and corresponding torchvision/torchaudio (pick your CUDA version https://pytorch.org/, pip install recommended)
**1. Install prerequisite if not yet met:**
```bash
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 --upgrade
```
(Or any other CUDA versions that your GPUs/driver support)
**2. Clone our repository:**
```bash
git clone https://github.com/hkchengrex/C2OT.git
```
**3. Install with pip (install pytorch first before attempting this!):**
```bash
cd C2OT
pip install -e .
```
(If you encounter the File "setup.py" not found error, upgrade your pip with pip install --upgrade pip)
## Demo
After installation, you can run our demo notebook at [moons.ipynb](moons.ipynb). More details to be added.
You can also run this demo on [Colab](https://colab.research.google.com/drive/1uhYPqnGlPoMTEqEgzpPvFQEcnr0faSBA?usp=sharing) without a local installation.
## Citation
```bibtex
@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={arXiv},
year={2025}
}
```
## Acknowledgement
Many thanks to:
- [torchcfm](https://github.com/atong01/conditional-flow-matching)