在有限預算下訓練材料基礎模型

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