[論文レビュー] Variational Wasserstein gradient flow

Jiaojiao Fan, Qinsheng Zhang+2 more|arXiv (Cornell University)|Dec 4, 2021

Advanced Neuroimaging Techniques and Applications被引用数 8

ひとこと要約

本稿では、密度に依存する目的関数を、パrametric関数上の変分形式に置き換えることで、高次元の経験的分布上での効率的かつスケーラブルな計算を可能にする、変分 Wasserstein 勾配フロー手法を提案する。f-発散の変分形式と入力凸ニューラルネットワーク（ICNN）を活用することで、高価な密度推定やヘッセ行列の行列式近似を回避し、GAN よりも安定した学習と優れたサンプル品質を実現した。CIFAR10 では、最先端の FID スコア（23.1）を達成した。

ABSTRACT

Wasserstein gradient flow has emerged as a promising approach to solve optimization problems over the space of probability distributions. A recent trend is to use the well-known JKO scheme in combination with input convex neural networks to numerically implement the proximal step. The most challenging step, in this setup, is to evaluate functions involving density explicitly, such as entropy, in terms of samples. This paper builds on the recent works with a slight but crucial difference: we propose to utilize a variational formulation of the objective function formulated as maximization over a parametric class of functions. Theoretically, the proposed variational formulation allows the construction of gradient flows directly for empirical distributions with a well-defined and meaningful objective function. Computationally, this approach replaces the computationally expensive step in existing methods, to handle objective functions involving density, with inner loop updates that only require a small batch of samples and scale well with the dimension. The performance and scalability of the proposed method are illustrated with the aid of several numerical experiments involving high-dimensional synthetic and real datasets.

研究の動機と目的

• 高次元空間における Wasserstein 勾配フローの計算不能性を、密度推定とヘッセ行列の計算に起因するものとして解決すること。
• 経験的分布上での Wasserstein 勾配フローを数値的に安定かつスケーラブルに計算する手法の開発。
• 密度依存の目的関数を、サンプルベースでの評価が可能な変分形式に置き換えること。
• 標準的な GAN と比較して、学習の安定性を向上させ、モード崩壊を回避すること。
• 合成データおよび実世界の高次元データセット（MNIST や CIFAR10 を含む）における本手法の有効性を実証すること。

提案手法

• f-発散に基づく変分表現を用いて目的関数を定式化し、パラメトリック関数族上での最大化を実行する。
• 入力凸ニューラルネットワーク（ICNN）を用いて、凸関数の勾配として最適輸送マップをパラメータライズする。
• 小さなバッチのサンプルに対する内側ループの更新により、明示的な密度評価やヘッセ行列の行列式近似の必要性を回避する。
• ステップごとに正則化された Wasserstein 距離と変分目的関数を最小化する形で、JKO スキームを確率的最適化により実装する。
• 初期分布（例：正規分布）を、勾配フローのダイナミクスに従って、目的分布へと段階的に変換するための、一連の輸送マップを用いる。
• 学習済みの輸送マップに沿ってサンプルを伝搬させることで、フロー全体に沿ったサンプリングと密度評価を可能にする。

実験結果

リサーチクエスチョン

• RQ1 Wasserstein 勾配フローにおける目的関数の変分形式は、高次元設定において明示的な密度推定を不要にすることができるか？
• RQ2 本手法は、Hessian 行列式近似を要する従来の JKO に基づく手法と比較して、計算効率とスケーラビリティにおいて優れているか？
• RQ3 変分形式は、モーメントマッチングや積分確率的度量に関する埋め込み不等式といった幾何学的・統計的性質を保持しているか？
• RQ4 本手法は、モード崩壊を回避しつつ、CIFAR10 などの画像生成ベンチマークで最先端のサンプル品質を達成できるか？
• RQ5 ステップサイズやフローのステップ数の変更に伴い、本手法の性能はどのように変化するか？

主な発見

• 本手法は、CIFAR10 で 23.1 の Fréchet Inception Distance（FID）スコアを達成し、GAN や他のフローに基づく生成モデルを上回った。
• ステップサイズ a=5.0 で学習を行うと、安定した収束が得られ、モード崩壊を回避できるが、より大きな値では不安定になり、サンプル品質が低下した。
• ヘッセ行列の行列式近似を必要とする従来の JKO に基づく手法と比較して、本手法は優れたスケーラビリティと計算効率を示した。
• 変分目的関数は、モーメントマッチングの性質と、積分確率的度量に関する埋め込み不等式を満たしており、理論的整合性を裏付けた。
• FID スコアは時間ステップに伴い収束する傾向を示し、フローのステップ数が増加しても安定した学習ダイナミクスを示した。
• CIFAR10 では、最先端のモデル（AE-OT-GAN や OTM と同等の）と比較して、競争力のある Inception Score（7.48 ± 0.12）を達成した。