[論文レビュー] Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent

Jason M. Altschuler, Sinho Chewi+2 more|arXiv (Cornell University)|Jun 15, 2021

Geometric Analysis and Curvature Flows被引用数 8

ひとこと要約

本稿は、ガウス分布の Wasserstein バリオセントラーや、エントロピー正則化バリオセントラーや幾何的中央値を計算する際、Bures–Wasserstein 多様体上でのリーマン的勾配降下法に対して次元に依存しない収束レートを確立する。著者らは、新たな測地的凸性に関する結果を証明することで、次元に依存しない収束保証を達成し、最適輸送距離の最適化における主要な理論的ギャップを解消する。

ABSTRACT

We study first-order optimization algorithms for computing the barycenter of Gaussian distributions with respect to the optimal transport metric. Although the objective is geodesically non-convex, Riemannian GD empirically converges rapidly, in fact faster than off-the-shelf methods such as Euclidean GD and SDP solvers. This stands in stark contrast to the best-known theoretical results for Riemannian GD, which depend exponentially on the dimension. In this work, we prove new geodesic convexity results which provide stronger control of the iterates, yielding a dimension-free convergence rate. Our techniques also enable the analysis of two related notions of averaging, the entropically-regularized barycenter and the geometric median, providing the first convergence guarantees for Riemannian GD for these problems.

研究の動機と目的

• 最適輸送バリオセントラーや最適輸送距離におけるリーマン的勾配降下法の理論的ギャップを解消する。特に、既存の収束レートが次元に指数関数的に依存する問題を解決する。
• ガウス分布の Wasserstein バリオセントラーや最適輸送距離におけるリーマン的勾配降下法に対して、次元に依存しない収束保証を提供する。
• 標準的なバリオセントラーよりも関連する2つの問題、すなわちエントロピー正則化バリオセントラーや幾何的中央値の解析を拡張し、これら問題におけるリーマン的勾配降下法の収束結果を初めて得る。
• より強い測地的凸性の制御を確立することで、リーマン的設定下での目的関数の非凸性を克服し、次元に依存しない収束を実現する。
• 実験的に、リーマン的勾配降下法が、理論的制限を受ける後続の手法であるユークリッド的勾配降下法や半定値計画法ソルバーと比較して優れていることを示す。

提案手法

• Wasserstein バリオセントラータスクを、正定値行列の錐に等長である Bures–Wasserstein 多様体上の最小化問題として定式化する。
• 収束の制御と次元に依存しない収束の実現のため、目的関数の新たな測地的凸性特性を導入し、証明する。
• 共分散行列の多様体上に直接リーマン的勾配降下法を適用し、ユークリッド空間への埋め込みではなく、多様体の内面的幾何を活用する。
• エントロピー正則化バリオセントラーや幾何的中央値問題を、同じ幾何的フレームワークを用いて分析し、標準的バリオセントラーよりも広い範囲の収束保証を拡張する。
• 指数写像や単位元における対数写像を含む、最適輸送とリーマン幾何学の道具を用いて収束挙動を分析する。
• 環境次元 d に依存しない理論的収束レートを提供する。多様体上の曲率と凸性の制御に依存する。

実験結果

リサーチクエスチョン

• RQ1 目的関数が測地的に非凸であるにもかかわらず、ガウス分布の Wasserstein バリオセントラーや最適輸送距離に対して、リーマン的勾配降下法が次元に依存しない収束を達成できるか？
• RQ2 Bures–Wasserstein 多様体のどのような幾何的性質と凸性特性が、高次元設定下での次元に依存しない収束を可能にするか？
• RQ3 同様の理論的枠組みを、これまでにリーマン的勾配降下法の収束保証が存在しなかったエントロピー正則化バリオセントラーや幾何的中央値問題に拡張可能か？
• RQ4 なぜ実験的にリーマン的勾配降下法はユークリッド的勾配降下法や SDP ソルバーを上回る性能を示すのか？この現象は理論的に説明可能か？
• RQ5 多様体の曲率と測地的構造が、収束の安定化と加速に果たす役割は何か？

主な発見

• 本稿は、ガウス分布の Wasserstein バリオセントラーや最適輸送距離に対して、Bures–Wasserstein 多様体上でのリーマン的勾配降下法の収束率が次元に依存しないという、初の結果を確立する。
• 収束の制御を強化するため、新たな測地的凸性に関する結果を証明し、次元 d に依存しない収束を可能にする。
• 収束保証をエントロピー正則化バリオセントラーや幾何的中央値問題に拡張し、これら問題におけるリーマン的勾配降下法の収束結果として初めての結果を提供する。
• 実験的結果から、リーマン的勾配降下法は、ステップサイズのチューニングが不要なまま、ユークリッド的勾配降下法や SDP ソルバーと比較してより速く収束することが示された。一方、後者の理論的境界は次元 d に指数関数的に依存する。
• 理論的分析により、リーマン的勾配降下法が、悪条件の行列や行列構造（例えば、非可換固有基底）に対して頑健であり、多様な合成データ分布にわたり安定した収束を示すことが確認された。
• 実験では固定ステップサイズ 1 で収束が達成されたが、ユークリッド的勾配降下法は慎重なチューニングを要し、依然として性能が劣ることから、内面的幾何の利点が顕著に現れている。