A Riemannian Approach to Low-Rank Optimal Transport
About
Low-rank optimal transport (OT) mitigates the quadratic scaling of classical solvers, yet existing approaches rely heavily on first-order mirror-descent updates that require careful hyperparameter tuning and ignore the optimization landscape's curvature. To address these limitations, we propose a unified Riemannian geometric framework for low-rank OT, modeling balanced and unbalanced rank-$r$ positive factored couplings as novel smooth embedded submanifolds of the positive orthant. By equipping these manifolds with the Fisher-Rao product metric, we derive tractable formulations for Riemannian projectors, retractions, and Hessian-vector products. Our cost-agnostic framework seamlessly extends to linear OT, Gromov-Wasserstein (GW), fused GW, and their unbalanced counterparts. For balanced OT, our geometric ingredients are computed via efficient conjugate-gradient and iterative Bregman updates. For the unbalanced OT, our operations elegantly reduce to closed-form scalings, completely eliminating inner iterative loops. In both regimes, per-iteration complexity scales linearly with dataset size, and we provide a rank-sufficiency certificate for global optimality verification. Extensive experiments across a range of problem sizes demonstrate that our regularization-free first- and second-order solvers achieve faster convergence and superior performance over existing state-of-the-art low-rank OT solvers.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Balanced Optimal Transport | Separated Gaussians (N((1, 1), I2) -> N(0, 0.1I2), n=m=5k) f*=2.921 | Ratio f/f*1.008 | 14 | |
| Balanced Linear Optimal Transport | Gaussian-mixture n=m=10k R2 | f/f* Ratio1 | 12 | |
| Balanced Gromov-Wasserstein | Anisotropic Gaussians (m=10k, n=20k, r=5) | Cost7.51e+3 | 6 | |
| Low-rank Optimal Transport | Balanced OT (n=m=10k, f*=0.287) rank r=50 | Ratio f/f*1.066 | 6 | |
| Low-rank Optimal Transport | Balanced OT (n=m=10k, f*=0.287) rank r=100 | Transport Ratio (f/f*)1.049 | 6 | |
| Partial-overlap Fused Gromov-Wasserstein | 3→2 Gaussian clusters (m=300, n=200, r=5, τ=5, α=0.25) | Cost0.062 | 6 | |
| Partial-overlap Fused Gromov-Wasserstein | 3→2 Gaussian clusters (m=300, n=200, r=5, τ=5, α=0.5) | Cost0.096 | 6 | |
| Balanced Gromov-Wasserstein | Anisotropic Gaussians (m=10k, n=20k, r=20) | Cost7.51e+3 | 6 | |
| Balanced Optimal Transport | Random Gaussian point clouds (n=20k, r=5, d=3) | Cost3.17 | 5 | |
| Unbalanced Optimal Transport | Gaussian clusters n=5k, εout=0.05 | Transport Cost0.201 | 5 |