cuRegOT: A GPU-Accelerated Solver for Entropic-Regularized Optimal Transport

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.

Optimal transport is a mathematical way to compare two probability distributions, and it is widely used in machine learning tasks that involve matching or aligning data. However, solving optimal transport problems accurately can be slow, especially when the data are large and the computation must be repeated many times. Most GPU-friendly optimal transport solvers rely on the Sinkhorn algorithm, which is easy to parallelize but can require many iterations to reach high accuracy. Our work develops cuRegOT, a GPU solver for entropic-regularized optimal transport based on a faster quasi-Newton method. The main challenge is that quasi-Newton methods use sparse matrix computations, and some of these computations are difficult to parallelize and can leave the GPU waiting. cuRegOT addresses this by reusing sparsity information across iterations, running useful Sinkhorn updates on the GPU while the CPU performs unavoidable sparse-matrix analysis, and using a fused CUDA kernel to reduce memory traffic. These designs help the CPU and GPU work together more efficiently. In experiments, cuRegOT is faster than common GPU Sinkhorn baselines, especially for large problems and high-accuracy solutions. This provides a practical path towards making more accurate optimal transport computations usable in large-scale machine learning pipelines.