A Truncated Newton Method for Optimal Transport

2025-04-02

AI Summary

This work addresses the entropy-regularized optimal transport (OT) problem by proposing a truncated Newton solver tailored for large-scale, high-dimensional settings. Methodologically, it abandons the conventional Lipschitz Hessian assumption and establishes, for the first time under milder conditions, local quadratic convergence. It further introduces adaptive Hessian approximation and GPU-accelerated parallelization to jointly ensure theoretical rigor and computational efficiency. Innovatively integrating truncated Newton optimization with numerical stabilization mechanisms, the method consistently outperforms mainstream approaches—including Sinkhorn and L-BFGS—across 24 benchmark experiments. It successfully solves high-dimensional OT problems with $n approx 10^6$, achieving high accuracy, low runtime (several times faster than state-of-the-art solvers), and strong robustness. The proposed framework delivers a scalable, reliable, and efficient new paradigm for large-scale OT computation.

Technology Category

• Search and Optimization: Non-convex Optimization
• Machine Learning: Optimization
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Security and Privacy: Large-scale security measurements
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Experiences and lessons learnt from Web-based algorithms and system deployments

Abstract

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 $ imes$ 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.

Problem

• Develops truncated Newton method for optimal transport problems
• Addresses GPU parallelization and high-dimensional scalability challenges
• Ensures fast convergence and numerical stability in practice

Innovation

• Truncated Newton method for OT
• GPU-parallel algorithm with scalability
• Quadratic convergence without Lipschitz Hessian