Treinando um Modelo de Base para Materiais com Orçamento Limitado
Training a Foundation Model for Materials on a Budget
August 22, 2025
Autores: Teddy Koker, Tess Smidt
cs.AI
Resumo
Os modelos de base para modelagem de materiais estão avançando rapidamente, mas seu treinamento permanece caro, muitas vezes colocando os métodos de ponta fora do alcance de muitos grupos de pesquisa. Apresentamos o Nequix, um potencial E(3)-equivariante compacto que combina um design simplificado do NequIP com práticas modernas de treinamento, incluindo normalização de camada equivariante de raiz quadrada média e o otimizador Muon, para manter a precisão enquanto reduz substancialmente os requisitos de computação. Construído em JAX, o Nequix possui 700 mil parâmetros e foi treinado em 500 horas de GPU A100. Nos benchmarks Matbench-Discovery e MDR Phonon, o Nequix ocupa a terceira posição geral, exigindo menos de um quarto do custo de treinamento da maioria dos outros métodos, e oferece uma velocidade de inferência uma ordem de magnitude mais rápida do que o modelo atualmente mais bem classificado. Disponibilizamos os pesos do modelo e o código totalmente reproduzível em 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