[論文レビュー] 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

ひとこと要約

本稿では、リーマン多様体上の非凸なmax-min問題をより取り扱いやすい形に再定式化することで、効率的にProjection Robust Wasserstein (PRW) 距離を計算するためのリーマン多様体ブロック座標降下法（RBCD）を提案する。この手法は、$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.

研究の動機と目的

• リーマン多様体上での非凸なmax-min定式化に起因する、Projection Robust Wasserstein (PRW) 距離の計算における計算課題に対処すること。
• 各反復でエントロピー正則化オペレーター理論問題を解く必要があるなど、RGASのような既存手法の高い反復コストを克服すること。
• 算術的複雑度と反復コストを低減することで、大規模な最適輸送問題に対してより効率的なアルゴリズムを開発すること。
• リーマン多様体上での提案されたRBCD手法の理論的収束保証を提供すること。
• 合成データおよび実世界のデータセットにおける実験的優位性を、実行時間とスケーラビリティの観点から、RGASおよび他のベースライン手法と比較して示すこと。

提案手法

• PRW計算のための正則化されたmax-min問題を、スティーフェル多様体上でのブロック座標降下に適した新しい形に再定式化する。
• リーマン多様体上でのブロック座標降下（RBCD）を用い、輸送計画$\pi$と射影行列$U$を、スティーフェル多様体$\mathrm{St}(d,k)$上で交互に最適化する。
• 更新の過程で制約$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 非凸なmax-min問題における$\epsilon$-停留点を達成するための、RBCDの理論的反復複雑度は何か？
• RQ5 提案手法は、MNISTを用いた数字分類などの実世界の機械学習タスクに効果的に応用可能か？

主な発見

• RBCDアルゴリズムは、$\epsilon$-停留点を求める反復複雑度として$O(\epsilon^{-3})$を達成し、RGASの$O(\epsilon^{-12})$ と比べて顕著な改善を示している。
• 合成データおよび実データセット（MNISTを含む）における数値実験から、RBCDはRGASよりも一貫して高速であることが示され、サンプル数が増えるほど顕著である。
• MNISTデータセットでは、RBCDは全数字ペア比較において、RGASと比較して実行時間を最大30％短縮したが、PRW距離の値は一貫していた。
• 全数字ペアにおいて、RBCDとRGASが算出したPRW距離は非常に一致しており、提案手法の信頼性を裏付けている。
• 適応的バージョンRABCDはさらに収束速度を向上させたが、RBCDおよびRABCDの両方が異なるパrameter設定においても頑健な性能を示した。
• 手法はパrameterチューニングに対して頑健であるが、最適な性能を得るためにはステップサイズおよび正則化パrameterの慎重な選定が必要である。