예산 내에서 재료 기반 파운데이션 모델 학습하기

초록

재료 모델링을 위한 기초 모델(foundation models)은 빠르게 발전하고 있지만, 그 훈련 비용이 여전히 높아 최신 기술을 많은 연구 그룹이 활용하기 어려운 상황입니다. 우리는 Nequix를 소개합니다. Nequix는 간소화된 NequIP 설계를 현대적인 훈련 기법과 결합한 컴팩트한 E(3)-등변(equivariant) 포텐셜로, 등변 제곱평균제곱근 레이어 정규화(equivariant root-mean-square layer normalization)와 Muon 옵티마이저를 포함하여 정확도를 유지하면서도 계산 요구량을 크게 줄였습니다. JAX로 구축된 Nequix는 70만 개의 매개변수를 가지며 500 A100-GPU 시간 동안 훈련되었습니다. Matbench-Discovery 및 MDR Phonon 벤치마크에서 Nequix는 전체적으로 3위를 차지하면서도 대부분의 다른 방법들보다 4분의 1 이하의 훈련 비용을 요구하며, 현재 최고 순위 모델보다 한 차원 빠른 추론 속도를 제공합니다. 우리는 모델 가중치와 완전히 재현 가능한 코드베이스를 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

예산 내에서 재료 기반 파운데이션 모델 학습하기

Training a Foundation Model for Materials on a Budget

초록

Support