[论文解读] Equivariant Graph Attention Networks for Molecular Property Prediction

Tuan Le, Frank Noé+1 more|arXiv (Cornell University)|Feb 20, 2022

Computational Drug Discovery Methods被引用 6

一句话总结

该论文提出了一种新型图神经网络架构——等变图注意力网络（EQGAT），通过利用笛卡尔坐标和一种空间-内容感知的注意力机制，实现对分子表示学习的旋转与平移等变性。该模型在小分子的量子性质预测任务上达到最先进性能，在蛋白质-配体结合亲和力等大生物分子任务上也表现出竞争力，同时计算效率高，推理速度优于基线等变图神经网络。

ABSTRACT

Learning and reasoning about 3D molecular structures with varying size is an emerging and important challenge in machine learning and especially in drug discovery. Equivariant Graph Neural Networks (GNNs) can simultaneously leverage the geometric and relational detail of the problem domain and are known to learn expressive representations through the propagation of information between nodes leveraging higher-order representations to faithfully express the geometry of the data, such as directionality in their intermediate layers. In this work, we propose an equivariant GNN that operates with Cartesian coordinates to incorporate directionality and we implement a novel attention mechanism, acting as a content and spatial dependent filter when propagating information between nodes. We demonstrate the efficacy of our architecture on predicting quantum mechanical properties of small molecules and its benefit on problems that concern macromolecular structures such as protein complexes.

研究动机与目标

• 开发一种图神经网络，通过在分子图中保持对三维几何结构信息的等变性，实现对旋转和平移的不变性。
• 通过一种新颖的注意力机制，将方向性（空间）和原子性（内容）信息整合到消息传递过程中，以提升分子性质预测性能。
• 在包括蛋白质和蛋白质-配体复合物在内的大分子体系上，实现可扩展且高效的模型学习，超越小分子的范畴。
• 在多样化的分子预测基准上实现最先进或具有竞争力的性能，同时保持较低的计算成本。

提出的方法

• 模型直接在笛卡尔坐标上运行，通过在消息传递框架中显式应用变换规则，确保对三维旋转和平移的等变性。
• 提出一种新型图注意力机制，计算依赖于相对位置（空间）和节点特征（内容）的边注意力权重，实现自适应、几何感知的消息传递。
• 采用标量和向量表示（0型和1型张量）以保持几何保真度，避免了张量场网络中复杂的基函数计算。
• 通过等变更新规则执行消息传递，确保特征在空间对称变换下保持其变换行为。
• 模型在消息传递神经网络（MPNN）框架内实现，采用基于注意力的聚合机制以增强表示学习能力。
• 通过直接在笛卡尔空间中操作，避免了昂贵的基函数构建，相比基于基函数的等变图神经网络，显著降低了浮点运算次数和内存占用。

实验结果

研究问题

• RQ1 一种同时编码空间与内容信息的注意力机制图神经网络，是否能在分子性质预测上优于标准图神经网络？
• RQ2 在笛卡尔空间中直接操作并显式实现等变性，是否在性能和效率上优于基于基函数的张量场网络？
• RQ3 等变图神经网络能否在蛋白质和蛋白质-配体复合物等大生物分子体系上实现有效扩展？
• RQ4 在多样化的分子基准上，EQGAT与最先进等变图神经网络（如PaiNN和E(n)-GNN）相比，在性能和推理速度方面表现如何？

主要发现

• EQGAT在小分子的量子力学性质预测任务上达到最先进性能，优于此前所有图神经网络在该基准上的表现。
• 在LBA（配体结合亲和力）数据集上，EQGAT在所有测试的图神经网络架构中表现最佳，尽管整体上3D卷积神经网络表现最优。
• EQGAT展现出卓越的计算效率，其在LBA、PSR和RSR数据集上的推理时间分别为23.67 ms、42.79 ms和54.80 ms，优于PaiNN和E(n)-GNN在所有基准上的表现。
• 该模型的推理速度优势归因于其直接在笛卡尔空间中的实现方式，避免了PaiNN中使用的内存密集型向量特征表示。
• PaiNN在生物大分子基准上表现欠佳，可能是因为其存储3D向量特征导致的高内存占用，且随图大小增长而性能显著下降。
• 作者指出，采用独立的配体和蛋白质编码器可能进一步提升性能，为未来研究指明了有前景的方向。