[论文解读] A new perspective on building efficient and expressive 3D equivariant graph neural networks

Weitao Du, Yuanqi Du+6 more|arXiv (Cornell University)|Apr 7, 2023

Machine Learning in Materials Science被引用 7

一句话总结

本文通过分析从局部到全局的几何表达能力，提出了一种基于3D同构层次的新型理论框架，用于设计高效且表达能力强的3D等变图神经网络。该工作提出了LEFTNet，通过整合局部子结构编码（LSE）与帧变换编码（FTE），在分子性质预测任务中实现了最先进性能，平均在标量属性上超越先前模型20%，在向量属性上提升16%。

ABSTRACT

Geometric deep learning enables the encoding of physical symmetries in modeling 3D objects. Despite rapid progress in encoding 3D symmetries into Graph Neural Networks (GNNs), a comprehensive evaluation of the expressiveness of these networks through a local-to-global analysis lacks today. In this paper, we propose a local hierarchy of 3D isomorphism to evaluate the expressive power of equivariant GNNs and investigate the process of representing global geometric information from local patches. Our work leads to two crucial modules for designing expressive and efficient geometric GNNs; namely local substructure encoding (LSE) and frame transition encoding (FTE). To demonstrate the applicability of our theory, we propose LEFTNet which effectively implements these modules and achieves state-of-the-art performance on both scalar-valued and vector-valued molecular property prediction tasks. We further point out the design space for future developments of equivariant graph neural networks. Our codes are available at \url{https://github.com/yuanqidu/LeftNet}.

研究动机与目标

• 通过分析几何同构的局部到全局特性，建立一个理论框架，用于评估3D等变GNN的表达能力。
• 识别出实现3D GNN高表达能力与计算效率的关键架构组件——LSE与FTE。
• 证明通过不变几何消息与等变聚合相结合，可有效捕捉全局3D结构。
• 为未来3D几何深度学习模型提供模块化的设计范式。

提出的方法

• 提出3D同构的局部层次结构——树、三角形与子图，用于分析GNN在捕捉几何结构方面的表达能力。
• 引入局部子结构编码（LSE）机制，利用微分几何中的标量化技术，实现对局部3D几何模式的不变表示。
• 开发帧变换编码（FTE），用于建模消息传递过程中局部坐标系的变换，从而实现全局几何表征。
• 设计LEFTNet为模块化架构，将LSE与FTE与等变张量更新相结合，以保持SE(3)对称性。
• 采用基于帧的消息聚合机制，使用张量积或求和操作，灵感来源于高阶体势能模型。
• 在纤维丛上应用标量化技术，实现对局部几何特征的不变编码，这对LSE至关重要。

实验结果

研究问题

• RQ1 不变标量在编码3D几何模式方面具有多大的表达能力？
• RQ2 在消息传递中保持等变性是否真正对捕捉全局3D几何结构至关重要？
• RQ3 如何系统性地聚合局部几何信息，以对称性保持的方式表示全局3D结构？
• RQ4 哪些架构组件能够同时实现3D等变GNN的高表达能力与计算效率？

主要发现

• LEFTNet在QM9数据集上达到最先进性能，平均在12项分子性质预测任务中超越先前SOTA结果20%。
• 在rMD17数据集上，使用100倍力加权损失时，LEFTNet在8种分子中的7种上优于所有基线模型，平均提升16%。
• 在QM9数据集的两种数据划分下，LEFTNet在12项属性中的10项达到第二名或最佳表现，展现出强鲁棒性与泛化能力。
• 消融实验确认LSE与FTE均为性能关键，两者对最终结果均有显著贡献。
• 该框架可推广至非分子场景，通过切空间帧定义，有望扩展至流形上的几何网格图。