Light Unbalanced Optimal Transport

This is the official Python implementation of the NeurIPS 2024 paper Light Unbalanced Optimal Transport (paper page on NeurIPS) by Milena Gazdieva, Arip Asadulaev, Evgeny Burnaev and Alexander Korotin.

The repository contains reproducible PyTorch source code for fast, simple, and effective solver which allows solving the continuous unbalanced entropic optimal transport (UEOT) problem in minutes on CPU. The algorithm can be used to approximate the UEOT plans.

Presentations

• Talk by Milena Gazdieva at AIRI meetup (5 December 2024, RU)
• Short Talk by Milena Gazdieva at NeurIPS conference website (EN)

Related repositories

• Repository for Light Schrödinger Bridge paper (ICLR 2024);
• Repository for ALAE model.

Citation

@inproceedings{
    gazdieva2024light,
    title={Light Unbalanced Optimal Tansport},
    author={Gazdieva, Milena and Asadulaev, Arip and Burnaev, Evgeny and Korotin, Alexander},
    booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems}
    year={2024}
}

Application to Unpaired Image-to-Image Translation Task

A popular testbed which is usually considered in OT/EOT papers is the unpaired image-to-image translation task. However, the usage of the balanced OT/EOT solvers in these tasks lead to failures in certain cases, e.g., when the attributes of objects from source and target distributions are not balanced or the distributions are contaminated with outliers.

This problem could be alleviated by using our unbalanced EOT solver. Our method offers a parameter $\tau\in \mathbb{R}_+$ which controls the level of unbalancedness of the learned plans. For small $\tau$, our solver truly learns the UEOT plans, while for large $\tau$, its results are close to those of balanced solvers. In contrast to other existing UEOT solvers, our ULightOT solver is fast, light and theoretically justified.

An example: unpaired Adult$\rightarrow$Young translation by our ULightOT solver ($\tau=100$) applied in the latent space of ALAE for 1024x1024 FFHQ images. Our ULightOT solver converges in several minutes on 4 cpu cores and is robust to the imbalance of gender classes in source and target distributions.

Repository structure

The implementation is CPU-based and tested with torch==2.0.0.

The experiments are issued in the form of pretty self-explanatory jupyter notebooks (notebooks/). For convenience, the majority of the evaluation output is preserved. Auxilary source code is moved to .py modules (src/). The code for launching ALAE model is located in ALAE/ folder. Checkpoints for classifies used to evaluate the performance of our solver in the unpaired image translation problem are moved to checkpoints/ folder.

• notebooks/ULightOT_gaussians.ipynb - experiment with Gaussian mixtures in 2D (demonstrates the ability of our solver to deal with class imbalance issue);
• notebooks/ULightOT_gaussians_outliers.ipynb - experiment with Gaussian mixtures with added outliers in 2D (demonstrates the ability of our solver to simultaneously hadle the class imbalance and outliers issues);
• notebooks/ULightOT_alae.ipynb - experiment with unpaired image-to-image translation using the latent space of ALAE autoencoder.

Credits

• Weights & Biases developer tools for machine learning;
• pytorch-fid repo to compute FID score;
• Inkscape for the awesome editor for vector graphics.