[论文解读] 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)和置换等变性的同时，实现了快速、精确的采样与密度评估。通过将原子坐标投影到通过学习得到的增强维度所构成的SE(3)不变基底上，该方法在玻尔兹曼分布学习任务上达到了最先进性能，采样速度比CNFs和扩散模型快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)和置换等变性，且无法实现快速采样与精确密度评估的问题。
• 通过确保可 tractable 采样与密度评估，实现归一化流在分子系统中的能量基训练。
• 仅使用笛卡尔坐标和能量函数，学习丙氨酸二肽的完整玻尔兹曼分布。
• 提供一种可扩展的、等变的流架构，可泛化至分子图和粒子系统，且无需依赖内坐标。
• 证明SE(3)等变耦合流在采样速度和似然度方面可与连续归一化流和扩散模型相媲美或超越，且采样速度显著更快。

提出的方法

• 该方法通过在增强输入空间中引入在SE(3)下变换的辅助变量，扩展了标准耦合流，从而实现等变性。
• 在每一层中，原子位置通过学习到的变换投影到SE(3)-不变基底上，随后在该基底上应用标准归一化流层（例如单调有理二次样条）。
• 输出随后被投影回原始的SE(3)-等变空间，从而保持流的可逆性与雅可比行列式计算的可 tractable 性。
• 通过避免像CNFs或扩散模型那样的迭代采样，该架构被设计为保持快速采样与密度评估。
• 该方法支持仅使用未归一化的玻尔兹曼密度进行最大似然和能量基学习，从而无需依赖MD采样进行训练。
• 评估了不同投影策略（向量投影、笛卡尔投影、球面投影）以优化不变性与表达能力。

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

实验结果

研究问题

• RQ1 能否在保持快速采样与精确密度评估的前提下，使耦合流具备SE(3)和置换等变性？
• RQ2 该流能否仅使用笛卡尔坐标和能量函数，学习丙氨酸二肽的完整玻尔兹曼分布？
• RQ3 所提出的流是否可在DW4和LJ13等系统上实现能量基训练，而这些系统因采样速度过慢，导致CNFs和扩散模型不可行？
• RQ4 在相同训练设置下，该流与SE(3) CNFs和扩散模型相比，在采样速度和似然度方面表现如何？
• RQ5 当有效样本量较大时，重要性重加权是否能提升低密度区域样本的质量？

主要发现

• 笛卡尔投影E-ACF在QM9-位置完整数据集上实现了-188.6的最低负对数似然（NLL），优于其他变体，并与最先进模型相当。
• 通过重要性重加权，球面投影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系统每样本推理时间为6.3–7.5秒，LJ13系统为42.7–45.0秒。
• 首次在完全等变流中实现了DW4和LJ13系统的能量基训练，因为该方法可实现快速密度评估与采样，无需依赖MD生成数据。
• 该方法是首个仅使用笛卡尔坐标、不依赖内坐标或分子图依赖关系，即学习到丙氨酸二肽完整玻尔兹曼分布的方法。

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.