[論文レビュー] Flow Annealed Importance Sampling Bootstrap

Laurence Illing Midgley, Vincent Stimper+3 more|arXiv (Cornell University)|Aug 3, 2022

Machine Learning in Materials Science被引用数 14

ひとこと要約

この論文では、正規化フローと冷却付き重要度サンプリング（AIS）を組み合わせた新規手法、Flow AIS Bootstrap（FAB）を紹介する。FABは、分子動力学（MD）シミュレーションにかかるコストを回避するため、未正規化されたターゲット密度のみを用いてフローモデルを訓練する。α=2における質量カバー型αダイバージェンスを最小化することで、重要度重みの分散を低減し、多次元分布の正確な近似が可能になる。特にアラニンジペプチドのボルツマン分布に対しても、MDサンプルを用いた最尤推定法に比べ100倍少ないターゲット評価回数で優れた結果を達成した。

ABSTRACT

Normalizing flows are tractable density models that can approximate complicated target distributions, e.g. Boltzmann distributions of physical systems. However, current methods for training flows either suffer from mode-seeking behavior, use samples from the target generated beforehand by expensive MCMC methods, or use stochastic losses that have high variance. To avoid these problems, we augment flows with annealed importance sampling (AIS) and minimize the mass-covering $α$-divergence with $α=2$, which minimizes importance weight variance. Our method, Flow AIS Bootstrap (FAB), uses AIS to generate samples in regions where the flow is a poor approximation of the target, facilitating the discovery of new modes. We apply FAB to multimodal targets and show that we can approximate them very accurately where previous methods fail. To the best of our knowledge, we are the first to learn the Boltzmann distribution of the alanine dipeptide molecule using only the unnormalized target density, without access to samples generated via Molecular Dynamics (MD) simulations: FAB produces better results than training via maximum likelihood on MD samples while using 100 times fewer target evaluations. After reweighting the samples, we obtain unbiased histograms of dihedral angles that are almost identical to the ground truth.

研究の動機と目的

• 高価なMDシミュレーションに依存する既存のフローモデル訓練手法の限界や、多次元分布におけるモード探索の欠落を是正すること。
• 重要度重みの分散を最小化しつつ、複雑で多次元なターゲット分布を正確に近似できる訓練目的関数の開発。
• 事前のMDサンプリングなしに、未正規化ターゲット密度のみを用いてボルツマンジェネレータを訓練可能にすること。
• 再利用可能なAISサンプルを活用する優先順位付きリプレイバッファを導入し、計算コストを低減すること。
• アラニンジペプチドを含む挑戦的なベンチマークで最先端の性能を示すこと。

提案手法

• FABは、α=2におけるαダイバージェンスを訓練目的関数として採用し、質量カバー型であり、推論時の重要度重みの分散を最小化する。
• AISを用い、フローを初期分布とし、ターゲットを分散最小の分布とする。これにより、αダイバージェンス損失を推定する。
• AISは、フローが近似が悪い領域からの高品質なサンプルを生成し、新たなモードの発見と訓練信号の質の向上を可能にする。
• 優先順位付きリプレイバッファを導入し、AISサンプルを再利用することで、計算コストを低減し、サンプル効率を向上させる。
• AISで推定されたα=2におけるαダイバージェンス目的関数を、確率的勾配降下法でエンドツーエンドにフローを訓練する。
• テスト段階では、訓練済みフローを用いた重要度サンプリングにより、ターゲット分布下での期待値の低分散推定が可能になる。

実験結果

リサーチクエスチョン

• RQ1 未正規化ターゲット密度のみを用いて、高価なMDシミュレーションに依存せずに正規化フローを効果的に訓練できるか？
• RQ2 α=2におけるαダイバージェンスを最小化することで、標準的なKLダイバージェンスと比較して、より優れたサンプル品質と低い重要度重み分散が達成できるか？
• RQ3 FABは、アラニンジペプチドのラマチャンドランプロットのような極めて多次元な分布のすべてのモードを発見し、正確にモデル化できるか？
• RQ4 MDシミュレーションから得たサンプルを用いた最尤推定法と比較して、FABの性能とサンプル効率はどのように異なるか？
• RQ5 リプレイバッファの使用が、訓練に必要なターゲット密度評価回数を顕著に削減できるか？

主な発見

• FABは、MDサンプルを一切使用せず、未正規化エネルギー関数のみを用いてアラニンジペプチドのボルツマン分布を学習し、最先端の結果を達成した。
• 最尤推定法に比べ、ターゲット評価回数が100倍以上少ないにもかかわらず、真の分布に近い近似を実現した。
• 再重み付け後、FABが生成した二面角のヒストограмグラムは、真値とほとんど同一であり、モードカバレッジの高さが裏付けられた。
• 32次元の「Many Well」問題において、有効サンプルサイズ（ESS）、対数尤度、KLダイバージェンスの観点で、他の手法を上回った。
• α=2を用いることで、ラマチャンドランプロットにおいてKLダイバージェンスが最小（2.51×10⁻³）かつESSが最大（52.2％）となり、他の手法を著しく上回った。
• リプレイバッファの使用により、計算コストが低減され、訓練の安定性が向上し、高品質なAISサンプルの効率的再利用が可能になった。