Fast Computation of Optimal Transport via Entropy-Regularized Extragradient Methods
- Resource Type
- Working Paper
- Authors
- Li, Gen; Chen, Yanxi; Chi, Yuejie; Poor, H. Vincent; Chen, Yuxin
- Source
- Subject
- Computer Science - Machine Learning
Computer Science - Data Structures and Algorithms
Computer Science - Information Theory
Mathematics - Optimization and Control
Statistics - Machine Learning
- Language
Efficient computation of the optimal transport distance between two distributions serves as an algorithm subroutine that empowers various applications. This paper develops a scalable first-order optimization-based method that computes optimal transport to within $\varepsilon$ additive accuracy with runtime $\widetilde{O}( n^2/\varepsilon)$, where $n$ denotes the dimension of the probability distributions of interest. Our algorithm achieves the state-of-the-art computational guarantees among all first-order methods, while exhibiting favorable numerical performance compared to classical algorithms like Sinkhorn and Greenkhorn. Underlying our algorithm designs are two key elements: (a) converting the original problem into a bilinear minimax problem over probability distributions; (b) exploiting the extragradient idea -- in conjunction with entropy regularization and adaptive learning rates -- to accelerate convergence.