[论文解读] Data-Augmentation for Graph Neural Network Learning of the Relaxed Energies of Unrelaxed Structures

Jason B. Gibson, Ajinkya C. Hire+1 more|arXiv (Cornell University)|Feb 28, 2022

Machine Learning in Materials Science被引用 5

一句话总结

本文提出了一种简单且基于物理原理的数据增强技术，可提升图神经网络（GNN）对未弛豫晶体结构的生成能预测性能——这类结构通常由晶体结构预测算法（CSPA）生成。通过在训练数据中加入经扰动的未弛豫结构版本，该方法在Nb-Sr-H氢化物测试集上将生成能预测的平均绝对误差（MAE）降低了66%（从251降至86 meV/atom），显著提升了模型在昂贵的第一性原理计算前筛选不稳定结构的能力。

ABSTRACT

Computational materials discovery has continually grown in utility over the past decade due to advances in computing power and crystal structure prediction algorithms (CSPA). However, the computational cost of the extit{ab initio} calculations required by CSPA limits its utility to small unit cells, reducing the compositional and structural space the algorithms can explore. Past studies have bypassed many unneeded extit{ab initio} calculations by utilizing machine learning methods to predict formation energy and determine the stability of a material. Specifically, graph neural networks display high fidelity in predicting formation energy. Traditionally graph neural networks are trained on large data sets of relaxed structures. Unfortunately, the geometries of unrelaxed candidate structures produced by CSPA often deviate from the relaxed state, which leads to poor predictions hindering the model's ability to filter energetically unfavorable prior to extit{ab initio} evaluation. This work shows that the prediction error on relaxed structures reduces as training progresses, while the prediction error on unrelaxed structures increases, suggesting an inverse correlation between relaxed and unrelaxed structure prediction accuracy. To remedy this behavior, we propose a simple, physically motivated, computationally cheap perturbation technique that augments training data to improve predictions on unrelaxed structures dramatically. On our test set consisting of 623 Nb-Sr-H hydride structures, we found that training a crystal graph convolutional neural networks, utilizing our augmentation method, reduced the MAE of formation energy prediction by 66\% compared to training with only relaxed structures. We then show how this error reduction can accelerates CSPA by improving the model's ability to filter out energetically unfavorable structures accurately.

研究动机与目标

• 为解决图神经网络在预测未弛豫晶体结构生成能方面表现不佳的问题，这类结构在晶体结构预测算法（CSPA）中十分常见。
• 克服在弛豫结构上预测精度高但在未弛豫结构上表现差的反向相关性，即在弛豫数据上训练的模型在未弛豫结构上表现欠佳。
• 开发一种计算成本低廉、具有物理依据的数据增强策略，以提升GNN对未弛豫几何构型的泛化能力。
• 通过实现更准确的机器学习驱动的不稳定结构预筛选，加速CSPA流程，减少对第一性原理计算的依赖。

提出的方法

• 该方法引入一种扰动技术，通过在弛豫晶体结构上施加小幅度、物理上合理的原子位移，生成未弛豫的结构变体。
• 将训练集中每个弛豫结构与其一个扰动变体配对，以极低的计算成本将训练数据量翻倍。
• 增强后的数据集每条原始条目包含一个弛豫结构和一个未弛豫结构，使模型能够更稳健地学习势能面（PES）。
• 使用晶体图卷积神经网络（CGCNN）在增强数据上进行训练以预测生成能，损失函数通过平均绝对误差（MAE）最小化。
• 在包含623个Nb-Sr-H氢化物结构（含弛豫与未弛豫构型）的保留测试集上评估模型性能。
• 采用受试者工作特征（ROC）曲线和真正率（TPR）分析，评估模型在不同能量阈值下筛选不稳定结构的能力。

实验结果

研究问题

• RQ1 一种简单且基于物理启发的数据增强技术，能否提升仅在弛豫数据上训练的图神经网络对未弛豫晶体结构的预测性能？
• RQ2 在弛豫与未弛豫结构上预测精度之间的反向相关性是否依然存在？能否通过数据增强加以缓解？
• RQ3 与仅在弛豫数据上训练的标准方法相比，数据增强在多大程度上降低了未弛豫结构生成能预测的平均绝对误差（MAE）？
• RQ4 在CSPA中，增强后的模型在筛选不稳定结构方面表现如何？需要何种能量阈值才能实现高TPR并最小化误报？

主要发现

• 在增强数据上训练使未弛豫结构的生成能预测MAE从251 meV/atom（CGCNN）和172 meV/atom（CGCNN-HD）分别降低至86 meV/atom和82 meV/atom。
• 与仅在弛豫结构上训练相比，增强数据训练的模型在未弛豫几何构型上的MAE降低了66%，显著提升了对未弛豫结构的泛化能力。
• 增强模型在39 meV/atom的过滤阈值下实现了1.0的真正率（TPR），仅错误分类130个结构，相比基线模型可将第一性原理计算需求减少5倍。
• 在TPR为0.7时，增强的CGCNN在6 meV/atom的阈值下仅错误分类54个结构，其准确性和计算效率均优于所有其他模型。
• ROC分析证实，增强模型显著优于仅在弛豫数据上训练的模型，尤其在高精度识别稳定结构方面表现突出。
• 本研究证明，即使采用简单的基于扰动的增强策略，也能显著提升GNN在未弛豫结构上的性能，从而在CSPA工作流中实现更有效的预筛选。