[논문 리뷰] The Design Space of E(3)-Equivariant Atom-Centered Interatomic Potentials

Ilyes Batatia, Simon Batzner외 7명|arXiv (Cornell University)|2022. 05. 13.

Machine Learning in Materials Science인용 수 98

한 줄 요약

본 논문은 Multi-ACE 프레임워크 내에서 ACE와 NequIP를 통합하고, 중요한 설계 선택을 식별하기 위한 어블레이션(ablation)을 수행하며, 더 간단하고 해석 가능하지만 정확한 모델로 BOTNet을 소개한다.

ABSTRACT

The rapid progress of machine learning interatomic potentials over the past couple of years produced a number of new architectures. Particularly notable among these are the Atomic Cluster Expansion (ACE), which unified many of the earlier ideas around atom density-based descriptors, and Neural Equivariant Interatomic Potentials (NequIP), a message passing neural network with equivariant features that showed state of the art accuracy. In this work, we construct a mathematical framework that unifies these models: ACE is generalised so that it can be recast as one layer of a multi-layer architecture. From another point of view, the linearised version of NequIP is understood as a particular sparsification of a much larger polynomial model. Our framework also provides a practical tool for systematically probing different choices in the unified design space. We demonstrate this by an ablation study of NequIP via a set of experiments looking at in- and out-of-domain accuracy and smooth extrapolation very far from the training data, and shed some light on which design choices are critical for achieving high accuracy. Finally, we present BOTNet (Body-Ordered-Tensor-Network), a much-simplified version of NequIP, which has an interpretable architecture and maintains accuracy on benchmark datasets.

연구 동기 및 목표

• ACE가 E(3) 등가성으로 확장되고 다층 MPNN 프레임워크에 통합될 수 있는 방법을 명확히 밝힌다.
• 디스크립터 기반의 ACE와 등가 MPNN을 Multi-ACE 설계 공간 아래에서 통합한다.
• 통합 프레임워크 내에서 NequIP의 어블레이션 연구를 통해 설계 선택을 체계적으로 조사한다.
• 높은 정확도에 필수적인 구성 요소를 식별하고 성능을 유지하면서 더 간단하고 해석 가능한 모델(BOTNet)을 도입한다.
• 설계 공간 내에서 신속한 실험을 가능하게 하는 실용적인 통찰과 모듈식 구현을 제공한다.

제안 방법

• 메시지 전달 방식으로 등가 ACE 레이어를 쌓아 구성하는 Multi-ACE 프레임워크를 정의한다.
• 연속 원소 임베딩을 갖는 등가적인 한 입자 기저를 도입하여 ACE를 확장한다.
• ACE 밀도 트릭과 곱 기저를 사용해 더 높은 차수의 특성을 형성한다.
• Wigner-D 기계를 사용해 O(3) 회전에 대해 기저 함수를 대칭화하여 불변/등가 특성을 얻는다.
• 선형/비선형 변환과 층별 상태 업데이트를 통해 MPNN 업데이트 및 읽기를 ACE 레이어와 연계한다.
• 비필수 비선형 활성화 없이도 정확성을 유지하는 간소화되고 해석 가능한 변형으로 BOTNet을 도입한다.

실험 결과

연구 질문

• RQ1 Multi-ACE 프레임워크에서 벤치마크 데이터셋에서 높은 정확도를 달성하는 데 가장 중요한 설계 구성요소는 무엇인가?
• RQ2 NequIP의 등가/비등가 구성요소를 어블레이션할 때 도메인 내(in-domain) 및 도메인 외(out-of-domain) 성능은 어떻게 달라지는가?
• RQ3 BOTNet과 같은 간소화된 모델이 해석성 및 모듈성을 향상시키면서 NequIP 수준의 정확도를 유지할 수 있는가?
• RQ4 정규화와 원소 임베딩이 등가 상호원자 포텐셜의 안정성과 외삽에 어떤 역할을 하는가?

주요 결과

• 디스크립터 기반 ACE와 등가 MPNN을 연결하는 통합된 Multi-ACE 프레임워크는 설계 선택을 체계적으로 탐구할 수 있게 한다.
• 어블레이션 연구는 최고 성능에 필수적인 NequIP 혁신과 어디에서 단순화가 가능한지 밝힌다.
• 정규화와 연속 원소 임베딩은 모델의 안정성 및 외삽 동작에서 중요한 역할을 한다.
• BOTNet은 벤치마크 데이터셋 전반에서 정확성을 유지하는 훨씬 단순하고 해석 가능한 아키텍처로 부상한다.
• 이 프레임워크는 가장 복잡한 비선형 활성화 중 일부 없이도 높은 정확도를 달성할 수 있음을 보여준다.