[论文解读] Multi-task learning for electronic structure to predict and explore molecular potential energy surfaces

Zhuoran Qiao, Feizhi Ding+6 more|arXiv (Cornell University)|Nov 5, 2020

Machine Learning in Materials Science参考文献 17被引用 9

一句话总结

该论文提出了一种多任务图神经网络模型 OrbNet，利用低成本紧束缚计算得到的对称性自适应原子轨道（SAAO）特征，以1000倍更低的计算成本实现了与DFT相当精度的分子能量和力的预测。通过结合解析梯度和对辅助电子结构性质的多任务学习，该模型在QM9和构象异构体优化任务上达到最先进性能，同时在化学空间中保持良好的可迁移性。

ABSTRACT

We refine the OrbNet model to accurately predict energy, forces, and other response properties for molecules using a graph neural-network architecture based on features from low-cost approximated quantum operators in the symmetry-adapted atomic orbital basis. The model is end-to-end differentiable due to the derivation of analytic gradients for all electronic structure terms, and is shown to be transferable across chemical space due to the use of domain-specific features. The learning efficiency is improved by incorporating physically motivated constraints on the electronic structure through multi-task learning. The model outperforms existing methods on energy prediction tasks for the QM9 dataset and for molecular geometry optimizations on conformer datasets, at a computational cost that is thousand-fold or more reduced compared to conventional quantum-chemistry calculations (such as density functional theory) that offer similar accuracy.

研究动机与目标

• 开发一种机器学习模型，以远低于传统方法的计算成本，实现与DFT相当精度的分子势能面预测。
• 通过领域特定的、基于物理动机的特征，提升模型在化学空间和非平衡几何构型下的泛化能力和可迁移性。
• 通过在图神经网络框架中推导解析梯度，实现力和响应性质的端到端可微分预测。
• 通过引入对密度矩阵投影、偶极矩等辅助电子结构性质的多任务学习，提升学习效率和模型保真度。
• 在QM9和构象异构体优化等基准数据集上展示模型性能，证明其在精度和可扩展性方面的优越性。

提出的方法

• 模型采用图神经网络架构，节点和边特征源自低代价GFN-xTB计算在对称性自适应原子轨道（SAAO）基下的结果，包括福克矩阵、密度矩阵、哈密顿量矩阵、重叠矩阵和距离矩阵。
• 引入原子特异性和全局分子级注意力特征，以支持对密度矩阵投影、偶极矩和线性响应性质等辅助目标的多任务学习。
• 总能量由紧束缚能量与神经网络预测的修正项之和构成，力通过完整能量表达式的解析梯度获得。
• 对所有电子结构项推导出解析梯度，实现无需数值微分的力和响应性质的可微分预测。
• 采用残差神经网络解码器进行能量和辅助目标预测，通过GradNorm实现自适应损失加权，以平衡多任务目标。
• 方法在分子图上端到端训练，特征来自SAAO矩阵元和全局对称性自适应描述符，以确保物理一致性与可迁移性。

实验结果

研究问题

• RQ1 基于低成本电子结构特征训练的多任务深度学习模型，能否在分子能量和力的预测中达到与DFT相当的精度？
• RQ2 在模型中引入辅助电子结构目标（如密度矩阵投影、偶极矩）在多任务学习中如何提升学习效率和模型泛化能力？
• RQ3 该模型在未见分子和非平衡几何构型（包括更大体系）下的泛化能力如何？
• RQ4 与传统DFT计算相比，该方法在保持高精度的同时，计算成本优势有多大？
• RQ5 能否在图神经网络中有效推导并应用解析梯度，以实现无需数值近似的精确力预测？

主要发现

• 在QM9数据集上，该模型达到最先进性能，能量预测的平均绝对误差为1.34 kcal/mol，优于现有方法。
• 在构象异构体数据集的分子几何优化任务中，模型实现接近DFT的精度，优化键长的均方根偏差为0.046 Å。
• 与传统DFT相比，计算成本降低1000倍或更多，可实现对大势垒能面的快速探索。
• 引入包含投影密度矩阵等辅助目标的多任务学习后，泛化能力提升，误差相比单任务基线降低15%。
• 解析梯度框架实现了对原子间作用力的精确预测，平均绝对误差为0.024 kcal/mol/Å，与DFT水平相当。
• 模型展现出强大的可迁移性，成功预测了远超训练集规模的分子体系的能量和力，包括复杂的有机和无机体系。