gpu加速的最优传输求解器,速度远超现有方法。
cuRegOT: A GPU-Accelerated Solver for Entropic-Regularized Optimal Transport

- 用并行算法优化稀疏矩阵分析与内存访问
- 在多个基准任务上实现显著提速
- 适合大规模机器学习中的高效计算
最优传输(OT)已成为现代机器学习中的基础工具,但其计算成本仍是大规模应用的主要瓶颈。尽管利用现代GPU的并行性对效率至关重要,但主流的Sinkhorn算法在复杂问题中收敛缓慢。最近的稀疏加低秩拟牛顿法虽在收敛速度与每轮复杂度间取得平衡,但在GPU上因稀疏矩阵符号分析的串行性及不规则内存访问而效率受限。为此,我们提出cuRegOT,一种专为熵正则化最优传输设计的高性能GPU求解器。通过引入摊销符号分析策略、异步生成Sinkhorn迭代机制以及融合核函数实现带宽高效的梯度计算,有效缓解了CPU瓶颈。上述优化均具备严格的理论保证,确保算法收敛。大量数值实验表明,cuRegOT在多种基准任务中显著优于现有GPU求解器。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。