[论文解读] Multi-task graph neural networks for simultaneous prediction of global and atomic properties in ferromagnetic systems

Massimiliano Lupo Pasini, Zhang, Pei+2 more|arXiv (Cornell University)|Feb 4, 2022

Machine Learning in Materials Science被引用 10

一句话总结

本文提出HydraGNN，一种多任务图神经网络，通过共享的卷积主干网络同时预测铁磁性FePt合金中的全局性质（如混合焓）和原子性质（如电荷转移、磁矩），并采用任务特定的输出头。该模型在预测精度上与单任务学习相当，同时相比训练独立模型将训练成本降低两倍以上，利用各性质间的物理相关性提升了数据效率。

ABSTRACT

We introduce a multi-tasking graph convolutional neural network, HydraGNN, to simultaneously predict both global and atomic physical properties and demonstrate with ferromagnetic materials. We train HydraGNN on an open-source ab initio density functional theory (DFT) dataset for iron-platinum (FePt) with a fixed body centered tetragonal (BCT) lattice structure and fixed volume to simultaneously predict the mixing enthalpy (a global feature of the system), the atomic charge transfer, and the atomic magnetic moment across configurations that span the entire compositional range. By taking advantage of underlying physical correlations between material properties, multi-task learning (MTL) with HydraGNN provides effective training even with modest amounts of data. Moreover, this is achieved with just one architecture instead of three, as required by single-task learning (STL). The first convolutional layers of the HydraGNN architecture are shared by all learning tasks and extract features common to all material properties. The following layers discriminate the features of the different properties, the results of which are fed to the separate heads of the final layer to produce predictions. Numerical results show that HydraGNN effectively captures the relation between the configurational entropy and the material properties over the entire compositional range. Overall, the accuracy of simultaneous MTL predictions is comparable to the accuracy of the STL predictions. In addition, the computational cost of training HydraGNN for MTL is much lower than the original DFT calculations and also lower than training separate STL models for each property.

研究动机与目标

• 开发一种数据高效的深度学习模型，用于在铁磁体系中同时预测全局与原子材料性质。
• 通过利用混合焓、电荷转移和磁矩等相互依赖性质之间的物理相关性，降低计算成本与数据需求。
• 设计一种统一的神经网络架构，以单个多头模型替代多个单任务模型，提升训练效率。
• 在覆盖FePt全组成范围（固定BCT结构与体积）的DFT生成数据集上验证该模型。
• 证明基于图神经网络的多任务学习可有效捕捉非平衡材料中复杂的物理关系。

提出的方法

• HydraGNN采用共享的图卷积神经网络（GCNN）主干网络，从以原子为节点、化学键为边的图结构中提取局部原子间特征。
• 在共享特征提取之后，使用任务特定的输出头分别预测不同物理性质：混合焓（全局）、原子电荷转移和原子磁矩（按原子）。
• 模型通过端到端的多任务学习（MTL）进行训练，所有损失函数的梯度联合反向传播，以提升泛化能力与特征共享效果。
• 该架构支持灵活的图卷积层聚合策略，可根据数据的物理特性进行定制化调整。
• 模型在固定体心四方（BCT）结构与体积的FePt DFT数据集上进行训练，覆盖从Fe富集到Pt富集的所有组成。
• 通过12次独立训练运行（随机权重初始化）确保统计稳健性并实现可靠的误差估计。

实验结果

研究问题

• RQ1 单一多任务图神经网络架构是否能有效预测铁磁材料中的全局与原子性质？
• RQ2 与单任务学习相比，采用共享表示的多任务学习在预测精度与数据效率方面有何提升？
• RQ3 材料性质间（如焓变、电荷转移、磁矩）的物理相关性在低数据条件下在多大程度上增强模型泛化能力？
• RQ4 训练多任务模型与为每个性质分别训练单任务模型相比，计算成本如何？
• RQ5 该模型是否能在不依赖显式数据增强或迁移学习的情况下，泛化至FePt的整个组成范围？

主要发现

• HydraGNN在混合焓预测上的测试均方根误差（RMSE）为1.02×10⁻² ± 1.16×10⁻³ eV/atom，与单任务学习（STL）模型相当。
• 电荷转移预测的RMSE为5.94×10⁻³ ± 4.39×10⁻⁴ eV/atom，磁矩预测的RMSE为8.77×10⁻³ ± 3.18×10⁻⁴ μB，与STL基线模型精度一致。
• 同时预测三项性质的多任务模型训练时间比分别训练三个单任务模型快2.2倍以上。
• 与原始DFT计算相比，该模型将计算成本降低了数百倍，实现了材料性质的快速筛选。
• 共享的卷积层有效学习到与局部原子间相互作用相关的特征，并在所有预测任务间实现有效迁移。
• 通过联合训练多个相关性质，提升了数据效率并降低了不确定性；但若引入过多无关性质，可能导致性能下降。