类型：a

基于Transformer生成原子嵌入提升晶体性质机器学习预测精度的研究

一、研究作者、机构与发表信息

本研究由复旦大学信息科学与工程学院的罗志杰（Luozhijie Jin）和复旦大学物理系的杜子健（Zijian Du）作为共同第一作者完成，通讯作者为复旦大学物理系岑燕（Yan Cen）和复旦大学信息科学与工程学院张浩（Hao Zhang）。其他合作者包括复旦大学乐舒（Le Shu）、山东建筑大学许元峰（Yuanfeng Xu）及复旦大学梅永丰（Yongfeng Mei）。该论文于2025年发表于*Nature Communications*，文章编号为16:1210，DOI为10.1038/s41467-025-56481-x。

二、学术背景与研究目的

在材料科学领域，加速新型晶体材料的发现对于推动清洁能源到信息处理等多种技术的发展至关重要。深度学习（deep learning，DL）和机器学习（machine learning，ML）的发展为材料性质的预测与发现提供了新的研究范式。近年来，研究者已提出多种图神经网络（graph neural network，GNN）模型用于晶体性质预测，如CGCNN（crystal graph convolutional neural network）、MEGNet（materials graph network）、ALIGNN（atomistic line graph neural network）等，在多种应用中取得了成功。

然而，在固体理论中，晶体性质由其组成原子的特征和空间拓扑排列决定，这些信息在深度学习算法中被封装为“原子嵌入”（atomic embedding）。原子嵌入是将原子性质数字化输入晶体模型的过程，其思想源于自然语言处理中的词嵌入技术。尽管合适的原子嵌入可以加速模型训练、提高预测精度并提供可解释信息，但目前材料信息学领域的多数研究将注意力集中在晶体模型架构设计上，对原子嵌入的研究较为稀少。传统方法普遍采用简单的0-1独热编码（one-hot embedding）作为原子嵌入算法，这通常产生稀疏的嵌入矩阵，不利于模型的信息提取。

为此，作者开发了自研的CrystalTransformer模型，基于Transformer架构生成通用原子嵌入（universal atomic embeddings，UAEs），称为CT-UAEs，为每个原子学习独特的“指纹”，捕捉其在材料中的角色与相互作用本质。该研究的主要目的在于：（1）提出一种不依赖预定义原子属性集的原子嵌入策略；（2）验证CT-UAEs在多种后端GNN模型和不同数据库上的预测精度提升效果；（3）探索CT-UAEs的可解释性与可迁移性。

三、研究流程与方法

总体框架。 本研究的核心思路是将材料性质预测流程分为前端模型（front-end model）和后端模型（back-end model）两部分。前端模型负责生成原子嵌入，后端模型为GNN模型（如CGCNN、MEGNet、ALIGNN），负责信息传递、节点特征更新和性质预测。作者提出的CrystalTransformer作为前端模型（方法i），有别于使用GNN模型进行预训练获取原子嵌入的传统方法（方法ii），以及基于已知原子性质人工构造特征的方法（方法iii）。CrystalTransformer直接从晶体数据库的化学信息中学习原子嵌入，不依赖于预定义的原子属性集，因而能够适应任意目标材料性质。

数据集。 研究使用了多个数据集。MP数据集为Materials Project 2018.6.1版本，包含69,239种材料，训练/验证/测试划分遵循60,000/5,000/4,239的分布；MP*数据集为2023.6.23版本，包含134,243种材料，按80%训练、10%验证、10%测试划分。目标性质包括形成能（formation energy，$E_f$）和PBE带隙（bandgap，$E_g$）。此外还使用了JARVIS数据集、MC3D数据集以及用于混合有机-无机钙钛矿（hybrid organic-inorganic perovskite，HOIP）研究的合并数据集（2,103种HOIP晶体），以及针对能量守恒原子间势（interatomic potential，IAP）模型的MPtrj数据集。

CrystalTransformer模型架构。 CrystalTransformer以原始Transformer算法为主体。模型输入包括两部分：原子输入（大小为batch × n × l，其中l为原子种类的独热编码）和坐标输入（大小为batch × n × d，d=3）。首先对坐标输入进行拓扑数据增强（平移和旋转变换），随后对两种输入分别进行线性变换嵌入至维度c。将嵌入后的原子特征和坐标特征沿特征维度拼接，得到形状为batch × n × 2c的特征矩阵m，再输入多头Transformer编码器。编码器由多层多头自注意力（multi-head self-attention）和前馈神经网络组成。最后选取输出序列的第一个token，经过线性层得到预测的材料性质。需要特别指出的是，模型的$W_a$和$b_a$矩阵（或$a W_a + b_a$的输出）即为原子信息的嵌入矩阵，这是CrystalTransformer模型最核心的部分。

多任务学习。 为提高嵌入的通用性，研究采用了多任务学习（multi-task learning，MTL）策略。模型在不同任务（如形成能、带隙、总能量、总磁化强度）上同时训练，损失函数为各任务损失的加权和。实验中涉及mt@2p（形成能+带隙）、mt@3p（增加总能量）、mt@4p（再增加总磁化强度）等变体。

实验设计。 实验分为以下几个层次：（1）在MP*数据集上预训练CrystalTransformer及对比的前端模型（CGCNN、ALIGNN、MEGNet），提取原子嵌入；（2）将这些嵌入转移至CGCNN等后端模型，在MP数据集上训练并比较MAE；（3）测试单任务嵌入与多任务嵌入的迁移效果（如$E_f$任务嵌入用于$E_g$预测）；（4）在JARVIS、MC3D和HOIP数据库上验证迁移性；（5）通过UMAP降维和K-means聚类进行可解释性分析；（6）通过CatBoost回归和SHAP算法进行反向训练实验，验证嵌入与原子物理化学性质之间的关联。

四、主要结果

CT-UAEs对预测精度的提升。 在MP数据集上，使用CT-UAEs的ct-CGCNN使形成能的MAE从0.083 eV/atom降至0.071 eV/atom（降低14%），带隙MAE从0.384 eV降至0.359 eV（降低7%）。ct-MEGNet的形成能MAE从0.051 eV/atom降至0.049 eV/atom（降低4%），带隙从0.324 eV降至0.304 eV（降低6%）。ct-ALIGNN的表现最为突出：形成能MAE从0.022 eV/atom降至0.018 eV/atom（降低18%），带隙从0.276 eV降至0.256 eV（降低7%）。相比之下，使用其他GNN前端嵌入的最佳结果（cg-CGCNN）仅分别改善4%和5%，表明CT-UAEs显著优于传统GNN生成的嵌入。

多任务嵌入的迁移性。 单任务交叉迁移实验中，$E_g$任务嵌入迁移至形成能预测时MAE降低6%，$E_f$任务嵌入迁移至带隙预测时降低0.2%。多任务嵌入（mt@2p）较单任务嵌入进一步提升性能：ctmt@2p-CGCNN实现$E_f$ MAE 0.068 eV/atom（较基线降低18%）和$E_g$ MAE 0.357 eV（降低7%）。引入总能量的mt@3p进一步将带隙MAE微降至0.356 eV，但加入总磁化强度的mt@4p使带隙MAE回升至0.367 eV，可能源于磁性与带隙之间的物理差异。

跨数据库迁移性。 在JARVIS数据集上，ct-CGCNN使形成能MAE从0.080 eV/atom降至0.066 eV/atom（降低17.5%），带隙从0.531 eV降至0.463 eV（降低12.8%）。在MC3D数据集上，ct-MEGNet使总能量MAE从5.029 eV降至4.687 eV（降低6.8%）。在MPtrj数据集训练的势能模型上，CHGNet加入CT-UAEs后力损失从0.284降至0.242（降低14.8%），展示了CT-UAEs在能量守恒势能模型中的适用性。

HOIP应用。 针对HOIP数据稀缺的挑战，研究合并了两个HOIP数据集（共2,103种晶体）。在预测形成能时，ct-CGCNN的MAE从0.054 eV/atom降至0.046 eV/atom（提升16%），ct-MEGNet的MAE从0.032 eV/atom降至0.021 eV/atom（提升34.38%），证实了CT-UAEs作为原子指纹在数据稀缺场景中的潜力。

可解释性分析。 UMAP将89×128维的CT-UAEs降至89×2，K-means聚类将元素分为三类（class a、b、c）。肘部图和轮廓系数图表明3或4类为最佳聚类数。聚类结果与元素周期表传统分类存在差异：class a包含IIA、IIIB和IVB族元素，class b包含大部分VB至VIIIB族过渡金属，class c包含IA、IB、IIB族及主族金属和非金属。对含氧化合物的统计分析显示：class a氧化物形成能集中在−2.5至−4.0 eV/atom（高化学稳定性），带隙集中在3–6 eV（宽禁带半导体）；class b氧化物带隙集中在0.5–2.5 eV（窄禁带半导体），且磁化分布最宽（部分元素磁化大于5 μB，表现出铁磁行为，Fe和Co属于此类）；class c氧化物的形成能和带隙居中。值得注意的是，不同多任务嵌入（mt@2p、mt@3p、mt@4p）之间的动态时间规整（dynamic time warping，DTW）分析显示嵌入的基本趋势保持高度相似，说明CrystalTransformer学习到的原子特征演化具有内在一致性。

反向训练验证。 CatBoost模型使用CT-UAEs作为输入，对原子半径、沸点、熔点、电导率和第一电离能等元素性质进行预测，$R^2$均超过0.78（第一电离能$R^2$达0.907），表明CT-UAEs与原子物理化学性质之间存在稳健的关联。SHAP分析进一步确定了不同性质对应的特定嵌入维度，揭示了嵌入向量中蕴含的“类基因”信息结构。

五、结论与研究价值

本研究成功提出并验证了一种基于Transformer架构的通用原子嵌入策略CT-UAEs。该策略直接利用CrystalTransformer从晶体数据库中学习原子嵌入，不仅有效提升了多种主流GNN模型（CGCNN、MEGNet、ALIGNN）在多个数据库（MP、JARVIS、MC3D）上的预测精度，还在数据稀缺的HOIP领域展现出显著的迁移学习优势。其科学价值在于证明了Transformer自注意力机制可以有效地从晶体结构与化学信息中学习具有物理意义的原子表征，而无需显式引入原子性质先验或图结构归纳偏置。其应用价值在于这种通用嵌入可以作为一种“即插即用”的原子指纹，为材料信息学中的跨模型、跨数据库迁移提供便捷工具，尤其有助于应对特定材料体系数据稀缺的挑战。

六、研究亮点

本研究的突出亮点包括：（1）提出了一种新型的基于Transformer的前端嵌入模型CrystalTransformer，突破了传统GNN嵌入的局限；（2）CT-UAEs在不同后端模型和不同数据库中均表现出一致而显著的精度提升，尤其是ALIGNN上18%的形成能改进和HOIP上MEGNet近34%的改进；（3）通过UMAP聚类分析将元素分为具有物理意义的三类，并通过对含氧化合物的统计验证揭示了嵌入中蕴含的原子特征与晶体性质之间的关联；（4）反向训练实验证实了CT-UAEs与原子基础物理化学性质（如第一电离能$R^2$=0.907）之间的强关联；（5）DTW分析表明不同多任务嵌入之间的基本趋势具有高度一致性，说明了嵌入的通用性和内在稳定性。此外，CrystalTransformer在不使用预定义图结构和边信息的情况下仍能达到与GNN模型相当的预测水平，其坐标信息分支在消融实验中被证明对嵌入质量有重要作用。