[论文解读] Spherical Message Passing for 3D Graph Networks

Yi Liu, Limei Wang+4 more|arXiv (Cornell University)|Feb 9, 2021

Advanced Graph Neural Networks参考文献 47被引用 31

一句话总结

提出球面信息传播（SMP）用于三维分子图，以高效利用距离、角度和扭转信息，进而得到 SphereNet，在若干分子基准测试上实现最先进或具有竞争力的性能，并具备更高的效率。

ABSTRACT

We consider representation learning of 3D molecular graphs in which each atom is associated with a spatial position in 3D. This is an under-explored area of research, and a principled message passing framework is currently lacking. In this work, we conduct analyses in the spherical coordinate system (SCS) for the complete identification of 3D graph structures. Based on such observations, we propose the spherical message passing (SMP) as a novel and powerful scheme for 3D molecular learning. SMP dramatically reduces training complexity, enabling it to perform efficiently on large-scale molecules. In addition, SMP is capable of distinguishing almost all molecular structures, and the uncovered cases may not exist in practice. Based on meaningful physically-based representations of 3D information, we further propose the SphereNet for 3D molecular learning. Experimental results demonstrate that the use of meaningful 3D information in SphereNet leads to significant performance improvements in prediction tasks. Our results also demonstrate the advantages of SphereNet in terms of capability, efficiency, and scalability. Our code is publicly available as part of the DIG library (https://github.com/divelab/DIG).

研究动机与目标

• 动机并形式化地完成使用球坐标中的相对三维信息对三维分子结构的完整识别。
• 开发 SMP，以实现包含 1 路几何信息的高效、可扩展的消息传递。
• 引入 SphereNet，将三维几何编码为具有物理意义的表示，并在标准分子基准上进行评估。
• 评估 SMP 相对于完整的 2-hop 方案在完备性与效率之间的权衡。

提出的方法

• 用球坐标系（SCS）定义3D分子图，并证明完备性需要距离、角度和扭转。
• 提出 SMP，其中边、节点和图的更新函数（φ^e、φ^v、φ^u）通过局部性驱动的 ρ^{p→e} 和 ρ^{e→v} 聚合整合3D信息。
• 使用基于边的1-hop信息结合参考平面结构来计算扭转角，以使复杂度保持为 O(nk^2)。
• 论证 SMP 在区分大多数3D结构方面近似完备，并通过应用物理意义的基函数展开为 SphereNet。
• 通过球贝塞尔函数和球谐函数定义 SphereNet 表示 Ψ(d,θ,φ)，以确保平移/旋转不变性并整合到消息传递块中。
• 将 SMP 与先前的方法（SchNet、DimeNet）在 SMP 框架内视为特殊情况进行关联。

实验结果

研究问题

• RQ1 是否能够在可扩展的GNN框架内，使用一组完整的相对几何特征（距离、角度、扭转）有效地识别3D分子图？
• RQ2 球面信息传播是否在保持表达能力的同时，提供一个有原则的、高效的替代边基2-hop完整表示的方法？
• RQ3 SphereNet 表示 Ψ(d,θ,φ) 是否在与现有基线相比的标准分子基准上带来改进的预测性能？
• RQ4 SMP/SphereNet 如何在大规模分子数据集上平衡准确性、效率和可扩展性？

主要发现

• SphereNet 在 OC20 IS2RE 的 energy MAE 和 EwT 各分割上均表现出竞争力或最佳性能，且对基线有显著的平均改进。
• 在 QM9 上，SphereNet 设立了 std. MAE 的新状态，且在多种性质上取得了强健结果。
• MD17 实验表明，在可比 WoFE 设置下，SphereNet 在力场 MAE 上优于 SchNet 和 DimeNet，并在使用更高的 WoFE 时接近 GemNet-T。
• 消融和效率分析表明，SMP 在准确性方面可与完整的 2-hop 表示相匹配，同时在运行时间上显著下降（在某些设置下约快4倍）。
• SphereNet 将训练复杂度降低到 O(nk^2)，同时维持近似完备的3D结构可辨别性，从而实现大规模分子学习。