动态熵最优传输的认证并行时间 Sinkhorn 方法

论文信息

标题: Certified Parallel-in-Time Sinkhorn for Dynamic Entropic Optimal Transport

作者: Xinyang Wen

发布日期: 2026-07-27

PDF 链接: 下载 PDF

引言：动态最优传输与时间并行化的需求

最优传输（Optimal Transport, OT）是衡量概率分布之间差异的强大工具，近年来在生成模型、流匹配（Flow Matching）等领域得到了广泛应用。特别是基于熵正则化的 Sinkhorn 算法，因其高效和可微性，已经成为求解大规模 OT 问题的标准方法。然而，在许多实际应用中，我们需要反复求解一系列相互关联的 OT 问题——例如，在基于 OT 的流匹配中，模型必须不断为新的小批量数据构建最优传输耦合，以获得更 “直” 的传输路径。传统的做法是顺序处理每一个问题，即每一帧都依次执行完整的 Sinkhorn 迭代，并在每次更新后进行一次全局同步。这种方式存在两个明显的计算浪费：第一，每个问题独立执行矩阵-向量乘法和通信操作，无法利用问题之间的相似性进行批量化加速；第二，对于变化缓慢或相关性较强的连续帧，很多迭代的初始状态与最终解的差距并不大，可以预见其收敛行为，但现有方法缺乏一种机制来安全地利用这种 “可预见性” 从而节省检查开销。

论文 TemporalSinkhorn 正是针对这一痛点，提出了一种时间维度上并行化的执行框架。它不再将连续帧视为孤立的问题，而是将它们打包成一个流，通过共享的更新和智能的审计调度，在不牺牲输出准确性的前提下，大幅缩短了整体运行时间。

核心方法：安全的批处理与遗忘引导的调度

TemporalSinkhorn 的设计理念是将预测与验证解耦。它的工作流程可以概括为：预测窗口—候选验证—共享修复—无误退役。整个框架围绕三个关键组件展开：

1. 中心化的局部证书（Centered Local Certificate）

在多 GPU 行分片（row-sharded）的设置下，判断一个候选解是否满足边际约束通常需要全局通信来计算列边际误差。为了减少通信压力，本文提出了一种仅需局部标量计算和单次打包通信的证书机制。具体而言，对于每一个未来候选解 $v_j$ ，选择一个已得到验证的锚点 $v_0$ ，并利用各 GPU 的局部列贡献 $p_{r,j}$ 、目标分布 $b_j$ 以及按行质量份额 $\alpha_r$ 计算出上界：

B_j = B_0 + \sum_{r=1}^R \left\| (p_{r,j} - p_{r,0}) - \alpha_r (b_j - b_0) \right\|_1

该值直观地衡量了候选解相对于锚点的偏移。三角不等式保证了 $\|\sum_r p_{r,j} - b_j\|_1 \le B_j$ 。这样，只要本地计算出的 $B_j$ 不超过预设的容忍度 $\tau$ ，就可以安全地接受该候选解，完全无需额外的全局残差检查。证书的关键优势在于：它的安全性完全独立于预测方法和遗忘率估计，即使预测完全错误，也不会输出一个不合格的解。这种确定性接受保证了方法的输出质量和顺序流的原始保证完全一致。

2. 共享的打包修复流（Packed Repair Stream）

证书只能为少数 “容易” 的帧提供立即退役，对于大多数未能通过证书检测的候选解，TemporalSinkhorn 将它们全部放入一个活跃打包流中，共享后续的 Sinkhorn 更新。也就是说，所有还需要迭代的候选解不再各自独立执行更新，而是以矩阵-矩阵乘法的形式一同推进。这种打包操作将重复的 GEMV（矩阵-向量乘）和标量通信转化为了更高计算密度的 GEMM（矩阵-矩阵乘），并显著减少了集体通信的轮次。当某些候选解在中间的审计点达到精度要求时，它们会被压缩移除，剩余的活跃列继续共享更新。这种设计使得即使证书只能接受零长度的前缀（即所有候选都需要修复），整个批处理仍然能够通过摊销通信和计算开销带来加速。

3. 遗忘率指导的审计调度（Forgetting-Guided Milestone Placement）

在共享修复流中，审计（即精确检查残差）仍然是必要的，但不需要在每一次迭代后都进行。论文观察到 Sinkhorn 迭代的投射残差 $\rho_t(v)$ 在局部常常以近似几何速率衰减，其衰减因子 $q$ 可以通过近期观测进行估计。据此可以预测出剩余的迭代深度并放置一个稀疏的审计里程碑。但预测仅仅是经济信号，绝不被当作停止准则。安全网是一个永久性的几何审计网格，以及最终的实际残差核查和常规 Sinkhorn 后备。因此，错误的预测只会浪费计算（过早审计），而不会导致不正确的输出。这项工作将检查开销显著降低，尤其是在核矩阵发生变化的复杂模式下，节省的计算量远大于增加的审计成本。

实验亮点与速度提升

论文在多种设定下进行了全面评估，以验证方法在效率与正确性上的双重承诺。实验使用了 4 块 A100 GPU 和一块消费级 RTX 4060 笔记本 GPU，分别代表了数据中心和本地部署路径。

在 4×A100 的受控合成流实验中，采用遗忘引导的里程碑调度相比每次迭代都进行审计的基线，在五个统计显著的场景中实现了 1.15× 到 1.47× 的加速，且零容忍违反。与目前最强的顺序软 c-变换预热基线相比，时间并行执行在所有 30 次配对测试中全部胜出，加速比达到 1.42× 到 3.55×。特别值得注意的是，在流匹配小批量耦合流测试中，时间执行比顺序 carry 快 3.05× 到 3.63×，并且在支持尺寸扩展到 4096 时仍然保持了 2.59× 到 2.76× 的加速，没有出现任何边际约束的违反。

消费级 GPU 上的表现同样亮眼：在 RTX 4060 上，固定核算子利用矩阵-矩阵乘法的批处理优势，相比顺序执行获得了 4.315× 的几何平均加速，而内存增量仅为 3.25 MiB。这一结果表明，即使在没有多卡通信的环境下，单纯摊销核读取和启动延迟也能带来显著收益。

更深入的分析揭示了模式（固定核 vs. 变化核）对调度策略的影响。对于廉价的固定核 GEMM，较粗的算术间隔审计（例如每 10 次迭代检查一次）效果更好，因为它减少了多余的审计和计算；而对于昂贵的变化核分组更新，利用一个稀疏的预测里程碑能够提前终止许多不必要的迭代，从而击败最好的静态几何网格。据此，论文提出了一种模式感知的控制器：固定核采用算术间隔 $h=10$ ，变化核采用预测策略。在独立测试种子下，该组合策略相对于逐迭代审计取得了 1.436× 的加速，相对于通用静态网格取得了 1.069× 的加速，且全部 10 对比较均获胜。

实践应用建议与未来方向

TemporalSinkhorn 为动态最优传输的高效求解提供了切实可行的系统设计。对于涉及反复求解 OT 的实践者，具有以下指导意义：

流匹配与生成模型加速：如果你的训练管道中包含大量基于 Sinkhorn 的耦合计算（例如 Mini-batch OT 流匹配），直接将顺序求解器替换为时间并行执行器，不需要任何超参数调整，就能获得数倍的求解速度提升。目前虽然尚未与训练管道完全集成，但论文提供的流匹配小批量实验已经充分证明了其潜力。
硬件自适应部署：对于仅有单 GPU 的消费级环境，可以利用固定核算子的打包 GEMM，在几乎不增加显存的情况下大幅加速。对于多 GPU 数据中心环境，可以启用行分片和变化核支持，并通过模式感知的调度策略最大化收益。
调度策略的自动选择：论文揭示了一个关键事实：没有一种通用的静态审计策略在所有模式下都是最优的。实践中，可以通过一个小型成本模型或在线测量，动态选择算术间隔、几何网格或预测里程碑。更重要的是，永远不要为了性能牺牲正确的停止条件——后验残差检查和常规 Sinkhorn 后备是不可或缺的。
未来工作方向：目前的工作仍然是一个单节点系统机制，尚未与端到端的训练流程（包括耦合采样、参数更新等）深度集成。接下来的关键一步是将 TemporalSinkhorn 嵌入到真实的流匹配或元最优传输训练中，测量总训练时间和生成质量。此外，在多节点、网络受限的场景下，打包通信的收益可能更加突出，但也可能暴露出滞后和资源竞争问题，需要针对性的优化。在数值安全方面，完整的 TF32 和 NCCL 归约兼容性保证仍需完善。最后，将当前的简单模式感知控制器升级为可以处理渐变核变化或混合模式的连续自适应控制器，也是一个富有价值的方向。

总结

TemporalSinkhorn 成功地将并行时间处理的思想引入了动态最优传输问题，其核心贡献在于分离了性能预测与正确性验证：利用遗忘率的经济信号调度审计以节省开销，同时永远依赖确定性证书和实际残差检查来保证输出的绝对安全。该设计不仅带来了可观的加速，更保持了与顺序基准完全相同的容限制约。通过模式感知的复合策略，它分别在固定核和变化核环境下建立了新的速度标杆。这一工作为处理最优传输流提供了强大的系统抽象，也为流匹配等前沿应用的大规模化铺平了计算基础。随着端到端集成和多节点部署的推进，TemporalSinkhorn 有望成为动态 OT 计算的基石工具之一。