✨ 要点🔬 技术摘要
想象一下,你正试图预测一座巨大的、隐形的乐高城堡会如何表现。如果你摇晃它,它会崩塌吗?它会像金属平底锅那样导热,还是像羊毛衫那样隔热?在材料科学的世界里,科学家们使用超级计算机来模拟这些微小的构建块(原子),以设计新的电池、药物和超强金属。但问题在于:计算这些原子如何相互作用最准确的方法,就像试图徒手拼凑一个百万块的拼图——这既耗时又极其昂贵。为了提高速度,科学家们创建了“机器学习原子间势函数”(MLIPs)。把这些模型想象成一种超级智能、能快进播放的视频游戏引擎。与其从头开始计算每一条物理规则,这个引擎已经通过数百万次先前的计算“学习”了原子通常是如何行为的,从而能够在一瞬间猜出结果。最近,科学家们开始在海量数据集上训练这些引擎,以创建“基础模型”——这是一种可以针对特定任务进行微调的通用“作弊码”。但这里有一个大问题:我们能否让这些作弊码既极其快速又极其准确,还是说我们必须二选一?
这篇论文讲述了一个团队决定打造“终极作弊码”的故事,他们证明了你并不需要非此即彼。他们开发了一系列新的“基础模型”,使用了两种特定的架构,分别叫做 NequIP 和 Allegro。想象一下,这些模型是两种不同类型的超级运动员:一种是细致入微的建筑师,会检查每一个角度(NequIP);而另一种则是闪电般的短跑选手,它能看到全局,而不纠结于细节(Allegro)。研究人员在超过 1 亿个原子结构的庞大库上训练了这些运动员,这个数据集如此巨大,在旧系统中训练可能需要数周时间。通过使用一些巧妙的软件技巧——比如“图编译”(在比赛开始前组织好数据)和“混合精度数学”(在处理简单部分时使用细节稍欠精细的计算器)——他们成功让这些模型的训练速度比以前快了 5 到 10 倍,将计算成本从数周缩减到了仅需几百个 GPU 小时。
结果具有变革意义。团队发现,这些新模型不仅速度快,而且在预测材料行为方面也是有史以来最准确的。他们在三个主要挑战上对其进行了测试:寻找新材料最稳定的形状(就像寻找完美的乐高结构)、预测材料的导热性能(这对电子设备至关重要),以及计算材料的拉伸或挤压程度。这些模型表现出色,甚至经常打破之前的纪录。然而,论文也指出了一个特定的“盲点”。虽然这些模型精通通用化学,但在处理过渡金属(如铁、铬和锰)时仍会遇到一些困难,因为这些棘手元素的训练数据有些不一致。这就像这些模型精通大多数语言,但会被几种语法规则变化莫测的方言搞糊涂。
作者还研究了这些模型如何扩展规模。他们展示了这些模型可以在数千个芯片上协同工作,模拟包含超过 1 亿个原子的系统。这意义重大,因为这意味着科学家现在可以以前所未有的原子级精度来模拟整粒沙子或微小的液滴,这在以前是不可能实现的。论文指出,未来更优质材料的关键不在于制造更大的模型,而在于清理训练数据——特别是针对那些棘手的过渡金属——并确保用于生成数据的“规则”保持一致。简而言之,他们构建了一个更快、更敏锐、更可靠的工具,用于发现未来的材料,并证明了通过正确的软件升级,我们可以鱼与熊掌兼得:速度与准确性,尽收囊中。
技术摘要:用于等变机器学习原子间势能的高速、高精度基础模型
问题陈述 机器学习原子间势能(MLIP)已成为计算材料科学中不可或缺的工具,特别是作为在大型多样化数据集(如 MPtrj、Alexandria、OMat24)上进行训练并用于下游微调的“基础模型”。然而,存在一个关键瓶颈:科学应用(如分子动力学 MD)要求兼顾高精度和高推理/训练速度。虽然等变架构(编码了旋转和反转等物理对称性)提供了卓越的精度和数据效率,但在超大规模数据集(数亿个结构)上训练这些架构在历史上是计算上难以承受的,通常需要数周的训练时间。此外,在向大规模数据集扩展时(此时数据效率的重要性降低,而原始性能变得更为关键),如何平衡模型规模、推理速度和精度之间相互竞争的需求,目前仍不明确。
方法论 作者展示了一系列基于 NequIP 和 Allegro 等变图神经网络架构的基础势能模型,这些模型是在规模不断扩大的无机材料数据集上训练的:DIRECT(186k)、MPtrj(1.58M)、MPA(10.5M)以及 OAM(113M 个结构)。
关键方法论进展包括:
基础设施加速: 该工作利用了 NequIP 框架的最新发展,包括训练时图编译(torch.compile)、用于多 GPU 训练的分布式数据并行(DDP)策略,以及优化的张量积内核(用于 NequIP 的 OpenEquivariance,用于 Allegro 的 cuEquivariance)。
混合精度训练: 作者利用混合精度策略(对于嵌入使用 float64,对于张量积使用 float32,并在中间计算中使用自动混合精度)将训练速度提升了高达 3 倍,且没有明显的精度损失。
两阶段训练: 一个优化的训练协议涉及初始阶段侧重于力误差权重,随后是一个增加能量误差权重、降低学习率并使用全精度的持续运行阶段。
模型变体: 训练了不同规模(Small, Medium, Large, Extra-Large)的模型,以探索参数量、精度和速度之间的权衡。OAM 数据集的训练涉及特定的协议:先在 OMat24(非平衡几何结构)上进行预训练,随后在 MPA 上进行微调,以确保与 Materials Project 设置的兼容性。
主要贡献
加速训练流水线: 图编译和加速内核的集成使得在仅需约 100–750 GPU 小时(使用 H100/H200 GPU)的适中计算成本下,即可训练出基于超过 1 亿个结构的高精度基础模型,这相比以往数周的训练需求实现了显著降低。
基准测试套件: 模型针对涵盖以下内容的社区基准进行了严格评估:
材料发现: 预测基态几何结构和凸包能量(Matbench Discovery)。
热导率: 通过声子力常数预测非平衡性质。
近平衡性质: 评估体积/剪切模量、热容和力幅比(MatCalc)。
可扩展性分析: 论文证明了 NequIP 架构(此前受限于迭代消息传递)现在可以通过支持 Kokkos 加速的 ML-IAP 接口在 LAMASS 中实现高效的多 GPU 推理,从而支持模拟高达 1 亿个原子的系统。
结果
精度: NequIP-OAM-XL 模型在各项基准测试中均达到了最先进水平(SOTA),其凸包能量 MAE 为 19.7 meV/atom,热导率对称相对平均误差(κ S R M E \kappa_{SRME} κ S R M E )为 0.129。
缩放法则(Scaling Laws):
模型规模: 对于中小到大型模型,精度通常遵循随模型规模提升的幂律。然而,XL 模型显示出饱和迹象,表明在极端规模下特定架构存在收益递减现象。
数据集规模: 热导率精度随数据集规模理想地扩展(NequIP 的幂律指数 α ≈ 0.24 \alpha \approx 0.24 α ≈ 0.24 ),证实了非平衡几何结构的价值。然而,材料发现的精度在从 MPA 扩展到 OAM 时提升有限。
误差分析: 材料发现精度未随更大规模数据集提升的原因归因于成分多样性 和过渡金属化合物描述的不一致性 。具体而言,多价态第一过渡金属(V, Cr, Mn, Fe)和锕系元素(Np, Pu)表现出显著较高的误差。作者将此归因于训练数据中 Hubbard U U U 校正的应用不一致(应用于氧化物/氟化物而非金属)以及相关电子系统的固有复杂性。
速度 vs. 精度: 训练好的模型定义了精度-速度权衡的帕累托前沿(Pareto front)。例如,NequIP-OAM-XL 实现了与其他领先模型(如 eSEN-30M-OAM)相当的综合性能评分,但推理速度快了约一个数量级。
过度软化(Over-softening): 在多样化数据集(包括非平衡结构)上训练的大型模型表现出极小的“过度软化”现象(即在高能状态下低估力),其力幅比(f / f D F T f/f_{DFT} f / f D F T )达到 0.98。
意义与主张 论文声称,训练高精度基础势能的主要障碍不再是计算成本,而是数据质量和多样性 。作者认为,未来材料发现 MLIP 的进展不仅取决于原始的模型缩放,更取决于:
一致的数据: 消除由计算设置不一致(例如 Hubbard U U U 的应用)引起的系统误差。
化学多样性: 将数据集从离子置换扩展到包含更广泛的化学和结构多样性,特别是针对复杂的过渡金属系统。
架构适应性: 开发能够根据局部化学环境动态调整径向和角度分辨率的模型,从而处理不同元素间迥异的复杂性,而不必统一增加模型容量。
这项工作确立了:凭借当前的基础设施,单个研究小组实现训练和部署能够模拟约 1 亿个原子的基础模型是可行的,这使高保真度大规模原子模拟的使用变得民主化。所提供的模型为针对特定化学空间的微调提供了稳健的起点。
每周获取最佳 applied physics 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。