[논문 리뷰] MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields

Ilyes Batatia, Dávid Péter Kovács외 3명|arXiv (Cornell University)|2022. 06. 15.

Machine Learning in Materials Science인용 수 269

한 줄 요약

MACE는 다체(고차) 등가 메시지 전달을 원자 간 포텐셜에 도입하여, rMD17, 3BPA, AcAc와 같은 벤치마크에서 최첨단 정확도와 빠른 학습을 가능하게 하는 두 겹 네트워크를 구현합니다.

ABSTRACT

Creating fast and accurate force fields is a long-standing challenge in computational chemistry and materials science. Recently, several equivariant message passing neural networks (MPNNs) have been shown to outperform models built using other approaches in terms of accuracy. However, most MPNNs suffer from high computational cost and poor scalability. We propose that these limitations arise because MPNNs only pass two-body messages leading to a direct relationship between the number of layers and the expressivity of the network. In this work, we introduce MACE, a new equivariant MPNN model that uses higher body order messages. In particular, we show that using four-body messages reduces the required number of message passing iterations to just two, resulting in a fast and highly parallelizable model, reaching or exceeding state-of-the-art accuracy on the rMD17, 3BPA, and AcAc benchmark tasks. We also demonstrate that using higher order messages leads to an improved steepness of the learning curves.

연구 동기 및 목표

• 화학 및 재료 과학을 위한 빠르고 정확한 힘 필드를 동기화합니다.
• 두 바디 MPNN의 한계를 넘어 더 높은 체 구성 메시지를 도입합니다.
• 효율적인 텐서 연산으로 확장 가능하고 병렬화 가능한 아키텍처를 개발합니다.
• 도전적인 벤치마크(rMD17, 3BPA, AcAc)에서 최첨단 정확도를 시연합니다.
• 체 구성 및 등가성(Equivariance)이 학습 곡선과 일반화에 어떤 영향을 미치는지 조사합니다.

제안 방법

• nu-바이더(interactions)까지 포함하는 계층적인 다체 메시지 구성을 도입하는 MACE를 소개합니다.
• 메시지를 이웃 집합의 합으로 표현하되 체 구성 순서를 증가시키며(식 7).
• 등가성을 강제하기 위해 텐서 곱과 Clebsch–Gordan 결합으로 고차 특징을 구성합니다(식 8–10).
• 메시지를 전파하기 위해 잔차 연결이 있는 선형 업데이트를 사용합니다(식 12).
• 레이어별 기여의 합으로 사이트 에너지를 얻기 위해 특징의 불변 부분을 읽어옵니다(식 13).
• 고차 메시징을 계산적으로 실행 가능하게 만들기 위해 효율적인 루프 텐서 수축을 구현합니다.

실험 결과

연구 질문

• RQ1 등가적 MPNN에서의 고차(다체) 메시지가 필요한 메시지 전달 층의 수를 감소시키면서도 정확도를 희생하지 않을 수 있는가?
• RQ2 체 구성의 증가와 등가성이 힘 필드 예측의 학습 곡선 및 데이터 효율성에 어떤 영향을 미치는가?
• RQ3 고차의 등가 메시지가 도메인 밖의 분자 구성에서 더 나은 외삽 및 일반화를 제공하는가?
• RQ4 기존 최첨단 등가 MPNN과 비교한 MACE의 속도 및 확장성은 벤치마크에서 어떤가?

주요 결과

• 고차 메시지는 필요한 층 수를 두 개로 줄이면서 벤치마크 작업에서 최첨단 정확도에 도달하거나 이를 능가합니다.
• 충분한 체 구성에서의 불변 메시지는 등가 모델의 성능에 근접할 수 있으며, 고차 메시지는 등가성을 갖출 때 최상의 결과를 냅니다.
• L=2인 MACE(두 레이어, 2L+M 등가 메시지)는 3BPA 외삽 테스트에서 경쟁 모델을 능가하며 벤치마크 전반에서 에너지 및 힘 예측이 우수합니다.
• MACE는 이전 최상위 모델 대비 학습 및 평가 속도에서 상당한 이점을 제공하며 정확도를 유지하거나 향상시킵니다.
• 학습 곡선은 고차 체인 정보를 도입하면 데이터 효율성의 기울기가 이동하고 가팔라져 학습 역학이 개선됨을 시사합니다.