[论文解读] GemNet: Universal Directional Graph Neural Networks for Molecules

Johannes Gasteiger, Florian Becker+1 more|arXiv (Cornell University)|Jun 2, 2021

Computational Drug Discovery Methods被引用 127

一句话总结

GemNet 确立了球面表示在旋转不变/协变分子预测中的普适性，引入带有定向边嵌入的几何两跳消息传递，并在多个分子动力学数据集上实现了最先进的力预测。

ABSTRACT

Effectively predicting molecular interactions has the potential to accelerate molecular dynamics by multiple orders of magnitude and thus revolutionize chemical simulations. Graph neural networks (GNNs) have recently shown great successes for this task, overtaking classical methods based on fixed molecular kernels. However, they still appear very limited from a theoretical perspective, since regular GNNs cannot distinguish certain types of graphs. In this work we close this gap between theory and practice. We show that GNNs with spherical representations are indeed universal approximators for predictions that are invariant to translation, and equivariant to permutation and rotation. We then discretize such GNNs via directed edge embeddings and two-hop message passing, and incorporate multiple structural improvements to arrive at the geometric message passing neural network (GemNet). We demonstrate the benefits of the proposed changes in multiple ablation studies. GemNet outperforms previous models on the COLL, MD17, and OC20 datasets by 34%, 41%, and 20%, respectively, and performs especially well on the most challenging molecules. Our implementation is available online.

研究动机与目标

• 解决常规 GNN 在分子图表达能力上的不足，并实现旋转不变与等变预测。
• 证明球面表示在旋转不变和等变任务中的普适性。
• 开发一个实用的、定向的、基于边的两跳消息传递框架，结合几何信息以提升准确性和样本效率。

提出的方法

• 证明球面表示对 SO(3)-不变和置换等变函数的普适性。
• 将球面表示离散化为有向边嵌入，并对边之间的两跳消息传递进行建模。
• 引入带有原子间方向和夹角（包括 φ 和 θ）的几何消息传递，使用球面傅里叶-贝塞尔基底。
• 结合对称消息传递和高效双线性层以提高稳定性和计算效率。
• 以预先设定的缩放因子稳定激活方差，而非采用标准归一化。
• 开发 GemNet 架构，包含三种交互形式（两跳几何消息传递、单跳几何消息传递、原子自相互作用）。

实验结果

研究问题

• RQ1 球面表示是否可以在分子图神经网络中实现对旋转不变和等变预测的普适性？
• RQ2 如何将定向（基于边）的表示离散化，以实现一个实用且表达力强的分子 GNN？
• RQ3 哪些架构创新（对称消息、双线性层、方差稳定化）最能提升分子动力学数据集中能量与力的预测准确性？
• RQ4 与最先进模型在 COLL、MD17 和 OC20 的力和能量预测相比，GemNet 的表现如何？

主要发现

• GemNet 在各数据集上实现显著的误差降低，在 COLL、MD17 和 OC20 上相较于前一代模型具有显著优势。
• 带有定向边嵌入的两跳几何消息传递显著提升性能，尤其在具有挑战性的非平面分子上。
• 对称消息传递、双线性层和激活方差稳定化对准确性有实质性贡献，消融研究证实了每个组件的益处。
• 直接力预测可以加速训练和推理，但在数据集和任务上有 MAE 的权衡。
• GemNet-T 和 GemNet-Q 变体在计算成本与准确性之间提供不同权衡，GemNet-Q 在更复杂的组合数据集上带来更大收益。