A Riemannian Approach to Low-Rank Optimal Transport

2026-06-10

AI Summary

This work addresses the limitations of conventional low-rank optimal transport methods, which rely on first-order gradient updates requiring careful hyperparameter tuning and neglect the intrinsic geometric structure of the underlying optimization problem, thereby compromising efficiency and stability. The authors propose the first unified Riemannian optimization framework that models both balanced and unbalanced low-rank couplings as smooth submanifolds within the positive orthant, equipped with a Fisher–Rao product metric. This formulation yields closed-form expressions for Riemannian projections, retractions, and Hessian-vector products, eliminating the need for inner iterative loops. The approach seamlessly accommodates linear optimal transport, Gromov–Wasserstein alignment, and their unbalanced variants, while providing a rank-sufficiency certificate for global optimality. Experiments demonstrate that the proposed first- and second-order solvers consistently outperform existing methods across problems of varying scales, with per-iteration complexity scaling only linearly in the data size.

Abstract

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.

Problem

• low-rank optimal transport
• optimization landscape curvature
• hyperparameter tuning
• computational scalability
• Riemannian geometry

Innovation

• Riemannian optimization
• low-rank optimal transport
• Fisher-Rao metric
• unbalanced optimal transport
• Gromov-Wasserstein