Обучение базовой модели для материалов при ограниченном бюджете
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 содержит 700 тысяч параметров и был обучен за 500 часов на GPU A100. На тестах 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