[论文解读] A Fast Proximal Point Method for Computing Exact Wasserstein Distance

Yujia Xie, Xiangfeng Wang+2 more|arXiv (Cornell University)|Feb 12, 2018

Geometric Analysis and Curvature Flows参考文献 39被引用 38

一句话总结

该论文提出IPOT，一种通过概率单纯形投影近似求解近端算子来快速计算精确Wasserstein距离的快速近端点方法。其在单次内迭代下实现线性收敛，效率与Sinkhorn相当，避免了生成模型中正则化引起的收缩问题，并生成比基于熵正则化的方法更清晰的Wasserstein中位数。

ABSTRACT

Wasserstein distance plays increasingly important roles in machine learning, stochastic programming and image processing. Major efforts have been under way to address its high computational complexity, some leading to approximate or regularized variations such as Sinkhorn distance. However, as we will demonstrate, regularized variations with large regularization parameter will degradate the performance in several important machine learning applications, and small regularization parameter will fail due to numerical stability issues with existing algorithms. We address this challenge by developing an Inexact Proximal point method for exact Optimal Transport problem (IPOT) with the proximal operator approximately evaluated at each iteration using projections to the probability simplex. The algorithm (a) converges to exact Wasserstein distance with theoretical guarantee and robust regularization parameter selection, (b) alleviates numerical stability issue, (c) has similar computational complexity to Sinkhorn, and (d) avoids the shrinking problem when apply to generative models. Furthermore, a new algorithm is proposed based on IPOT to obtain sharper Wasserstein barycenter.

研究动机与目标

• 解决像Sinkhorn这样的正则化最优传输方法的局限性，这些方法在中等正则化下性能下降，或在小正则化下出现数值不稳定性。
• 开发一种稳健计算精确Wasserstein距离的方法，无需对正则化参数进行精细调优。
• 实现生成建模和Wasserstein中位数计算等应用中最优传输计划的准确且稳定计算。
• 在保持精确性的同时，实现与Sinkhorn相当的计算效率，并避免生成模型中的收缩问题。

提出的方法

• 基于Bregman散度提出一种用于最优传输的非精确近端点方法（IPOT），采用广义近端点迭代。
• 通过将投影映射到概率单纯形来近似每次迭代中的近端算子，从而实现高效且数值稳定的更新。
• 采用两步迭代方案：使用当前传输计划更新对偶变量，然后将结果投影回概率单纯形以保持约束。
• 该方法在理论上保证即使仅使用一次内迭代，也能线性收敛至精确最优传输解。
• 通过将近端点框架扩展至多个分布，将IPOT应用于计算Wasserstein中位数。
• 使用固定的小正则化参数 β ≈ 0.001，对应有效 ε ≈ 2×10⁻⁵，确保高精度且无数值不稳定性。

实验结果

研究问题

• RQ1 能否设计一种近端点方法，以高效方式计算精确Wasserstein距离，而无需依赖熵正则化？
• RQ2 通过单纯形投影对近端算子进行非精确评估，是否仍能保证收敛至精确解？
• RQ3 IPOT能否通过避免正则化引起的收缩问题，在生成建模中优于基于Sinkhorn的方法？
• RQ4 在计算Wasserstein中位数方面，IPOT与最先进方法相比，在清晰度和准确性上表现如何？
• RQ5 IPOT的计算复杂度是否与Sinkhorn相当，同时保持精确性和对参数选择的鲁棒性？

主要发现

• 即使仅使用一次内迭代，IPOT仍能线性收敛至精确Wasserstein距离，展现出超越理论边界的强大经验性能。
• 该方法避免了生成模型中的收缩问题：当Sinkhorn使用中等ε值时无法覆盖MNIST中的所有数字，而IPOT在生成样本中成功恢复了全部十个数字。
• IPOT生成的Wasserstein中位数显著比基于Sinkhorn的方法更清晰，后者因熵正则化而产生模糊伪影。
• IPOT的计算复杂度与Sinkhorn几乎无法区分，每轮迭代的运行时间和内存使用量相似。
• IPOT对参数选择具有鲁棒性：在不同β值下均保持高性能，而Sinkhorn则需要仔细调整ε。
• 在中位数计算中，IPOT在视觉质量上优于Cuturi、Solomon和Benamou等人方法，生成的数字重建更清晰、更准确。