[论文解读] Graph Nets for Partial Charge Prediction

Yuanqing Wang, Josh Fass+3 more|arXiv (Cornell University)|Sep 17, 2019

Machine Learning in Materials Science参考文献 13被引用 16

一句话总结

该论文提出了一种图神经网络（Graph Nets），用于在分子中实现高精度且快速的局部电荷预测，其基于可微分、置换等变的图神经网络架构，并在从密度泛函理论（DFT）获得的电荷数据上进行训练。该模型在与DFT参考值对比时达到0.0223 e的均方根误差（RMSE），相比传统方法提速超过500倍，且在不同大小的分子中保持高精度，通过二次规划方法强制实现全局电荷守恒。

ABSTRACT

Atomic partial charges are crucial parameters for Molecular Dynamics (MD) simulations, molecular mechanics calculations, and virtual screening, as they determine the electrostatic contributions to interaction energies. Current methods for calculating partial charges, however, are either slow and scale poorly with molecular size (quantum chemical methods) or unreliable (empirical methods). Here, we present a new charge derivation method based on Graph Nets---a set of update and aggregate functions that operate on molecular topologies and propagate information thereon---that could approximate charges derived from Density Functional Theory (DFT) calculations with high accuracy and an over 500-fold speed up.

研究动机与目标

• 开发一种快速、准确的预测小分子中原子局部电荷的方法，以替代计算缓慢的量子化学方法或不可靠的经验方法。
• 解决分子图学习中针对原子级回归任务缺乏可扩展、可微分模型的问题。
• 通过确保预测局部电荷之和等于分子总电荷，实现物理一致性。
• 通过提供可微分、端到端可训练的架构，实现与力场优化的集成。
• 在保持预测精度与体系大小无关的前提下，实现对大规模分子数据集的可扩展性。

提出的方法

• 模型采用完整的全连接图网络架构，其学习更新函数和聚合函数作用于节点、边和全局属性。
• 节点特征通过原子元素的一次性编码初始化，并经由前馈网络生成初始节点嵌入。
• 边特征通过键级（对共振结构取平均）初始化，并通过独立的前馈网络处理。
• 通过将最终电荷预测表述为可计算的二次规划问题，强制实现全局电荷守恒，使原子电荷之和与分子总电荷一致。
• 在来自ChEMBL的350,259个分子数据集上进行训练，目标为DFT计算得到的电荷，使用ε=4的介电常数。
• 框架基于TensorFlow 2.0实现，支持分布式并行推理，代码已以MIT许可证开源。

实验结果

研究问题

• RQ1 可微分、基于图的深度学习模型是否能在局部电荷预测中达到与DFT相当的精度，同时显著快于量子化学方法？
• RQ2 随着分子尺寸增大，模型性能如何变化，特别是在预测误差方面？
• RQ3 图神经网络架构能否以可微分且可扩展的方式强制实现全局电荷守恒？
• RQ4 模型是否能在不依赖基于规则的指纹的前提下，对包括重原子和卤素原子在内的多种元素表现出良好的泛化能力？
• RQ5 由于其可微分性和端到端训练能力，该模型是否可被集成到力场优化流程中？

主要发现

• 在包含299,811个原子的保留测试集上，模型整体RMSE为0.0223 e，95%置信区间为[0.0221, 0.0225] e。
• 对于氢原子，模型R²达0.9744，RMSE为0.0144 e，表明对这一常见元素具有极高预测精度。
• 模型在不同大小的分子中保持一致性能，大型体系（最多63个原子）的误差无显著增加。
• 与Bleiziffer等人提出的随机森林基线相比，本模型在精度和可微性方面均更优，支持与力场参数的联合优化。
• 该方法比AM1-BCC快约500倍，单分子推理时间约为0.03秒，支持高通量虚拟筛选。
• 模型在碳（R²=0.9932，RMSE=0.0217 e）和硫（R²=0.9935，RMSE=0.0524 e）上表现优异，但对碘的性能下降（R²=0.2819，RMSE=0.1948 e），可能由于训练数据有限。