A Truncated Newton Method for Optimal Transport
This provides a more efficient solver for optimal transport problems, which is incremental but addresses practical bottlenecks in computational speed and scalability for applications in machine learning and data science.
The paper tackles the challenge of developing a fast and scalable optimal transport solver by introducing a truncated Newton algorithm for entropic-regularized optimal transport, achieving high precision orders of magnitude faster than existing alternatives on 24 problem sets and scaling to problems with up to 1 million points.
Developing a contemporary optimal transport (OT) solver requires navigating trade-offs among several critical requirements: GPU parallelization, scalability to high-dimensional problems, theoretical convergence guarantees, empirical performance in terms of precision versus runtime, and numerical stability in practice. With these challenges in mind, we introduce a specialized truncated Newton algorithm for entropic-regularized OT. In addition to proving that locally quadratic convergence is possible without assuming a Lipschitz Hessian, we provide strategies to maximally exploit the high rate of local convergence in practice. Our GPU-parallel algorithm exhibits exceptionally favorable runtime performance, achieving high precision orders of magnitude faster than many existing alternatives. This is evidenced by wall-clock time experiments on 24 problem sets (12 datasets $\times$ 2 cost functions). The scalability of the algorithm is showcased on an extremely large OT problem with $n \approx 10^6$, solved approximately under weak entopric regularization.