[論文レビュー] Variational inference via Wasserstein gradient flows

Marc Lambert, Sinho Chewi+3 more|arXiv (Cornell University)|May 31, 2022

Markov Chains and Monte Carlo Methods被引用数 12

ひとこと要約

本稿では、ガウス測度のBures–Wasserstein空間におけるワサーティン勾配フローを用いた、新たな変分推論手法を提案する。これにより、複雑な後方分布の原理的近似が可能となる。対数凸性の下で強い理論的保証を達成し、ODEに基づく最適化により混合ガウス分布によるマルチモーダルなターゲットの有効な近似が可能である。

ABSTRACT

Along with Markov chain Monte Carlo (MCMC) methods, variational inference (VI) has emerged as a central computational approach to large-scale Bayesian inference. Rather than sampling from the true posterior $π$, VI aims at producing a simple but effective approximation $\hat π$ to $π$ for which summary statistics are easy to compute. However, unlike the well-studied MCMC methodology, algorithmic guarantees for VI are still relatively less well-understood. In this work, we propose principled methods for VI, in which $\hat π$ is taken to be a Gaussian or a mixture of Gaussians, which rest upon the theory of gradient flows on the Bures--Wasserstein space of Gaussian measures. Akin to MCMC, it comes with strong theoretical guarantees when $π$ is log-concave.

研究の動機と目的

• 特にターゲット後方分布が対数凸である場合に、ガウス分布および混合ガウス分布近似における変分推論（VI）の理論的保証の欠如に取り組む。
• 確率測度のBures–Wasserstein空間における勾配フローに基づく、原理的で整合性のあるVIフレームワークを構築する。
• 時間発展する後方分布近似をODEでモデル化することで、変分カルマンフィルタのアイデアを連続時間動的システムに拡張する。
• ヘッセ行列の計算を必要とせず、平均および共分散パラメータの学習にスケーラブルで最適化ベースのアルゴリズムを提供する。
• 平均場VIを改善し、混合ガウス分布近似によって相関構造およびマルチモーダル構造を捉える。

提案手法

• 提案分布と真の後方分布のKLダイバージェンスを最小化する変分推論を、ガウス測度のBures–Wasserstein幾何における勾配フローとして定式化する。
• 部分積分およびワサーティン勾配フロー理論を用いて、各ガウス成分の平均および共分散パラメータの連続時間ODEを導出する。
• K個のガウス成分を有する粒子ベースの表現を用い、各成分の平均および共分散行列の下三角チャレスキー因子が結合されたODEで時間発展を記述する。
• 4次ルンゲ＝クッタ法を用いてODEを数値的に統合し、全平均およびベクトル化されたチャレスキー因子を含む状態ベクトルを用いる。
• 共分散行列の正定値性を保つために、共分散そのものを直接進めるのではなく、下三角チャレスキー因子を進める。
• 数値積分のために、サグマポイントを用いたコロナチャージュールールを用いてODE内の期待値を計算する。

実験結果

リサーチクエスチョン

• RQ1 Bures–Wasserstein空間におけるワサーティン勾配フローは、従来の変分推論に対する理論的根拠とスケーラビリティを備えた、ガウス分布および混合ガウス分布近似の代替手段として有効であるか？
• RQ2 ターゲット後方分布が対数凸である場合に、本手法はラングジュアン拡散のようなMCMC手法に類する強い収束保証を達成するか？
• RQ3 本手法は、ガウス成分の数を増やすことで、マルチモーダルな後方分布をどれほどよく近似できるか？
• RQ4 アルゴリズムは初期値にどれほど敏感か？実務において、悪い初期値をどのように緩和できるか？
• RQ5 ヘッセ行列の情報が不要な状況でも、本手法は相関構造や複数のモードを効果的に捉えることができるか？

主な発見

• 本手法は、対数凸ターゲットの下で理論的収束保証を達成し、MCMCの非漸近的解析が知られている分野を、変分推論へと拡張した。
• 二モーダルなターゲットに対して、20個のガウス成分を用いた混合分布により、両方のモードを的確に捉えた。等高線プロットでは、真の後方分布と密接に一致していることが示された。
• KLダイバージェンスは約30ステップ（ステップサイズ0.1で300ステップ）後に安定化し、最適化プロセスの収束が確認された。
• ガウス粒子の数を増やすことで近似品質が向上し、4〜6個の非等価な重みを持つモードを持つターゲットにおいて、より良いKLダイバージェンスと正確な密度推定が達成された。
• 初期値の選択が収束に大きく影響する：あるモードの近くに初期化された粒子はその周辺に留まり、モードから等距離に初期化された粒子は平均方向に収束し、共分散が増大する傾向を示した。これは、複数の初期粒子を用いる必要性を示唆している。
• 部分積分を用いてヘッセ行列を必要としない更新式を導出することで、Hessian計算を回避し、効率的で安定した最適化を実現した。