[논문 리뷰] UMA: A Family of Universal Models for Atoms

Brandon M. Wood, Misko Dzamba외 16명|ArXiv.org|2025. 06. 30.

Machine Learning in Materials Science인용 수 20

한 줄 요약

UMA는 분자, 재료, 촉매 전반에 걸쳐 약 5억 개의 원자 구조를 학습한 범용 머신러닝 인터원자 포텐셜의 가족을 제시합니다. 모델은 용량을 효율적으로 확장하기 위해 Mixture of Linear Experts (MoLE)를 사용하며, 특정 작업별 파인튜닝 없이도 경쟁력 있거나 최첨단 성능을 달성합니다.

ABSTRACT

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.

연구 동기 및 목표

• 다양한 화학 영역(재료, 분자, 촉매)에서도 빠르고 정확한 DFT 대리 모델의 필요성을 제시한다.
• 하나의 대규모 모델이 파인튜닝 없이도 작업 간 일반화가 가능하다는 점을 보여준다.
• 추론 속도를 손실 없이 용량을 증가시킬 수 있는 확장 가능한 아키텍처(MoLE)를 개발한다.
• 속도와 에너지 보존 정확도 사이의 균형을 맞추기 위한 두 단계 학습 절차를 제안한다.
• 커뮤니티 전반의 활용과 검증을 가능하게 하는 코드, 가중치, 데이터의 공개를 추진한다.

제안 방법

• 전체 전하, 스핀, DFT 작업 입력을 확장한 eSEN 기반 등각 그래프 신경망 아키텍처를 채택한다.
• 출력이 선형 전문가의 조밀한 조합으로 이루어지는 MoLE을 도입해 부드러운 에너지 표면과 회전 등각성을 보존한다.
• 소형 MLP를 통해 시스템 수준 임베딩으로 전문가 가중치 α를 계산하고 미리 계산된 W* = Σk αkWk를 사용해 추론 속도를 유지한다.
• 두 단계 스케줄로 학습한다: 먼저 힘을 직접 예측하고, 그다음 자동 미분을 통해 에너지 보존 및 응력을 미세조정한다.
• 정확도 유지를 위해 프리트레이닝은 BF16, 미세튜닝 시 FP32로 전환하고, 큰 MoLE 구성을 확장하기 위해 메모리/그래프 병렬성을 활용한다.
• 다양한 데이터셋(재료, 분자, 촉매)으로 약 5억 개의 원자 구조에 대해 학습하고, 서로 다른 DFT 설정에서 다중 작업 학습을 위한 에너지 참조 스킴을 채용한다.

실험 결과

연구 질문

• RQ1 단일의 비 파인튜닝 모델이 재료, 분자, 촉매에 걸친 다양한 DFT 작업에서 경쟁력 있는 정확도를 달성할 수 있는가?
• RQ2 모델 크기, 데이터 양, 컴퓨트 비용이 UMA 구성의 최적화를 어떻게 결정하는가?
• RQ3 다중 작업 MLIPs에서 MoLE 아키텍처가 밀집 모델에 비해 얻는 이점은 무엇이며, 특히 MD와 같은 장기간 시뮬레이션에 대해 어떤가?
• RQ4 단일 모델이 다양한 작업과 데이터셋 전반에서 에너지 보존 및 매끄러운 포텐셜 에너지 표면을 유지할 수 있는가?

주요 결과

• UMA는 작업별 파인튜닝 없이도 재료, 촉매, 분자, 분자 결정구조 및 MOF에 대해 경쟁력 또는 최첨단 성능을 달성한다.
• MoLE은 밀집 모델에 비해 대략 2.5배 적은 활성 매개변수로 비슷한 손실을 달성하는 등 상당한 효율성을 제공한다(예: UMA-M).
• UMA-S, UMA-M, UMA-L은 Matbench Discovery 및 흡착에너지 벤치마크에서 강력한 성과를 보이며 촉매에서 AdsorbML 성공률이 25% 개선된다.
• 단일 모델이 긴 MD 롤아웃을 우수한 추론 속도와 메모리 사용으로 처리하여 80GB GPU 단일 노드에서 100k+ 원자까지의 시뮬레이션과 다중 GPU 확장을 가능하게 한다.
• 두 단계 학습은 효율적인 학습과 에너지 보존을 보장하며, BF16 프리트레이닝에 이은 FP32 미세조정으로 구성된다.
• 경험적 스케일링 법칙은 성능이 모델 크기와 데이터에 대해 로그-선형으로 확장된다는 것을 시사하며, 컴퓨트 최적화 및 추론 최적화 학습 전략의 방향을 잡아준다.