[논문 리뷰] Distillation of atomistic foundation models across architectures and chemical domains

John L. A. Gardner, Daniel F. Thomas du Toit외 9명|ArXiv.org|2025. 06. 12.

Machine Learning in Materials Science인용 수 3

한 줄 요약

본 논문은 아키텍처에 구애받지 않는 증류 프로토콜을 제시하여 대형 원자적 기초 모델에서 더 작고 빠른 학생 MLIP로 지식을 전이하고, 합성 데이터를 통해 다양한 화학 영역에서 상당한 속도 향상(10배 이상에서 100배 이상)을 달성합니다. 이 접근법은 짧은 미세 조정 세트를 사용하여 다수의 아키텍처로 증류함으로써 보통의 하드웨어에서도 정확하고 확장 가능한 MD 시뮬레이션을 가능하게 합니다.

ABSTRACT

Machine-learned interatomic potentials have transformed computational research in the physical sciences. Recent atomistic `foundation' models have changed the field yet again: trained on many different chemical elements and domains, these potentials are widely applicable, but comparably slow and resource-intensive to run. Here we show how distillation via synthetic data can be used to cheaply transfer knowledge from atomistic foundation models to a range of different architectures, unlocking much smaller, more efficient potentials. We demonstrate speed-ups of $> 10 imes$ by distilling from one graph-network architecture into another, and $> 100 imes$ by leveraging the atomic cluster expansion framework. We showcase applicability across chemical and materials domains: from liquid water to hydrogen under extreme conditions; from porous silica and a hybrid halide perovskite solar-cell material to modelling organic reactions. Our work shows how distillation can support the routine and computationally efficient use of current and future atomistic foundation models in real-world scientific research.

연구 동기 및 목표

• 원자적 기초 모델(FMs)로부터 더 작고 빠른 학생 MLIP로 지식을 화학 도메인 전반에 걸쳐 이전하기 위한 일반적 증류 프로토콜을 입증한다.
• 다수의 MLIP 아키텍처로 증류하고 합성 데이터 라벨링을 활용하여 아키텍처에 구애받지 않는 적용 가능성을 보여준다.
• 대표 시스템에서의 메모리 사용량 및 MD 안정성 등 계산 효율성과 정확성의 트레이드오프를 평가한다.
• MD 기반 진단 및 벤치마크를 통해 증류된 모델이 본질적 물리적 특성을 보존하는지 검증한다.
• 원자적 FM을 보통의 하드웨어로도 접근 가능하게 하기 위한 실용적 함의를 강조한다.

제안 방법

• 양자역학 라벨을 가진 도메인 특정한 소규모 구조 세트에서 기존의 원자적 FM을 미세 조정한다.
• MD 시뮬레이션 없이 rattle-relax-repeat 확장으로 대규모 합성 데이터를 생성하기 위해 미세 조정된 FM을 사용한다.
• 합성 데이터에 대해 작은, 빠른 학생 MLIP 아키텍처를 학습시켜 FM의 예측 및 라벨을 근사하도록 한다.
• 합성 데이터로 훈련된 모델을 DFT 테스트 세트와 대조하고 MD 시뮬레이션에서 구조/열역학적 특성을 비교한다.
• 다양한 아키텍처(TensorNet, PaiNN, ACE) 간의 속도 향상 및 규모 확장을 시연하고 ACE/EDDP 계열 내에서도 확장성을 보인다.
• ASE 계산기와 augment-atoms를 사용해 엔드투엔드 워크플로우를 가능하게 하는 아키텍처-독립적 호환성을 보여준다.

실험 결과

연구 질문

• RQ1 고용량의 원자적 FM에서 작은 빠른 학생 모델로의 지식 전이가 서로 다른 아키텍처 간에 합성 데이터 증류를 통해 가능한가?
• RQ2 정확도 대비 DFT 라벨에 대해 어느 정도의 속도향상과 메모리 효율성을 달성할 수 있는가?
• RQ3 증류된 MLIP가 다양한 화학 영역에서 MD의 핵심 구조적 및 동적 특성을 재현하는가?
• RQ4 반응성 및 고에너지 구성에 대한 증류의 실용적 한계와 도메인 경계는 무엇인가?
• RQ5 아키텍처, 컷오프 반경, 미세 조정 데이터의 양에 따라 증류 결과가 어떻게 달라지는가?

주요 결과

• 그래프 네트워크 FM에서 다른 그래프 네트워크 아키텍처로의 전이 시 10배 이상 속도 향상을 달성하고, ACE 프레임워크를 활용할 때는 100배 이상 속도 향상을 달성한다.
• 증류된 모델(TensorNet, PaiNN, ACE)은 DFT 라벨에 대한 미세 조정된 FM에 근접한 포스 MAE를 달성하면서 MD 속도 이점을 크게 얻는다.
• 증류된 모델은 단일 GPU에서 안정적인 MD를 가능하게 하고 FM의 메모리 한계를 넘어 더 큰 시스템 크기로 확장할 수 있다.
• 물, 수소, 실리카, MAPI, 용매 내 유기 반응 등 도메인 전반에서 증류된 모델은 일부 지표에서 교사 모델과 비교하여 핵심 구조적·동적 특성을 재현하거나 더 잘 재현한다.
• 제거 연구에서 합성 데이터 스케일링이 FM-DFT 정확도를 향상시키고, 증류된 모델은 FM보다 작은 컷오프에서도 큰 정확도 저하 없이 작동할 수 있다.
• 이 접근법은 미세 조정을 위한 도메인 데이터(<50개의 DFT 표기 구조) 정도가 필요하며 완전 자동화가 오픈 소스 도구로 가능하다.