在有限预算下训练材料基础模型

Training a Foundation Model for Materials on a Budget

August 22, 2025

作者: Teddy Koker, Tess Smidt

cs.AI

摘要

材料建模的基础模型正在快速发展，但其训练成本依然高昂，往往使得许多研究团队难以触及最先进的方法。我们推出了Nequix，一种紧凑的E(3)等变势能模型，它结合了简化的NequIP架构与现代训练实践，包括等变均方根层归一化和Muon优化器，在保持精度的同时大幅降低了计算需求。基于JAX构建的Nequix拥有70万参数，仅用500个A100 GPU小时完成训练。在Matbench-Discovery和MDR Phonon基准测试中，Nequix总体排名第三，而所需训练成本不到大多数其他方法的四分之一，并且其推理速度比当前排名第一的模型快一个数量级。我们在https://github.com/atomicarchitects/nequix上发布了模型权重及完全可复现的代码库。

English

Foundation models for materials modeling are advancing quickly, but their training remains expensive, often placing state-of-the-art methods out of reach for many research groups. We introduce Nequix, a compact E(3)-equivariant potential that pairs a simplified NequIP design with modern training practices, including equivariant root-mean-square layer normalization and the Muon optimizer, to retain accuracy while substantially reducing compute requirements. Built in JAX, Nequix has 700K parameters and was trained in 500 A100-GPU hours. On the Matbench-Discovery and MDR Phonon benchmarks, Nequix ranks third overall while requiring less than one quarter of the training cost of most other methods, and it delivers an order-of-magnitude faster inference speed than the current top-ranked model. We release model weights and fully reproducible codebase at https://github.com/atomicarchitects/nequix