Machine Learning Interatomic Potentials: Advancing Open-Source Software for Efficient and Scalable Molecular Simulation

본 논문은 재설계된 모듈형 API, 고성능 공변 백엔드, eSEN 아키텍처 및 개선된 정전기 처리와 같은 고급 기능을 통해 기계 학습 원자 간 퍼텐셜의 효율성, 확장성 및 유연성을 향상시키는 차세대 오픈소스 소프트웨어인 mlip v2 를 소개합니다.

핵심

수십억 개의 작고 움직이는 기어 (원자) 로 구성된 복잡한 기계가 어떻게 작동하는지 시뮬레이션해 보려고 상상해 보세요. 가장 정확한 그림을 얻으려면 양자 물리학의 법칙을 사용해야 하지만, 그렇게 하는 것은 슈퍼컴퓨터로 시뮬레이션 1 초를 완료하는 데 수년이 걸리는 방식으로 기어 하나하나의 경로를 계산해 보려는 것과 같습니다. 이는 실용적일 만큼 너무 느립니다.

이제 **기계 학습 원자 간 퍼텐셜 (MLIPs)**이 등장합니다. 이것들을 "현명한 단축키"로 생각하세요. 이들은 그 느리고 완벽한 물리 계산 결과들로 훈련된 AI 모델입니다. 일단 훈련되면, 슈퍼컴퓨터와 거의 동일한 정확도로 원자의 움직임을 거의 즉시 예측할 수 있지만, 그 시간은 그보다 훨씬 짧습니다.

그러나 지금까지 이러한 현명한 단축키를 사용하는 것은 고장 난 조향 장치와 한 특정 도시에서만 작동하는 지도로 고성능 레이싱 카를 운전해 보려는 것과 같았습니다. 도구들은 분산되어 있었고, 확장하기 어려우며 경직되어 있었습니다.

이 논문은 이러한 시뮬레이션을 구동하는 소프트웨어 툴킷의 주요 업그레이드인 mlip v2를 소개합니다. 그들이 무엇을 구축했는지 간단히 설명해 보겠습니다:

1. 새로운 엔진실 (소프트웨어 프레임워크)

저자들은 소프트웨어의 "엔진실"을 완전히 재설계했습니다.

• 과거의 방식: 모든 도구가 특정 손잡이에 접착된 도구 상자를 상상해 보세요. 손잡이를 바꾸려면 도구를 부러뜨려야 했습니다.
• 새로운 방식 (mlip v2): 데이터 처리, 훈련, 시뮬레이션 등 모든 도구가 고품질 레고 블록처럼 서로 맞물리는 모듈식 시스템을 구축했습니다. 전체 구조를 깨뜨리지 않고도 조각들을 쉽게 교체할 수 있습니다. 이로 인해 과학자들이 자신의 특정 필요에 맞춰 소프트웨어를 커스터마이징하기가 훨씬 쉬워졌습니다.

2. 터보차저 (e3j 백엔드)

이러한 시뮬레이션에서 가장 큰 병목 현상 중 하나는 3 차원 모양과 관련된 복잡한 수학 ( "공변 연산"이라고 함) 을 수행하는 것입니다.

• 유사성: 3 차원 물체를 머릿속에서 회전시켜 보려고 상상해 보세요. 수백만 개의 원자에 대해 이를 수행하는 것은 지칠 대로 지치게 합니다.
• 해결책: 그들은 e3j라는 새로운 고속 엔진을 통합했습니다. 이는 3 차원 수학에 특화된 터보차저를 소프트웨어에 장착한 것과 같습니다. 논문은 이것이 현대 컴퓨터 칩 (GPU 및 TPU) 에서 소프트웨어를 최대 3 배 더 빠르게 실행하게 만든다고 보여줍니다.

3. 새로운 초능력

이 업데이트는 단순히 속도를 높인 것뿐만 아니라, 소프트웨어에 이전에는 없던 새로운 능력을 부여했습니다:

• "전문가" 시스템 (전문가 혼합):

  • 문제: 물 (water) 에서 복잡한 약물까지 모든 유형의 분자에 대해 하나의 거대한 뇌를 훈련시키는 것은 어렵습니다. 종종 혼란에 빠집니다.
  • 해결책: 그들은 eSEN이라는 아키텍처를 도입했는데, 이는 전문가 팀처럼 작동합니다. 모든 것을 알고자 하는 하나의 뇌 대신, 시스템은 서로 다른 문제를 모델 내의 서로 다른 "전문가"로 라우팅합니다. 이를 통해 압도되지 않고 방대하고 복잡한 데이터셋에서 학습할 수 있습니다.

• 전기 이해 (정전기):

  • 문제: 원자는 종종 전하를 띱니다. 이전 모델들은 총 전하가 변하는 시스템을 처리하는 데 어려움을 겪어 부정확한 예측을 초래했습니다.
  • 해결책: 새로운 버전은 시스템의 총 전하를 명시적으로 "듣습니다". 이는 AI 에게 항상 "북쪽" (총 전하) 을 아는 나침반을 주는 것과 같아서, 배터리나 소금물과 같은 이온이 포함된 전하 시스템을 훨씬 더 정확하게 모델링할 수 있게 합니다.

• 곡선 감지 (헤세안 레이블):

  • 문제: 원자의 움직임 (힘) 을 아는 것은 언덕의 경사를 아는 것과 같습니다. 하지만 공이 어떻게 굴러가며 진동하는지 예측하려면 언덕의 곡률도 알아야 합니다.
  • 해결책: 이제 소프트웨어는 이 "곡률" (헤세안이라고 함) 을 예측하도록 훈련될 수 있습니다. 이는 AI 가 에너지 지형의 모양을 더 잘 이해하도록 도와주어, 분자가 어떻게 진동하고 반응하는지에 대한 예측을 더 정확하게 만듭니다.

• 경로 찾기 (전이 상태 탐색):

  • 문제: 화학 반응이 일어날 때, 반응물은 다른 쪽으로 가기 위해 고에너지 "산길" (전이 상태) 을 통과해야 합니다. 이 길을 찾는 것은 건초더미에서 바늘을 찾는 것과 같습니다.
  • 해결책: 그들은 NEB(Nudged Elastic Band) 라는 내장 도구를 추가했는데, 이는 시작점과 끝점 사이에 원자들의 고무줄을 자동으로 당겨서 그 산길을 효율적으로 찾습니다.

• 숨 쉴 공간 (NPT 앙상블):

  • 문제: 실제 세계에서는 액체와 고체가 압력이나 온도가 변할 때 팽창하고 수축합니다. 이전 시뮬레이션들은 종종 용기 크기를 고정했는데, 이는 현실적이지 않습니다.
  • 해결책: 새로운 소프트웨어는 이제 뜨거운 공기 속에서 팽창하는 실제 풍선처럼 압력을 일정하게 유지하기 위해 용기 크기가 변하는 시스템을 시뮬레이션할 수 있습니다 (NPT).

4. 결과

저자들은 거대한 분자 데이터셋으로 이미 훈련된 사전 훈련된 모델 ( "뇌") 을 출시했는데, 이는 사용 준비가 되어 있습니다. 그들은 이러한 모델을 테스트했고 에너지, 힘, 심지어 원자의 전하를 예측하는 데 매우 정확하다는 것을 발견했습니다.

요약하자면: 저자들은 원자를 시뮬레이션하기 위한 강력하지만 투박한 도구를 가져와서 세련되고 모듈식이며 번개처럼 빠른 플랫폼으로 변모시켰습니다. 그들은 새로운 "근육" (속도), 새로운 "감각" (전하 및 곡률 인식), 그리고 새로운 "도구" (반응 경로 찾기) 를 추가하여, 이전에는 너무 어렵거나 느려서 모델링하기 어려웠던 복잡하고 현실적인 화학 시스템을 시뮬레이션할 수 있게 만들었습니다. 이 소프트웨어는 오픈 소스이므로 누구나 다운로드하여 즉시 사용할 수 있습니다.

기술 요약

기술 요약: mlip v2 – 효율적이고 확장 가능한 분자 시뮬레이션을 위한 오픈소스 소프트웨어 발전

문제 제기

기계 학습 원자간 퍼텐셜 (MLIPs) 은 전자 구조 방법 (예: 밀도 범함수 이론, DFT) 에 비해 계산 비용의 일부로 원자 단위 시뮬레이션에서 거의 ab initio 수준의 정확도를 달성할 수 있는 경로를 제공합니다. 그러나 단편화된 도구, 제한된 확장성, 그리고 효율적인 시뮬레이션, 응용 프로그램의 사용 편의성, 그리고 빠른 방법론적 혁신을 지원하기 어려운 비유연한 소프트웨어 설계로 인해 보다 광범위한 채택이 저해되고 있습니다. mlip 라이브러리의 첫 번째 버전 (v1) 은 MLIP 훈련 및 배포를 위한 통합된 JAX 기반 프레임워크를 확립했지만, 1 세대 시스템으로 설계되었습니다. 이로 인해 구성 가능성, 엔드투엔드 파이프라인에 대한 제어, 그리고 복잡한 정전기학, 전이 상태 탐색, 대규모 다중 데이터셋 훈련과 같은 고급 과학적 기능을 처리하는 능력에 한계가 있었습니다.

방법론 및 아키텍처

본 논문은 타겟팅된 API 개편과 새로운 고성능 구성 요소의 통합을 통해 이러한 한계를 해결하는 mlip 라이브러리의 포괄적인 재설계인 mlip v2를 제시합니다.

1. 통합 소프트웨어 프레임워크 및 API 재설계

• 통합 Graph 클래스: 라이브러리는 더 이상 사용되지 않는 jraph.GraphsTuple 을 새로운 통합 Graph 클래스로 대체합니다. 이는 모델 입력, 출력 및 중간 잠재 특징을 위한 핵심 데이터 구조로 작용하여 모든 모델 구성 요소 (Graph $\to$ Graph) 간의 인터페이스를 표준화합니다. 이는 보관된 프로젝트에 대한 의존성을 제거하고 더 깨끗한 구성 및 확장을 용이하게 합니다.
• 모듈형 아키텍처: 이 설계는 데이터 처리, 훈련, 추론과 같은 핵심 구성 요소를 최소한의 명확하게 정의된 인터페이스를 통해 분리합니다. 이를 통해 다중 데이터셋 훈련 및 멀티헤드 파인튜닝을 포함한 워크플로우의 유연한 사용자 정의가 가능해집니다.
• 마이그레이션 전략: 내부 리팩토링에도 불구하고 라이브러리는 기존 사용자를 위한 중단 없는 변경을 최소화하기 위해 핵심 워크플로우에 익숙한 인터페이스를 유지하며, 포괄적인 마이그레이션 가이드를 지원합니다.

2. 고성능 백엔드 (e3j)

다양한 하드웨어 전반에 걸친 런타임을 최적화하기 위해 mlip v2 는 등변 (equivariant) 연산을 위한 새로운 오픈소스 고성능 백엔드인 e3j를 통합합니다.

• 구현: e3j 는 Pallas(TPU 용) 와 CUDA(GPU 용) 를 모두 사용하여 등변 연산을 위한 전용 커널을 제공합니다.
• 대상 모델: 클리브시 - 고르단 텐서 곱 (Clebsch-Gordan Tensor Products) 에 의존하는 MACE 및 NequIP와 같은 모델을 특히 가속화하며, 이러한 모델들은 종종 계산 병목 현상이 됩니다.
• 성능: 벤치마크에 따르면 v1 구현 대비 최대 3 배의 런타임 속도 향상이 나타났습니다.

3. 확장된 과학적 기능

프레임워크는 MLIP 응용 범위를 넓히기 위해 여러 새로운 방법론을 도입합니다.

• Mixture-of-Experts(MoE) 를 활용한 eSEN 아키텍처: 라이브러리는 MoE 공식을 활용하는 eSEN 아키텍처를 통합합니다. 이는 효율적인 추론을 유지하면서 대규모 다양성 데이터셋에서의 확장 가능한 훈련을 가능하게 합니다. 라우팅 메커니즘은 추론 시 전문화된 전문가들을 단일 밀집 커널로 결합할 수 있게 합니다.
• 고급 정전기학 및 전하 모델링:
  • 부분 전하 예측: 모든 모델이 원자 부분 전하 예측을 지원합니다.
  • 장거리 상호작용: 발산을 방지하기 위한 소프트코어 정규화를 포함하여 장거리 정전기학을 처리하기 위해 수정된 쿨롱 상호작용 항 (PhysNet 공식 따름) 이 구현되었습니다.
  • 전체 전하 조건부 처리: 다양한 전체 전하를 가진 시스템의 정확도를 향상시키기 위해 모델은 전체 시스템 전하의 임베딩을 원자 번호 임베딩과 연결하여 포함합니다.
• 헤시안 (Hessian) 레이블 훈련: 라이브러리는 에너지의 2 차 도함수 (헤시안) 를 이용한 훈련을 지원합니다. 계산 비용을 관리하기 위해 모든 원자 좌표에 대해 선택된 힘 성분만 미분하는 서브샘플링 전략 (Vector-Jacobian Products) 을 사용합니다. 이는 전체 헤시안 역전파의 prohibitive 한 비용 없이 곡률 정보에 기반한 파운데이션 모델 훈련을 가능하게 합니다.
• 전이 상태 탐색: 전이 상태 위치를 찾기 위해 ASE 와 인터페이스하는 Nudged Elastic Band(NEB) 방법 (클라이밍 이미지 변형 포함) 을 구현한 커스텀 엔진이 통합되었습니다.
• NPT 앙상블 시뮬레이션: 라이브러리는 랑제빈 적분기와 결합된 JAX 기반 몬테카를로 (MC) 바로스탯을 통해 등온 - 등압 (NPT) 시뮬레이션을 지원합니다. 이 접근법은 포텐셜 에너지 변화에 기반한 메트로폴리스 기준을 사용하여 다른 바로스탯 (예: Berendsen 또는 Parrinello-Rahman) 에서 필요한 비싼 응력 평가를 피합니다.

4. 멀티헤드 파인튜닝

대규모 데이터셋에서 사전 훈련된 모델을 하류 작업 (예: 특정 화학 또는 이론 수준) 에 맞게 전문화하면서도 파국적 망각을 방지할 수 있는 멀티헤드 파인튜닝을 위한 통합 프레임워크가 도입되었습니다. 이는 공유 등변 백본과 데이터셋별 리드아웃 헤드 및 원자 에너지 테이블을 결합하여 달성됩니다.

결과 및 검증

저자들은 OMOL25 데이터셋의 큐레이션된 부분집합 (특히 약 176 만 개의 구조를 포함하는 SPICE2 부분집합) 에서 훈련된 새로운 라이브러리와 사전 훈련된 모델 (MACE, NequIP, ViSNet, eSEN) 에 대한 광범위한 검증을 제공합니다.

• 정확도: 사전 훈련된 모델은 SPICE2 의 일곱 가지 분자 부분집합에서 평가되었습니다. eSEN 아키텍처는 대부분의 부분집합에서 에너지와 힘 모두에 대해 가장 낮은 평균 절대 오차 (MAE) 를 달성했습니다.
• 물리적 충실도: MLIPAudit를 사용한 평가 결과, 모든 아키텍처가 결합 길이 분포, 고리 평면성, 그리고 기준 기하학적 안정성에 대해 거의 완벽한 점수를 달성했습니다. eSEN 이 전체 점수 (0.716) 에서 가장 높았으며, ViSNet(0.699) 이 그 뒤를 이었습니다.
• 정전기학 및 전하: 전체 전하 임베딩을 가진 모델은 그렇지 않은 모델에 비해 전체적으로 전하를 띤 시스템의 에너지 예측 정확도가 크게 향상되었습니다. 부분 전하 예측은 모든 부분집합에서 정확했습니다.
• 헤시안 훈련: 통제된 연구는 헤시안 레이블로 훈련하는 것이 에너지와 힘만으로 훈련된 기준 모델에 비해 예측된 진동 주파수의 오차를 크게 감소시켰음을 보여주었습니다.
• NPT 검증: JAX 기반 NPT 적분기는 온도, 등온 압축률, 그리고 방사상 분포 함수 측면에서 기준 ASE 구현 (Berendsen 및 Parrinello-Rahman) 과 excellent 한 일치를 보였으며, 2.2 배에서 4.0 배의 속도 향상을 제공했습니다.
• 런타임: 벤치마크는 e3j 통합과 최적화된 백엔드가 MACE 및 NequIP 모델 전반에 걸쳐 일관된 속도 향상을 가져왔음을 확인했으며, 라이브러리가 단일 장치에서 배치 시뮬레이션을 지원함을 입증했습니다.

중요성 및 주장

본 논문은 mlip v2를 ML 기반 분자 시뮬레이션을 위한 확장 가능하고 적응적인 기반으로 위치시킵니다. 그 주요 중요성은 다음을 통해 ML 연구와 실제 응용 간의 격차를 해소하는 데 있습니다:

1. 스택 통합: 데이터 처리, 모델 훈련, 분자 시뮬레이션을 연결하는 단일 확장 가능 프레임워크를 제공합니다.
2. 확장성 강화: MoE 공식과 고성능 백엔드 (e3j) 를 통해 대규모 다양성 데이터셋에서의 효율적인 훈련을 가능하게 합니다.
3. 적용 범위 확대: 전하 종, 전이 상태, 그리고 실제 열역학적 조건 (NPT) 하의 응축상 환경을 포함하여 복잡하고 반응성이 있으며 비평형 시스템의 모델링을 가능하게 하는 기능을 도입합니다.
4. 오픈소스 접근성: 적용 연구자와 방법론 개발자 모두의 진입 장벽을 낮추기 위해 사전 훈련된 모델과 포괄적인 문서와 함께 Apache 2.0 라이선스로 라이브러리를 공개합니다.

저자들은 라이브러리가 소프트웨어 인프라의 최첨단 상태를 크게 발전시켰지만, 제시된 결과는 라이브러리로 달성 가능한 성능을 나타내는 것이지 아키텍처 간의 결정적인 벤치마크가 아니라고 강조합니다. 서로 다른 모델 패밀리 간에 비교 가능한 하이퍼파라미터 설정을 정의하기 어렵기 때문입니다.