[论文解读] Interior-Point Methods Strike Back: Solving the Wasserstein Barycenter Problem

Dongdong Ge, Haoyue Wang+2 more|arXiv (Cornell University)|May 30, 2019

Markov Chains and Monte Carlo Methods被引用 13

一句话总结

该论文提出了一种新型内点法——基于矩阵的自适应交替内点法（MAAIPM），通过利用块矩阵结构来加速牛顿迭代，从而高效求解Wasserstein中位数问题。该方法在MNIST和Fashion-MNIST基准测试中，相比最先进的第一类方法和正则化方法，实现了高达10倍的速度提升，且在速度和可扩展性方面优于Gurobi和MOSEK，达到高精度解。

ABSTRACT

Computing the Wasserstein barycenter of a set of probability measures under the optimal transport metric can quickly become prohibitive for traditional second-order algorithms, such as interior-point methods, as the support size of the measures increases. In this paper, we overcome the difficulty by developing a new adapted interior-point method that fully exploits the problem's special matrix structure to reduce the iteration complexity and speed up the Newton procedure. Different from regularization approaches, our method achieves a well-balanced tradeoff between accuracy and speed. A numerical comparison on various distributions with existing algorithms exhibits the computational advantages of our approach. Moreover, we demonstrate the practicality of our algorithm on image benchmark problems including MNIST and Fashion-MNIST.

研究动机与目标

• 为解决传统二阶方法（如内点法，IPM）在大规模Wasserstein中位数问题上的计算不可行性。
• 通过开发一种非正则化、高精度的内点法，克服熵正则化方法固有的精度与速度权衡。
• 设计一种自适应算法，联合优化支撑点与传输方案，提升收敛性并避免陷入局部极小值。
• 利用内点法中正规方程的块矩阵结构，降低迭代复杂度并加速牛顿步长。
• 在真实图像数据集（如MNIST和Fashion-MNIST）上展示所提方法的实际可行性与可扩展性。

提出的方法

• MAAIPM采用预测-校正IPM框架，结合Mehrotra算法，并自适应选择步长与中心化参数。
• 提出两种基于块矩阵的加速算法，通过利用问题固有的低秩与块稀疏结构，高效求解牛顿方程。
• 对于自由支撑中位数问题，MAAIPM通过无约束二次规划交替更新支撑点，并应用“跳跃”启发式策略以逃离局部极小值。
• 采用两阶段策略：初期对支撑点进行粗粒度更新，随着收敛逼近，转为频繁的细粒度更新。
• 将中位数问题建模为具有O(Nm + ∑mi)个约束和O(m∑mi + m)个变量的线性规划问题，随后应用具有结构感知分解的IPM。
• 该算法保持二次收敛性，时间复杂度接近N和m′的线性关系，内存使用量保持在正则化方法的常数倍以内。

实验结果

研究问题

• RQ1 通过利用问题特定的矩阵结构，能否使内点法在大规模Wasserstein中位数问题上实现可扩展性？
• RQ2 非正则化IPM方法在解的精度与计算效率方面，是否优于熵正则化方法和第一类方法？
• RQ3 在自由支撑中位数计算中，采用带‘跳跃’启发式的自适应支撑点更新能否改善收敛性并避免陷入局部极小值？
• RQ4 在真实图像数据集上，MAAIPM与Gurobi、MOSEK、BADMM和IBP等最先进求解器相比，在解的质量与运行时间方面表现如何？
• RQ5 问题规模（N, m, mi）对所提方法的迭代次数与计算时间有何影响？

主要发现

• 在所有测试方法中，MAAIPM以第二短的计算时间达到最小的目标函数值，在速度和可扩展性方面均优于Gurobi和MOSEK。
• 在MNIST和Fashion-MNIST数据集上，MAAIPM在最短时间内计算出最清晰、最锐利的中位数图像，200张数字8图像的计算时间低于250秒。
• 当N=5000时，MAAIPM仅需3098.23秒，而MOSEK耗时超过20,000秒，展现出卓越的可扩展性。
• 该方法的迭代次数几乎不随问题规模变化，表明其时间复杂度接近N和m′的线性关系。
• 在自由支撑情形下，MAAIPM以第二快的速度达到最低目标函数值，而IBP在ε=0.01时虽最快但解的质量较低。
• MAAIPM的内存使用效率高，保持在IBP的常数倍以内，且在大规模问题上显著低于Gurobi和MOSEK。