[논문 리뷰] Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics

Albert Musaelian, Simon Batzner외 5명|arXiv (Cornell University)|2022. 04. 11.

Machine Learning in Materials Science인용 수 21

한 줄 요약

Allegro는 엄격히 지역적이며 E(3)-등각 신경망 interatomic potential로 소분자 벤치마크에서 최첨단 정확도를 달성하는 동시에 확장 가능한 병렬 시뮬레이션과 분포 외 데이터에 대한 강한 전달 가능성을 가능하게 한다.

ABSTRACT

A simultaneously accurate and computationally efficient parametrization of the energy and atomic forces of molecules and materials is a long-standing goal in the natural sciences. In pursuit of this goal, neural message passing has lead to a paradigm shift by describing many-body correlations of atoms through iteratively passing messages along an atomistic graph. This propagation of information, however, makes parallel computation difficult and limits the length scales that can be studied. Strictly local descriptor-based methods, on the other hand, can scale to large systems but do not currently match the high accuracy observed with message passing approaches. This work introduces Allegro, a strictly local equivariant deep learning interatomic potential that simultaneously exhibits excellent accuracy and scalability of parallel computation. Allegro learns many-body functions of atomic coordinates using a series of tensor products of learned equivariant representations, but without relying on message passing. Allegro obtains improvements over state-of-the-art methods on the QM9 and revised MD-17 data sets. A single tensor product layer is shown to outperform existing deep message passing neural networks and transformers on the QM9 benchmark. Furthermore, Allegro displays remarkable generalization to out-of-distribution data. Molecular dynamics simulations based on Allegro recover structural and kinetic properties of an amorphous phosphate electrolyte in excellent agreement with first principles calculations. Finally, we demonstrate the parallel scaling of Allegro with a dynamics simulation of 100 million atoms.

연구 동기 및 목표

• 메시지 전달 방식의 접근법과 동등한 또는 이를 상회하는 정확도를 갖춘 엄격히 지역적이고 등각 불변(interatomic) 포텐셜을 개발한다.
• 대규모 시스템에서의 품질 저하 없이 확장성 및 병렬화 가능성을 입증한다.
• 분포 외 및 더 높은 온도 데이터에 대한 강한 전달 가능성을 보인다.
• 복잡한 전해질의 MD 시뮬레이션에서 구조적 및 운동학적 특성을 재현함으로써 물리적 현실성을 검증한다.

제안 방법

• 임의 깊이의 E(3)-등각 네트워크 Allegro를 도입하고 스칼라(불변) 및 텐서(등각) 잠재 공간을 활용한다.
• 크기-확장적이고 로컬한 에너지 기여를 보장하기 위해 전체 에너지를 종-종/중심-쌍 에너지로 분해한다.
• 텐서-곱 층을 사용하여 이웃 구형조화 함수의 가중합으로 정보가 포함된 환경을 통해 정보를 전파하고 계산 비용을 줄이기 위한 밀도-트릭을 적용한다.
• 두 체 간의 스칼라 특징을 임베딩한 후 불변 및 등각 잠재 상태를 레이어 간에 반복적으로 업데이트한 다음 쌍 에너지를 예측하는 MLP를 적용한다.
• 훈련 안정성 및 시스템 크기와 구성의 변화에 따른 크기 확장성을 개선하기 위해 내부 특징과 타깃을 정규화한다.

실험 결과

연구 질문

• RQ1 엄격히 로컬하고 등각 신경망이 메시지 전달 포텐셜과 동등하거나 더 나은 최첨단 정확도에 도달할 수 있는가?
• RQ2 로컬 Allegro 모델은 분포 외 구성 및 더 높은 온도에 얼마나 잘 일반화되는가?
• RQ3 역학 동안 에너지 보전과 정확성을 유지하면서 매우 큰 시스템에 대해 Allegro가 확장 가능한가?
• RQ4 복잡한 전해질의 MD 시뮬레이션에서 구조적 및 운동학적 특성을 충실히 재현하는가?

주요 결과

Molecule	Energy (meV)	Forces (meV/Å)
Aspirin	6.2	2.3
Azobenzene	2.8	2.3
Benzene	0.3	0.2
Ethanol	0.9	2.1
Malonaldehyde	1.5	3.6
Naphthalene	1.2	0.9
Paracetamol	2.9	4.9
Salicylic acid	1.8	2.9
Toluene	1.6	1.8
Uracil	0.6	1.8

• Allegro는 개정된 MD-17 벤치마크에서 힘 구성 성능 MAE가 최첨단에 도달하고, 여러 소분자 목표에서 경쟁력 있거나 우수한 성능을 보여준다.
• Allegro의 단일 텐서-곱 층은 QM9 벤치마크에서 기존의 딥 MPNN 및 트랜스포머를 능가할 수 있다.
• Allegro는 분포 외 데이터에 대한 강한 전달 가능성을 보여주며 다른 로컬 MLIP보다 뛰어나다.
• Allegro를 이용한 분자 역학은 Li3PO4의 구조적 및 운동학적 특성을 1st-principles 계산과 훌륭히 일치시키며 재현한다.
• 모델은 1억 원자까지의 시스템을 병렬로 시뮬레이션할 수 있을 만큼 확장 가능하여 실용적 대규모 적용성을 보여준다.