📊 statistics

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

本文介绍了 cuRegOT，这是一种高性能的 GPU 加速求解器，用于熵正则化最优传输问题，它通过新颖的算法和架构优化克服了现有方法的局限性，在各类基准测试中实现了显著的加速效果并提供了严格的收敛性保证。

原作者： Yixuan Qiu

发布于 2026-05-12

📖 1 分钟阅读☕ 轻松阅读

原作者： Yixuan Qiu

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象你是一位物流经理，试图将一堆沙子从一个地点（“源”）移动到另一个地点（“目的地”）。你的目标是用尽可能少的燃料（成本）移动每一粒沙子。在数学和机器学习领域，这被称为最优传输（Optimal Transport）。它是一种强大的工具，用于比较不同的数据组，例如匹配照片中的面部或翻译语言。

然而，为海量数据解决这个“运沙”难题极其缓慢且计算成本高昂。这就像试图用一把铲子一粒一粒地移动整座山。

问题：旧铲子与新卡车

多年来，解决这一问题的标准方法是使用一种名为Sinkhorn的算法。将 Sinkhorn 想象为一个非常有条理、并行化的工人团队。他们可以同时工作（这对现代计算机芯片即 GPU 来说非常有利），但他们有点固执。在困难的情况下，他们需要很长时间才能完成任务，缓慢地来回 shuffle。

最近，数学家们开发了一种更聪明、更快的方法，称为SPLR（一种拟牛顿法）。这就像一辆高科技卡车，它了解地形并能走捷径。它能更快地收敛到解。但有一个陷阱： 这辆“卡车”有一个沉重、缓慢的引擎部件，只能在老式的 CPU（计算机的主脑）上运行，而无法在快速的 GPU（显卡）上运行。具体来说，它在移动之前需要执行复杂的“地图分析”（符号分析）。这种分析是一次一步完成的，导致强大的 GPU 坐以待毙，等待 CPU。

解决方案：cuRegOT

本文的作者构建了cuRegOT，这是一种新的软件工具，旨在让这辆“智能卡车”在现代 GPU 上全速运行。他们不仅仅是编写代码；他们利用三个巧妙的技巧重新设计了工作流程：

1. “复用地图”策略（摊销符号分析）

类比： 想象你在城市中导航。每走一步，旧方法都迫使你停下来，拿出地图，在再次移动之前从头重绘整个路线。这很慢。
cuRegOT 的修复： 作者意识到，“地图”（问题的结构）在每一步之间变化不大。因此，他们决定每 10 步绘制一次地图，并在接下来的 9 步中复用它，只更新具体的数字（如交通状况），同时保持道路布局不变。
结果： 这阻止了 CPU 成为瓶颈。GPU 得以持续工作，而无需等待 CPU 每次都重绘地图。

2. “支线任务”策略（协作式 CPU-GPU）

类比： 当 CPU 忙于绘制那张地图（这需要时间）时，GPU 只是坐在那里，无所事事。
cuRegOT 的修复： 作者建立了一个系统，当 CPU 绘制地图时，GPU 不会等待。相反，它在后台开始执行一种不同的、更简单的计算（使用较旧的 Sinkhorn 方法）。这就像一名工人在等待蓝图时，开始准备材料。
结果： 当 CPU 完成地图时，GPU 已经准备好了一份“备用计划”。系统随后快速检查哪个计划更好并选出胜者。这隐藏了等待时间，加快了整个过程。

3. “一体化”工具（融合内核）

类比： 想象一名工厂工人必须走到仓库去拿螺丝，走回桌子使用它，再走回去拿螺母，依此类推。这种来回走动（内存访问）浪费了大量时间。
cuRegOT 的修复： 他们构建了一个定制的“超级工具”（融合 CUDA 内核），一次性抓取螺丝、螺母和说明，完成工作，并在单次行程中存放结果。
结果： 这极大地减少了移动数据所花费的时间，而这通常是 GPU 上最大的速度杀手。

验证：它有效吗？

作者使用以下数据将 cuRegOT 与现有的最佳工具（如 POT 和 OTT-JAX 包中的工具）进行了测试：

合成数据： 具有不同形状和大小的虚构问题。
真实数据： 来自著名 CIFAR-10 数据集的图像（例如区分猫和狗的图片）。

发现：

速度： cuRegOT 始终比其他工具更快地解决问题。
精度： 当任务需要极高水平的准确性（使解“恰到好处”）时，这种优势变得更大。
规模： 随着问题变大（更多数据点），cuRegOT 进一步拉开差距，证明其可扩展性适用于大规模任务。
安全性： 他们在数学上证明了他们的捷径（复用地图和运行支线任务）不会破坏数学原理。解保证收敛到正确答案，就像原始较慢的方法一样。

总结

cuRegOT 是一个高性能引擎，用于解决复杂的数据匹配难题。它将一个聪明但 CPU 密集型算法进行优化，使其能够在强大的 GPU 上流畅运行，方法是复用工作、在 CPU 思考时保持 GPU 忙碌，并简化数据移动。其结果是，该工具解决大规模问题的速度显著快于当前的行业标准。

技术摘要：cuRegOT

问题陈述

最优传输（OT）是现代机器学习中用于比较概率测度的基础工具，广泛应用于域适应、生成建模和计算机视觉。然而，由于标准线性规划求解器具有超立方复杂度，在大规模上求解离散 OT 问题在计算上是不可行的。虽然熵正则化（Cuturi, 2013）将问题转化为可通过 Sinkhorn 算法求解的平滑无约束对偶优化问题，但后者在具有挑战性的问题上往往收敛缓慢。

最近的进展引入了二阶和拟牛顿方法（例如 SNS、SSNS 以及稀疏加低秩或 SPLR 算法），这些方法通过利用海森矩阵的结构提供了更优越的收敛速率。然而，当这些方法在 GPU 上部署时，面临一个关键瓶颈：求解大型稀疏线性系统的需求。具体而言，稀疏 Cholesky 分解所必需的符号分析和重排序步骤本质上是串行的，难以在 SIMT 架构上并行化，通常由 CPU 执行。这造成了同步瓶颈，导致 GPU 处于空闲状态，严重阻碍了原本强大的拟牛顿求解器的效率。

方法论

本文提出了cuRegOT，这是一个专为熵正则化 OT 设计的高性能 GPU 求解器，基于 SPLR 算法构建。为了克服 CPU-GPU 同步瓶颈和内存带宽限制，作者提出了三种具体的算法和架构优化：

1. 摊销符号分析

标准的 SPLR 在每次迭代中重新评估海森近似（ $H_\Omega$ ）的稀疏模式，从而触发昂贵的基于 CPU 的符号分析。作者观察到，稀疏模式（由传输计划 $T$ 的前 $k$ 个条目决定）在迭代过程中演变缓慢。

策略：求解器仅每隔 $S$ 次迭代（例如 $S=10$ ）执行一次基于 CPU 的符号分析和矩阵重排序。在中间的迭代中，重用已建立的稀疏模式，仅更新非零元素的数值。
理论保证：作者证明，重用稀疏模式不会扩大海森近似的特征值范围。因此，条件数保持有界，确保了拟牛顿方法的收敛性质得以保留。

2. 协同 CPU-GPU 计算（异步候选生成）

即使进行了摊销，在周期性的基于 CPU 的符号分析阶段，GPU 仍然处于空闲状态。

策略：作者引入了一种异步机制，当 CPU 执行符号分析时，GPU 同时针对当前对偶变量执行轻量级的 Sinkhorn 迭代。这生成了一个“候选”迭代点（ $x^s_k$ ）。
选择：一旦 CPU 完成分析且 GPU 完成主要的拟牛顿步骤（生成 $x^q_k$ ），求解器将评估两个候选点的目标函数和梯度范数。选择带来更大改进的迭代点。
优势：这种方法掩盖了 CPU 延迟，并在不增加挂钟时间的情况下提供了对解空间的额外探索。理论分析证实，插入这些候选选择步骤保留了原始算法的全局收敛性。

3. 用于梯度评估的融合 CUDA 内核

朴素的梯度计算涉及多次遍历代价矩阵 $M$ 和传输计划 $T$ ，导致 GPU 上出现严重的内存带宽瓶颈。

策略：设计了一个融合的 CUDA 内核，在单次数据遍历中执行逐元素计算、行归约和列归约。
实现：该内核利用网格步长循环（grid-stride loop），并在寄存器和共享内存中最大化数据重用。它采用线程束级（warp-level）洗牌指令进行行归约（零共享内存使用），并使用共享内存缓冲区进行列归约，从而最小化全局内存写入和原子竞争。

主要贡献

cuRegOT 库：一个高性能 GPU 求解器，用于熵正则化 OT，成功将 SPLR 拟牛顿方法适配到 GPU 硬件。
算法创新：
- 一种摊销符号分析策略，通过重用稀疏模式减少 CPU 开销。
- 一种异步候选生成机制，利用 CPU 分析期间的 GPU 空闲时间以加速收敛。
- 一个融合梯度内核，显著减少了内存 I/O 开销。
理论保证：严格的证明表明，所提出的修改（模式重用和候选选择）保持了原始 SPLR 算法的全局收敛性和线性收敛速率。

实验结果

作者在合成数据和 CIFAR-10 数据集上，将 cuRegOT 与最先进的基于 GPU 的求解器进行了评估，包括 POT（Sinkhorn, Greenkhorn）、OTT-JAX（Sinkhorn, Anderson 加速）和 AccSinkhorn。

性能：cuRegOT 始终比所有基线实现显著更快的收敛，特别是在需要高精度的情况下。
可扩展性：cuRegOT 的性能优势随着问题规模的增加而提升（测试规模高达 $6400 \times 4800$ 和 $5000 \times 5000$ ），验证了其在大规模应用中的可扩展性。
消融研究：分别移除摊销分析（"-A"）或候选生成（"-S"）的实验显示收敛变慢，证实了两个组件对于求解器性能都是必不可少的。
基线：虽然 Anderson 加速在某些情况下为标准 Sinkhorn 提供了加速，但与 cuRegOT 相比，它仍面临收敛到高精度缓慢的问题。Greenkhorn 被发现对于实际比较而言过于缓慢。

意义

本文解决了在现代硬件上部署最优传输高级优化方法的关键差距。通过弥合拟牛顿方法的理论收敛速度与 GPU 执行（特别是关于稀疏线性求解器）的实际局限性之间的效率差距，cuRegOT 使得高效求解大规模、高精度的 OT 问题成为可能。这项工作表明，通过精心设计的算法——特别是摊销串行 CPU 任务和融合内存受限操作——二阶方法可以在 GPU 加速器上变得可行，并优于一阶基线。