[论文解读] Efficient and Accurate Physics-aware Multiplex Graph Neural Networks for 3D Small Molecules and Macromolecule Complexes

Shuo Zhang, Yang Liu+1 more|arXiv (Cornell University)|Jun 6, 2022

Machine Learning in Materials Science被引用 6

一句话总结

该论文提出PaxNet，一种物理感知的多路图神经网络，通过使用独立的消息传递方案分别建模局部和非局部分子相互作用，实现了对三维小分子及大分子复合物中标量和矢量性质的高效且精确预测。与最先进基线相比，PaxNet在小分子上将预测误差降低15%，内存使用减少73%；在大分子复合物上将推理时间减少85%，内存使用减少33%。

ABSTRACT

Recent advances in applying Graph Neural Networks (GNNs) to molecular science have showcased the power of learning three-dimensional (3D) structure representations with GNNs. However, most existing GNNs suffer from the limitations of insufficient modeling of diverse interactions, computational expensive operations, and ignorance of vectorial values. Here, we tackle these limitations by proposing a novel GNN model, Physics-aware Multiplex Graph Neural Network (PaxNet), to efficiently and accurately learn the representations of 3D molecules for both small organic compounds and macromolecule complexes. PaxNet separates the modeling of local and non-local interactions inspired by molecular mechanics, and reduces the expensive angle-related computations. Besides scalar properties, PaxNet can also predict vectorial properties by learning an associated vector for each atom. To evaluate the performance of PaxNet, we compare it with state-of-the-art baselines in two tasks. On small molecule dataset for predicting quantum chemical properties, PaxNet reduces the prediction error by 15% and uses 73% less memory than the best baseline. On macromolecule dataset for predicting protein-ligand binding affinities, PaxNet outperforms the best baseline while reducing the memory consumption by 33% and the inference time by 85%. Thus, PaxNet provides a universal, robust and accurate method for large-scale machine learning of molecules. Our code is available at https://github.com/zetayue/Physics-aware-Multiplex-GNN.

研究动机与目标

• 为解决现有GNN在建模多样化分子相互作用（尤其是3D结构）方面的局限性。
• 通过分离局部与非局部相互作用建模，降低计算成本，避免昂贵的角度相关计算。
• 实现对矢量性质（如偶极矩和受力）的精确预测，而不仅限于标量性质。
• 在资源受限条件下，提升对蛋白质-配体等大分子复合物的效率与可扩展性。
• 开发一种统一、稳健且精确的大规模分子表征学习框架。

提出的方法

• PaxNet采用多路图结构，其中每个'路'代表一种不同类型的相互作用（如局部共价、非局部非共价），支持按相互作用类型独立进行消息传递。
• 引入受物理启发的设计，将局部（一跳和两跳）与非局部（全局）消息传递分离，以更准确高效地建模相互作用。
• 通过原子位置和消息权重导出的几何向量，直接预测矢量性质（如偶极矩）。
• 在消息传递方案之间应用可学习注意力机制，动态加权不同相互作用类型的贡献。
• 采用改进的消息传递方案，通过使用物理感知近似替代角度计算，避免O(Nk²)复杂度。
• 通过为每个原子关联一个向量并利用消息传递更新该向量，支持标量和矢量分子性质的端到端学习。

实验结果

研究问题

• RQ1 是否可通过为局部与非局部相互作用分别设计独立消息传递的多路GNN框架，提升对3D分子性质的预测精度？
• RQ2 与统一处理所有相互作用的传统GNN相比，分离相互作用建模在计算成本上如何降低？
• RQ3 GNN模型在多大程度上能以比仅预测标量的模型更高的保真度预测矢量性质（如偶极矩）？
• RQ4 PaxNet如何在显著减少内存和推理时间的同时，保持在大分子复合物上的高性能？
• RQ5 各组件（如注意力机制、路间通信、多跳消息传递）对整体模型性能的贡献如何？

主要发现

• 在QM9数据集上，PaxNet将偶极矩预测的平均绝对误差（MAE）降低至0.0108 D，优于最佳基线（0.021 D）超过50%。
• 在QM9上，PaxNet相比最佳基线模型实现15%的更低预测误差和73%的更少内存使用。
• 在PDBbind数据集上，PaxNet相比最佳基线将内存消耗减少33%，推理时间减少85%，同时实现更优性能。
• 消融实验表明，移除注意力模块后，标量目标的MAE最多增加4%，矢量目标最多增加3%，证明其重要性。
• 移除路间通信或消息传递方案后，标量目标性能最多下降28%，矢量目标性能最多下降102%，证实集成设计的必要性。
• 在局部消息传递中使用两跳几何信息（θ）相比基线，使U₀性能提升64%，μ性能提升111%，凸显更丰富几何编码的价值。