🤖 AI

FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU

FlashSinkhorn 是一种面向 I/O 的 GPU 熵最优传输求解器，它利用 FlashAttention 风格的融合与分块技术大幅减少高带宽内存流量，在实现大规模点云任务可扩展优化的同时，相比最先进基线模型实现了高达 161 倍的加速。

原作者： Felix X. -F. Ye, Xingjie Li, An Yu, Ming-Ching Chang, Linsong Chu, Davis Wertheimer

发布于 2026-05-22

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Felix X. -F. Ye, Xingjie Li, An Yu, Ming-Ching Chang, Linsong Chu, Davis Wertheimer

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

想象一下，你试图匹配两群庞大的人群。一群人站在场地的一侧（“源”），另一群站在对面（“目标”）。你的目标是找出最高效的方式将每个人配对，使得所有人需要行走的总距离最小化。这是一个经典的数学问题，称为最优传输。

在现代机器学习中，我们通常会给这个匹配过程添加一点“模糊性”，以使数学处理更容易。这被称为熵最优传输。为了解决它，计算机使用一种称为Sinkhorn 迭代的方法，这就像一场“热土豆”游戏：计算机在两组人群之间不断来回传递信息，反复优化匹配，直到找到最佳解。

问题：交通堵塞

论文指出，虽然这种方法在小规模人群上表现良好，但当人群变得巨大（例如数万人）时，它会遭遇巨大的障碍。

将计算机的内存想象成一座城市：

HBM（高带宽内存）：这是城市的主干道。它容量巨大，能容纳大量数据，但访问速度较慢。
SRAM（片上内存）：这是位于计算机处理器内部的一个极小但超快的私人办公室。它速度极快，但容量非常小。

旧有的解决此类匹配问题的方法，就像一辆送货卡车，每次需要检查一对人时，都必须从主干道（HBM）开到办公室（SRAM）再返回。由于存在数百万种可能的配对，这辆卡车始终被困在主干道的交通堵塞中，不断来回搬运数据。计算机花费在等待数据上的时间，比实际进行数学计算的时间还要多。

解决方案：FlashSinkhorn

作者开发了一种新工具，名为FlashSinkhorn。他们意识到，这个匹配问题背后的数学形式，与Transformer（即支撑像你现在正在对话的 AI 聊天机器人这样的技术）所使用的数学完全相同。

在 Transformer 中，有一个巧妙的技巧叫做FlashAttention，它解决了类似的交通堵塞问题。FlashAttention 不再让卡车来回奔波，而是将整个“块”（一小批数据）加载到快速办公室中，在那里完成所有必要的计算，仅将最终结果写回主干道。

FlashSinkhorn采用了同样的“基于块”的策略，并将其应用于匹配问题：

不再需要完整地图：它不再试图将每一对可能连接的完整地图写入内存（那会太大而无法容纳），而是按需逐个计算小块连接。
“办公室”策略：它将当前批次的计算保留在快速、小型的办公室（SRAM）中。它直接在那里更新“匹配分数”，而无需将庞大的中间列表写回缓慢的主干道。
流式处理：它像传送带一样流式处理数据，在过程中进行计算并丢弃繁重的中间步骤，从而保持主干道畅通。

结果：速度与规模

论文在强大的 GPU（具体为 A100）上对此进行了测试。结果令人瞩目：

速度：与现有的最佳在线方法相比，初始计算速度最快提升了32 倍，完整流程（包括从错误中学习）的速度最快提升了161 倍。
内存：旧方法在尝试匹配 30,000 人的群体时会崩溃（内存不足），而 FlashSinkhorn 可以轻松处理50,000 人，因为它从未试图一次性存储整个地图。
实际应用：他们展示了它在真实任务中的有效性，例如比较大型数据集（如数千张图像）以及解决数据顺序被打乱的复杂回归问题。

核心结论

FlashSinkhorn就像将一辆困在交通堵塞中的送货卡车升级为一架高速无人机。它并不改变目的地（数学答案依然精确），但改变了数据移动的方式。通过将繁重的计算保留在计算机的快速“办公室”内，而仅将最终结果使用缓慢的“主干道”传输，它使得解决大规模匹配问题变得切实可行且快速，将原本需要数小时或导致计算机崩溃的任务，缩短为仅需数秒。

技术摘要：FlashSinkhorn：面向 GPU 的 IO 感知熵最优传输

问题陈述

熵最优传输（EOT）通过 Sinkhorn 迭代求解，是现代机器学习中用于比较概率测度的基础工具，应用范围从生成建模到分布对齐。然而，在 GPU 上将 EOT 扩展到大规模点云（例如 $n \approx 70,000$ ）仍然是一个持续存在的瓶颈。

现有的 GPU 求解器面临两个主要效率问题：

张量化实现：这些方法会实例化稠密的 $n \times m$ 代价矩阵，导致二次方的高带宽内存（HBM）流量（$O(nm) $次读/写）以及内存容量问题，随着$ n$ 的增长会引发内存溢出（OOM）错误。
在线后端：虽然这些方法通过即时评估交互（例如通过 KeOps 或 OTT-JAX）避免了存储稠密矩阵，但它们依赖于通用的分块 Map-Reduce 归约内核。这些内核融合度有限，无法利用张量核心的 GEMM 结构，导致计算利用率低下和运行时过长，尽管内存使用量有所减少。

核心系统瓶颈在于，当前方法要么受限于内存（由于过量的 HBM 流量），要么受限于计算但指令效率低下，无法扩展到下游流水线所需的重复大规模求解。

方法论

作者提出了 FlashSinkhorn，这是一种针对平方欧氏距离代价的熵最优传输（EOT）的 IO 感知 GPU 求解器。该方法将稳定化的对数域 Sinkhorn 更新与 Transformer 中的注意力机制进行直接类比，利用 FlashAttention 的原则来优化内存访问模式。

1. 注意力形式重构

核心理论洞察（命题 1）是，对于平方欧氏距离代价，稳定化的对数域 Sinkhorn 更新可以重写为有偏点积得分矩阵的逐行 LogSumExp（LSE）归约。

具体而言，更新 $f_i \leftarrow -\varepsilon \text{LSE}_j [(g_j - C_{ij})/\varepsilon + \log b_j]$ 等价于计算矩阵 $S_X = (QK^\top + \mathbf{1}(\hat{g} + \delta))/\varepsilon$ 的 LSE，其中 $Q=\sqrt{2}X$ ， $K=\sqrt{2}Y$ ，且 $\delta, \gamma$ 为预计算的 log 权重偏差。
这一重构将 Sinkhorn 迭代转化为与缩放点积注意力中使用的完全相同的归一化统计量，从而使得应用 IO 感知的流式技术成为可能。

2. 流式内核与分块

FlashSinkhorn 使用融合的 Triton 内核实现这一重构，这些内核采用两级内存模型（慢速 HBM 和快速片上 SRAM）：

无实例化： $n \times m$ 代价矩阵永远不会在 HBM 中实例化。相反， $Q$ 和 $K$ 的块被流式传输。
在线 LSE：内核在处理得分矩阵的块时，在片上（SRAM）维护运行的逐行最大值和 sum-exp 统计量。这使得能够以显著更少的 HBM 访问次数计算精确的 LSE 归约。
融合：整个更新过程（得分计算、偏差添加、LSE 归约和潜在更新）被融合到单个内核传递中。这消除了对 HBM 的中间写入，并减少了内核启动开销。
行驻留遍历：与通常在外部循环中遍历键/值的 FlashAttention 不同，FlashSinkhorn 采用行驻留遍历（将查询行块固定在 SRAM 中并流式传输所有键块），以更好地匹配 Sinkhorn 更新的逐行归约特性。

3. 微分与二阶优化

本文将这种 IO 感知策略扩展到微分和二阶方法：

传输矩阵应用：重心投影和梯度计算被表示为类注意力操作（加权平均），允许通过相同的流式内核进行计算，而无需实例化传输矩阵 $P^*$ 。
Hessian-向量积（HVP）：作者推导了一个流式 HVP 预言机，将 Hessian-向量积分解为一系列传输向量和传输矩阵乘积。这避免了实例化 $O(n^2 d^2)$ 的 Hessian 张量或 $O(nm) $的传输矩阵，将工作内存减少到$ O((n+m)d)$。

主要贡献

注意力形式 Sinkhorn 更新：作者将稳定化的 Sinkhorn 更新表示为有偏点积 LSE 归约，使得 EOT 能够进行精确的 FlashAttention 风格流式处理。
IO 感知 EOT 求解器：他们提供了用于流式 Sinkhorn 迭代的融合 Triton 内核，以及用于微分的无传输矩阵流式算子，大幅减少了 HBM IO 量。
可扩展的二阶优化：他们引入了流式 HVP，使得基于 EOT 的损失函数能够使用可扩展的 Newton-CG 和信赖域方法，而这些方法此前因内存限制在大规模下不可行。
大规模评估：论文在 A100 GPU 上的合成基准测试和下游任务（如最优传输数据集距离 OTDD 和打乱回归）中展示了显著的速度提升和内存效率。

实验结果

在 NVIDIA A100-80GB GPU 上的评估表明，FlashSinkhorn 相比最先进基线（使用 KeOps/张量化后端的 GeomLoss 和 OTT-JAX）表现出显著的性能提升：

速度：
- 前向传播：相比 KeOps 最高加速 32 倍，相比 OTT-JAX 加速 5.1 倍。
- 端到端（前向 + 后向）：在高维度（ $d=512$ ）下，相比 KeOps 最高加速 161 倍。
- HVP：根据维度不同加速 3–25 倍，使得 Newton-CG 在其他方法超时的规模下变得可行。
内存效率：
- FlashSinkhorn 保持 线性内存复杂度 $O((n+m)d)$ ，而张量化基线则按 $O(nm) $或$ O(n^2)$ 扩展。
- 它成功处理了 $n=50,000$ 且维度高达 $d=1024$ 的点云，而张量化方法在此规模下内存溢出，在线基线（KeOps）则超出时间限制。
下游任务：
- OTDD：FlashSinkhorn 将 OTDD 距离计算和梯度流扩展到 $n=60,000$ （MNIST $\leftrightarrow$ Fashion-MNIST），峰值内存低于 1 GB，而张量化基线在 $n > 20,000$ 时发生 OOM。
- 鞍点逃逸：在打乱线性回归中，流式 HVP 使得能够监测最小 Hessian 特征值以检测鞍点，允许混合 Adam-Newton 优化器逃逸鞍点并快速收敛。

意义与主张

本文主张 FlashSinkhorn 使得大规模 EOT 在内存受限 regime 下变得可行。通过将 Sinkhorn 更新重写以匹配 FlashAttention 的 IO 感知结构，作者弥合了最优传输与现代 Transformer 加速技术之间的差距。

其意义在于：

系统级优化：收益归因于内核级专业化（融合与分块），而非对 Sinkhorn 迭代本身的算法更改。
启用新工作负载：内存占用和运行时的减少使得 EOT 能够用于此前被认为计算上不可行的下游任务（如数据集距离和隐式微分）。
可复现性：作者发布了开源实现，以促进可扩展最优传输的进一步研究。

作者谦逊地指出，该方法目前仅适用于具有特定结构的代价（平方欧氏距离和归一化后的余弦距离）。没有这种结构的代价（例如带平方根的原始欧氏距离或学习到的神经代价）被确定为未来的工作。