Zatom-1: A Multimodal Flow Foundation Model for 3D Molecules and Materials

이 논문은 분자와 물질을 모두 다루는 3D 화학 모델링을 위해 생성 및 예측 작업을 통합한 최초의 오픈소스 기반 모델인 'Zatom-1'을 제안하며, 멀티모달 흐름 매칭을 통해 기존 특화 모델 대비 우수한 성능과 빠른 추론 속도를 달성하고 도메인 간 전이 학습 효과를 입증했습니다.

핵심

ZATOM-1: 화학의 '만능 요리사'를 만나다

이 논문은 ZATOM-1이라는 새로운 인공지능 모델을 소개합니다. 이 모델은 분자 (약물 같은 작은 입자) 와 재료 (금속, 결정 같은 큰 구조) 를 모두 이해하고, 새로운 것을 만들어낼 수 있는 **초능력을 가진 '만능 요리사'**라고 생각하시면 됩니다.

기존의 AI 들은 "약물 요리만 하는 셰프"이거나 "금속 요리만 하는 셰프"처럼 한 가지 일만 잘했습니다. 하지만 ZATOM-1 은 모든 재료를 다룰 수 있는 천재 셰프입니다.

이 모델을 쉽게 이해할 수 있도록 몇 가지 비유로 설명해 드릴게요.

---

1. 왜 이 모델이 필요한가요? (기존의 문제점)

지금까지 과학자들은 새로운 약을 만들거나 더 강한 재료를 개발할 때, AI 에게 "이런 모양의 분자를 만들어줘"라고 요청하거나 "이 분자의 성질이 뭐야?"라고 물어보곤 했습니다.

• 문제: 기존 AI 들은 '생성 (만들기)'과 '예측 (알기)'을 따로 배웠습니다. 마치 요리사가 레시피를 외우는 것만 배우고, 실제로 요리를 하거나 맛을 보는 법은 따로 배운 것과 비슷합니다.
• 결과: 서로 다른 일을 하려면 모델을 따로 만들어야 했고, 지식이 공유되지 않아 비효율적이었습니다.

2. ZATOM-1 의 핵심 아이디어: "흐름 맞추기" (Flow Matching)

ZATOM-1 은 **'흐름 맞추기 (Flow Matching)'**라는 기술을 사용합니다. 이를 비유하자면 다음과 같습니다.

• 비유: imagine you have a jar of mixed-up Lego bricks (source) and you want to build a specific castle (target).
  • 기존 AI 들은 레고 조각을 하나씩 붙여가며 (autoregressive) 천천히 성을 만들거나, 잠긴 상자에 숨겨진 레고 (latent space) 를 꺼내야 하는 복잡한 과정을 거쳤습니다.
  • ZATOM-1 의 방식: 이 모델은 **"레고 조각들이 어떻게 움직여야 성으로 변할지"**를 직접 계산합니다. 마치 흐르는 강물이 바다로 향하듯, 무질서한 원자 (레고) 들이 자연스럽게 원하는 모양 (분자나 재료) 으로 흐르게 유도하는 것입니다.
  • 장점: 이 방식은 매우 빠르고, 별도의 복잡한 사전 준비 (잠재 공간 인코더) 없이 바로 3D 공간에서 원자 위치와 종류를 동시에 예측할 수 있습니다.

3. ZATOM-1 의 두 가지 주요 능력

이 모델은 두 가지 큰 일을 동시에 합니다.

① 새로운 것 창조하기 (생성 능력)

• 상황: "새로운 약을 만들어줘" 또는 "더 강한 배터리 재료를 찾아줘."
• ZATOM-1 의 역할: 원자 (Atom) 들의 종류와 3D 위치를 마치 점토를 빚듯이 자유롭게 조합하여, 실제로 존재할 수 있는 새로운 분자나 결정 구조를 만들어냅니다.
• 성과: 기존 모델들보다 약 10 배 더 빠르며, 만들어진 분자가 물리적으로 타당한지 (예: 원자들이 겹치지 않는지) 확인하는 테스트에서 압도적인 성적을 냈습니다.

② 속성 예측하기 (예측 능력)

• 상황: "이 분자가 얼마나 단단할까?", "전기를 얼마나 잘 통할까?"
• ZATOM-1 의 역할: 앞서 만들어낸 '만들기' 경험을 바탕으로, 새로운 분자의 성질을 즉시 예측합니다.
• 신기한 점 (전이 학습): ZATOM-1 은 분자 (작은 것) 와 재료 (큰 것) 를 함께 배우면서 서로의 지식을 공유합니다. 마치 "요리하는 법을 배운 셰프가 그 지식을 바탕으로 음식의 맛을 더 잘 예측하는 것"과 같습니다.
  • 결과: 재료를 배운 경험이 분자 예측의 정확도를 높여주었고, 반대의 경우도 마찬가지였습니다. 이는 과학계에서 처음 있는 일입니다.

4. 왜 이것이 혁신적인가요?

1. 단 하나의 모델로 모든 것: 분자와 재료를 구분하지 않고 하나의 모델로 처리합니다.
2. 압도적인 속도: 기존 모델들이 10,000 개의 분자를 만드는 데 몇 시간이 걸렸다면, ZATOM-1 은 4 분도 안 걸립니다. (NVIDIA A100 GPU 기준)
3. 확장성: 모델의 크기를 키우면 성능이 꾸준히 좋아집니다. 마치 더 큰 뇌를 가진 천재가 더 많은 일을 해내는 것처럼 예측 가능합니다.
4. 오픈 소스: 이 모델의 코드와 가중치는 모두 공개되어 있어, 누구나 무료로 사용할 수 있습니다.

5. 결론: 화학의 미래를 바꾸다

ZATOM-1 은 단순히 "빠른 AI"가 아닙니다. 이는 화학 연구의 패러다임을 바꾸는 도구입니다.

• 과거: 새로운 약이나 재료를 찾려면 수천 년의 실험과 시행착오가 필요했습니다.
• 미래: ZATOM-1 같은 AI 가 수초 만에 수만 가지 후보를 생성하고, 가장 유망한 것만 실험실로 보내면 됩니다.

이 기술은 새로운 암 치료제, 더 효율적인 태양전지, 차세대 배터리 등을 훨씬 빠르게 발견하게 해줄 것입니다. 마치 화학이라는 거대한 미로에서 길을 잃지 않고 가장 빠른 출구를 찾아주는 나침반과 같은 존재입니다.

---

한 줄 요약:

ZATOM-1 은 분자와 재료를 동시에 '만들고' '이해'하는 초고속 AI 요리사로, 기존에 불가능했던 속도와 정확도로 과학적 발견을 가속화합니다.

기술 요약

ZATOM-1: 3D 분자와 물질을 위한 멀티모달 흐름 기초 모델 (Multimodal Flow Foundation Model) 기술 요약

본 논문은 3D 분자 (비주기적) 와 물질 (주기적 결정) 의 생성 및 예측을 통합한 최초의 엔드 - 투 - 엔드 오픈소스 기초 모델인 ZATOM-1을 소개합니다. 기존 AI 접근법이 특정 도메인 (분자 또는 물질) 이나 단일 작업 (생성 또는 예측) 에 최적화되어 있어 표현 공유와 전이 학습이 제한되었던 문제를 해결합니다.

1. 문제 정의 (Problem)

• 분리된 모델링: 기존 연구들은 분자 생성 (De novo generation) 과 물질 생성, 그리고 물성 예측 (Property prediction) 을 별도의 모델로 처리하는 경우가 많았습니다. 이로 인해 도메인 간 지식 공유가 어렵고, 모델 개발 비용이 중복되었습니다.
• 복잡한 아키텍처: 많은 3D 생성 모델이 잠재 공간 (Latent space) 을 사용하는 잠재 확산 (Latent Diffusion) 모델을 기반으로 하여, 인코더/디코더 학습이 필요하고 추론 속도가 느리며 불안정할 수 있습니다.
• 스케일링 한계: 과학적 기초 모델 (Scientific Foundation Models) 의 경우, 파라미터 수 증가에 따른 성능 향상이 예측 가능하지 않거나, 도메인 간 전이 학습 (Transfer Learning) 이 효과적으로 이루어지지 않았습니다.

2. 방법론 (Methodology)

ZATOM-1 은 Transformer 아키텍처를 기반으로 하며, **멀티모달 흐름 매칭 (Multimodal Flow Matching)**을 학습 목표로 사용합니다.

2.1. 데이터 표현 (Tokenization)

원자 단위 (Atomic granularity) 에서 3D 시스템을 표현하기 위해 5 가지 통합 모달리티를 사용합니다:

1. 원자 유형 (Atom types, $A$): 이산적 (Discrete) 모달리티 (원소 종류).
2. 3D 좌표 (3D coordinates, $X$): 연속적 (Continuous) 모달리티 (Angstrom 단위).
3. 분수 좌표 (Fractional coordinates, $F$): 단위 셀 내 상대적 위치.
4. 단위 셀 길이 (Unit cell lengths, $L_{len}$): 결정 격자의 크기.
5. 단위 셀 각도 (Unit cell angles, $L_{ang}$): 결정 격자의 각도.

• 비주기적 분자의 경우 격자 관련 입력은 마스킹되며, 주기적 물질의 경우 3D 좌표는 회전 불변성을 위해 마스킹됩니다.

2.2. 아키텍처: Trunk-based Flow Transformer (TFT)

• Standard Transformer: 기존에 복잡한 그래프 신경망 (GNN) 이나 수동으로 설계된 생성 사전 지식 (Priors) 에 의존하던 것과 달리, 표준 Transformer 아키텍처를 사용하여 모든 모달리티를 공유 잠재 표현으로 학습합니다.
• 조건부 흐름 매칭 (Conditional Flow Matching):
  • 소스 분포 (노이즈) 에서 타겟 분포 (데이터) 로 샘플을 이동시키는 속도장 (Velocity field) 을 학습합니다.
  • 이산 (Discrete) 및 연속 (Continuous) 동시 처리: 원자 유형은 이산 흐름 모델 (Discrete Flow Models) 프레임워크를, 좌표 및 격자 파라미터는 유클리드 확산/흐름 매칭을 통해 동시에 학습합니다.
  • 비자기회귀 (Non-autoregressive): 한 번에 전체 구조를 생성하므로 생성 속도가 매우 빠릅니다.

2.3. 학습 단계 (Two-Stage Training)

1. 생성적 사전 학습 (Generative Pretraining):
   • MP20 (물질) 과 QM9/GEOM-Drugs (분자) 데이터셋을 혼합하여 학습합니다.
   • 목적: 3D 분자와 물질의 유효한 구조를 생성하는 능력을 배우고, 풍부한 표현 (Embeddings) 을 획득합니다.
   • 손실 함수: 연속 모달리티에 대한 MSE 와 이산 모달리티에 대한 교차 엔트로피를 가중치 ($\lambda_{discrete}$) 와 함께 결합합니다.
2. 예측적 파인튜닝 (Predictive Finetuning):
   • 사전 학습된 Trunk 가중치를 고정 (Freeze) 하거나 일부만 미세 조정합니다.
   • 하위 작업 (물성, 에너지, 힘 예측) 을 위한 보조 헤드 (Heads) 를 추가하여 학습합니다.
   • 제로샷 (Zero-shot) 전이: 분자 데이터로만 학습된 모델이 물질 예측에, 혹은 그 반대로 적용될 수 있는 전이 학습 효과를 검증합니다.

3. 주요 기여 (Key Contributions)

1. 첫 번째 통합 기초 모델: 3D 분자와 물질을 동시에 생성하고 예측할 수 있는 최초의 오픈소스 엔드 - 투 - 엔드 기초 모델을 제안했습니다.
2. 성능 및 효율성:
   • 생성: QM9, GEOM-Drugs, LeMat-GenBench(물질) 에서 SOTA(최고 수준) 또는 경쟁력 있는 결과를 달성했습니다.
   • 속도: 잠재 확산 (Latent Diffusion) 기반 모델 대비 10 배 이상 빠른 추론 속도와 3 배 빠른 학습 속도를 달성했습니다 (예: 300M 파라미터 모델이 4 분 이내에 10,000 개 샘플 생성).
   • 스케일링: 모델 크기가 커질수록 성능이 예측 가능하게 향상되는 스케일링 법칙을 보였습니다.
3. 긍정적 전이 학습 (Positive Transfer):
   • 분자와 물질을 함께 사전 학습한 것이 도메인 간 예측 정확도를 향상시킵니다. 특히, 물질 생성 학습이 분자 물성 예측 정확도를 높이는 것을 실험적으로 증명했습니다.
4. 간소화된 설계: 사전 학습된 오토인코더 (Autoencoder) 나 잠재 공간이 필요 없어, 모델 구조가 단순하고 학습/추론이 안정적입니다.

4. 실험 결과 (Results)

• 물질 생성 (LeMat-GenBench):
  • MatterGen, PLaID++ 등 기존 모델들과 비교하여 유효성 (Validity), 고유성 (Uniqueness), 안정성 (Stability) 측면에서 경쟁력 있는 성능을 보였습니다.
  • 특히, 완화 (Relaxation) 과정 없이도 높은 안정성을 가진 결정을 생성했습니다.
• 분자 생성 (QM9, GEOM-Drugs):
  • QM9 에서 생성된 분자의 약 99% 가 PoseBusters(물리적 타당성 검사) 를 통과하여 SOTA 를 기록했습니다.
  • GEOM-Drugs(대형 약물 유사 분자) 에서도 94% 이상의 유효성을 달성했습니다.
• 물성 예측:
  • QM9 의 다양한 물성 예측에서 단일 작업 (Single-task) 및 다중 작업 (Multi-task) 학습 기반 SOTA 모델들을 능가하거나 경쟁하는 성능을 보였습니다.
  • 분자와 물질 데이터를 함께 학습한 모델이 개별 도메인만 학습한 모델보다 더 나은 예측 성능을 보였습니다.
• 에너지 및 힘 예측:
  • MPtrj(물질) 와 OMol25(분자) 데이터셋에서 힘 (Force) 예측 시 Orb-v1 과 같은 SOTA 모델과 경쟁하거나 상회하는 성능을 보였습니다.

5. 의의 및 결론 (Significance)

• 과학적 머신러닝 (SciML) 의 패러다임 전환: 생성 모델링을 단순한 생성 도구가 아닌, 강력한 자기지도 학습 (Self-supervised) 사전 학습 전략으로 활용하여 화학적 표현을 학습하는 새로운 방식을 제시했습니다.
• 효율성과 확장성: 복잡한 인과적 (Autoregressive) 모델이나 잠재 공간 기반 모델의 단점을 극복하고, 표준 Transformer 아키텍처를 통해 계산 효율성과 확장성을 동시에 확보했습니다.
• 미래 지향성: ZATOM-1 은 신약 개발, 신소재 발견, 촉매 설계 등 다양한 화학 및 재료 과학 분야에서 단일 모델로 활용될 수 있는 기반을 마련했습니다. 코드와 가중치는 오픈소스로 공개되어 연구 커뮤니티의 발전을 촉진할 것으로 기대됩니다.

요약하자면, ZATOM-1 은 3D 화학 시스템의 생성과 예측을 통합하여 속도와 정확도를 모두 혁신한 최초의 범용 기초 모델로, 과학적 발견의 가속화에 중요한 기여를 할 것으로 평가됩니다.