[论文解读] SE(3) Equivariant Graph Neural Networks with Complete Local Frames

Weitao Du, He Zhang+5 more|arXiv (Cornell University)|Oct 26, 2021

Machine Learning in BioinformaticsBiochemistry, Genetics and Molecular Biology被引用 22

一句话总结

该论文提出ClofNet，一种计算高效的SE(3)-等变图神经网络，通过使用通过叉积构建的完整局部坐标系，实现对几何张量的忠实、无参考系表示。通过将高阶几何量投影到正交基上，该模型在多体动力学建模和分子构象生成任务中均实现了最先进性能，具备优异的样本效率和低等变性误差。

ABSTRACT

Group equivariance (e.g. SE(3) equivariance) is a critical physical symmetry in science, from classical and quantum physics to computational biology. It enables robust and accurate prediction under arbitrary reference transformations. In light of this, great efforts have been put on encoding this symmetry into deep neural networks, which has been shown to improve the generalization performance and data efficiency for downstream tasks. Constructing an equivariant neural network generally brings high computational costs to ensure expressiveness. Therefore, how to better trade-off the expressiveness and computational efficiency plays a core role in the design of the equivariant deep learning models. In this paper, we propose a framework to construct SE(3) equivariant graph neural networks that can approximate the geometric quantities efficiently. Inspired by differential geometry and physics, we introduce equivariant local complete frames to graph neural networks, such that tensor information at given orders can be projected onto the frames. The local frame is constructed to form an orthonormal basis that avoids direction degeneration and ensure completeness. Since the frames are built only by cross product operations, our method is computationally efficient. We evaluate our method on two tasks: Newton mechanics modeling and equilibrium molecule conformation generation. Extensive experimental results demonstrate that our model achieves the best or competitive performance in two types of datasets.

研究动机与目标

• 解决SE(3)-等变深度学习在3D多体系统中表达能力与计算效率之间的平衡挑战。
• 克服现有等变图神经网络依赖不完整径向坐标系所导致的方向退化问题。
• 在无参考系条件下，实现对几何量（如向量、二面角力）的无损表示，且保持等变性。
• 构建一个支持任意阶张量的框架，无需高维提升或复杂张量运算即可保持群等变性。
• 在物理建模和分子构象生成任务中，展示出强大的泛化能力和数据效率。

提出的方法

• 通过粒子向量的叉积构造局部完整坐标系，在三维空间中形成正交基，确保完备性并避免方向退化。
• 通过标量化模块将几何张量（标量、向量、高阶张量）投影到这些局部坐标系上，将其转换为标量系数以用于消息传递。
• 在标量系数上使用标准图神经网络进行消息传递，由于坐标系的等变性，模型保持SE(3)等变性。
• 仅通过叉积操作构建坐标系，相比依赖Clebsch-Gordan分解或球谐函数的方法，计算效率更高。
• 通过递归应用叉积并保持正交性，将坐标系构造方法扩展至处理高阶张量输入。
• 将该框架集成到基于分数的生成建模设置中，用于分子构象生成，采用退火Langevin动力学进行采样。

实验结果

研究问题

• RQ1 能否设计一种计算高效的SE(3)-等变图神经网络，避免高维提升或张量分解的高昂开销？
• RQ2 与仅依赖径向或不完整坐标系的方法相比，使用完整局部坐标系是否能提升几何表示的表达能力和准确性？
• RQ3 该方法在物理系统中的分布外和低数据场景下，泛化能力如何？
• RQ4 该模型是否能在保持强等变性的同时，在动力学建模和分子构象生成任务中均实现最先进性能？
• RQ5 与替代的等变函数空间表示相比，通过叉积构造坐标系在效率和稳定性方面表现如何？

主要发现

• ClofNet在POS数据集上表现最佳，插值MSE为0.138，外推MSE为2.502，显著优于Radial Field（7.665）和EGNN（6.449）。
• 模型保持了极低的等变性误差（$\Delta_{EQ} = 4.13 \cdot 10^{-6}$），表明SE(3)等变性得到出色的数值保持。
• 在GEOM-QM9和GEOM-Drugs数据集上，ClofNet生成了高质量的分子构象，RMSD较低，可视化结果展示了最佳对齐结构。
• 即使仅使用0.4%的训练数据，ClofNet仍取得具有竞争力的结果，凸显其卓越的样本效率。
• 消融实验确认，完整局部坐标系的使用对建模复杂几何量（如二面角力）至关重要。
• 通过仅使用叉积操作构建坐标系，避免了昂贵的张量积分解，验证了该框架的计算高效性。