🤖 AI Summary
This work addresses the computational challenges of large-scale entropy-regularized optimal transport, which is notoriously difficult to solve efficiently due to its high complexity. Existing GPU-based methods struggle to balance convergence speed and parallel efficiency. To overcome this limitation, we propose a high-performance GPU solver that synergistically combines the strengths of quasi-Newton methods and the Sinkhorn algorithm. Our approach introduces several key innovations: amortized sign analysis, asynchronous Sinkhorn iterations, and fused gradient kernels, complemented by sparse-plus-low-rank approximations, asynchronous computation, and optimized memory access patterns. While preserving theoretical convergence guarantees, the proposed method consistently outperforms state-of-the-art GPU solvers across multiple benchmark tasks, achieving substantial acceleration.
📝 Abstract
Optimal transport (OT) has emerged as a fundamental tool in modern machine learning, yet its computational cost remains a significant bottleneck for large-scale applications. While harnessing the massive parallelism of modern GPU hardware is critical for efficiency, the de facto standard Sinkhorn algorithm, despite its ease of parallelization, often suffers from slow convergence in challenging problems. More recently, the sparse-plus-low-rank quasi-Newton method offers a balance between convergence rate and per-iteration complexity; however, its efficiency on GPUs is severely hindered by the serial nature of sparse matrix symbolic analysis and irregular memory access patterns. To bridge this gap, we present cuRegOT, a high-performance GPU solver tailored for entropic-regularized OT. We introduce a suite of algorithmic and architectural optimizations, including an amortized symbolic analysis strategy to mitigate CPU bottlenecks, an asynchronous Sinkhorn iterates generation mechanism, and a fused kernel for bandwidth-efficient gradient evaluation. These strategies are backed by rigorous theoretical guarantees ensuring algorithmic convergence. Extensive numerical experiments demonstrate that cuRegOT achieves significant speedups over state-of-the-art GPU-based solvers across a variety of benchmark tasks.