[논문 리뷰] Ab-Initio Potential Energy Surfaces by Pairing GNNs with Neural Wave Functions

Nicholas Gao, Stephan Günnemann|arXiv (Cornell University)|2021. 10. 11.

Machine Learning in Materials Science참고 문헌 47인용 수 5

한 줄 요약

이 논문은 변분 몬테카를로(Variational Monte Carlo, VMC)를 사용하여 여러 분자 기하구조에 대해 동시에 슈뢰딩거 방정식을 푸는 데에 사용할 수 있는 새로운 아키텍처인 PESNet을 소개한다. 그래프 신경망(GNN)과 신경파면함수를 결합한 PESNet은 단일 학습 프로세스에서 다양한 구조에 대한 파면함수 파arameter를 예측함으로써, 최신 기술인 FermiNet에 비해 최대 40배 빠른 학습 속도를 달성하면서도 정확도는 유지하거나 초월한다.

ABSTRACT

Solving the Schrödinger equation is key to many quantum mechanical properties. However, an analytical solution is only tractable for single-electron systems. Recently, neural networks succeeded at modeling wave functions of many-electron systems. Together with the variational Monte-Carlo (VMC) framework, this led to solutions on par with the best known classical methods. Still, these neural methods require tremendous amounts of computational resources as one has to train a separate model for each molecular geometry. In this work, we combine a Graph Neural Network (GNN) with a neural wave function to simultaneously solve the Schrödinger equation for multiple geometries via VMC. This enables us to model continuous subsets of the potential energy surface with a single training pass. Compared to existing state-of-the-art networks, our Potential Energy Surface Network PESNet speeds up training for multiple geometries by up to 40 times while matching or surpassing their accuracy. This may open the path to accurate and orders of magnitude cheaper quantum mechanical calculations.

연구 동기 및 목표

• 변분 몬테카를로(Variational Monte Carlo, VMC) 시뮬레이션에서 각 분자 기하구조에 대해 별도의 신경파면함수를 학습시키는 데 드는 높은 계산 비용을 해결한다.
• 재학습 없이도 연속적인 위치 에너지 표면(Potential Energy Surface, PES)을 단일 신경망으로 모델링할 수 있도록 한다.
• 물리적 대칭성(이동, 회전, 반사)을 파면함수의 파rameter화 과정에 통합하여 일반화 능력과 불변성을 향상시킨다.
• 기존의 신경파면함수 모델보다 정확도를 유지하거나 초월하면서도 학습 시간을 극적으로 줄이는 확장성 있고 효율적인 프레임워크를 개발한다.

제안 방법

• 분자의 핵의 기하구조를 기반으로 신경파면함수(WFModel)의 파arameter를 예측하기 위해 메타그래프 신경망(MetaGNN)을 사용한다.
• 메타그래프 신경망을 학습시켜, 각 입력 기하구조에 대해 변분 에너지를 최소화하는 파면함수를 생성할 수 있도록 한다.
• GNN 입력에 등변좌표계(equivariant coordinate system)를 설계하여 물리적 대칭성(이동, 회전, 반사)에 대한 불변성을 강제한다.
• 전자 구성을 확률적 샘플링을 통해 최적화하는 변분 몬테카를로 프레임워크를 활용한다.
• 전체 PESNet 아키텍처를 단일 프로세스에서 엔드 투 엔드로 학습시켜, 여러 분자 기하구조에 동시에 학습을 수행할 수 있도록 한다.
• GNN 출력에 의해 조건화되는 공통된 파면함수 아키텍처(FermiNet 기반 등)를 사용하여 연속적인 PES 영역에 걸쳐 일반화 능력을 확보한다.

실험 결과

연구 질문

• RQ1 단일 신경망 아키텍처가 동시에 여러 분자 기하구조에 대해 정확한 전자 파면함수를 예측할 수 있는가?
• RQ2 GNN에 의해 조건화된 파면함수 모델은 기존의 기하구조 별 학습 방식에 비해 학습 시간을 얼마나 줄일 수 있는가?
• RQ3 GNN 입력에 물리적 대칭성을 통합함으로써 예측된 위치 에너지 표면의 일반화 능력과 불변성은 얼마나 향상되는가?
• RQ4 재학습 없이도 새로운, 볼 수 없는 기하구조에 대해 연속적인 PES 영역에서 일반화가 가능한가?
• RQ5 다중 구조적 위치 에너지 표면을 모델링할 때, 학습 효율성과 예측 정확도 사이의 상충 관계는 어떻게 나타나는가?

주요 결과

• PESNet은 H₄⁺, H₄, H₁₀, N₂, 사이클로부타디엔 등 다양한 시스템에서 FermiNet 대비 최대 47배 빠른 학습 속도를 기록했다.
• H₄⁺의 경우, PESNet은 A100 GPU에서 20 GPU 시간으로 학습을 완료했으며, FermiNet의 127시간에 비해 빠르면서도 정확도는 유지하거나 초월했다.
• H₄에서는 PauliNet과 DeepErwin에 비해 정확도와 학습 효율성에서 모두 PESNet이 뛰어나 24시간 대비 118시간을 기록했고, N₂의 경우 89시간 대비 4196시간으로 뚜렷한 성능 향상을 보였다.
• 모델은 단일 학습 프로세스에서 연속적인 위치 에너지 표면의 부분 영역을 포괄할 수 있어 기하구조 수에 비례하는 선형 확장성을 확보했다.
• PESNet은 복잡한 분자인 사이클로부타디엔을 포함한 모든 테스트 시스템에서 최신 기술 수준의 정확도를 달성했으며, 계산 비용을 크게 줄였다.
• 등변좌표계를 통해 물리적 대칭성을 통합함으로써 모델의 안정성과 일반화 능력이 향상되었으며, 예측 성능에 영향을 주지 않았다.