[論文レビュー] Efficient Neural Network Approaches for Conditional Optimal Transport with Applications in Bayesian Inference

Zheyu Oliver Wang, Ricardo Baptista+3 more|arXiv (Cornell University)|Oct 25, 2023

Gaussian Processes and Bayesian Inference被引用数 5

ひとこと要約

本稿では、ベイズ推論における効率的な条件付き最適輸送（COT）のための2つの深層学習ベースの手法—PCP-Map と COT-Flow—を提案する。単調な輸送写像をニューラルネットワーク（部分的凸ネットワークおよび正則化されたノーマライジングフロー）でパラメータ化することにより、複雑な条件付き事後分布からの高速なサンプリングと密度推定が可能となり、高次元の逆問題において、収束速度および数値的精度の面で最先端の手法を上回る。

ABSTRACT

We present two neural network approaches that approximate the solutions of static and dynamic $\unicode{x1D450}\unicode{x1D45C}\unicode{x1D45B}\unicode{x1D451}\unicode{x1D456}\unicode{x1D461}\unicode{x1D456}\unicode{x1D45C}\unicode{x1D45B}\unicode{x1D44E}\unicode{x1D459}\unicode{x0020}\unicode{x1D45C}\unicode{x1D45D}\unicode{x1D461}\unicode{x1D456}\unicode{x1D45A}\unicode{x1D44E}\unicode{x1D459}\unicode{x0020}\unicode{x1D461}\unicode{x1D45F}\unicode{x1D44E}\unicode{x1D45B}\unicode{x1D460}\unicode{x1D45D}\unicode{x1D45C}\unicode{x1D45F}\unicode{x1D461}$ (COT) problems. Both approaches enable conditional sampling and conditional density estimation, which are core tasks in Bayesian inference$\unicode{x2013}$particularly in the simulation-based ($\unicode{x201C}$likelihood-free$\unicode{x201D}$) setting. Our methods represent the target conditional distribution as a transformation of a tractable reference distribution. Obtaining such a transformation, chosen here to be an approximation of the COT map, is computationally challenging even in moderate dimensions. To improve scalability, our numerical algorithms use neural networks to parameterize candidate maps and further exploit the structure of the COT problem. Our static approach approximates the map as the gradient of a partially input-convex neural network. It uses a novel numerical implementation to increase computational efficiency compared to state-of-the-art alternatives. Our dynamic approach approximates the conditional optimal transport via the flow map of a regularized neural ODE; compared to the static approach, it is slower to train but offers more modeling choices and can lead to faster sampling. We demonstrate both algorithms numerically, comparing them with competing state-of-the-art approaches, using benchmark datasets and simulation-based Bayesian inverse problems.

研究の動機と目的

• 高次元のベイズ推論における条件付き最適輸送（COT）問題の計算的に不成立な性質に対処すること。
• スケーラブルで微分可能なニューラルネットワークによる COT 写像のパラメータ化を構築し、高速なサンプリングと密度推定を可能にすること。
• 測度輸送における敵対的学習の限界を克服し、学習確率の最大化による訓練を用いること。
• 複数の条件変数にわたるアモアタイズド条件付きサンプリングを可能にし、1回の推論コストを低減すること。
• 尤度や事前分布が不成立な状況においても、実用的で再現可能な条件付き事後分布の近似フレームワークを提供すること。

提案手法

• PCP-Map は、部分的入力凸ニューラルネットワーク（PICNN）の凸ポテンシャルの勾配として条件付き輸送写像をパラメータ化し、凸ポテンシャルによる単調性と一意性を保証する。
• COT-Flow は、正則化されたニューラルODEの流れとして条件付き輸送をモデル化し、ニューラルネットワークでパラメータ化された速度場を通じて保存的力学を強制する。
• 両手法とも、観測されたターゲットサンプルの尤度最大化により学習を行い、敵対的最適化を回避し、エンドツーエンドのバックプロパゲーションを可能にする。
• PCP-Map は凸最適化を用いてサンプリングを行うが、COT-Flow はODE統合を用いるため、学習が完了した後はより高速なサンプリングが可能である。
• 多様なデータセットや問題設定において堅牢な構成を特定するため、2段階のハイパーパramータサーチヒューリスティックを採用する。
• これらの手法は、確率的 Lotka-Volterra 方程式および1次元浅水域方程式を含むベンチマーク逆問題に適用され、高次元の条件付けがなされている。

Figure 2: Posterior samples in log scale and MAP point quality comparisons between proposed approaches and ABC with ${\bf x}^{*}=(0.01,0.5,1,0.01)^{\top}$ . Left : posterior samples generated by proposed approaches trained on 50k samples. Middle : posterior samples generated by proposed approaches t

実験結果

リサーチクエスチョン

• RQ1 ニューラルネットワークは、単調性や一意性といった望ましい理論的性質を保ちつつ、条件付き最適輸送写像を効率的にパラメータ化できるか？
• RQ2 敵対的学習と比較して、尤度に基づく輸送写像の学習は、ベイズ推論における収束速度およびサンプル品質の面でどのように異なるか？
• RQ3 静的（PCP-Map）と動的（COT-Flow）なアプローチの間で、訓練時間とサンプリング速度のトレードオフはどのようなものか？
• RQ4 従来のグリッドベースや多項式近似では失敗する高次元問題において、これらの手法はどの程度スケーラブルに拡張可能か？
• RQ5 ハイパーパramータの選択に対してこれらの手法はどの程度頑健か？また、単純なヒューリスティックによって多様な応用分野における構成選択が改善可能か？

主な発見

• PCP-Map は、14個の有効パラメータと3,500件の測定値を持つ高次元の逆問題において、約715秒で収束した。
• PCP-Map は、静的定式化であるにもかかわらず、アモアタイズド CP-Flow 法よりも収束が速かった。
• COT-Flow は、学習済みの速度場が整った後、ODE統合のおかげで PCP-Map よりも高速なサンプリングを実現した。
• 1次元浅水域方程式の問題において、両手法とも SMC-ABC を上回る数値的精度を示したが、ABC は強い理論的保証を持っているにもかかわらずである。
• PCP-Map は、特徴幅、コンテキスト幅、深さといった多くのハイパーパramータにおいて一貫した性能を示し、3つの主要なハイパーパラメータのみを必要とした。
• 尤度や事前分布の評価が不成立な状況においても、これらの手法は効果的であり、SDE やレガシーコードを含む実世界の逆問題における実用性を示した。

Figure 3: Posterior samples log scale and MAP point quality comparisons between proposed approaches and ABC with ${\bf x}^{*}=(0.02,0.02,0.02,0.02)^{\top}$ . Left : posterior samples generated by proposed approaches trained on 50k samples. Middle : posterior samples generated by proposed approaches