[論文レビュー] Entropic Wasserstein Gradient Flows

Gabriel Peyré|arXiv (Cornell University)|Feb 22, 2015

Geometric Analysis and Curvature Flows参考文献 48被引用数 5

ひとこと要約

本稿では、計算コストの高いウォッサーシュタイン距離の代わりにエントロピー正則化を用いてKullback-Leibler発散に置き換えることで、ウォッサーシュタイン勾配流れの計算を高速化する手法を提案する。この手法により、Dykstraのアルゴリズムと線形フィルタリング（例：ガウス畳み込み）を用いた並列処理可能な数値スキームが可能となり、グリッド上ではほぼ線形時間計算量を達成し、複雑な領域上でも効率的な熱拡散近似が実現される。

ABSTRACT

This article details a novel numerical scheme to approximate gradient flows for optimal transport (i.e. Wasserstein) metrics. These flows have proved useful to tackle theoretically and numerically non-linear diffusion equations that model for instance porous media or crowd evolutions. These gradient flows define a suitable notion of weak solutions for these evolutions and they can be approximated in a stable way using discrete flows. These discrete flows are implicit Euler time stepping according to the Wasserstein metric. A bottleneck of these approaches is the high computational load induced by the resolution of each step. Indeed, this corresponds to the resolution of a convex optimization problem involving a Wasserstein distance to the previous iterate. Following several recent works on the approximation of Wasserstein distances, we consider a discrete flow induced by an entropic regularization of the transportation coupling. This entropic regularization allows one to trade the initial Wasserstein fidelity term for a Kulback-Leibler divergence, which is easier to deal with numerically. We show how KL proximal schemes, and in particular Dykstra's algorithm, can be used to compute each step of the regularized flow. The resulting algorithm is both fast, parallelizable and versatile, because it only requires multiplications by a Gibbs kernel. On Euclidean domains discretized on an uniform grid, this corresponds to a linear filtering (for instance a Gaussian filtering when $c$ is the squared Euclidean distance) which can be computed in nearly linear time. On more general domains, such as (possibly non-convex) shapes or on manifolds discretized by a triangular mesh, following a recently proposed numerical scheme for optimal transport, this Gibbs kernel multiplication is approximated by a short-time heat diffusion.

研究の動機と目的

• ウォッサーシュタイン勾配流れのインパルシット・オイラー時刻離散化における各ステップの計算コストの高さに対処する。
• 勾配流れスキームにおける正確なウォッサーシュタイン距離計算の代替として、数値的に効率的な手法を開発する。
• 非線形拡散および群衆移動PDEのスケーラブルで安定な数値近似を、複雑な領域上でも可能にする。
• エントロピー正則化を活用して、ウォッサーシュタイン適合項をKullback-Leibler発散に変換し、数値的取り扱いやすさを向上させる。
• 行列乗算のみを必要とする、汎用的かつ並列処理可能なアルゴリズムを設計する。ただし、Gibbsカーネル $e^{-c/ u}$ との乗算を対象とする。

提案手法

• JKO（Jordan-Kinderlehrer-Otto）勾配流れスキームの各ステップにおける最適輸送問題にエントロピー正則化を適用し、ウォッサーシュタイン距離をKullback-Leibler発散に置き換える。
• KL発散に基づく一次の近接スキームを用いて各時刻ステップを計算し、Dykstraのアルゴリズムによる効率的最適化を可能にする。
• Gibbsカーネル $e^{-c/ u}$ を均一グリッド上での線形フィルタ（例：ガウスフィルタ）として表現し、高速畳み込みを用いてほぼ線形時間計算量を達成する。
• 非均一または多様体に離散化された領域では、Gibbsカーネル乗算を短時間熱拡散を用いて近似し、既存の数値スキームを活用する。
• 重み付き幾何平均を用いた確率測度の重み付き幾何平均を用いて、合成関数の近接作用素を定式化し、複数の周辺分布間でモジュラーな計算を可能にする。
• 関数形（例：和、積、周辺化）のさまざまな場合について、近接作用素の閉形式表現を導出し、計算効率を保証する。

実験結果

リサーチクエスチョン

• RQ1 エントロピー正則化は、ウォッサーシュタイン勾配流れの解法を安定的かつ数値的に効率的に高速化するために効果的に適用可能か？
• RQ2 ウォッサーシュタイン距離の部分問題を解く際の計算ボトルネックを、PDEの弱解への収束を保ちつつどのように軽減できるか？
• RQ3 三角形メッシュや多様体などの非ユークリッド領域上では、Gibbsカーネル乗算をどの程度効率的に近似できるか？
• RQ4 得られたアルゴリズムは、群衆移動や非線形拡散の多様な応用にわたり、並列処理可能で汎用的であるか？
• RQ5 エントロピー正則化は、勾配流れ近似の精度および収束速度にどのような影響を与えるか？

主な発見

• エントロピー正則化により、ウォッサーシュタイン距離の部分問題がKullback-Leibler発散の最小化問題に変換され、元の線形計画問題よりも数値的に取り扱いやすくなる。
• Gibbsカーネル演算に高速線形フィルタリング（例：ガウス畳み込み）を用いることで、均一グリッド上ではほぼ線形時間計算量を達成する。
• 非凸形状や三角形メッシュなどの複雑な領域では、Gibbsカーネル乗算が短時間熱拡散を用いて近似され、精度と安定性が保持される。
• 本手法により、多孔質媒体方程式や空間的に変化する拡散係数を持つ方程式を含む、非線形PDEのJKO型勾配流れの効率的計算が可能になる。
• 数値実験により、複雑な幾何形状上での群衆移動シミュレーションおよび変数係数を有する非線形拡散のシミュレーションにおいて、本手法の有効性が示された。
• KLに基づく近接スキームにDykstraのアルゴリズムを用いることで、すべてのテスト設定において強固でスケーラブルな収束が保証された。