[論文レビュー] Annealed Flow Transport Monte Carlo

Michael Arbel, Alexander Matthews+1 more|arXiv (Cornell University)|Feb 15, 2021

Model Reduction and Neural Networks参考文献 73被引用数 4

ひとこと要約

本稿では、複雑なターゲット分布下での正規化定数および期待値の推定を改善するために、順次モンテカルロ法（SMC）と正規化フロー（NF）を組み合わせた、新たなモンテカルロ法であるアニールドフロー輸送（AFT）を提案する。AFTは、逐次的に粒子を次々とアニールドターゲットへと輸送するNFを学習することで、分散を低減し、モード崩壊を緩和する。これにより、温度ステップを減らしても標準的なSMCに比べて優れた性能を達成し、連続時間極限において制御された拡散過程に理論的に収束することが示された。

ABSTRACT

Annealed Importance Sampling (AIS) and its Sequential Monte Carlo (SMC) extensions are state-of-the-art methods for estimating normalizing constants of probability distributions. We propose here a novel Monte Carlo algorithm, Annealed Flow Transport (AFT), that builds upon AIS and SMC and combines them with normalizing flows (NFs) for improved performance. This method transports a set of particles using not only importance sampling (IS), Markov chain Monte Carlo (MCMC) and resampling steps - as in SMC, but also relies on NFs which are learned sequentially to push particles towards the successive annealed targets. We provide limit theorems for the resulting Monte Carlo estimates of the normalizing constant and expectations with respect to the target distribution. Additionally, we show that a continuous-time scaling limit of the population version of AFT is given by a Feynman--Kac measure which simplifies to the law of a controlled diffusion for expressive NFs. We demonstrate experimentally the benefits and limitations of our methodology on a variety of applications.

研究の動機と目的

• 分布間の大きな差異がある場合に生じる、アニールド重要度サンプリング（AIS）およびSMCにおける高分散推定器の問題に対処すること。
• 直接単純な事前分布からフローを学習するのではなく、中間のアニールドステップでフローを学習することで、正規化フローに基づく輸送写像におけるモード崩壊やテールの低減を克服すること。
• 再サンプリング、MCMC、重要度サンプリングの強みを有するSMCと、表現力のある正規化フローを組み合わせ、より正確で安定した推論を実現すること。
• 提案されたアルゴリズム下でのモンテカルロ推定量について、弱大数法則および中心極限定理を含む理論的保証を確立すること。
• 連続時間極限において、AFTが制御された拡散過程に簡略化されるフェインマン＝カック測度に収束することを示すこと。

提案手法

• AFTは、πₖ ∝ π₀^(1−βₖ)π_K^βₖ で定義される一連のアニールドターゲットを用い、0 = β₀ < β₁ < ... < β_K = 1 とし、単純な初期分布π₀からターゲットπ_K = πへと橋渡しを行う。
• 各ステップkにおいて、(T_{k+1})#πₖ と π_{k+1} の間のKLダイバージェンスを最小化するように正規化フローT_{k+1}を学習し、πₖの粒子をπ_{k+1}をよりよく近似するように輸送する。
• フローによる輸送後、π_{k+1}/πₖ の比率を用いて粒子を再重み付けし、π_{k+1}の重要度サンプリング近似を構築する。
• その後、π_{k+1}を不変分布とするMCMCカーネルを用いて粒子を再サンプリングおよび変異化させ、多様性を維持し、フロー近似誤差を是正する。
• Unadjusted Langevin Algorithm（ULA）カーネルを用いて連続時間スケーリング極限に拡張し、フェインマン＝カック測度に収束することを示した。
• 表現力のあるNFが正確な輸送を可能にする場合、極限測度は制御された拡散過程の法則に対応する。

実験結果

リサーチクエスチョン

• RQ1 正規化フローを効果的に用いることで、重要度重みの分散を低減させ、SMCサンプラーの効率と正確性を向上させることができるか？
• RQ2 直接事前分布からではなく、中間のアニールドステップで輸送写像を学習することで、モード崩壊を緩和し、サンプリング品質を向上させることができるか？
• RQ3 学習ベースの輸送スキームをSMCに適用した場合、モンテカルロ推定量にどのような理論的保証を確立できるか？
• RQ4 AFTの連続時間極限は、制御された拡散過程のような確率過程とどのように関係するか？
• RQ5 より少ない温度ステップで、AFTは標準的なSMCに比べて推定の正確性と収束速度において優れているか？

主な発見

• AFTは、ベンチマーク問題において温度ステップを減らしても標準的なSMCに比べて低い推定誤差を達成し、より高いサンプル効率を示した。
• 理論的分析により、学習された輸送写像の性質にもかかわらず、AFT推定量に対して弱大数法則および中心極限定理が確立された。
• ULAカーネルを用い、再サンプリングを行わない連続時間極限において、AFTはフェインマン＝カック測度に収束し、フローが十分に表現力を持つ場合、これは制御された拡散過程の法則に簡略化される。
• バイナリズドMNISTおよび空間 Cox プロセスにおける実験から、AFTはベースラインのSMCおよび変分推論に比べ、より正確な正規化定数推定とより良い事後分布近似を提供した。
• 中間のアニールドターゲットを活用することで、直接正規化フローを訓練する際に生じるモード崩壊やテールの低減の影響を効果的に低減した。
• アブレーションスタディにより、特に高次元または挑戦的なターゲット分布において、各遷移における十分なHMC更新回数がAFTの性能にとって不可欠であることが確認された。