Entrenamiento de un Modelo Fundacional para Materiales con un Presupuesto Limitado
Training a Foundation Model for Materials on a Budget
August 22, 2025
Autores: Teddy Koker, Tess Smidt
cs.AI
Resumen
Los modelos base para el modelado de materiales están avanzando rápidamente, pero su entrenamiento sigue siendo costoso, lo que a menudo sitúa los métodos de vanguardia fuera del alcance de muchos grupos de investigación. Presentamos Nequix, un potencial E(3)-equivariante compacto que combina un diseño simplificado de NequIP con prácticas modernas de entrenamiento, incluyendo la normalización de capas equivariante de raíz cuadrática media y el optimizador Muon, para mantener la precisión mientras reduce sustancialmente los requisitos de cómputo. Desarrollado en JAX, Nequix tiene 700K parámetros y fue entrenado en 500 horas de GPU A100. En los benchmarks Matbench-Discovery y MDR Phonon, Nequix ocupa el tercer lugar en general mientras requiere menos de un cuarto del costo de entrenamiento de la mayoría de los otros métodos, y ofrece una velocidad de inferencia un orden de magnitud más rápida que el modelo actualmente mejor clasificado. Publicamos los pesos del modelo y un código base completamente reproducible en 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