[논문 리뷰] DPA-1: Pretraining of Attention-based Deep Potential Model for Molecular Simulation

Duo Zhang, Hangrui Bi외 4명|arXiv (Cornell University)|2022. 08. 17.

Machine Learning in Materials Science인용 수 27

한 줄 요약

DPA-1은 대규모 다원소 데이터에 대해 주의(attention) 기반의 Deep Potential 모델을 사전학습하고 다운스트림 PES 작업에 맞춰 미세조정하여 우수한 전이 가능성과 샘플 효율성을 달성합니다. 또한 주기적 경향성과 일치하는 해석 가능한 원소 유형 임베딩을 드러냅니다.

ABSTRACT

Machine learning assisted modeling of the inter-atomic potential energy surface (PES) is revolutionizing the field of molecular simulation. With the accumulation of high-quality electronic structure data, a model that can be pretrained on all available data and finetuned on downstream tasks with a small additional effort would bring the field to a new stage. Here we propose DPA-1, a Deep Potential model with a novel attention mechanism, which is highly effective for representing the conformation and chemical spaces of atomic systems and learning the PES. We tested DPA-1 on a number of systems and observed superior performance compared with existing benchmarks. When pretrained on large-scale datasets containing 56 elements, DPA-1 can be successfully applied to various downstream tasks with a great improvement of sample efficiency. Surprisingly, for different elements, the learned type embedding parameters form a $spiral$ in the latent space and have a natural correspondence with their positions on the periodic table, showing interesting interpretability of the pretrained DPA-1 model.

연구 동기 및 목표

• 다수 원소에 걸친 대규모 사전 학습을 활용할 수 있는 확장 가능하고 정확한 PES 모델을 제시합니다.
• 에너지 보존성, 매끄러움, 병렬 확장성을 보존하는 주의 기반 Deep Potential(DPA-1)을 개발합니다.
• AlMgCu, SSE, HEA 시스템에서 사전학습-다운스트림 전이 효율성을 입증합니다.
• 학습된 원소 유형 임베딩의 해석 가능성과 그것의 주기표와의 관계를 보여줍니다.

제안 방법

• 원소 유형 정보를 가진 로컬 임베딩과 이웃 원자들에 대한 자기 주의(self-attention) 기작을 도입합니다.
• 변환된 상대 좌표 행렬과 유형 임베딩을 통해 로컬 디스크립터를 구성하고, 주의 기반 집계로 학습 가능한 디스크립터를 구축합니다.
• 에너지/힘 손실의 공동 최적화를 위해 Adam을 사용하고, 에너지 항과 힘 항의 균형을 위한 스케줄링 방식을 적용합니다.
• 최종 에너지 바이어스를 조정하고 매개변수를 선택적으로 고정시키는 방식으로 사전학습 모델을 미세조정하며, 특히 유형 임베딩을 중심으로 합니다.
• OC2M(56 원소)에서 사전학습하고 AlMgCu, SSE, HEA 데이터셋을 포함한 다운스트림 작업에서 평가합니다.

실험 결과

연구 질문

• RQ1 사전학습된 주의 기반 PES 모델이 다양한 조성(단일, 이진, 삼원)과 서로 다른 데이터셋에 일반화될 수 있나요?
• RQ2 사전학습이 스크래치 학습에 비해 다운스트림 작업의 데이터 및 계산 요구를 감소시키나요?
• RQ3 학습된 원소 유형 임베딩이 주기표와 관련된 물리적으로 해석 가능한 구조를 보이나요?
• RQ4 대규모 시뮬레이션에서 모델이 에너지 보존성, 매끄러움 및 확장성을 유지할 수 있나요?

주요 결과

• DPA-1은 제한된 조성 데이터에서 전이할 때 AlMgCu, SSE, HEA에 대해 검증 정확도에서 DeepPot-SE를 능가합니다.
• OC2M으로의 사전학습은 강력한 샘플 효율성을 가능하게 하며, 삼원 데이터의 필요성을 크게 낮춥니다(예: AlMgCu에서 삼원 샘플을 90% 감소).
• OC2M에서 사전학습된 DPA-1은 OC2M 검증에서(0.681 eV, 0.076 eV/Å) 다른 최신 모델들과 경쟁력 있는 에너지/힘 MAE를 달성합니다.
• OC2M-사전학습 DPA-1의 원소 유형 임베딩은 주기표 그룹화와 대략 연관된 나선형으로 배열되어 해석 가능한 물리학을 시사합니다.
• DPA-1은 일부 기준 모델과 달리 학습 비용이 이차적으로 증가하지 않고 다양한 원소가 있는 데이터셋으로 확장됩니다.
• 사전학습 방식은 더 나은 샘플 효율성을 제공하고 MD에 적합한 에너지 보존의 매끄러운 힘을 유지합니다.