发表机构
University at Albany; IBM T. J. Watson Research Center（奥尔巴尼大学; IBM T. J. Watson 研究中心）
机构由 AI 辅助整理，请以论文原文为准。AI 中文总结
提出FlashSinkhorn 2，通过粗质心与块稀疏细阶段耦合，在单GPU上高效求解大规模熵最优传输，实现亿级粒子问题的高精度求解。
AI 中文摘要
用于熵最优传输（EOT）的流式GPU求解器（如FlashSinkhorn）虽然避免了存储稠密核，但在每次Sinkhorn迭代中仍会评估所有$n\ imes m$个点对。我们提出FlashSinkhorn 2（FS2），一种针对低维点云上平方欧氏代价的求解器，通过耦合两个阶段，在单个GPU上将大规模离散EOT问题求解至指定的边际残差。粗阶段在单元质心上求解，将势提升至每个点，当采样边际检查拒绝该提升时，继续在质心上进行，从而替代大多数点级更新。随后，块稀疏细阶段消除粗更新无法处理的质心误差。其Morton序块支持筛选和融合张量核心执行，由块质量设定的阈值限制了每个被省略瓦片对每一行和每一列的贡献。在合成基准测试中，FS2在所有32个问题上达到目标残差，在10个问题上优于GeomLoss多尺度方法。在一张A100上，FS2解决了来自宇宙学$N$-体模拟的两个$1.34\ imes10^8$粒子测度之间的离散EOT，熵模糊等于平均粒子间距，全粒子边际残差低于0.01，耗时不到2.5小时。据我们所知，这是数小时内达到此精度的最大离散EOT问题。为可复现性，我们在该https URL发布了开源实现。
英文摘要
Streaming GPU solvers for entropic optimal transport (EOT), such as FlashSinkhorn, avoid storing the dense kernel but still evaluate all $n\times m$ point pairs in every Sinkhorn iteration. We present \textbf{FlashSinkhorn~2} (FS2), a solver for squared-Euclidean cost on low-dimensional point clouds that solves large discrete EOT problems to a prescribed marginal residual on a single GPU by coupling two stages. A coarse stage solves on cell centroids, lifts the potentials to every point and, when a sampled marginal check rejects the lift, continues on the centroids, replacing most point-level updates. A block-sparse fine stage then removes the centroid error that coarse updates cannot. Its Morton-ordered blocks support screening and fused tensor-core execution, and a threshold set by the block masses bounds each omitted tile's contribution to every row and column. On synthetic benchmarks, FS2 reaches the target residual on all 32 problems and GeomLoss multiscale on 10. On one A100, FS2 solves discrete EOT between two $1.34\times10^8$-particle measures from a cosmological $N$-body simulation, at an entropic blur equal to the mean interparticle distance, to an all-particle marginal residual below 0.01 in under 2.5 hours. To our knowledge, it is the largest discrete EOT problem solved to this accuracy within hours. For reproducibility, we release an open-source implementation at https://github.com/ot-triton-lab/flash-sinkhorn
