SE(3) Equivariant Augmented Coupling Flows

要約

正規化フローを結合すると、高速サンプリングと密度評価が可能になり、物理システムの確率的モデリングに最適なツールになります。
ただし、標準的な結合アーキテクチャでは、原子のデカルト座標上で動作するフローに SE(3) および物理システムの順列不変性を与えることができません。
この研究では、追加の拡張次元に沿って座標分割を実行することにより、SE(3) と順列等分散を保存する結合フローを提案しています。
各層で、フローは原子の位置を学習した SE(3) 不変基底にマッピングします。そこで、元の基底に戻る前に、単調有理 2 次スプラインなどの標準的なフロー変換を適用します。
重要なことは、私たちのフローは高速なサンプリングと密度評価を維持しており、重要度サンプリングを通じてターゲット分布に関する期待の不偏推定値を生成するために使用できます。
DW4、LJ13、および QM9 位置データセットでトレーニングすると、私たちのフローは等変連続正規化フローや拡散モデルと競合し、1 桁以上高速なサンプリングが可能になります。
さらに、私たちの知る限り、原子のデカルト位置をモデル化するだけでアラニン ジペプチドの完全なボルツマン分布を学んだのは私たちが初めてです。
最後に、DW4 および LJ13 粒子システムのエネルギー関数のみを使用して、そのボルツマン分布から近似的にサンプリングするようにフローをトレーニングできることを示します。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google