[论文解读] A Riemannian Block Coordinate Descent Method for Computing the Projection Robust Wasserstein Distance

Minhui Huang, Shiqian Ma+1 more|arXiv (Cornell University)|Dec 9, 2020

Stochastic Gradient Optimization Techniques参考文献 29被引用 8

一句话总结

该论文提出了一种黎曼块坐标下降（RBCD）方法，以高效计算投影鲁棒Wasserstein（PRW）距离，将定义在Stiefel流形上的非凸极大极小问题重新表述为更易处理的形式。该方法实现了 $O(\epsilon^{-3})$ 的迭代复杂度，显著优于先前RGAS算法的 $O(\epsilon^{-12})$ 复杂度，并在大规模数据集（包括MNIST）上表现出更优的效率。

ABSTRACT

The Wasserstein distance has become increasingly important in machine learning and deep learning. Despite its popularity, the Wasserstein distance is hard to approximate because of the curse of dimensionality. A recently proposed approach to alleviate the curse of dimensionality is to project the sampled data from the high dimensional probability distribution onto a lower-dimensional subspace, and then compute the Wasserstein distance between the projected data. However, this approach requires to solve a max-min problem over the Stiefel manifold, which is very challenging in practice. The only existing work that solves this problem directly is the RGAS (Riemannian Gradient Ascent with Sinkhorn Iteration) algorithm, which requires to solve an entropy-regularized optimal transport problem in each iteration, and thus can be costly for large-scale problems. In this paper, we propose a Riemannian block coordinate descent (RBCD) method to solve this problem, which is based on a novel reformulation of the regularized max-min problem over the Stiefel manifold. We show that the complexity of arithmetic operations for RBCD to obtain an $ε$-stationary point is $O(ε^{-3})$. This significantly improves the corresponding complexity of RGAS, which is $O(ε^{-12})$. Moreover, our RBCD has very low per-iteration complexity, and hence is suitable for large-scale problems. Numerical results on both synthetic and real datasets demonstrate that our method is more efficient than existing methods, especially when the number of sampled data is very large.

研究动机与目标

• 解决由于在Stiefel流形上采用非凸极大极小公式化而带来的计算投影鲁棒Wasserstein（PRW）距离的挑战。
• 克服现有方法（如RGAS）在每次迭代中需求解熵正则化OT问题所带来的高昂计算成本。
• 通过降低算术复杂度和每次迭代的计算成本，开发一种更高效的大型最优传输问题算法。
• 为所提出的RBCD方法在黎曼流形上的收敛性提供理论保证。
• 在合成数据集和真实世界数据集上，通过运行时间和可扩展性指标，证明RBCD在性能上优于RGAS及其他基线方法。

提出的方法

• 将PRW计算的正则化极大极小问题重新表述为一种新颖形式，使其适用于在Stiefel流形上进行块坐标下降。
• 应用黎曼块坐标下降（RBCD）方法，交替优化流形 $\mathrm{St}(d,k)$ 上的传输计划 $\pi$ 和投影矩阵 $U$。
• 利用黎曼优化技术在更新过程中处理流形约束 $U^\top U = I_k$。
• 提出一种自适应变体RABCD，采用线搜索选择步长，以进一步提升收敛速度。
• 利用目标函数 $f(\pi, U) = \sum_{i,j} \pi_{ij} \|U^\top x_i - U^\top y_j\|^2$ 的结构，实现子问题的高效求解。
• 集成黎曼共轭梯度或信赖域方法求解子问题，实现低每次迭代成本。

实验结果

研究问题

• RQ1 黎曼块坐标下降方法是否能在计算PRW距离方面，实现优于RGAS等现有方法的迭代复杂度？
• RQ2 所提出的RBCD方法是否在保持精度的同时，显著降低大规模数据集上的计算成本？
• RQ3 在不同数据维度和样本规模下，RBCD算法与RGAS相比在收敛速度和可扩展性方面表现如何？
• RQ4 RBCD在Stiefel流形上求解非凸极大极小问题时，达到 $\epsilon$-驻点的理论迭代复杂度是多少？
• RQ5 所提出的方法是否能有效应用于真实世界机器学习任务，如使用MNIST进行数字分类？

主要发现

• RBCD算法在寻找 $\epsilon$-驻点时达到 $O(\epsilon^{-3})$ 的迭代复杂度，显著优于RGAS的 $O(\epsilon^{-12})$ 复杂度。
• 在合成数据集和真实数据集（包括MNIST）上的数值实验表明，RBCD在样本数量增加时始终比RGAS更快。
• 在MNIST数据集上，RBCD在所有数字对比较中运行时间最多减少了30%，且PRW距离值保持一致。
• RBCD与RGAS计算出的PRW距离在所有数字对之间高度一致，验证了所提方法的可靠性。
• 自适应变体RABCD进一步提升了收敛速度，且RBCD与RABCD在不同参数设置下均表现出稳健性能。
• 该方法对参数调优具有鲁棒性，但最优性能仍需仔细选择步长和正则化参数。