CHGNet：一种用于电荷信息原子模拟的预训练通用神经网络势

作者：Bowen Deng、Peichen Zhong、Kyujung Jun、Janosh Riebesell、Kevin Han、Christopher J. Bartel 与 Gerbrand Ceder。研究机构主要为加州大学伯克利分校材料科学与工程系、劳伦斯伯克利国家实验室材料科学部，以及剑桥大学卡文迪许实验室和明尼苏达大学化学工程与材料科学系。本文于2023年9月14日在线发表在《自然·机器智能》（Nature Machine Intelligence）期刊上。

研究背景

大规模原子模拟是凝聚态材料计算探索中的核心工具，尤其在研究反应活性、界面反应、部分无序结构中的输运以及复杂材料技术应用中的非均匀现象时具有不可替代的价值。然而，传统经典力场通常将电荷视为预先分配给每个原子的静态属性，难以准确描述电子态与离子重排之间的耦合。相比之下，基于密度泛函理论（DFT）的从头算分子动力学（AIMD）能够以量子力学精度显式计算电子结构，但其计算复杂度随系统规模呈三次方增长，严重限制了其在长时间、大尺度模拟中的应用。机器学习原子间势（MLIP）近年来在桥接高精度电子结构计算与高效经典势之间显示出巨大潜力，尤其是基于图神经网络（GNN）的MLIP在能量和力的预测精度上已达到先进水平。然而，已有MLIP大多仅以元素种类作为化学身份标签，无法区分同一元素的不同价态，例如Mn²⁺、Mn³⁺和Mn⁴⁺在配位偏好、局域畸变能力和化学行为上存在本质差异。因此，将电子自由度，尤其是原子电荷信息，显式引入MLIP成为该领域亟待解决的问题。本文提出利用磁矩（magnetic moment, magmom）作为原子电荷的代理变量，开发并预训练了一种名为CHGNet（Crystal Hamiltonian Graph Neural Network）的通用图神经网络势能模型，旨在实现电荷信息约束下的原子模拟。

研究流程

CHGNet的核心架构建立在图神经网络之上。模型首先将周期性晶体结构转换为原子图，通过截断半径搜索每个原子的近邻原子，并以原子间距离构建边信息。此外，利用辅助的键图（bond graph）捕获三体角度信息。节点特征由原子序数嵌入获得，边特征由可训练的平滑径向贝塞尔函数（SmoothRBF）展开距离得到，角度特征则通过傅里叶基函数展开。模型包含多个交互块（interaction block），每个交互块中原子特征、键特征和角度特征通过带权消息传递（weighted message passing）进行更新。与以往GNN模型不同，CHGNet在第三个交互块之后对原子特征施加磁矩投影约束，即在中间特征上增加磁矩预测损失，从而将电荷信息正则化注入潜在空间。经过该约束后的特征再进入最后一个交互块，最终通过全连接层预测体系总能量。力（force）和应力（stress）通过能量对原子坐标和晶格应变张量的自动微分获得。因此，CHGNet可以仅以原子核位置和元素类型为输入，同时输出能量、力、应力和原子磁矩，并由此推断原子电荷态。

在数据方面，研究团队从Materials Project数据库中提取了2022年9月版本中的全部GGA及GGA+U计算任务轨迹数据，经过严格筛选后构建了Materials Project Trajectory（MPTrj）数据集。该数据集包含1,580,395个原子构型，每个构型均对应总能量、原子受力、应力以及7,944,833个原子磁矩。数据覆盖89种元素、超过146,000种无机化合物。训练集、验证集和测试集按材料编号以8:1:1的比例划分，确保同一材料的所有轨迹帧只出现在一个集合中。模型在训练时使用Huber损失函数，同时拟合能量、力、应力和磁矩，力的权重为1.0，应力和磁矩的权重各为0.1。优化器为Adam，初始学习率为10⁻³，并通过余弦退火调度器逐步衰减至10⁻⁵。模型总可训练参数量为400,438个。

主要结果

在MPTrj测试集上，CHGNet在包含磁矩损失训练后，能量平均绝对误差为30 meV/atom，力平均绝对误差为77 meV/Å，应力平均绝对误差为0.348 GPa，磁矩平均绝对误差为0.032 μB。相比之下，未加入磁矩损失时，能量误差为33 meV/atom，力误差为79 meV/Å，应力误差为0.351 GPa。这表明磁矩约束不仅没有损害能量和力的预测精度，反而对能量和力的精度有一定提升。此外，研究者将CHGNet提交至MatBench Discovery基准测试，用于评估其对未知无机晶体结构稳定性的预测能力。在文章撰写时，CHGNet在该基准上的高通量稳定晶体发现性能优于同期其他八个模型。为验证CHGNet在分子动力学（MD）中的可靠性，研究者对多种锂超离子导体进行了模拟。结果显示，CHGNet预测的室温离子电导率和活化能与已报道的AIMD结果在误差范围内一致，并且能够正确区分由DFT识别的快离子导体与慢离子导体。

在电荷约束验证方面，研究者以钠超离子导体Na₄V₂(PO₄)₃为例，构建了2,268个原子的超胞，并随机移除一半钠离子以生成Na₂V₂(PO₄)₃结构。CHGNet在对该结构进行弛豫后，成功将216个钒离子区分为V³⁺和V⁴⁺两组，尽管模型事先并未被告知任何电荷分布信息。主成分分析显示，所有钒离子的潜在空间特征向量明显聚成两类，分别对应不同价态。这说明磁矩标签的引入对潜在空间特征产生了强正则化作用，使能量、力和应力的预测建立在电荷感知的特征基础之上。

在LiₓMnO₂体系中，研究者使用预训练CHGNet对部分脱锂的正交LiMnO₂（o-LMO）超胞进行了1,100 K下的1.5 ns电荷信息MD模拟。模拟的X射线衍射图谱显示，体系随时间发生由正交相向尖晶石相的结构转变。在约0.8 ns时，Mn²⁺在体系中的比例显著增加，同时体系势能下降，尖晶石特征峰增强。后续DFT计算验证了CHGNet对Mn离子磁矩的预测精度。重要的是，模拟发现Mn²⁺在四面体位的形成时间远长于Mn离子在四面体和八面体位之间的单次跳跃时间尺度，表明四面体Mn²⁺的形成与长程尖晶石有序化相关，而非单纯由离子跳跃驱动。

在LiₓFePO₄相图研究中，CHGNet被用作能量计算器生成两种不同的团簇展开模型：一种考虑Li⁺/空位和Fe²⁺/Fe³⁺的电荷自由度，另一种仅考虑Li⁺/空位。半巨正则蒙特卡洛模拟结果表明，包含电荷自由度的相图在中间锂浓度处出现类共析转变，与实验相符；而缺少电荷自由度的相图仅显示单一的混溶隙，无法再现实验中的中间固溶体行为。该结果直接证明了显式包含电子自由度对正确描述混合价氧化物有限温度相稳定性的必要性。

在石榴石型锂超离子导体Li₃La₃Te₂O₁₂体系中，CHGNet通过零样本测试精确复现了少量过量锂（δ = 1/48）对锂离子扩散的激活效应。对于化学计量比的Li₃结构，CHGNet模拟得到的活化能为0.591 eV，而在Li₃.₁₂₅结构中，活化能降至0.158 eV，与AIMD计算结果高度一致。这表明CHGNet能够捕捉由锂离子间强相互作用引起的复杂扩散行为。

结论与价值

CHGNet通过将磁矩作为电荷信息代理变量引入图神经网络势函数，实现了电子自由度和原子自由度在同一模型中的耦合建模。其预训练基于规模巨大的MPTrj数据集，覆盖了除稀有气体和锕系元素以外的绝大多数元素化学空间，具备通用势函数的广泛应用潜力。在方法论层面，磁矩约束作为一种潜在空间正则化手段，在不依赖显式电荷输入的前提下推断原子价态，为MLIP中的电子信息编码提供了新思路。在应用层面，CHGNet可广泛用于混合价态过渡金属氧化物的高通量筛选、离子导体中的扩散机制研究、以及带电耦合相变过程的长时间模拟。相较于AIMD，CHGNet在保持近DFT精度的前提下大幅提高了计算效率，能够模拟纳秒级时间尺度的过程，为传统方法无法触及的电子-离子耦合动力学研究提供了可能。研究的局限性在于：磁矩对于非磁性离子的电荷推断存在模糊性，且形式价态的分配仍需依赖外部化学知识。未来可通过引入电子局域函数、电极化或瓦尼尔函数等更丰富的电荷表示方法进一步提升模型能力。

研究亮点

本研究的主要亮点在于：首次在通用MLIP中引入磁矩作为电荷信息显式约束；构建了迄今规模最大的包含磁矩标签的DFT轨迹数据集MPTrj；通过长时MD直接观察到LiₓMnO₂中电荷歧化与长程结构有序化的耦合；并通过LiₓFePO₄相图比较实验清楚地表明，缺失电子自由度会导致错误的相平衡预测；同时验证了CHGNet对石榴石型快离子导体中锂离子扩散的精准预测能力。总体上，这项工作为电子自由度参与的大规模原子模拟提供了强有力的工具和方法学基础。