[論文レビュー] SE(3) Equivariant Augmented Coupling Flows

Laurence I. Midgley, Vincent Stimper+4 more|arXiv (Cornell University)|Aug 20, 2023

Computer Graphics and Visualization TechniquesComputer Science被引用数 3

ひとこと要約

本稿では、分子系における回転・並進および置換対称性を保ちながら、高速かつ正確なサンプリングと密度推定を可能にするSE(3)不変拡張カップリングフロー（E-ACF）を提案する。原子座標を学習された拡張次元によるSE(3)不変基底に投影することで、ボルツマン分布学習において最先端の性能を達成し、CNF や拡散モデルと比較してサンプリング速度が10倍以上速く、DW4 や LJ13 システムにおけるエネルギーに基づく学習を可能にする。

ABSTRACT

Coupling normalizing flows allow for fast sampling and density evaluation, making them the tool of choice for probabilistic modeling of physical systems. However, the standard coupling architecture precludes endowing flows that operate on the Cartesian coordinates of atoms with the SE(3) and permutation invariances of physical systems. This work proposes a coupling flow that preserves SE(3) and permutation equivariance by performing coordinate splits along additional augmented dimensions. At each layer, the flow maps atoms' positions into learned SE(3) invariant bases, where we apply standard flow transformations, such as monotonic rational-quadratic splines, before returning to the original basis. Crucially, our flow preserves fast sampling and density evaluation, and may be used to produce unbiased estimates of expectations with respect to the target distribution via importance sampling. When trained on the DW4, LJ13, and QM9-positional datasets, our flow is competitive with equivariant continuous normalizing flows and diffusion models, while allowing sampling more than an order of magnitude faster. Moreover, to the best of our knowledge, we are the first to learn the full Boltzmann distribution of alanine dipeptide by only modeling the Cartesian positions of its atoms. Lastly, we demonstrate that our flow can be trained to approximately sample from the Boltzmann distribution of the DW4 and LJ13 particle systems using only their energy functions.

研究の動機と目的

• カートェシアン座標上で動作し、SE(3)および置換対称性を強制し、高速なサンプリングと密度推定を可能にする正規化フローの欠如に対処すること。
• トレーサブルなサンプリングと密度推定を保証することで、分子系における正規化フローのエネルギーに基づく学習を可能にすること。
• 内部座標や分子グラフ依存性に依存せずに、カートェシアン座標とエネルギー関数のみを用いてアラニンジペプチドの完全なボルツマン分布を学習すること。
• 内部座標に依存せず、分子グラフや粒子系に一般化可能なスケーラブルな不変フロー構造を提供すること。
• SE(3)不変カップリングフローが連続正規化フローおよび拡散モデルと同等またはそれを上回る性能を達成するが、サンプリングがはるかに高速であることを示すこと。

提案手法

• 標準カップリングフローを拡張し、SE(3)に作用する補助変数を含む拡張入力空間を導入することで、不変性を実現する。
• 各レイヤーで、原子位置が学習された変換を用いてSE(3)不変基底に投影され、標準的な正規化フロー層（例：単調な有理関数スプライン）が適用される。
• 出力は元のSE(3)不変空間に再投影され、フローの可逆性とヤコビアン行列式の取り扱いやすさが保持される。
• 反復的サンプリングを回避することで、CNF や拡散モデルとは異なり、高速なサンプリングと密度推定を実現する。
• 最大尤度学習およびエネルギーに基づく学習が可能であり、未正規化ボルツマン密度のみを用いてMDサンプルを必要としない。
• 不変性と表現力の最適化を図るため、異なる投影戦略（ベクトル投影、カートェシアン投影、球面投影）を評価する。

Figure 2: Inter-atomic distances for samples from the train-data and Spherical-proj E-ACF .

実験結果

リサーチクエスチョン

• RQ1 カップリングフローをSE(3)および置換対称性に不変に保ちつつ、高速なサンプリングと正確な密度推定を実現できるか？
• RQ2 カートェシアン座標とエネルギー関数のみを用いて、アラニンジペプチドの完全なボルツマン分布を学習できるか？
• RQ3 CNF や拡散モデルが遅いサンプリングのため不適切なDW4 や LJ13 システムにおいて、提案手法のエネルギーに基づく学習が可能か？
• RQ4 同じ学習設定下で、提案手法のサンプリング速度と尤度性能がSE(3)不変CNF や拡散モデルと比較してどうなるか？
• RQ5 有効サンプルサイズが大きい場合、特に分布の低密度領域において、重要度再重み付けがサンプル品質を向上させるか？

主な発見

• カートェシアン投影によるE-ACFは、QM9-位置データセット全体でNLL = -188.6を達成し、他のバリエーションや最先端モデルを上回る性能を示した。
• 球面投影E-ACFモデルは、重要度再重み付けにより、800 KにおけるアラニンジペプチドのKullback-Leibler発散（KLD）を1.88×10⁻³から1.39×10⁻³に低下させ、低密度領域のサンプル品質を顕著に向上させた。
• 300 Kにおけるアラニンジペプチドでは、球面投影E-ACFが再重み付け後にKLD = 1.97×10⁻³を達成し、内部座標フロー（KLD = (7.57±3.80)×10⁻³）を上回った。
• E-ACFモデルは、SE(3)連続正規化フローおよび拡散モデルと比較して10倍以上高速にサンプリング可能であり、DW4では1サンプルあたり6.3–7.5秒、LJ13では42.7–45.0秒の推論時間であった。
• MDシミュレーションデータを必要とせず、完全に不変なフローとして初めてDW4およびLJ13におけるエネルギーに基づく学習が可能になった。
• 本手法は、内部座標や分子グラフ依存性に依存せず、カートェシアン座標のみを用いてアラニンジペプチドの完全なボルツマン分布を学習した最初の手法である。

Figure 3: Ramachandran plots, i.e. marginal distribution of the dihedral angles $\phi$ and $\psi$ (see Sec. C.3.2 ), obtained with MD (ground truth) and various models.