This is the repository for the paper "Low-rank optimal transport through factor relaxation with latent coupling" which computes an LC-factorization as an interpretable means of decomposing a coupling matrix for general OT costs (W, GW, FGW) and relaxations (balanced, semi-relaxed, and unbalanced).
- `FRLC.py` contains `FRLC_opt` and `FRLC_LR_opt` to run the FRLC algorithm and return a decomposed plan
$(\mathbf{Q}, \mathbf{R}, \mathbf{T})$ - `FRLC_demo.ipynb is a demonstration notebook which shows users how to run the method, with plotting code in `plotting.py.
- `objective_grad.py contain helper functions for the gradients of `FRLC.py`
- `util.py contains helper functions for `FRLC.py`, as well as `low_rank_distance_factorization` which can be used to compute an approximate factorization of the distance matrix, if it is too large to be computed or stored.
If you have any questions about how to run this method, encounter bugs, or have any feedback on it, feel free to send an email to ph3641@princeton.edu or open a GitHub discussion. If this algorithm has been helpful for your work, or has been relevant to your paper, also feel free to cite this work!
@inproceedings{
halmos2024lowrank,
title={Low-Rank Optimal Transport through Factor Relaxation with Latent Coupling},
author={Peter Halmos and Xinhao Liu and Julian Gold and Benjamin Raphael},
booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
year={2024},
url={https://openreview.net/forum?id=hGgkdFF2hR}
}
Development of this repository is ongoing. The original files and experiments from the submission are also available on OpenReview for reproducibility.