[论文解读] Fast Computation of Wasserstein Barycenters

Marco Cuturi, Randal Douc|arXiv (Cornell University)|Oct 16, 2013

Topological and Geometric Data Analysis参考文献 33被引用 462

一句话总结

该论文提出两种高效算法，通过使用熵正则化平滑最优传输问题，实现对经验概率测度的Wasserstein中位数的快速计算，借助矩阵缩放实现快速梯度计算。核心贡献是一种可扩展的方法，相较于精确次梯度方法，实现了数量级的加速，同时保持了高精度的中位数估计，在图像可视化与约束聚类任务中得到验证。

ABSTRACT

We present new algorithms to compute the mean of a set of empirical probability measures under the optimal transport metric. This mean, known as the Wasserstein barycenter, is the measure that minimizes the sum of its Wasserstein distances to each element in that set. We propose two original algorithms to compute Wasserstein barycenters that build upon the subgradient method. A direct implementation of these algorithms is, however, too costly because it would require the repeated resolution of large primal and dual optimal transport problems to compute subgradients. Extending the work of Cuturi (2013), we propose to smooth the Wasserstein distance used in the definition of Wasserstein barycenters with an entropic regularizer and recover in doing so a strictly convex objective whose gradients can be computed for a considerably cheaper computational cost using matrix scaling algorithms. We use these algorithms to visualize a large family of images and to solve a constrained clustering problem.

研究动机与目标

• 解决由于重复求解大规模最优传输问题，导致使用精确次梯度方法计算Wasserstein中位数在计算上不可行的问题。
• 通过熵平滑降低计算成本，使Wasserstein中位数在数据分析中具备实际可用性。
• 开发高效算法，在自由支持与均匀权重约束下计算中位数，适用于约束聚类等应用。
• 展示Wasserstein中位数在图像族可视化与平衡聚类问题求解中的实用性。

提出的方法

• 使用熵正则化平滑Wasserstein距离，将非光滑、非凸的次梯度问题转化为严格凸、可微的优化问题。
• 利用Sinkhorn-Knopp算法通过矩阵缩放高效计算平滑对偶最优传输问题的梯度。
• 提出两种基于次梯度的算法：一种用于固定支撑中位数，另一种用于自由支撑中位数，均利用平滑对偶目标函数。
• 仅通过矩阵-向量乘积计算梯度，与求解完整的原始-对偶最优传输问题相比，显著降低计算成本。
• 将算法应用于图像族与人口普查数据聚类，通过约束质心权重实现平衡分配。
• 利用平滑对偶公式化，即使在大规模数据集上也能实现快速收敛。

实验结果

研究问题

• RQ1 熵正则化是否可用于使大规模经验测度的Wasserstein中位数计算变得可行？
• RQ2 平滑对偶公式化如何实现基于次梯度的中位数优化中的高效梯度计算？
• RQ3 具有均匀权重的约束Wasserstein中位数是否相比标准k-means能提升聚类的平衡性？
• RQ4 在强制实施均匀权重约束时，聚类准确率与计算成本之间的权衡如何？
• RQ5 Wasserstein中位数是否能有效总结复杂图像族，同时保持几何结构？

主要发现

• 在57,647个空间位置与48个质心的约束中位数计算中，所提算法耗时仅12.5秒，而使用精确最优传输求解器则需约1小时。
• 无约束k-means算法在1.55秒内收敛，而约束均匀权重中位数计算耗时12.5秒，表明计算开销适中且可接受。
• 具有均匀权重的Wasserstein中位数产生更平衡的聚类结果，每个质心均捕获总质量的相等份额，而标准k-means可能使质量集中在少数质心上。
• 嵌套椭圆的可视化结果表明，2-Wasserstein中位数能保持结构特征，而欧氏、RKHS与Jeffrey质心方法无法生成有意义的均值。
• 平滑对偶公式化通过矩阵缩放实现梯度计算，将每次次梯度评估的成本从求解大规模传输问题降低为简单的矩阵-向量运算。
• 该方法成功推广至涉及多个Wasserstein距离的变分问题，如Wasserstein传播。