[논문 리뷰] Rotation Invariant Graph Neural Networks using Spin Convolutions

Muhammed Shuaibi, Adeesh Kolluru외 5명|arXiv (Cornell University)|2021. 06. 17.

Machine Learning in Materials Science참고 문헌 34인용 수 51

한 줄 요약

SpinConv은 per-edge 로컬 프레임을 사용하고 남은 회전 자유도에 대한 스핀 컨볼루션을 통해 회전 불변 그래프 신경망으로 원자 구조를 모델링하여 에너지와 힘 추정에서 OC20에서 최첨단 결과를 달성합니다.

ABSTRACT

Progress towards the energy breakthroughs needed to combat climate change can be significantly accelerated through the efficient simulation of atomic systems. Simulation techniques based on first principles, such as Density Functional Theory (DFT), are limited in their practical use due to their high computational expense. Machine learning approaches have the potential to approximate DFT in a computationally efficient manner, which could dramatically increase the impact of computational simulations on real-world problems. Approximating DFT poses several challenges. These include accurately modeling the subtle changes in the relative positions and angles between atoms, and enforcing constraints such as rotation invariance or energy conservation. We introduce a novel approach to modeling angular information between sets of neighboring atoms in a graph neural network. Rotation invariance is achieved for the network's edge messages through the use of a per-edge local coordinate frame and a novel spin convolution over the remaining degree of freedom. Two model variants are proposed for the applications of structure relaxation and molecular dynamics. State-of-the-art results are demonstrated on the large-scale Open Catalyst 2020 dataset. Comparisons are also performed on the MD17 and QM9 datasets.

연구 동기 및 목표

• 회전 불변 GNN으로 원자 힘과 에너지를 DFT에 버금가는 효율적이고 정확하게 예측하는 동기를 부여합니다.
• 각 에지에 대한 로컬 좌표계를 개발하여 각도 정보를 손실 없이 포착할 수 있도록 합니다.
• 남은 회전 자유도에 대한 스핀 컨볼루션을 제안하여 회전 불변한 에지 메시지를 달성합니다.
• 에너지 중심 경로(E를 통한 -∂E/∂x에 의한 힘)와 힘 중심 경로(직접 힘 예측) 두 가지를 Present하여 에너지 보존과 직접 힘 추정의 균형을 맞춥니다.
• OC20에서의 최첨단 성능과 MD17 및 QM9에서의 경쟁력 있는 결과를 입증합니다.

제안 방법

• 원자 구조를 노드로, 이웃 간 에지를 컷오프 내의 그래프로 표현합니다.
• 에지 메시지 h_st를 공유 업데이트로 반복적으로 계산하며 스핀 컨볼루션, 거리 인코딩, 원자번호 임베딩을 포함합니다.
• 출발점 s와 목표점 t로 정의된 로컬 프레임을 사용하여 이웃들을 s 주위의 구에 투사하고, 그다음 st 축을 따라 회전 자유도에 대해 스핀 컨볼루션을 적용합니다.
• 학습 목표를 위해 에너지 중심 경로(E를 통해 -∂E/∂x로 힘을 구하는 방식) 또는 힘 중심 경로(힘을 직접 예측하는 방식) 중 하나를 채택합니다.
• 스핀 컨볼루션에 대해 그리드 기반 또는 구-함수(구면 조화) 기반의 구 표현을 채택하고, 회전 불변을 보장하기 위해 풀링을 수행합니다.
• 임베딩 블록(혼합-전문가 스타일)을 도입하여 원자번호 정보를 메시지 업데이트에 융합합니다.
• Robustness를 높이고 힘 블록에서 근사 회전 등가를 유도하기 위해 학습 중 전체 구조를 회전하는 것도 선택적으로 적용합니다.

실험 결과

연구 질문

• RQ1 스핀 컨볼루션 기반의 회전 불변 GNN이 원자 시스템의 에너지 및 힘 필드를 정확히 학습할 수 있는가?
• RQ2 에너지 중심 SpinConv와 힘 중심 SpinConv 변형 간 에너지 보존 및 예측 정확도 측면의 차이는 무엇인가?
• RQ3 Representation 선택(그리드 기반 대 구면 조화) 및 모델 크기가 OC20, MD17, QM9 성능에 미치는 영향은 무엇인가?
• RQ4 다중 무작위 회전을 통한 회전 견고성 유도가 힘 예측을 얼마나 향상시키는가?
• RQ5 IS2RS에서 SpinConv의 성능은 도메인 내/도메인 외(흡착종/촉매) 분할에서 어떤 차이가 있는가?

주요 결과

• 힘 중심 학습을 가진 SpinConv이 OC20 S2EF 작업에서 테스트된 모델들 중 가장 낮은 에너지 MAE와 힘 MAE를 달성했습니다.
• 그리드 기반 스핀 컨볼루션은 학습이 더 빠르고 구면 조화에 비해 경쟁력 있는 정확도를 달성하며, 두 표현 모두 유사한 정확도를 보입니다.
• 에너지 중심 SpinConv은 에너지 정확도를 개선하지만 S2EF 지표 및 IS2RS 결과에서 힘 중심 변형에 비해 전체적으로 뒤처집니다.
• OC20 IS2RE에서 SpinConv은 특히 이완된 에너지 예측을 위한 직접 회귀 설정에서 기존 접근법을 능가합니다.
• OC20 IS2RS에서 SpinConv은 ADwT에서 최고를 달성하지만 DimeNet++-large가 AFbT에서 최고를 달성하며, SpinConv도 추론 속도에서 더 빠릅니다(약 263.2h 대 814.6h).
• MD17 및 QM9 전반에 걸쳐 SpinConv은 상태--of-더 강력한 방법과 경쟁적 또는 우수한 성능을 보이며, 특히 힘 중심 형식화와 적절한 표현 크기를 활용할 때 그렇습니다.