Technical Summary: cuRegOT

Problem Statement

Optimal Transport (OT) is a fundamental tool in modern machine learning for comparing probability measures, widely used in domain adaptation, generative modeling, and computer vision. However, solving the discrete OT problem at scale is computationally prohibitive due to the super-cubic complexity of standard linear programming solvers. While entropic regularization (Cuturi, 2013) transforms the problem into a smooth, unconstrained dual optimization problem solvable via the Sinkhorn algorithm, the latter often suffers from slow convergence on challenging problems.

Recent advances have introduced second-order and quasi-Newton methods (e.g., SNS, SSNS, and the Sparse-Plus-Low-Rank or SPLR algorithm) that offer superior convergence rates by exploiting the structure of the Hessian matrix. However, these methods face a critical bottleneck when deployed on GPUs: the requirement to solve large sparse linear systems. Specifically, the symbolic analysis and reordering steps necessary for sparse Cholesky decomposition are inherently sequential, difficult to parallelize on SIMT architectures, and typically executed on the CPU. This creates a synchronization bottleneck where the GPU remains idle, severely hindering the efficiency of otherwise powerful quasi-Newton solvers.

Methodology

The paper presents cuRegOT, a high-performance GPU solver tailored for entropic-regularized OT, built upon the SPLR algorithm. To overcome the CPU-GPU synchronization bottleneck and memory bandwidth limitations, the authors propose three specific algorithmic and architectural optimizations:

1. Amortized Symbolic Analysis

Standard SPLR re-evaluates the sparsity pattern of the Hessian approximation ($H_\Omega$) at every iteration, triggering expensive CPU-based symbolic analysis. The authors observe that the sparsity pattern (determined by the top-$k$ entries of the transport plan $T$) evolves slowly across iterations.

• Strategy: The solver performs the CPU-based symbolic analysis and matrix reordering only once every $S$ iterations (e.g., $S=10$). For the intervening iterations, the established sparsity pattern is reused, updating only the numerical values of the non-zero elements.
• Theoretical Guarantee: The authors prove that reusing the sparsity pattern does not expand the eigenvalue range of the Hessian approximation. Consequently, the condition number remains bounded, ensuring that the convergence properties of the quasi-Newton method are preserved.

2. Collaborative CPU-GPU Computing (Asynchronous Candidate Generation)

Even with amortization, the GPU remains idle during the periodic CPU-based symbolic analysis phases.

• Strategy: The authors introduce an asynchronous mechanism where, while the CPU performs symbolic analysis, the GPU simultaneously executes lightweight Sinkhorn iterations on the current dual variables. This generates a "candidate" iterate ($x^s_k$).
• Selection: Once the CPU completes the analysis and the GPU finishes the primary quasi-Newton step (generating $x^q_k$), the solver evaluates the objective function and gradient norms for both candidates. The iterate yielding greater improvement is selected.
• Benefit: This approach masks CPU latency and provides additional exploration of the solution space without increasing wall-clock time. Theoretical analysis confirms that inserting these candidate selection steps preserves the global convergence of the original algorithm.

3. Fused CUDA Kernel for Gradient Evaluation

Naive gradient computation involves multiple passes over the cost matrix $M$ and transport plan $T$, leading to severe memory bandwidth bottlenecks on GPUs.

• Strategy: A fused CUDA kernel is designed to perform element-wise computation, row reduction, and column reduction in a single pass over the data.
• Implementation: The kernel utilizes a grid-stride loop and maximizes data reuse in registers and shared memory. It employs warp-level shuffle intrinsics for row reductions (zero shared memory usage) and shared memory buffers for column reductions, minimizing global memory writes and atomic contention.

Key Contributions

1. cuRegOT Library: A high-performance GPU solver for entropic-regularized OT that successfully adapts the SPLR quasi-Newton method to GPU hardware.
2. Algorithmic Innovations:
   • An amortized symbolic analysis strategy that reduces CPU overhead by reusing sparsity patterns.
   • An asynchronous candidate generation mechanism that utilizes GPU idle time during CPU analysis to accelerate convergence.
   • A fused gradient kernel that significantly reduces memory I/O overhead.
3. Theoretical Guarantees: Rigorous proofs demonstrating that the proposed modifications (pattern reuse and candidate selection) maintain the global convergence and linear convergence rate of the original SPLR algorithm.

Experimental Results

The authors evaluated cuRegOT against state-of-the-art GPU-based solvers, including POT (Sinkhorn, Greenkhorn), OTT-JAX (Sinkhorn, Anderson acceleration), and AccSinkhorn, across synthetic data and the CIFAR-10 dataset.

• Performance: cuRegOT consistently achieves significantly faster convergence than all baselines, particularly when high precision is required.
• Scalability: The performance advantage of cuRegOT increases with problem size (tested up to $6400 \times 4800$ and $5000 \times 5000$), validating its scalability for large-scale applications.
• Ablation Study: Experiments removing the amortized analysis ("-A") or the candidate generation ("-S") individually showed slower convergence, confirming that both components are essential to the solver's performance.
• Baselines: While Anderson acceleration provided speedups for standard Sinkhorn in some cases, it still faced slow convergence to high precision compared to cuRegOT. Greenkhorn was found to be too slow for practical comparison.

Significance

The paper addresses a critical gap in the deployment of advanced optimization methods for Optimal Transport on modern hardware. By bridging the efficiency gap between the theoretical convergence speed of quasi-Newton methods and the practical limitations of GPU execution (specifically regarding sparse linear solvers), cuRegOT enables the efficient solution of large-scale, high-precision OT problems. The work demonstrates that with careful algorithmic design—specifically amortizing serial CPU tasks and fusing memory-bound operations—second-order methods can be made viable and superior to first-order baselines on GPU accelerators.