[论文解读] E(n) Equivariant Normalizing Flows for Molecule Generation in 3D

Victor García Satorras, Emiel Hoogeboom+3 more|arXiv (Cornell University)|May 19, 2021

Machine Learning in Materials Science被引用 17

一句话总结

该论文提出E(n)等变归一化流（E-NFs），一种基于似然的生成模型，通过将E(n)等变图神经网络整合到连续时间归一化流中，学习生成3D分子。该方法在QM9数据集及3D粒子系统上实现了最先进的对数似然性能，是首个具备精确E(n)等变性的基于似然的3D分子生成模型。

ABSTRACT

This paper introduces a generative model equivariant to Euclidean symmetries: E(n) Equivariant Normalizing Flows (E-NFs). To construct E-NFs, we take the discriminative E(n) graph neural networks and integrate them as a differential equation to obtain an invertible equivariant function: a continuous-time normalizing flow. We demonstrate that E-NFs considerably outperform baselines and existing methods from the literature on particle systems such as DW4 and LJ13, and on molecules from QM9 in terms of log-likelihood. To the best of our knowledge, this is the first likelihood-based deep generative model that generates molecules in 3D.

研究动机与目标

• 开发一种深度生成模型，能够生成具有欧几里得变换（旋转与平移）精确不变性的3D分子结构。
• 将E(n)等变消息传递神经网络整合到连续时间归一化流框架中，实现可逆且等变的生成。
• 在3D分子与粒子系统上实现最先进的对数似然性能，包括QM9以及DW4和LJ13等基准数据集。
• 建立一个基于似然的3D分子生成模型，弥补尽管3D生成研究日益受到关注但该领域仍存在空白的现状。

提出的方法

• 该模型使用E(n)等变图神经网络作为连续归一化流中的神经网络组件，确保对旋转和平移的等变性。
• 通过求解常微分方程（ODE）对由GNN学习得到的向量场进行积分，定义流，实现可逆变换。
• 通过似然最大化进行端到端训练，对数密度的变化通过雅可比矩阵梯度的迹来计算。
• 通过构建基于等变GNN的流向量场来处理原子位置与特征，从而保持精确的E(n)等变性。
• 在3D粒子系统与分子上进行训练，使用潜在空间采样实现生成。
• 该方法支持生成与似然评估，使与其它基于似然的模型可直接比较。

实验结果

研究问题

• RQ1 能否设计一种归一化流模型，使其在3D分子生成中对欧几里得对称性（E(n)）具有精确等变性？
• RQ2 与非等变基线相比，E(n)等变性在3D分子与粒子系统上的对数似然与生成质量方面有何提升？
• RQ3 是否能够实现基于似然的模型在3D分子生成中达到最先进性能，同时保持精确的几何不变性？
• RQ4 将等变GNN整合到连续归一化流框架中，是否可行用于可逆生成建模？

主要发现

• E-NFs在QM9数据集上实现了最先进的对数似然性能，优于现有3D分子生成方法。
• 在DW4与LJ13等基准3D粒子系统上，该模型在对数似然方面表现出色。
• E-NFs是首个保持精确E(n)等变性的基于似然的深度生成模型，用于3D分子生成。
• 将E(n)等变GNN整合到连续归一化流中，实现了高质量生成与精确似然计算的统一。