OmniMol: Transferring Particle Physics Knowledge to Molecular Dynamics with Point-Edge Transformers

본 논문은 포인트-에지 트랜스포머 아키텍처를 활용하고 고에너지 물리학으로부터의 지식 전이를 통해 최소한의 미세 조정과 독보적으로 빠른 추론으로 탁월한 성능을 달성하는 소분자를 위한 최첨단 기계학습 기반 원자간 퍼텐셜인 OmniMol을 소개합니다.

핵심

두 가지 매우 다른 세계를 상상해 보세요: 하나는 입자 물리학의 혼란스럽고 초고속 세계 (과학자들이 원자를 부딪혀 무엇이 튀어나오는지 관찰하는 곳) 이고, 다른 하나는 정교하고 끈적한 분자 화학의 세계 (원자들이 서로 붙어 약, 소재, 그리고 생명체를 형성하는 곳) 입니다.

오랫동안 이 두 분야의 과학자들은 각자의 세계를 이해하기 위해 완전히 다른 도구를 사용해 왔습니다. 하지만 이 논문에서 저자들은 OmniMol이라는 새로운 도구를 소개하며, 이미 구축한 '기초 모델 (foundation model)'을 활용하여 입자 물리학 전문가들에게 화학을 이해하도록 가르칩니다.

다음은 그들이 어떻게 수행했는지와 무엇을 발견했는지에 대한 간단한 요약입니다:

1. "마스터 셰프" 비유

Omnilearned라고 불리는 원래 모델을 **입자 제트 (particle jets)**로 수년 동안 요리해 온 마스터 셰프로 생각해 보세요.

• 재료: 입자 물리학에서 "제트"는 충돌에서 튀어나오는 (양성자와 중성자 같은) 아원자 입자들의 분무입니다.
• 기술: 이 셰프는 이러한 분무 속의 패턴을 인식하는 법을 배웠습니다. 그들은 입자들이 어떻게 상호작용하고, 어떻게 군집을 이루며, 다음에 어떤 일이 일어날지 예측하는 법을 알고 있습니다. 그들은 10 억 개의 서로 다른 입자 분무로 훈련되었습니다.

이제 저자들은 질문합니다: 이 같은 셰프가 분자 요리를 할 수 있을까요?

• 새로운 재료: 아원자 입자 대신 "재료"는 분자 내의 원자들 (탄소, 산소, 수소 등) 입니다.
• 도전 과제: 원자들은 아원자 입자들과 다르게 행동하지만, 유사한 구조를 공유합니다: 그들은 특정 유형을 가진 공간 내의 점들일 뿐입니다.

2. "범용 번역기" (아키텍처)

이를 작동시키기 위해 그들은 처음부터 새로운 셰프를 만들지 않았습니다. 기존 "마스터 셰프"(Omnilearned) 를 가져와 새로운 도구 세트를 제공했습니다:

• Point-Edge Transformer (PET): 셰프가 음식 접시를 바라본다고 상상해 보세요. 이 도구는 한 번에 하나의 재료만 보는 것이 아니라, 모든 재료를 한 번에 보고 각 재료가 다른 모든 재원과 어떻게 관련되는지 보게 합니다.
• "물리적 편향 (Physics Bias)": 이것이 비결입니다. 이 모델에는 "이 두 입자/원자는 서로 가까이 있으므로 서로에게 더 많은 주의를 기울여야 한다"는 규칙이 내장되어 있습니다. 이는 모델이 노이즈에 혼란을 느끼지 않고 가장 중요한 관계들에 집중하도록 돕습니다.

3. 실험: 파인튜닝 (Fine-Tuning)

저자들은 이 입자 훈련 모델을 가져와 oMol(수백만 개의 분자 컬렉션) 이라는 데이터 세트를 사용하여 화학에 대한 "단기 집중 과정"을 제공했습니다.

• 목표: 그들은 이 모델이 **Machine-Learned Interatomic Potential (MLIP)**로 작동하기를 원했습니다. 쉬운 말로, 이 모델은 임의의 원자 그룹에 대해 두 가지를 예측해야 한다는 뜻입니다:
  1. 에너지: 그들을 붙잡고 있는 "접착제"는 얼마나 많은가?
  2. 힘: 하나의 원자를 밀면, 그것은 얼마나 강하게 밀어내겠는가?

4. 결과: 빠르고 놀라울 정도로 훌륭함

이 논문은 몇 가지 흥미로운 사실을 발견했습니다:

• "Few-Shot" 초능력: 일반적으로 컴퓨터에게 화학을 가르치려면 방대한 양의 데이터가 필요합니다. 하지만 OmniMol 은 입자 물리학의 "지식"으로 시작했기 때문에 화학을 매우 빠르게 배웠습니다. 상대적으로 적은 양의 새로운 데이터 (약 10 만 개의 분자 등) 로도 수백만 개로 훈련된 모델만큼 잘 수행했습니다. 이는 맛과 열에 대한 기본 원리를 이미 이해하고 있는 마스터 셰프가 몇 가지 레시피만으로 새로운 요리를 배울 수 있는 것과 같습니다.
• 속도: OmniMol 은 놀라울 정도로 빠릅니다. 다른 모델들이 분자의 움직임을 계산하는 데 오랜 시간이 걸릴 수 있는 반면, OmniMol 은 눈 깜짝할 사이에 해냅니다. 저자들은 계산 시간 1 시간당 OmniMol 이 일부 경쟁 모델들보다 3 배 더 많은 분자를 시뮬레이션할 수 있다고 지적합니다.
• 트레이드오프: 그들이 엄청난 양의 데이터 (수백만 개의 분자) 를 가졌을 때, 입자 물리학 지식으로 시작하는 이점은 다소 약화되었습니다. 이는 "입자 물리학 지식"이 강력한 출발점 역할을 하지만, 처음부터 모델을 훈련할 만큼 충분한 시간과 데이터가 있다면 그 출발점의 중요성은 줄어든다는 것을 시사합니다.

5. 큰 그림

이 논문은 OmniMol이 한 과학 분야 (입자 물리학) 를 위해 구축된 "기초 모델"이 완전히 다른 분야 (화학) 로 성공적으로 이전된 첫 사례라고 결론지었습니다.

그들은 한 분야에서 공간 내 점들이 어떻게 상호작용하는지 이해하는 똑똑한 모델이 있다면, 이를 다른 분야에서 공간 내 점들이 어떻게 상호작용하는지 이해하도록 적응시켜 시간과 컴퓨팅 파워를 절약할 수 있음을 증명했습니다.

요약하자면: 저자들은 고에너지 입자 충돌로 훈련된 초지능 AI 를 가져와, 입자가 아닌 원자를 이해하도록 그 두뇌를 조정했고, 그 결과 특히 데이터가 부족할 때 분자의 거동을 예측하는 데 있어 매우 빠르고 정밀한 도구가 되었습니다.

기술 요약

기술 요약: OmniMol

문제 정의

머신러닝 (ML) 은 입자 물리학과 분자 화학을 포함한 복잡한 물리 시스템의 표현 및 시뮬레이션을 혁신적으로 변화시켰습니다. 이러한 분야들은 에너지 규모에서 극적으로 다르지만, 위상 공간 내의 가변 크기의 입자 (또는 원자) 집합이라는 근본적인 데이터 구조를 공유하며, 이는 효과적으로 구조화된 점 구름 (point clouds) 을 형성합니다.

주요 과제는 효율적인 머신러닝 기반 원자 간 전위 (MLIPs) 개발입니다. 밀도 범함수 이론 (DFT) 과 같은 전통적인 방법은 계산 비용이 매우 높아 대규모 및 장기 분자 역학 (MD) 시뮬레이션을 제한합니다. MLIPs 는 이러한 비용의 일부만으로 전위 에너지 표면과 힘을 근사하는 것을 목표로 합니다. 그러나 견고한 MLIPs 를 훈련하려면 일반적으로 방대한 데이터셋과 상당한 계산 자원이 필요합니다. 본 논문은 입자 물리학 (특히 입자 제트) 의 점 구름으로 사전 훈련된 기반 모델이 분자 역학으로 전이될 수 있으며, 이는 저데이터 환경에서 최적화를 가속화하고 정확도를 향상시킬 수 있다는 가설을 제기합니다.

방법론

아키텍처: 포인트-에지 트랜스포머 (PET)

OmniMol 은 고에너지 물리학 (HEP) 에서 입자 제트 분류 및 생성을 위해 설계된 기반 모델인 Omnilearned를 적응시켜 구축되었습니다. 핵심 아키텍처는 $k$-최근접 이웃에 대한 국소 어텐션과 전역 모든-대-모든 트랜스포머 블록을 결합한 **포인트-에지 트랜스포머 (PET)**입니다.

주요 아키텍처 구성 요소는 다음과 같습니다:

1. 입력 임베딩: 원자들은 위치 정보 ($\vec{r}$), 이산 원자 번호 ($Z$), 그리고 전하, 스핀과 같은 추가 특징을 결합한 토큰 공간에 임베딩됩니다.
2. 국소 어텐션 블록: 각 원자에 대해 $K$-최근접 이웃 ($K=15$, 제트의 경우 $K=10$과 비교) 을 사용하여 국소 이웃이 구성됩니다. 거리 항, 거리의 역제곱, 그리고 학습된 원자 임베딩 함수를 포함한 쌍별 물리 특징이 계산됩니다. 이러한 특징들은 작은 국소 트랜스포머에 의해 처리되어 국소 임베딩 벡터를 생성합니다.
3. 상호작용 편향이 있는 전역 어텐션: 전역 자기 어텐션 메커니즘은 쌍별 물리 특징에서 유도된 명시적 편향을 통합합니다. 어텐션 로짓은 $A^*_{ij} = A_{ij} + B_{ij}$로 수정되며, 여기서 $B_{ij}$는 MLP 에 임베딩된 편향 항입니다. 이 "상호작용 행렬 어텐션 편향"은 쌍별 물리 사전 지식을 트랜스포머에 직접 주입하여 표현력을 희생하지 않으면서 네트워크를 물리적으로 의미 있는 이웃으로 유도합니다.
4. 출력 헤드: Omnilearned 의 생성 헤드는 두 가지 작업에 재사용됩니다:
   • 힘 예측: 원자별 힘을 예측하는 치환 공변 (permutation-equivariant) 헤드.
   • 에너지 예측: 전체 분자 에너지를 산출하기 위해 합산되는 원자별 에너지 보정을 예측하는 헤드. 이는 확장적 사전 지식을 보존합니다.

불변성 및 보존 제약 조건

물리적 제약을 충족하기 위해 저자는 두 가지 요구 사항을 다룹니다:

• 에너지 보존: 힘은 직접 예측되지 않고 에너지 출력의 역전파를 통해 계산됩니다 ($\vec{F}_i = \nabla_{\vec{r}_i} E$). 이는 정확한 에너지 보장을 하지만 훈련 중 계산 비용을 증가시킵니다 (이중 역전파 필요). 따라서 이 제약은 "소형" 모델 변형에만 적용됩니다.
• 회전 공변성: 표준 아키텍처는 원시 좌표 차이가 MLP 에 입력되기 때문에 본질적으로 공변적이지 않습니다. 이를 해결하기 위해 저자는 "공변적이고 보존적인" 변형을 도입합니다. 이 버전은 쌍별 특징에서 직접적인 좌표 차이 항을 제거하고 대신 국소 블록에 인접한 원자 간의 벡터가 형성하는 각도의 코사인 (각도 정보) 을 통합합니다. 이 수정은 좌표 항을 제거함으로써 발생하는 성능 손실을 크게 회복하면서 공변성을 유지합니다.

훈련 및 파인튜닝 전략

모델은 oMol 데이터셋 (특히 oMol-25, oMol-4M, oMol-100M, oMol-140M 하위 집합) 에서 파인튜닝됩니다. 두 가지 전략이 탐구됩니다:

1. LoRA (저랭크 적응): 사전 훈련된 PET 백본 가중치는 고정됩니다. 저랭크 어댑터는 트랜스포머 본체 행렬 ($W_Q, W_K, W_V, W_O, W_{MLP}$) 에만 도입되며, 분자 입력 인코더, 편향 MLP, 작업 헤드의 훈련과 병행됩니다. 또한 학습된 임베딩을 수정하기 위해 "임베딩 적응" 레이어가 추가됩니다.
2. 전체 파인튜닝: 본체 및 입력 인코더의 모든 가중치가 고정 해제되어 훈련되며, 작업 헤드는 처음부터 훈련됩니다.

훈련 목적 함수는 에너지와 힘에 대한 평균 절대 오차 (MAE) 의 합을 최소화하며, 힘은 더 무겁게 가중치 ($\lambda_F = 10$) 가 부여됩니다.

주요 결과

oMol 성능

• 전체 훈련: 대규모 데이터셋 (oMol-4M 및 oMol-100M/140M) 에서 훈련될 때, OmniMol 은 최첨단 MLIPs 와 경쟁력 있는 성능을 달성합니다. 예를 들어, oMol-140M 에서 OmniMol-large 모델은 원자당 에너지 MAE 1.04 meV 및 힘 MAE 13.59 meV/Å를 달성합니다.
• 저데이터 환경: 훈련 데이터가 제한될 때 가장 큰 개선이 관찰됩니다. 10 만 개의 분자만으로도 파인튜닝하거나 oMol-4M 에서 매우 적은 에포크 (2 회 통과) 로 파인튜닝할 때, 사전 훈련된 OmniMol 변형은 처음부터 훈련된 모델보다 훨씬 뛰어난 성능을 보입니다.
  • 10 만 개 하위 집합에서, 사전 훈련은 중간 모델의 에너지 MAE 를 최대 **29.4%**까지, 힘 MAE 를 **26.9%**까지 개선했습니다.
  • oMol-4M 에서 단 2 에포크만 훈련했을 때, 중간 모델은 사전 훈련되지 않은 대응 모델에 비해 에너지 MAE 에서 54.6%, 힘 MAE 에서 **56.9%**의 개선을 보였습니다.
• 공변/보존 변형: 공변적이고 보존적인 모델 변형은 저데이터 환경에서 (특히 힘에 대해) 성능이 크게 개선되지만, 데이터셋 크기가 증가함에 따라 이러한 이점은 감소합니다.

확장성 및 추론 속도

• 확장성: OmniMol 은 모델 크기에 따라 깔끔한 멱법칙 (power-law) 확장성을 따르며, 10 억 개 파라미터까지 포화 징후가 나타나지 않아 최근 트랜스포머 기반 MLIPs 에 대한 연구 결과와 일치합니다.
• 추론 속도: 대규모 파라미터 수에도 불구하고, 트랜스포머를 위한 하드웨어 최적화 덕분에 OmniMol 은 고유하게 빠른 추론 속도를 보여줍니다. 약 100 개의 원자로 구성된 시스템에서 A100 GPU 를 사용할 때, OmniMol-medium 은 비교 가능한 그래프 신경망 (GNN) 베이스라인 (eSEN-md-d 및 AllScAIP-md) 보다 약 3 배 빠르며, 경쟁력 있는 정확도 (AllScAIP-md 대비 원자당 에너지 오차가 약 0.7 meV/Å만 높음) 를 유지합니다.

중요성 및 주장

본 논문은 과학적 점 구름 기반 모델에 대한 학제 간 전이의 첫 번째 증명을 제시한다고 주장합니다. 고에너지 물리학 입자 제트로 사전 훈련된 모델을 분자 역학에 적응시킴으로써, 저자는 다음을 증명합니다:

1. 교차 영역 전이의 실현 가능성: 입자 물리학을 위해 구축된 기반 모델이 분자 화학으로 효과적으로 전이될 수 있으며, 이는 근본적인 점 구름 구조가 극도로 다른 물리 규모에 걸쳐 학습 가능한 특징을 공유함을 시사합니다.
2. 유도 편향의 학습 가속화: 사전 훈련은 강력한 유도 편향으로 작용합니다. 데이터가 부족할 때 공변성이 도움이 되는 것과 유사하게, 사전 훈련의 "쓴 교훈"은 훈련 데이터가 제한될 때 빠른 최적화와 향상된 정확도를 가능하게 합니다.
3. 효율성: 아키텍처 전이는 고유하게 빠른 추론 속도를 가능하게 하며, 이는 소분자 신약 개발과 같은 설계 공간의 빠른 탐색이 필요한 응용 분야에 중요합니다.

저자는 연구가 MLIPs 에 초점을 맞추고 있지만, 점 구름 기반 모델에 관한 교훈은 상호작용하는 무질서한 집합으로 기술되는 시스템이 존재하는 과학 분야 전반에 걸쳐 광범위한 유용성을 가질 수 있다고 결론지었습니다. 그들은 모든 환경에서 기존 모든 방법보다 보편적으로 우월하다고 주장하지는 않지만, 저데이터 시나리오와 추론 속도에서의 특정 이점을 강조합니다.