[論文レビュー] Learning High Dimensional Wasserstein Geodesics

Liu Shu, Shaojun Ma+3 more|arXiv (Cornell University)|Feb 5, 2021

Seismic Imaging and Inversion Techniques参考文献 42被引用数 5

ひとこと要約

本稿では、ラグランジュ乗数とKKT条件を用いて最適輸送問題をミニマックスサドルポイント問題に定式化することで、高次元 Wasserstein geodesic を計算する深層学習ベースの手法を提案する。この手法は、輸送写像と速度場を深層ニューラルネットワークでパrameter化し、明示的な正則化や次元削減を伴わずに、サンプルベースで双方向に学習可能である。これにより、Wasserstein 距離、最適写像、geodesic 路径を同時に推定できる。

ABSTRACT

We propose a new formulation and learning strategy for computing the Wasserstein geodesic between two probability distributions in high dimensions. By applying the method of Lagrange multipliers to the dynamic formulation of the optimal transport (OT) problem, we derive a minimax problem whose saddle point is the Wasserstein geodesic. We then parametrize the functions by deep neural networks and design a sample based bidirectional learning algorithm for training. The trained networks enable sampling from the Wasserstein geodesic. As by-products, the algorithm also computes the Wasserstein distance and OT map between the marginal distributions. We demonstrate the performance of our algorithms through a series of experiments with both synthetic and realistic data.

研究の動機と目的

• 従来のPDEベースの手法が次元の呪いに苦しむ高次元空間における Wasserstein geodesic の計算という課題に対処すること。
• 空間離散化や正則化子を回避するスケーラブルでサンプルベースのアルゴリズムを開発し、高次元設定への応用を可能にすること。
• 統一されたフレームワークを用いて、分布間の Wasserstein 距離、最適輸送写像、geodesic 路径を同時に計算すること。
• geodesic 沿いの自然なサンプリングメカニズムを提供することで、機械学習、ロボット工学、制御分野における実用的応用を可能にすること。
• 距離や写像推定にのみ焦点を当て、全 geodesic 軌道を扱わない既存のOT手法の限界を克服すること。

提案手法

• ラグランジュ乗数を用いて動的最適輸送問題をミニマックス問題に定式化し、サドルポイント問題に変換する。
• KKT条件を適用して探索空間を縮小し最適化を簡素化することで、より取り扱いやすい定式化を得る。
• 速度場と双対ポテンシャルを深層ニューラルネットワークでパラメータ化し、高次元における柔軟な関数近似を可能にする。
• 前向きおよび逆向きの輸送経路を交互に最適化する双方向でサンプルベースの学習アルゴリズムを設計する。
• 前処理と双方向ダイナミクスを組み込み、高次元空間における学習の安定性と収束性を向上させる。
• ミニバッチサンプリングを用いてネットワークをエンドツーエンドで学習させ、明示的な正則化制約なしにスケーラブルな推論と geodesic サンプリングを可能にする。

実験結果

リサーチクエスチョン

• RQ1 空間離散化や明示的な正則化に依存せずに、深層学習ベースのアプローチが高次元 Wasserstein geodesic を計算できるか？
• RQ2 geodesic 計算の過程で、最適輸送写像と Wasserstein 距離を同時に推定する方法は何か？
• RQ3 提案された双方向学習戦略は、高次元OT問題における安定性と収束性をどの程度向上させるか？
• RQ4 この手法は、geodesic 路径に沿った現実的な中間サンプルを生成できるか？画像間変換やデータ生成への応用が可能か？
• RQ5 複雑で高次元な分布（例：画像、手書き数字）を含む実世界データにおいて、この手法はどの程度の性能を示すか？

主な発見

• 提案手法は、次元削減を伴わず、高次元空間（例：28×28 MNIST ディジット、256×256 画像）における Wasserstein geodesic を成功裏に計算した。
• 画像の色移動やディジット変換タスクにおいて、ソース分布とターゲット分布の間の視覚的に現実的な中間サンプルを生成した。
• 画像変換タスク（例：夏から秋への森の風景変換）において、生成画像は真値に近く、パレット分布の類似性が高く維持された。
• POT ライブラリとの比較により、Wasserstein 距離と最適輸送写像の推定が高精度でなされた。
• 双方向学習スキームにより、安定性と収束性が向上し、複雑な高次元設定でも信頼性のある geodesic サンプリングが可能になった。
• 一般の凸コスト関数（Lp ノルムを含む）に対して、サンプルベースの geodesic 計算で最先端の性能を達成し、前向きおよび逆向きの両方の輸送をサポートした。