[論文レビュー] Rigid Body Flows for Sampling Molecular Crystal Structures

Jonas Köhler, Michele Invernizzi+2 more|arXiv (Cornell University)|Jan 26, 2023

Computer Graphics and Visualization Techniques被引用数 10

ひとこと要約

本論文は、3D空間での複数の剛体（分子）の位置と姿勢をサンプリングするために特化した正規化フローを導入し、SO(3)密度を滑らかにかつ正確にモデルする二重カバーのクォータニオン表現を用います。 tetrahedral system と TIP4P の Ice XI で Boltzmann generator スタイルのサンプリングを実演し、分子結晶における自由エネルギー差の推定を効率化します。

ABSTRACT

Normalizing flows (NF) are a class of powerful generative models that have gained popularity in recent years due to their ability to model complex distributions with high flexibility and expressiveness. In this work, we introduce a new type of normalizing flow that is tailored for modeling positions and orientations of multiple objects in three-dimensional space, such as molecules in a crystal. Our approach is based on two key ideas: first, we define smooth and expressive flows on the group of unit quaternions, which allows us to capture the continuous rotational motion of rigid bodies; second, we use the double cover property of unit quaternions to define a proper density on the rotation group. This ensures that our model can be trained using standard likelihood-based methods or variational inference with respect to a thermodynamic target density. We evaluate the method by training Boltzmann generators for two molecular examples, namely the multi-modal density of a tetrahedral system in an external field and the ice XI phase in the TIP4P water model. Our flows can be combined with flows operating on the internal degrees of freedom of molecules and constitute an important step towards the modeling of distributions of many interacting molecules.

研究の動機と目的

• 多くの相互作用する剛体が結晶中で位置と姿勢を同時に効率的に、尤度ベースでサンプリングできるよう動機付け、実現する。
• 回転多様体上で滑らかで可逆な流れをクォータニオンの二重カバーを介して開発し、物理密度を保ち、厳密な密度計算をサポートする。
• 外部姿勢と内部自由度を独立に扱えるようにし、内部DOFフロー・モデルとの組み込みを可能にする。
• 多モード回転密度と異なるサイズ・温度での氷XI相のサンプリングの Boltzmann ジェネレーターを構築してアプローチを実証する。

提案手法

• 各剛体を (x0, R, Psi) で表現。x0 は平行移動、R は SO(3) の回転、Psi は内部自由度。
• SO(3) の回転を二重カバーを使って単位クォータニオン3-スフィア S3 にリフトし、S3 上で滑らかで反転同値性を有するフローを可能にする。
• S3 上で反転対称の結合フローを、二つのクラス（対称化されたモビウス変換と射影凸勾配写像）を用いて構築する（解析的な逆と体積変化公式を持つ）。
• 外見アーギュメントを維持しつつ、位置とクォータニオンを更新するカップリングベースのフロー F を定義し、二重カバー射影の下で厳密な密度を保証する。
• 訓練は Boltzmann Generators として説明: バイアス付き MD データ上の最尤推定と逆KL最小化または LFEPを組み合わせ、自由エネルギー差を推定する。
• 二つの実験設定を示す: (i) 外部場中のテトラヘドロン体の多モード回転密度、(ii) TIP4P 水での Ice XI のサイズと温度を変えたサンプリング。

Figure 1: Flow on SO(3) via the $S^{3}$ double cover: we first transform a system of Cartesian coordinates $\bm{x}$ via $T$ into the triple $(\bm{x}_{0},\bm{R},\bm{\Psi})$ , containing the global translation $\bm{x}_{0}$ , the global rotation $\bm{R}$ and its fixed inner degrees of freedom $\bm{\Psi

実験結果

リサーチクエスチョン

• RQ1 クォータニオン二重カバー上の滑らかで反転同値性を持つフローは、剛体の多モード回転密度を正確に表現できるか？
• RQ2 剛体正規化フローは分子結晶をどれくらいうまくサンプリングし、MDを基盤とする参照と比較して自由エネルギー差（ΔF）の正確な推定をどれくらい可能にするか？
• RQ3 内部DOFフローと組み合わせて結晶中の相互作用分子をモデル化できるか？
• RQ4 結晶系の訓練効率と密度精度に対する S3 経由の SO(3) 表現の影響は何か？

主な発見

• 射影凸勾配写像と反転同値性を持つ対称化モビウス変換は、多モード回転密度を忠実に再構成でき、外部場下のテトラヘドロン試験でアフィンなクォータニオン層を上回る。
• TIP4P水の Ice XI サンプリングでは、流れベースの LFEP が分子あたりの ΔF を、さまざまな系サイズと対象温度で MBAR 参照結果とほぼ一致する。
• 基底フローで写像されたエネルギーと目標 MD 分布との重なりが高く、リウェイトにより RDF を高忠実度で回復する。
• 一つの MD 軌道を訓練と LFEP 評価に用いることで自由エネルギー差の効率的推定を実現し、階層的な MD シミュレーションの必要性を減らす。
• 外部姿勢と内部 DOF の独立性をサポートし、より大きな結晶系向けの置換対称性を考慮したスケーラブルなフロー設計と互換性がある。

Figure 2: Tetrahedron in an external field: a) each colored bead is attracted with the same force towards the external point $\bm{c}$ according to the potential defined in Eq. ( 15 ). b) Density of rotational degree of freedom. Rotations are represented as unit-quaternions $\bm{q}=(x,y,z,w)$ . First