[論文レビュー] Neural Wasserstein Gradient Flows for Maximum Mean Discrepancies with Riesz Kernels

Fabian Altekrüger, Johannes Hertrich+1 more|arXiv (Cornell University)|Jan 27, 2023

Generative Adversarial Networks and Image Synthesis被引用数 5

ひとこと要約

本稿では、滑らかでないリーツ核を伴う最大平均差分（MMD）汎関数に対して、生成的ニューラルネットワークを用いて輸送計画および速度計画を近似することで、写像に依存せず、特異的で絶対連続な測度の両方を扱えるニューラルウォッサースタイン勾配フローを提案する。この手法は、前向きおよび後向きスキームの両方が解析的解へ収束することを示しており、滑らかでない核を伴うMMDフローおよびバリセンターの計算に有効である。

ABSTRACT

Wasserstein gradient flows of maximum mean discrepancy (MMD) functionals with non-smooth Riesz kernels show a rich structure as singular measures can become absolutely continuous ones and conversely. In this paper we contribute to the understanding of such flows. We propose to approximate the backward scheme of Jordan, Kinderlehrer and Otto for computing such Wasserstein gradient flows as well as a forward scheme for so-called Wasserstein steepest descent flows by neural networks (NNs). Since we cannot restrict ourselves to absolutely continuous measures, we have to deal with transport plans and velocity plans instead of usual transport maps and velocity fields. Indeed, we approximate the disintegration of both plans by generative NNs which are learned with respect to appropriate loss functions. In order to evaluate the quality of both neural schemes, we benchmark them on the interaction energy. Here we provide analytic formulas for Wasserstein schemes starting at a Dirac measure and show their convergence as the time step size tends to zero. Finally, we illustrate our neural MMD flows by numerical examples.

研究の動機と目的

• 滑らかでないリーツ核を伴うMMD汎関数のウォッサースタイン勾配フローを計算するニューラルネットワークベースの手法を開発すること。特に、特異的でない連続的測度への移行を可能にすること。
• 従来のJKOスキームやパーティクルフローの限界を克服するため、輸送マップに依存せず、生成的ネットワークを用いて輸送計画および速度計画を近似すること。
• 目的測度が解析的に与えられておらず、核が凸的でも滑らかでもない状況においても、MMDフローおよびバリセンターの計算を可能にすること。
• 絶対連続測度の範囲を越えて、MMDに基づく生成的モデリングに向けた数値的に安定かつスケーラブルなフレームワークを提供すること。
• 非連続的測度の進化を含む複雑なフローの収束解析および数値実験を通じて、ニューラルスキームの有効性を検証すること。

提案手法

• リーツ核を伴うMMD汎関数のジョルダン＝キンダーレーラー＝オットー（JKO）後向きスキームを、生成的ニューラルネットワークを用いて近似する。このネットワークは、MMD汎関数の輸送計画の分解を学習する。
• ウォッサースタイン勾配降下フローに基づく前向きスキームを用い、速度計画についても生成的ネットワークを用いて近似することで、特異的初期測度からの進化を可能にする。
• MMDにおける相互作用エネルギーおよびポテンシャルエネルギー項を、サンプルベースの近似に適した二重積分としてモデル化し、勾配計算を効率化する。
• ウォッサースタイン距離およびリーツ核の構造に適合した損失関数を用いてニューラルネットワークを訓練することで、最適輸送ダイナミクスと整合性を保つ。
• MMDにリーツ核を適用するとスライス表現が得られることを活用し、1次元のソーティングにより高速な勾配計算が可能になるため、トレーニング効率が向上する。
• 提案手法を用いて、異なる部分多様体にサポートを持つ測度間のMMDバリセンターを計算し、複雑な標的密度からのサンプリングも実行する。

実験結果

リサーチクエスチョン

• RQ1 滑らかでないリーツ核を伴うMMD汎関数のウォッサースタイン勾配フローを、特に特異的測度から連続的測度への移行が生じる状況において、ニューラルネットワークが効果的に近似できるか。
• RQ2 初期測度が特異的である場合、ニューラル後向きおよび前向きスキームは、古典的パーティクルフローと比較して収束性および安定性に優れているか。
• RQ3 測度の絶対連続性を仮定しない状況でも、生成的ネットワークが輸送計画および速度計画をどれほど正確にモデル化できるか。
• RQ4 時間ステップサイズが0に近づくにつれて、提案手法がMMDフローの解析的解に収束するか。
• RQ5 目的測度が異なる部分多様体にサポートを持つ場合、ニューラルスキームがMMDバリセンターを正しく計算できるか。

主な発見

• 時間ステップサイズが0に近づくにつれて、ニューラル後向きおよび前向きスキームが、リーツ核を伴うMMDフローの解析的解に収束することが確認され、理論的整合性が裏付けられた。
• 本手法は測度の移行を効果的に扱える。例えば、ディラック測度は次元に応じて絶対連続測度または特異的でない非ディラック測度に進化するが、これは従来のパーティクルフローではモデル化できない。
• 数値実験の結果、初期測度が特異的である場合、両方のニューラルスキームが古典的パーティクルフローを上回り、特に非滑らかなダイナミクスを捉える能力に優れていることが示された。
• リーツ核の使用により、スライシングを介した高速勾配計算が可能になり、ニューラルネットワークのトレーニングが高速かつ安定した最適化を可能にした。
• 異なる部分多様体（例えば、円と正方形の境界）にサポートを持つ測度間のMMDバリセンターを、正確に理論的バリセンターに収束させる形で計算できた。
• MNISTの例では、初期測度を特異的ではなく一様分布にした場合でも、すべての手法で同等の性能を示した。これは、ニューラルスキームが初期測度の種別に対して頑健であることを示している。