UMA: A Family of Universal Models for Atoms

메타 FAIR 은 분자, 물질, 촉매 등 다양한 화학 도메인의 5 억 개 3 차원 원자 구조로 학습된 'UMA(원자를 위한 범용 모델)'를 공개하여, 미세 조정 없이도 다양한 응용 분야에서 전문 모델과 동등하거나 더 뛰어난 성능을 발휘하면서도 속도와 정확도를 동시에 확보하는 새로운 AI 모델을 제시했습니다.

핵심

이 논문은 메타 FAIR 연구소가 발표한 **'UMA (Universal Models for Atoms, 원자를 위한 범용 모델)'**라는 획기적인 인공지능 기술에 대한 내용입니다.

이걸 아주 쉽게 비유해서 설명해 드릴게요.

🌟 핵심 비유: "만능 요리사" vs "전문 요리사"

과거에 과학자들은 새로운 약을 만들거나 배터리 소재를 개발할 때, 컴퓨터로 원자 수준에서 시뮬레이션을 돌렸습니다. 이때 사용하는 전통적인 방법 (DFT) 은 **정밀하지만 너무 느린 '수제 요리'**였습니다. 한 접시 (시뮬레이션) 를 만드는 데 몇 시간이 걸려서, 수천 가지 재료를 실험해 보려면 몇 년이 걸릴 수도 있었죠.

그래서 등장한 것이 기존의 AI 모델들이었습니다. 이들은 **'전문 요리사'**처럼 특정 재료 (예: 금속만, 또는 유기분자만) 에만 능숙했습니다. 금속 요리사는 유기분자 요리를 못 하고, 유기분자 요리사는 금속 요리를 못 했죠.

UMA 는 바로 이 모든 것을 해결한 '만능 천재 요리사'입니다.

---

🚀 UMA 가 왜 특별한가요?

1. 방대한 경험 (데이터)

UMA 는 약 5 억 개의 서로 다른 3D 원자 구조를 배웠습니다. 이는 지금까지 만들어진 어떤 모델보다도 훨씬 많은 양입니다. 마치 요리사가 전 세계의 모든 식재료를 맛보고 조합해 본 것과 같습니다. 덕분에 약, 배터리, 반도체, 촉매 등 어떤 분야든 막힘없이 다룰 수 있게 되었습니다.

2. 똑똑한 구조 (MoLE: 전문가 혼합 시스템)

이 모델의 가장 큰 특징은 **'MoLE (Mixture of Linear Experts)'**라는 기술입니다.

• 비유: UMA 는 거대한 도서관 (모델) 안에 14 억 권의 책 (파라미터) 을 가지고 있습니다. 하지만 매번 모든 책을 다 읽을 필요는 없죠.
• 작동 원리: 질문이 들어오면, UMA 는 그 질문에 가장 적합한 '전문가 (Expert)' 몇 명만 골라내서 답을 냅니다. 예를 들어, "금속 재료"에 대한 질문이면 금속 전문가만, "약 분자"에 대한 질문이면 분자 전문가만 활성화됩니다.
• 효과: 이렇게 하면 모델의 총 크기는 거대해져서 정확도가 높아지지만, 실제로 계산할 때는 필요한 부분만 쓰므로 속도는 매우 빠릅니다. 마치 거대한 회사에 수천 명의 직원이 있지만, 업무마다 필요한 팀원만 투입해서 일하는 것과 같습니다.

3. 한 번에 모든 것 (Zero-shot 학습)

기존에는 특정 작업을 하려면 모델을 다시 훈련시켜야 (Fine-tuning) 했지만, UMA 는 한 번 훈련된 모델로 바로 다양한 작업을 수행할 수 있습니다. 약을 설계하든, 배터리를 개발하든, 별도의 추가 학습 없이도 기존에 특화된 모델들보다 잘하거나 비슷하게 잘합니다.

---

📊 실제로 얼마나 잘하나요?

이 논문에서는 UMA 를 여러 분야에서 테스트했습니다.

• 재료 과학: 새로운 금속 합금을 찾는 대회 (Matbench Discovery) 에서 최고 점수를 받았습니다.
• 촉매 (화학 반응): 화학 반응을 일으키는 물질을 찾는 데 있어, 기존 최고 모델보다 25% 더 정확하게 성공적인 결과를 냈습니다.
• 약물 설계: 약이 몸속에서 어떻게 작용하는지 예측할 때, 고가의 실험실 장비 (DFT) 를 대체할 만큼 정확한 결과를 보여줍니다.
• 속도: 일반 GPU 하나에서 1,000 개의 원자로 이루어진 분자를 하루에 1.4 나노초 동안 시뮬레이션할 수 있습니다. 이는 기존 방식보다 수천 배 빠른 속도입니다.

💡 결론: 과학의 미래를 여는 열쇠

UMA 는 "하나의 모델로 모든 원자 세계를 이해할 수 있다"는 것을 증명했습니다.

• 이전: "이걸 연구하려면 이 전용 모델을 써야 해."
• 이제: "UMA 하나면 모든 걸 해결할 수 있어."

이 기술은 신약 개발, 친환경 에너지, 반도체 등 우리 삶에 중요한 분야들의 연구 속도를 획기적으로 높여줄 것입니다. 연구팀은 이 모델의 코드와 데이터를 모두 공개하여, 전 세계 과학자들이 함께 더 발전된 AI 를 만들 수 있도록 장을 마련했습니다.

한 줄 요약:

"UMA 는 5 억 개의 원자 경험을 바탕으로, 어떤 재료든 빠르고 정확하게 예측할 수 있는 '만능 원자 시뮬레이션 AI'입니다."

기술 요약

UMA (Atoms 를 위한 범용 모델) 기술 요약

이 논문은 Meta FAIR 연구팀이 제안한 UMA (Universal Models for Atoms) 에 대한 연구 결과를 다루고 있습니다. UMA 는 화학 및 재료 과학 분야에서 원자 시뮬레이션 속도와 정확도를 획기적으로 개선하기 위해 설계된 일련의 범용 머신러닝 상호작용 포텐셜 (MLIP) 모델 패밀리입니다.

1. 문제 정의 (Problem)

• DFT 의 한계: 밀도 범함수 이론 (DFT) 은 현대 계산 화학 및 재료 과학의 기초이지만, 계산 비용이 매우 높아 ($O(n^3)$) 대규모 시스템이나 장시간 시뮬레이션에 적용하기 어렵습니다.
• 기존 MLIP 의 한계: 기존 머신러닝 상호작용 포텐셜 (MLIP) 은 DFT 를 빠르게 근사할 수 있지만, 대부분 특정 도메인 (예: 분자, 재료, 촉매 중 하나) 에만 특화되어 있어 일반화 능력이 부족했습니다. 또한, 데이터와 모델 크기를 확장할 때 정확도와 효율성 사이의 균형을 맞추는 것이 어려웠습니다.
• 다중 작업 학습의 어려움: 서로 다른 DFT 설정 (예: VASP/PBE vs ORCA/ωB97M-V) 을 사용하는 다양한 도메인 데이터를 통합하여 하나의 모델로 학습시키는 것은 기술적으로 난해한 과제였습니다.

2. 방법론 (Methodology)

2.1 대규모 데이터셋 통합

UMA 는 약 5 억 개 (500M) 의 고유한 3 차원 원자 구조로 구성된 방대한 데이터셋으로 학습되었습니다. 이는 다음과 같은 여러 공개 데이터셋을 통합하여 구성되었습니다:

• OMat24: 무기 재료 (Materials)
• OMol25: 분자 (Molecules)
• OC20/OC22/OC25: 촉매 및 흡착제 (Catalysis)
• OMC25: 분자 결정 (Molecular Crystals)
• ODAC25: 금속 유기 골격체 (MOFs)
  이 통합 데이터셋은 거의 모든 화학 공간 (방사성 원소 제외) 을 커버하며, 다양한 화학 환경과 상호작용을 학습할 수 있게 합니다.

2.2 아키텍처: 선형 전문가 혼합 (Mixture of Linear Experts, MoLE)

UMA 의 핵심 혁신은 MoLE (Mixture of Linear Experts) 아키텍처 도입입니다.

• 개념: 기존의 희소 (Sparse) MoE 와 달리, MoLE 는 선형 전문가들의 밀집된 (Dense) 혼합을 사용합니다.
• 장점:
  • 정확도 유지: 에너지 표면의 연속성을 보장하고 회전 공변성 (Rotational Equivariance) 을 유지합니다.
  • 효율성: 전체 모델 파라미터 수는 크지만 (예: UMA-M 은 14 억 개), 실제 추론 시 활성화되는 파라미터는 적습니다 (UMA-M 의 경우 약 5 천만 개).
  • 사전 계산 가능: 시스템 구성 (원소 종류, 전하, 스핀 등) 이 고정된 시뮬레이션 (예: 분자 동역학) 에서 전문가 가중치를 미리 합쳐 (Pre-merge) 추론 속도를 비 MoE 모델과 동일하게 유지할 수 있습니다.
• 입력 확장: 기존 eSEN 아키텍처에 전하 (Charge), 스핀 (Spin), 그리고 DFT 작업 설정 (Task) 정보를 임베딩하여 통합했습니다.

2.3 학습 전략

• 2 단계 학습:
  1. 1 단계 (Pre-training): 직접 힘 (Direct Force) 예측을 통해 BF16 정밀도로 빠르게 학습.
  2. 2 단계 (Fine-tuning): 자동 미분 (Autograd) 을 사용하여 에너지 보존 (Conservative) 특성을 갖춘 힘과 응력을 예측하도록 FP32 정밀도로 미세 조정.
• 스케일링 법칙 (Scaling Laws): 계산 비용, 데이터 크기, 모델 크기 간의 경험적 스케일링 법칙을 분석하여 주어진 컴퓨팅 예산 내에서 최적의 모델 크기를 결정했습니다.

3. 주요 기여 (Key Contributions)

1. 범용성 달성: 단일 모델 (Fine-tuning 없이) 이 재료, 분자, 촉매, 분자 결정, MOF 등 다양한 화학 도메인에서 전문화된 모델과 동등하거나 더 나은 성능을 발휘함을 증명했습니다.
2. MoLE 아키텍처 제안: MLIP 에 적용하여 모델 용량을 대폭 늘리면서도 추론 속도와 메모리 효율성을 유지하는 새로운 아키텍처를 제시했습니다.
3. 대규모 데이터와 스케일링 연구: 5 억 개 원자 시스템 데이터셋을 활용한 학습과 이를 위한 최적의 컴퓨팅/모델 스케일링 전략을 제시했습니다.
4. 오픈 소스 및 접근성: 코드, 가중치, 관련 데이터를 공개하여 커뮤니티의 연구 가속화를 도모했습니다.

4. 결과 (Results)

UMA 모델은 UMA-S (소형), UMA-M (중형), UMA-L (대형) 세 가지 크기로 제공되며, 다양한 벤치마크에서 SOTA(State-of-the-Art) 성능을 기록했습니다.

• 정확도:
  • 재료 (Matbench Discovery): UMA-M 이 최고 F1 점수를 기록했습니다.
  • 촉매 (AdsorbML): 흡착 에너지 계산 성공률이 기존 SOTA 대비 25% 향상되었습니다.
  • 분자 (OMol25): 약물 설계에 중요한 리간드 변형 에너지 (Ligand-strain energy) 예측에서 DFT 기준에 근접하는 정확도를 보였습니다.
  • 분자 결정: 결정 구조 예측 (CSP) 블라인드 테스트에서 격자 에너지 예측 오차가 3 kJ/mol 이하로 매우 정확했습니다.
• 속도 및 효율성:
  • UMA-S: 단일 80GB GPU 에서 1,000 개의 원자로 구성된 시스템을 초당 16 단계 (약 1.4 ns/day) 로 시뮬레이션할 수 있으며, 10 만 개 이상의 원자를 메모리에 로드할 수 있습니다.
  • MoLE 효과: MoLE 를 적용한 모델은 비 MoLE 모델보다 더 적은 활성 파라미터로 동일한 손실 (Loss) 을 달성할 수 있었습니다 (약 2.5 배 효율성 향상).
• 일반화: 특정 작업에 대한 미세 조정 (Fine-tuning) 없이도 다양한 도메인에서 우수한 성능을 발휘하여 범용 모델로서의 가능성을 입증했습니다.

5. 의의 및 결론 (Significance)

UMA 는 화학 및 재료 과학 분야에서 단일 범용 모델 (Universal Model) 이 실제 연구 응용에 충분한 정확도와 속도를 제공할 수 있음을 보여준 획기적인 연구입니다.

• 과학적 발견 가속화: DFT 의 계산 비용을 획기적으로 줄임으로써, 약물 발견, 에너지 저장, 반도체 제조 등 다양한 분야에서 대규모 시뮬레이션과 고속 스크리닝이 가능해졌습니다.
• AI 모델 패러다임 전환: 언어 및 비전 모델에서와 마찬가지로, 데이터와 모델 크기의 스케일링이 MLIP 분야에서도 성능 향상의 핵심임을 입증했습니다.
• 지속 가능한 발전: 오픈 소스 리소스를 통해 전 세계 연구자들이 더 강력하고 정교한 AI 모델을 개발할 수 있는 기반을 마련했습니다.

결론적으로, UMA 는 계산 화학의 새로운 표준을 제시하며, 단일 모델이 다양한 화학적 과제를 해결할 수 있는 "범용 원자 모델"의 시대를 열었다고 평가할 수 있습니다.