Nat. Commun. | 通用原子嵌入，晶体材料性质预测迈出新一步

今天为大家介绍的是来自复旦大学张浩团队的一篇论文。通过机器学习加速新型晶体材料的发现对于推动从清洁能源到信息处理的各种技术至关重要。用于预测材料性质的机器学习模型需要嵌入原子信息，而传统方法在提高预测准确性方面效果有限。在此，作者提出了一种称为通用原子嵌入(universal atomic embeddings，UAEs)的原子嵌入策略，以实现其作为原子指纹的广泛应用，并基于所提出的CrystalTransformer模型生成了UAE张量。通过在广泛使用的材料数据库上进行实验，作者的基于CrystalTransformer的UAEs(ct-UAEs)被证明能够准确捕捉复杂的原子特征，在使用Materials Project数据库的形成能作为目标时，CGCNN的预测准确率提高了14%，ALIGNN提高了18%。作者还展示了ct-UAEs在各种数据库之间良好的迁移性。基于多任务ct-UAEs的聚类分析，周期表中的元素可以进行分类，在原子特征和目标晶体性质之间建立合理联系。将ct-UAEs应用于预测杂化钙钛矿数据库中的形成能后，作者实现了准确性的提高，MEGNET提升了34%，CGCNN提升了16%，展示了它们作为原子指纹解决数据稀缺性挑战的潜力。

深度学习(DL)和机器学习(ML)的发展为各种研究领域创造了研究方法。在材料科学中，这一发展正在引领材料性质的发现，这对传统方法而言可能是一项具有挑战性的任务。许多DL算法和模型已被提出，如晶体图卷积神经网络(CGCNN)、材料图网络(MEGNET)、原子线图神经网络(ALIGNN)、改进的晶体图卷积神经网络(iCGCNN)、OrbNet等。它们在各种应用中取得了成功，例如从多保真度数据中学习性质、发现稳定的无铅杂化有机-无机钙钛矿、映射晶体结构相、设计材料微观结构等。

在固态理论中，晶体或其他凝聚系统中组成原子的特征和空间拓扑排列决定了它们的性质，这些在DL算法中被巧妙地封装到"原子嵌入"的实体中。具体来说，原子嵌入是将原子特性以数字方式输入晶体模型的过程，这一理念源自自然语言处理技术，其中词嵌入转变了文本数据表示的方式。适当的原子嵌入可以加速模型训练，提高预测准确性，并且可以从中导出一些可解释的信息。目前，材料信息学领域的大部分注意力集中在晶体模型架构的设计上，以提高性质预测的准确性，而关于原子嵌入的研究很少。通常，简单采用0-1嵌入作为原子嵌入算法是常见的做法，这通常会生成不利于模型信息提取的稀疏嵌入矩阵。

近年来，在化学分子性质和结构预测领域，开发了大量基于Transformer的训练方法和预测模型，如OrbNet、3D-Transformer等，这些被认为能够充分利用Transformer架构在处理原子交互和捕捉三维结构方面的优势，能够有效表示原子之间的复杂交互。受这些进展的启发，作者开发了自制的CrystalTransformer模型，基于transformer架构生成称为ct-UAEs的通用原子嵌入，它为每个原子学习独特的"指纹"，捕捉它们在材料中的角色和交互的本质。

图 1

大体上，在深度学习模型预测材料的性质（如形成能和带隙）时，首先将每个原子嵌入为特征。这种嵌入过程是CGCNN、ALIGNN和MEGNET等GNNs模型的内在过程。然后进行更深层次的特征提取过程，包括信息传输和聚合、节点特征更新等，以预测晶体性质。在这种情况下，这些GNNs被称为后端模型，而获取原子嵌入的方法被称为前端模型。本质上，原子嵌入的参数可以使用预训练参数进行迁移，或者基于预定义的性质进行构建，这在图1b、c所示的方法（I、II、III）的前端模型中实现。

如图1a所示，对于前端模型，作者使用了作者提出的CrystalTransformer来生成原子嵌入（方法I）。其他预训练的原子嵌入使用了GNN模型（如图1b所示的方法II）。而有些则使用基于已知原子性质的人工构建特征，如基于自编码器的方法（如图1c所示的方法III）。CrystalTransformer模型直接从晶体数据库中的化学信息学习原子嵌入。与方法III相比，方法III通过处理预定义的原子性质集来生成原子嵌入，作者提出的ct-UAE可以适应任何期望的材料性质，而不依赖于预定义的原子属性。

为了检查从不同模型获得的原子嵌入张量，作者使用了MP和MP*数据集来研究形成能(Ef)和PBE带隙(Eg)，这些是评估材料化学稳定性和电子性能的关键特性。MP代表2018.6.1版本的MP数据集，包含69,239种具有各种性质的材料。MP*表示2023.6.23版本，包含134,243种材料。对于训练、验证和测试集的划分，作者遵循了与之前工作相同的分布方式，即MP数据集的60,000个样本用于训练，5,000个用于验证，4,239个用于测试。而对于MP*数据集，其材料及其性质被分为80%训练集、10%验证集和10%测试集。

值得注意的是，材料数据库（如MP）中固体能带结构的带隙被定义为导带最小值(CBM)和价带最大值(VBM)特征值之间的差值，这些值是通过使用Perdew-Burke-Ernzerhof(PBE)参数化的交换关联(xc)方程解Kohn-Sham(KS)方程得到的。在半导体和绝缘体中，这些PBE带隙EPBE g并不等于它们的基本带隙EG，而是相差一个称为交换关联能导数间断Δxc的项，导致与EG相比，EPBE g被大幅低估，幅度可达40-50%。然而，由于KS方程是基于动能和带电粒子（电子和离子）之间的库仑势能构建的，当使用特定的交换关联泛函时，KS方程的特征值应该能够捕捉相互作用系统内的主要物理相互作用。因此，如果在深度学习模型中使用PBE带隙作为目标，所得到的原子嵌入应该包含原子性质和结构信息，因为在使用PBE类型的交换关联泛函构建KS哈密顿量时，EPBE g已经包含了这些信息。

前端模型（如CrystalTransformer、CGCNN、ALIGNN和MEGNET）首先在扩展的MP*数据集上进行预训练，专注于带隙能量Eg和形成能Ef的预测任务。随后，提取的原子嵌入被整合到CGCNN后端模型中，并在原始MP数据集上进行训练，这形成了CT-CGCNN、CG-CGCNN、ALI-CGCNN等模型。

表 1

表1展示了比较性MAE分析，用于评估归因于前端原子嵌入的相对性能提升，在表1中表示为N-CGCNN（N表示上述前端模型）。如表1所列，在不同模型预训练的原子嵌入中，使用ct-UAEs（CT-CGCNN）的性能最佳，Ef和Eg的MAE分别减少了14%和7%，同时也分别比最佳GNN前端嵌入（在此情况下为CG-CGCNN）在这两种性质上表现更好，分别提高了4%和5%。这些模型预测的形成能与目标形成能的对比列在图2a-c中。

图 2

此外，如表1所列，使用CrystalTransformer生成的原子嵌入（ct-UAEs）可以增强CGCNN、MEGNET和ALIGNN等GNN模型在MP数据集上的性能。采用CrystalTransformer生成的嵌入（ct-UAEs）的CGCNN模型（在表1中表示为CT-CGCNN）显示出形成能Ef的MAE值显著降低，从0.083 eV/原子减少到0.071 eV/原子，降低了14%，而带隙Eg则从0.384 eV减少到0.359 eV，降低了7%。对于MEGNET（在表1中表示为CT-MEGNET）也可以观察到类似的降低，Ef从0.051 eV/原子减少到0.049 eV/原子，降低了4%，而带隙Eg从0.324 eV减少到0.304 eV，降低了6%。ALIGNN（在表1中表示为CT-ALIGNN）在Ef预测准确性方面也有所提高，从0.022 eV/原子减少到0.018 eV/原子，降低了18%，而带隙Eg从0.276 eV减少到0.256 eV，降低了7%。

为了进一步研究ct-UAEs在不同性质上的表现，作者将任务生成的嵌入应用于不同任务。例如，将Ef任务生成的原子嵌入应用于带隙预测，将Eg任务生成的嵌入应用于形成能任务。结果列在表2中，表示为CTEf-CG和CTEg-CG。在带隙任务上训练的嵌入，当转移到形成能任务时，导致准确性有明显提高，MAE从0.083降低到0.078 eV/原子，减少了6%。此外，虽然在形成能这样的简单任务上训练，该嵌入在更具挑战性的带隙预测上也将MAE降低了0.2%。

表 2

进一步的实验关注多任务生成的嵌入(MT)。如表2所列，从两种性质（形成能和带隙）训练的嵌入，表示为MT@2p，与单任务生成的嵌入相比产生了更好的性能。当转移到CGCNN模型(CTMT@2p-CGCNN)时，该模型对Ef达到0.068 eV/原子的MAE，对Eg达到0.357 eV的MAE，优于基准CGCNN（Ef减少18%，Eg减少7%），也优于使用单任务嵌入的CGCNN变体，Ef减少4%，带隙减少0.5%。

额外引入了包含总能量和总磁化的多任务变体(MT@3p和MT@4p)。当引入额外包含总能量性质的MT@3p时，带隙的MAE减少了0.2%，而形成能几乎保持不变。然而，在MT@4p中引入磁化会导致带隙预测的MAE从0.357略微增加到0.367 eV，这可能是由于这两种性质之间的物理差异造成的。

表 3

然后，作者使用不同的训练策略来评估模型的性能，结果列在表3中。使用从CrystalTransformer或ct-UAEs冻结的预训练嵌入的CTfreeze-CGCNN，对形成能Ef达到0.073 eV/原子的MAE，对带隙Eg达到0.358 eV的MAE。然而，当将坐标嵌入与ct-UAEs（化学信息）一起整合到CGCNN框架中(CTchem+coords-CGCNN)时，MAE从仅使用原子嵌入模型的0.071 eV/原子增加到0.085 eV/原子。同样，带隙Eg的MAE从0.359 eV恶化到0.395 eV。

通用原子嵌入的能力和可迁移性在不同数据库和任务上得到了进一步测试。每个数据库按8:1:1的比例分为训练集、验证集和测试集。对于Jarvis数据集，结果如表1所示。CT-CGCNN模型在预测形成能Ef和带隙能量Eg方面都表现出改进。形成能和带隙的MAE分别从0.080 eV/原子减少到0.066 eV/原子，降低了17.5%，以及从0.531 eV减少到0.463 eV，降低了12.8%。

嵌入进一步在MC3D数据集上进行评估。选择总能量(E)等性质作为任务，结果如表1所示。CGCNN的MAE从5.558 eV降低到5.341 eV，表明有3.9%的改进。对于ALIGNN模型，MAE几乎保持不变。而对于MEGNET模型，MAE从5.029 eV减少到4.687 eV，显示出6.8%的改进。

此外，作者还研究了ct-UAE在能量守恒的原子间势(IAP)模型上的适用性，这些模型基于MPtrj数据集进行训练。作者在矢量和标量目标上训练了ct-UAE，即力、应力和能量。为了进行基准测试，作者在M3GNet提出的MP-RELAX数据集上重新训练了CHGNet、M3GNet和MACE模型。值得注意的是，将ct-UAEs添加到CHGNet中导致力的损失显著减少，从0.284减少到0.242（降低了14.8%），同时应力损失从1.496减少到1.437，能量损失从0.460略微降低到0.457。对于M3GNet，ct-UAE导致总损失（能量、力、应力）从2.1236略微减少到2.1234，能量损失从0.3597减少到0.3595，表明性能略有提高。然而，对于MACE，ct-UAE并未导致损失减少。

编译|黄海涛

审稿|王梓旭

Jin, L., Du, Z., Shu, L., Cen, Y., Xu, Y., Mei, Y., & Zhang, H. (2025). Transformer-generated atomic embeddings to enhance prediction accuracy of crystal properties with machine learning. Nature Communications, 16(1), 1210.