用CUDA实现快速稳定最优传输，小正则化下也高效
Fast Log-Domain Sinkhorn Optimal Transport with Warp-Level GPU Reductions
- 在对数域用线程包级归约与共享内存分块加速
- 8192×8192问题上比POT快12倍，仅用256MB显存
- 适合需要高精度、大规模最优传输的科研与工程应用
基于熵正则化最优传输的Sinkhorn算法已成为机器学习中的基础工具，但现有实现要么在小正则化参数下数值不稳，要么因深度学习框架开销过大。我们提出FastSinkhorn，一种轻量级原生CUDA实现，结合线程包级混洗归约与共享内存分块，在不牺牲数值稳定性的情况下实现高GPU利用率。求解器完全运行于对数域，可稳定处理小至ε=10⁻⁴的正则化参数，标准域方法在此时已失效。在n=m=8192的密集最优传输问题上，相比广泛使用的POT库提升12倍，比GPU加速的PyTorch基线快5.9倍，显存占用仅256MB。我们在图像颜色迁移、3D点云匹配及收敛性分析中验证了该求解器的有效性，表明经过精细数值处理的原生CUDA内核可为大规模最优传输计算提供实用高效的基石。
原文摘要 · Abstract (English)
Entropic regularized optimal transport (OT) via the Sinkhorn algorithm has become a fundamental tool in machine learning, yet existing implementations either suffer from numerical instability for small regularization parameters or incur significant overhead from deep learning frameworks. We present FastSinkhorn, a lightweight, native CUDA implementation of the log-domain Sinkhorn algorithm that combines warp-level shuffle reductions with shared-memory tiling to achieve high GPU utilization without sacrificing numerical stability. Our solver operates entirely in the log-domain, enabling robust computation for regularization parameters as small as epsilon = 10^{-4} where standard-domain methods fail. On dense OT problems with n = m = 8192, our implementation achieves 12x speedup over the widely-used POT library and 5.9x speedup over GPU-accelerated PyTorch baselines, while consuming only 256 MB of GPU memory. We validate our solver on image color transfer, 3D point cloud matching, and convergence analysis, demonstrating that native CUDA kernels with careful numerical treatment provide a practical and efficient foundation for large-scale optimal transport computation.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。