This repository contains code for the experiments of NeurIPS 2024 paper Optimal Flow Matching: Learning Straight Trajectories in Just One Step by Nikita Kornilov, Petr Mokrov, Alexander Gasnikov and Alexander Korotin.
We propose novel Optimal Flow Matching (OFM) method which extends conventional Flow Matching algorithm for specific "optimal" parameterization of transport ODE vector field.
Pros. of our method:
- OFM recovers straight trajectories between source and target distributions in just one Flow Matching step.
- OFM solves quadratic Optimal Transport problem between source and target distributions.
Cons. of our method:
- OFM is based on rather complex training procedure (with auxiliary optimization step).
- OFM relies on Input Convex NNs which have questionable scalability.
The notebooks for Illustrative 2D Gaussian->Eight Gaussians experiments:
.\ofm_notebooks\toy8gau_*.ipynbThe script for Wasserstein-2 benchmark:
.\ofm_scripts\bench_script.pyBefore running this scripts place the checkpoints from the W2 benchmark github repository folder in the folder .\benchmarks\Mix3toMix10.
The notebooks for ALAE experiment:
.\ofm_notebooks\alae_*.ipynbBefore running ALAE notebooks, have a look at the ./ALAE folder (e.g. some packages need to be installed).
If you find this repository or the ideas presented in our paper useful, please consider citing our paper.
@inproceedings{
kornilov2024optimal,
title={Optimal Flow Matching: Learning Straight Trajectories in Just One Step},
author={Nikita Maksimovich Kornilov and Petr Mokrov and Alexander Gasnikov and Alexander Korotin},
booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
year={2024},
url={https://openreview.net/forum?id=kqmucDKVcU}
}