[論文レビュー] Smooth Normalizing Flows

Jonas Köhler, Andreas Krämer+1 more|arXiv (Cornell University)|Oct 1, 2021

Protein Structure and Dynamics被引用数 15

ひとこと要約

本論文は、コンpact区間および超トーラス上での確率分布のモデリングを目的とした、C∞-スムーズな混合変換を用いたスムーズな正規化フローを提案する。分子系の内部座標に不可欠な分野である。逆関数定理を用いた微分可能な逆変換を可能にすることで、双方向学習、力マッチング、分子動力学シミュレーションへの応用が可能となり、滑らかで高精度なポテンシャルを実現し、力の挙動が良好で、高いサンプリング効率を達成する。

ABSTRACT

Normalizing flows are a promising tool for modeling probability distributions in physical systems. While state-of-the-art flows accurately approximate distributions and energies, applications in physics additionally require smooth energies to compute forces and higher-order derivatives. Furthermore, such densities are often defined on non-trivial topologies. A recent example are Boltzmann Generators for generating 3D-structures of peptides and small proteins. These generative models leverage the space of internal coordinates (dihedrals, angles, and bonds), which is a product of hypertori and compact intervals. In this work, we introduce a class of smooth mixture transformations working on both compact intervals and hypertori. Mixture transformations employ root-finding methods to invert them in practice, which has so far prevented bi-directional flow training. To this end, we show that parameter gradients and forces of such inverses can be computed from forward evaluations via the inverse function theorem. We demonstrate two advantages of such smooth flows: they allow training by force matching to simulation data and can be used as potentials in molecular dynamics simulations.

研究の動機と目的

• 分子内部座標のモデリングに不可欠なコンパクト区間および超トーラス上でのC∞-スムーズな正規化フローの欠如を解消する。
• 逆関数定理による勾配計算を通じて、解析的でない逆関数を持つフローの双方向学習を可能にする。
• 分子動力学シミュレーションに適した、滑らかで微分可能なポテンシャルの開発を目的とする。
• サポートの事前知識を必要としたり、不連続な力を生じる既存手法の制限を克服する。

提案手法

• コンパクト区間および超トーラス向けに、学習可能な形状パラメータを有するバーム関数を用いたC∞-スムーズな混合変換のクラスを提案する。
• ブラックボックス根の探索法（例：マルチビンサーチ）を用いて、逆方向へのフローの逆算を実行する。
• 逆関数定理を用いて逆関数の勾配を計算し、根の探索ステップを逆伝播可能にする。
• 最大尤度推定（MLE）、力マッチング（FM）、逆Kullback-Leibler発散のハイブリッド損失を用いてフローを学習する。
• 速度Verlet積分器を用いて、学習済みフローを分子動力学シミュレーションにおけるポテンシャルとして適用する。
• OpenMMを用いてエネルギーと力の評価を行い、サンプリング効率とエネルギー保存の観点から検証する。

実験結果

リサーチクエスチョン

• RQ1 C∞-スムーズな正規化フローを、コンパクト区間および超トーラス上で構築可能であり、分子内部座標の滑らかなモデリングを可能にするか？
• RQ2 逆関数定理を用いて、解析的でない逆関数を持つフローに対して勾配を逆伝播可能か？
• RQ3 密度推定と組み合わせた力マッチングにより、分子系における学習ポテンシャルの精度が向上するか？
• RQ4 滑らかなフローは、古典的力場と同等の安定性とエネルギー保存性を示す分子動力学シミュレーションを可能にするか？
• RQ5 スムーズなフローは、スプラインベースやアフィンフローに比べ、サンプリング効率および力の精度で優れるか？

主な発見

• スムーズなフローはアラニンジペプチドにおいて42%のサンプリング効率を達成し、スプラインベースのフロー（38%）やアフィンフロー（<1%）を上回った。
• 力マッチングとMLEを用いて学習したスムーズなフローは、逆Kullback-Leibler発散が0.0003にまで低下し、密度マッチングがほぼ完全であることを示した。
• スムーズなフローを用いたMDシミュレーションにおけるエネルギーのフラクチュエーションは、Amber ff99SB-ILDN力場と同等であり、数値的安定性を示した。
• スムーズなフローから計算された力は、スプラインベースのフローに比べて顕著に精度が高く、テストデータにおける力の残差が低かった。
• 逆関数定理による勾配計算により、解析的でない逆関数を持つフローの双方向学習が可能となり、先行研究の主な制限を克服した。
• スムーズなフローのアーキテクチャは、滑らかで整合性のある力とエネルギーを生成し、力の不連続性が生じないため、分子動力学シミュレーションへの直接適用が可能となった。