UMA: A Family of Universal Models for Atoms

The ability to quickly and accurately compute properties from atomic simulations is critical for advancing a large number of applications in chemistry and materials science including drug discovery, energy storage, and semiconductor manufacturing. To address this need, Meta FAIR presents a family of Universal Models for Atoms (UMA), designed to push the frontier of speed, accuracy, and generalization. UMA models are trained on half a billion unique 3D atomic structures (the largest training runs to date) by compiling data across multiple chemical domains, e.g. molecules, materials, and catalysts. We develop empirical scaling laws to help understand how to increase model capacity alongside dataset size to achieve the best accuracy. The UMA small and medium models utilize a novel architectural design we refer to as mixture of linear experts that enables increasing model capacity without sacrificing speed. For example, UMA-medium has 1.4B parameters but only ~50M active parameters per atomic structure. We evaluate UMA models on a diverse set of applications across multiple domains and find that, remarkably, a single model without any fine-tuning can perform similarly or better than specialized models. We are releasing the UMA code, weights, and associated data to accelerate computational workflows and enable the community to continue to build increasingly capable AI models.

한 줄 요약

UMA는 5억 개 이상의 3D 원자 구조를 학습한 대규모 원자 모델로, Mixture of Linear Experts(MoLE) 아키텍처를 통해 정확도와 속도를 동시에 달성한다.

핵심 기여도

UMA-medium 모델은 1.4B 파라미터를 사용하지만, 구조당 약 50M만 활성화되어 처리 속도 저하 없이 모델 용량 확장 가능.
5억 개의 3D 원자 구조로 학습한 최대 규모의 모델로, Matbench Discovery에서 최고 F1 점수 달성.
MoLE 아키텍처를 도입해, 25% 성공률 향상(AdsorbML 벤치마크)을 기록.
학습 코드, 가중치, 데이터를 공개하여 연구 및 산업 활용 확대.

핵심 아이디어

UMA는 DFT 계산을 대체할 수 있는 MLIP(머신러닝 인터아토믹 포텐셜)로, 다양한 화학 도메인(분자, 재료, 촉매)에서 일반화된 성능을 제공하는 것이 목표이다. 기존 모델은 특정 도메인에 최적화되어 있었으나, UMA는 대규모 데이터와 MoLE 아키텍처를 통해 단일 모델로 여러 도메인에서 경쟁력을 유지한다. 핵심 아이디어는 **모델 크기와 데이터셋 크기의 경험적 스케일링 법칙**을 통해 최적의 성능을 도출하는 것이다. MoLE는 **혼합 전문가 모델**로, 각 구조에 필요한 파라미터만 활성화하여 처리 속도를 유지하면서도 모델 용량을 확장한다.

기술적 접근법

**UMA 아키텍처**: eSEN 기반의 equivariant graph neural network를 사용.
**입력**: 3D 원자 위치, 원자 번호, 총 전하, 스핀, DFT 설정.
**출력**: 총 에너지, 원자별 힘, 스트레스.
**MoLE(Mixture of Linear Experts)**: 8개의 전문가 중 필요한 전문가만 활성화.
**데이터**: 5억 개의 3D 원자 구조(OMat, OC20, OMol25 등).
**훈련 전략**: 두 단계 훈련(예비 훈련 + 정제 단계).
**하이퍼파라미터**: UMA-medium: 1.4B 파라미터, 활성 파라미터 50M.
**스케일링 법칙**: FLOP 기반의 모델 크기와 데이터 크기의 관계를 분석.

주요 결과

**Matbench Discovery**: UMA-medium이 최고 F1 점수 달성.
**AdsorbML**: UMA-L이 25% 성공률 향상.
**OC20 S2EF**: 에너지 예측 오차 80% 감소.
**UMA-L**: 700M 파라미터 모델로, 5.6B 파라미터 MoLE 모델과 유사한 성능.
**모델 성능**: UMA-S, UMA-M은 고주파 관련 속성(κ_SRME, 자유 에너지 MAE)에서 우수한 성능.

의의 및 한계

UMA는 DFT 계산을 대체할 수 있는 MLIP으로, **단일 모델이 다양한 화학 도메인에서 경쟁적인 성능**을 보이는 것은 학술적·실용적으로 의미가 크다. 특히, **MoLE 아키텍처**는 모델 크기 확장과 속도 유지의 균형을 맞추는 데 기여하며, 대규모 데이터 학습을 가능하게 한다. 그러나 **UMA-L 모델이 HEA 데이터셋에서 과적합**을 보이는 점은 데이터셋의 한계를 드러낸다. 또한, **모델 크기가 커질수록 MoLE와 dense 아키텍처의 성능 차이는 줄어들며**, 데이터셋 크기의 한계가 성능 향상에 제약이 될 수 있다.

실용적 활용

UMA는 **의약품 발견, 에너지 저장, 반도체 제조** 등에서 DFT 계산을 대체할 수 있는 MLIP로 활용 가능하다. 특히, **AdsorbML과 OC20 S2EF**와 같은 촉매 설계, **Matbench Discovery**와 같은 재료 발굴 분야에서 실용적 가치가 높다. 공개된 코드와 데이터는 연구자들이 보다 정확하고 빠른 MLIP 모델을 개발하는 데 기여할 것으로 기대된다.