[论文解读] Improved Complexity Bounds in Wasserstein Barycenter Problem

Darina Dvinskikh, Daniil Tiapkin|arXiv (Cornell University)|Oct 9, 2020

Geometric Analysis and Curvature Flows参考文献 22被引用 6

一句话总结

本文提出了两种计算 m 个大小为 n 的离散概率测度的 Wasserstein barycenter 的新算法，实现了更优的收敛速率。第一种算法采用带有定制范数的镜像逼近法，实现 $ϪO(mn^2ackslash sqrt{n}/\varepsilon)$ 的复杂度，且在高精度下无数值不稳定性；第二种算法利用面积凸性与对偶外推法，达到目前已知最佳的 $ϪO(mn^2/\varepsilon)$ 收敛速率，在理论和实践中均优于现有的熵正则化方法，尤其在高精度场景下表现更优。

ABSTRACT

In this paper, we focus on computational aspects of the Wasserstein barycenter problem. We propose two algorithms to compute Wasserstein barycenters of $m$ discrete measures of size $n$ with accuracy $\e$. The first algorithm, based on mirror prox with a specific norm, meets the complexity of celebrated accelerated iterative Bregman projections (IBP), namely $\widetilde O(mn^2\sqrt n/\e)$, however, with no limitations in contrast to the (accelerated) IBP, which is numerically unstable under small regularization parameter. The second algorithm, based on area-convexity and dual extrapolation, improves the previously best-known convergence rates for the Wasserstein barycenter problem enjoying $\widetilde O(mn^2/\e)$ complexity.

研究动机与目标

• 解决熵正则化方法（如 IBP）在需要高精度解时出现的数值不稳定性问题，以应对 Wasserstein barycenter 问题。
• 开发一种方法，其收敛速率与加速版 IBP 相当或更优，且不依赖正则化，从而避免在小正则化参数下性能下降的问题。
• 实现优于现有方法的理论复杂度界限，尤其适用于高精度应用场景。
• 为迭代 Bregman 投影（IBP）提供一种数值稳定的替代方案，确保在小正则化参数下仍保持性能。

提出的方法

• 第一种算法采用针对 Wasserstein barycenter 问题量身定制的近端函数的镜像逼近法，实现无需正则化的稳定收敛。
• 第二种算法将 barycenter 问题重新表述为鞍点问题，并在面积凸性条件下应用对偶外推法，从而实现更快的收敛速度。
• 两种方法均避免使用熵正则化，从而消除了小正则化参数下带来的数值不稳定性。
• 通过先进优化技术（包括面积凸性和 Nesterov 类外推法）对算法进行分析，推导出更优的收敛速率。
• 利用基代价矩阵的结构和测度数量，推导出理论复杂度界限，实现对 $n$ 和 $\varepsilon$ 更优的依赖关系。
• 通过在 MNIST、notMNIST 和高斯分布上的实验，验证了算法收敛于真实 barycenter 的性能。

实验结果

研究问题

• RQ1 我们能否设计一种 Wasserstein barycenter 算法，在不依赖熵正则化的情况下，仍能保持高精度下的稳定性和准确性？
• RQ2 在不使用正则化的情况下，是否可能实现优于 $\widetilde{O}(mn^2\sqrt{n}/\varepsilon)$ 的 barycenter 问题收敛速率？
• RQ3 所提出的结合面积凸性的对偶外推法与现有基于 IBP 的方法相比，在收敛速度和数值稳定性方面表现如何？
• RQ4 与使用小正则化参数的 IBP 相比，基于镜像逼近的算法在实际性能上表现如何？

主要发现

• 基于镜像逼近的算法实现了 $\widetilde{O}(mn^2\sqrt{n}/\varepsilon)$ 的复杂度，与加速版 IBP 的最优收敛速率相当，但高精度下无数值不稳定性。
• 对偶外推算法实现了更优的复杂度界限 $\widetilde{O}(mn^2/\varepsilon)$，优于所有先前的 Wasserstein barycenter 方法。
• 在 MNIST 和 notMNIST 上的实验表明，所提算法在小正则化参数下产生的 barycenter 比 IBP 更精确，而 IBP 在此时已出现不稳定。
• 在高斯测度上，两种所提算法均比正则化 IBP 更快、更准确地收敛到真实 barycenter，尤其在高精度下优势明显。
• 尽管对偶外推法的理论复杂度更优，但镜像逼近算法在早期迭代中表现出更优的实际收敛行为，因其子问题代价更低。
• 结果表明，如 IBP 这类熵正则化方法在正则化参数 $\gamma$ 过小时会变得数值不稳定，而所提方法则保持鲁棒性。