[论文解读] Importance Sparsification for Sinkhorn Algorithm

Mengyu Li, Jun Yu+2 more|arXiv (Cornell University)|Jun 11, 2023

Bayesian Modeling and Causal Inference被引用 5

一句话总结

本文提出Spar-Sink，一种新颖的重要性稀疏化方法，通过利用自然上界构造采样概率，对熵正则化最优传输（OT）和非平衡OT（UOT）的Sinkhorn算法进行加速。该方法将每轮迭代的计算成本从$O(n^2)$降低至$\tildewide{O}(n)$，在温和条件下实现一致估计器，并在合成数据和真实世界的心动超声图像数据（包括心脏周期分析与生成建模）上，展现出优于现有方法的速度与精度表现。

ABSTRACT

Sinkhorn algorithm has been used pervasively to approximate the solution to optimal transport (OT) and unbalanced optimal transport (UOT) problems. However, its practical application is limited due to the high computational complexity. To alleviate the computational burden, we propose a novel importance sparsification method, called Spar-Sink, to efficiently approximate entropy-regularized OT and UOT solutions. Specifically, our method employs natural upper bounds for unknown optimal transport plans to establish effective sampling probabilities, and constructs a sparse kernel matrix to accelerate Sinkhorn iterations, reducing the computational cost of each iteration from $O(n^2)$ to $\widetilde{O}(n)$ for a sample of size $n$. Theoretically, we show the proposed estimators for the regularized OT and UOT problems are consistent under mild regularity conditions. Experiments on various synthetic data demonstrate Spar-Sink outperforms mainstream competitors in terms of both estimation error and speed. A real-world echocardiogram data analysis shows Spar-Sink can effectively estimate and visualize cardiac cycles, from which one can identify heart failure and arrhythmia. To evaluate the numerical accuracy of cardiac cycle prediction, we consider the task of predicting the end-systole time point using the end-diastole one. Results show Spar-Sink performs as well as the classical Sinkhorn algorithm, requiring significantly less computational time.

研究动机与目标

• 为解决大规模最优传输（OT）与非平衡OT（UOT）问题中Sinkhorn算法的高计算成本问题。
• 开发一种稀疏近似方法，在显著降低计算复杂度的同时保持估计精度。
• 实现对真实世界应用（如心动超声图像分析与生成建模）中OT与UOT的高效且精确计算。
• 在温和正则性条件下，建立所提方法的理论一致性。
• 在合成与真实数据集上，实现比现有最先进方法更优的速度与估计误差表现。

提出的方法

• 提出Spar-Sink方法，利用未知最优传输计划的自然上界推导出有效的采样概率，用于核矩阵的稀疏化。
• 基于这些重要性概率进行逐元素采样，构建稀疏核矩阵，将每轮Sinkhorn迭代的计算成本从$O(n^2)$降低至$\tildewide{O}(n)$。
• 在Sinkhorn算法中应用稀疏化核矩阵，高效近似熵正则化OT与UOT的解。
• 在生成建模中使用子采样参数$s = 10s_0(n)$，在色彩迁移中使用$s = 8s_0(n)$，以平衡精度与速度。
• 采用鲁棒变体Robust-NysSink用于OT框架中的对比实验，并将Spar-Sink集成到一种新型自编码器变体——Spar-Sink自编码器（SSAE）中。
• 理论分析证明，在温和正则性条件下，所提估计器对正则化OT与UOT均具有一致性。

实验结果

研究问题

• RQ1 通过自然上界实现的重要性稀疏化，能否在不牺牲估计精度的前提下显著降低Sinkhorn算法的计算成本？
• RQ2 在OT与UOT问题中，Spar-Sink相较于Nys-Sink与Robust-NysSink等现有变体在速度与精度方面表现如何？
• RQ3 Spar-Sink能否有效加速心动超声图像数据中的心脏周期分析，实现从心室舒张末期准确预测收缩末期？
• RQ4 将Spar-Sink集成到生成模型（SSAE）中，是否能保持或提升相比原始Sinkhorn自编码器（SAE）的性能？
• RQ5 在温和正则性条件下，Spar-Sink估计器对正则化OT与UOT的理论一致性如何？

主要发现

• 在色彩迁移任务中，Spar-Sink将计算传输计划的CPU时间从60.45秒（Sinkhorn）降低至3.15秒，视觉效果近乎一致。
• 在心动超声图像数据上，Spar-Sink在显著更短时间内实现了与经典Sinkhorn相当的收缩末期时间预测精度。
• 在生成建模中，SSAE的FID得分（23.65 ± 0.06）低于SAE（24.72 ± 0.13），且训练时间缩短至一半（64.81秒 vs. 125.87秒）。
• 该方法实现了每轮迭代$\tildewide{O}(n)$的计算成本，相较于标准Sinkhorn的$O(n^2)$有显著提升。
• 合成数据实验验证了Spar-Sink在估计误差与计算速度方面均优于主流竞争方法。
• 理论分析证实，在温和正则性条件下，Spar-Sink对正则化OT与UOT的估计器具有一致性。