[論文レビュー] A Fast Proximal Point Method for Computing Exact Wasserstein Distance

Yujia Xie, Xiangfeng Wang+2 more|arXiv (Cornell University)|Feb 12, 2018

Geometric Analysis and Curvature Flows参考文献 39被引用数 38

ひとこと要約

この論文は、確率単体への射影を用いて近似的にプロキシマル作用素を解くことで、正確なウォッサーシュタイン距離を高速に計算するIPOTを提案する。1つの内部反復で線形収束を達成し、Sinkhornと同等の効率性を発揮し、生成モデルにおける正則化由来の収縮を回避し、エンタロピー正則化手法よりもシャープなウォッサーシュタイン重心を生成する。

ABSTRACT

Wasserstein distance plays increasingly important roles in machine learning, stochastic programming and image processing. Major efforts have been under way to address its high computational complexity, some leading to approximate or regularized variations such as Sinkhorn distance. However, as we will demonstrate, regularized variations with large regularization parameter will degradate the performance in several important machine learning applications, and small regularization parameter will fail due to numerical stability issues with existing algorithms. We address this challenge by developing an Inexact Proximal point method for exact Optimal Transport problem (IPOT) with the proximal operator approximately evaluated at each iteration using projections to the probability simplex. The algorithm (a) converges to exact Wasserstein distance with theoretical guarantee and robust regularization parameter selection, (b) alleviates numerical stability issue, (c) has similar computational complexity to Sinkhorn, and (d) avoids the shrinking problem when apply to generative models. Furthermore, a new algorithm is proposed based on IPOT to obtain sharper Wasserstein barycenter.

研究の動機と目的

• Sinkhornのような正則化最適輸送手法の限界、特に中程度の正則化では性能低下が生じるか、小さな正則化では数値的不安定性を示す問題を解決すること。
• 正則化パラメータの微調整を必要とせず、正確なウォッサーシュタイン距離を堅牢に計算する手法を開発すること。
• 生成モデルやウォッサーシュタイン重心計算などの応用分野において、最適輸送計画の正確かつ安定した計算を可能にすること。
• 正確性を保ちつつ、Sinkhornと同等の計算効率を達成し、生成モデルにおける収縮問題を回避すること。

提案手法

• Bregman散発に基づく一般化プロキシマル点反復を用いた最適輸送のための不正確プロキシマル点法（IPOT）を提案する。
• 各反復で確率単体への射影を用いてプロキシマル作用素を近似的に計算し、効率的かつ数値的に安定した更新を実現する。
• 2段階の反復スキームを採用：まず現在の輸送計画を用いて双対変数を更新し、次に制約を満たすように確率単体に再投影する。
• 理論的に、1つの内部反復でも正確な最適輸送解への線形収束が保証される。
• 複数の分布に対してプロキシマル点フレームワークを適用することで、ウォッサーシュタイン重心の計算へとIPOTを拡張する。
• 固定された小さな正則化パrameter β ≈ 0.001 を用い、これにより効果的な ε ≈ 2×10⁻⁵ が得られ、数値的不安定性を避ける高精度な計算が可能になる。

実験結果

リサーチクエスチョン

• RQ1 エンタロピー正則化に依存せずに、正確なウォッサーシュタイン距離を効率的に計算できるプロキシマル点法を設計できるか？
• RQ2 単体への射影によるプロキシマル作用素の不正確評価でも、依然として正確な解への収束が保証されるか？
• RQ3 正則化による収縮問題を回避できるため、IPOTは生成モデルにおいてSinkhornに基づく手法を上回る性能を発揮できるか？
• RQ4 シャープネスと正確性の観点から、IPOTはウォッサーシュタイン重心計算において最先端の手法と比較してどうか？
• RQ5 IPOTの計算複雑度はSinkhornと同等であり、正確性とパrameter選択へのロバスト性を維持できるか？

主な発見

• IPOTは、1つの内部反復でも正確なウォッサーシュタイン距離への線形収束を示し、理論的境界をはるかに超える強力な経験的性能を発揮する。
• 生成モデルにおける収縮問題を回避する：中程度の ε を用いたSinkhornではMNISTのすべての10桁の数字をカバーできないが、IPOTは生成サンプルで全10桁を回復する。
• エンタロピー正則化によるぼやけたアーチファクトを引き起こすSinkhornに基づく手法と比較して、IPOTははるかにシャープなウォッサーシュタイン重心を生成する。
• IPOTの計算複雑度はSinkhornとほとんど区別できない。反復ごとの実行時間とメモリ使用量が類似している。
• IPOTはパrameter選択に対してロバストである：異なる β 値に対しても高い性能を維持するが、Sinkhornは ε の微調整が必須である。
• 重心計算において、IPOTはCuturi、Solomon、およびBenamouらの手法と比較して優れた視覚的品質を達成し、より明確で正確な数字再構成が可能である。