[論文レビュー] Fast Computation of Wasserstein Barycenters

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

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

ひとこと要約

本稿では、最適輸送問題にエントロピー正則化を導入することで滑らかにし、行列スケーリングを用いた高速勾配計算を可能にする2つの効率的なアルゴリズムを提案する。主な貢献は、正確な部分勾配法と比較して、オーダー・オブ・マグニチュードの高速化を達成するスケーラブルな手法であり、画像可視化および制約付きクラスタリングタスクでの検証がなされた。

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アルゴリズムを用い、行列スケーリングにより滑らか化された双対最適輸送問題の勾配を効率的に計算する。
• 固定サポート重心と自由サポート重心の両方を扱う2つの部分勾配ベースのアルゴリズムを提案し、両者とも滑らか化された双対目的関数を活用する。
• 勾配計算に行列-ベクトル積のみを用いることで、大規模なプライマル-双対最適輸送問題を解く必要がなくなり、計算コストを著しく削減する。
• 画像ファミリーおよびインコグナイトデータのクラスタリングにアルゴリズムを適用し、重心の重みに制約を課してバランスの取れた割り当てを実現する。
• 滑らか化された双対定式化により、大規模データセットに対しても実用的な高速収束が可能になる。

実験結果

リサーチクエスチョン

• RQ1 エントロピー正則化を用いることで、大規模な経験的測度に対するWasserstein重心の計算が実行可能になるか？
• RQ2 滑らか化された双対定式化は、部分勾配ベースの重心最適化における効率的勾配計算をどのように可能にするか？
• RQ3 均一な重み制約を課した制約付きWasserstein重心は、標準的なk-meansと比較してクラスタリングのバランスを改善できるか？
• RQ4 均一な重み制約を課した場合の、クラスタリング精度と計算コストのトレードオフは何か？
• RQ5 Wasserstein重心は、幾何的構造を保持しながら、複雑な画像ファミリーを効果的に要約できるか？

主な発見

• 48個の重心と57,647個の空間的位置を持つ制約付き重心の計算において、提案手法は12.5秒で完了したのに対し、正確な最適輸送ソルバーを用いた場合約1時間かかった。
• 非制約k-meansは1.55秒で収束したが、均一重み制約付きのWasserstein重心計算は12.5秒で完了し、わずかだが許容できるオーバーヘッドであった。
• 均一重み付きWasserstein重心は、各重心が合計質量の均等な割合を捉えるなど、よりバランスの取れたクラスタリングを実現した。これに対して標準k-meansは少数の重心に質量が集中する傾向がある。
• ネストされた楕円の可視化では、2-Wasserstein重心が構造的特徴を保持している一方、ユークリッド、RKHS、ジェフリー重心法は意味のある平均を生成できなかった。
• 滑らか化された双対定式化により、行列スケーリングを用いた勾配計算が可能となり、各部分勾配評価のコストが大規模な輸送問題の解法から単純な行列-ベクトル積にまで低下した。
• 本手法は、Wasserstein伝搬など、複数のWasserstein距離を含む変分問題へも成功裏に一般化された。