[논문 리뷰] Generalizing Neural Wave Functions

Nicholas Gao, Stephan Günnemann|arXiv (Cornell University)|2023. 02. 08.

Machine Learning in Materials ScienceMaterials Science인용 수 3

한 줄 요약

이 논문은 Globe를 소개하며, GNN의 공간 메시지 전파를 통해 局소 전자 구조를 학습함으로써 다양한 분자 간에 일반화할 수 있도록 하는 그래프-학습된 궤도 임베딩 방법이다. 또한 Moon은 크기 일관성 있는 신경 파동함수로, 이전 방법보다 최대 4.5배 빠른 수렴 속도를 달성하면서도 시스템 크기에 따라 에너지가 덧셈적으로 스케일링되는 정확한 특성을 유지하여, 하나의 모델로 서로 다른 분자의 슈뢰딩거 방정식을 동시에 해결할 수 있도록 한다.

ABSTRACT

Recent neural network-based wave functions have achieved state-of-the-art accuracies in modeling ab-initio ground-state potential energy surface. However, these networks can only solve different spatial arrangements of the same set of atoms. To overcome this limitation, we present Graph-learned orbital embeddings (Globe), a neural network-based reparametrization method that can adapt neural wave functions to different molecules. Globe learns representations of local electronic structures that generalize across molecules via spatial message passing by connecting molecular orbitals to covalent bonds. Further, we propose a size-consistent wave function Ansatz, the Molecular orbital network (Moon), tailored to jointly solve Schrödinger equations of different molecules. In our experiments, we find Moon converging in 4.5 times fewer steps to similar accuracy as previous methods or to lower energies given the same time. Further, our analysis shows that Moon's energy estimate scales additively with increased system sizes, unlike previous work where we observe divergence. In both computational chemistry and machine learning, we are the first to demonstrate that a single wave function can solve the Schrödinger equation of molecules with different atoms jointly.

연구 동기 및 목표

• 기존의 신경 파동함수는 같은 분자의 다양한 기하구조에 대해서만 일반화 가능하지만, 다른 분자 조성에는 일반화하지 못하는 한계를 극복하기 위해.
• 동일한 신경 파동함수로 서로 다른 원자 조성을 가진 여러 분자의 슈뢰딩거 방정식을 동시에 해결할 수 있도록 하기 위해.
• 신경 파동함수에서 크기 일관성을 확보하여 에너지가 시스템 크기에 따라 덧셈적으로 스케일링되도록 하여, 양자화학에서 물리적 정확성에 필수적인 조건을 충족시키기 위해.
• 공간 메시지 전파와 핵 잠금점(anchoring points)을 활용한 파동함수 아키텍처 설계를 통해 변분 몽테카를로 학습의 수렴 속도를 가속화하기 위해.

제안 방법

• Globe는 분자 궤도를 3차원 공간에 임bedding하기 위해 그래프 신경망에서 공간 메시지 전파를 사용하여, 다양한 분자 간에 국소적인 전자 구조를 학습한다.
• 분자 궤도를 공유 결합에 연결함으로써 궤도 임베딩을 학습함으로써, 다른 수의 궤도를 가진 분자 간에도 일반화 가능하게 한다.
• 파동함수 아키텍처인 Moon은 국소 전자-핵 상호작용에 초점을 맞추고 핵을 메시지 전파의 잠금점으로 사용함으로써 크기 일관성을 강제한다.
• 파동함수를 궤도 가중치가 Globe에 의해 재매arameter화된 결정식들의 합으로 구성함으로써, 다양한 분자 간의 공동 학습을 가능하게 한다.
• 이 방법은 두 수준의 신경망을 사용한다: 하나는 궤도 가중치를 재매arameter화하기 위한 것(Globe), 다른 하나는 최종 파동함수를 계산하기 위한 것(Moon).
• 학습은 에너지 최소화를 통한 변분 몽테카를로 방법으로 수행되며, Globe의 재매arameter화 과정은 단일 스텝당 한 번만 계산되어 샘플링 및 에너지 평가와 분리된다.

Figure 1: Globe steps. 1) Atom embeddings are obtained by message passing between atoms. 2) Orbital locations and embeddings are determined by core orbitals and molecular bonds. 3) Orbital embeddings are updated via a unidirectional message passing. For clarity, we omitted messages from each atom to

실험 결과

연구 질문

• RQ1 다른 원자 조성을 가진 여러 분자의 슈뢰딩거 방정식을 하나의 신경 파동함수로 동시에 풀 수 있는가?
• RQ2 다양한 파라미터화가 가능하고 확장 가능한 신경 파동함수 아키텍처에서 크기 일관성을 어떻게 확보할 수 있는가?
• RQ3 GNN에서의 공간 메시지 전파가 서로 다른 궤도 수를 가진 분자 간에 궤도 임베딩의 일반화를 가능하게 하는가?
• RQ4 개별 학습 대비 다수 분자에 대한 공동 학습이 수렴 속도를 가속화하는가?
• RQ5 다양한 분자 간 일반화된 신경 파동함수의 에너지 추정치가 시스템 크기에 따라 덧셈적으로 스케일링되는가?

주요 결과

• Moon은 이전 방법보다 최대 4.5배 적은 학습 스텝 수로 수렴하면서도 유사하거나 더 낮은 에너지 추정치를 달성하여 최적화 수렴 속도가 빠르다는 것을 입증한다.
• Moon의 에너지 추정치는 시스템 크기에 따라 덧셈적으로 스케일링되며, 이는 이전 방법이 분리되는 스케일링을 보이는 것과 대비된다. 이는 크기 일관성이 확인됨을 의미한다.
• Globe는 추가적인 자가일관장(스elf-consistent field, SCF) 계산 없이도 더 큰 분자 시스템으로의 전이가 가능하여, 직접적으로 더 큰 사슬에서 평가가 가능하다.
• 수소 사슬에서 Globe와 Moon은 시스템 크기가 증가함에 따라 원자당 에너지가 일정하게 유지되며, 이는 강건한 광범위성과 물리적 일관성을 나타낸다.
• 에틸렌과 사이클로부타디엔과 같은 더 큰 분자에서도 Globe와 Moon은 사전학습 없이도 수렴을 달성하는 반면, 다른 방법들은 이를 위해 사전학습이 필수적임을 보였다.
• Moon은 파rameter 수가 60만 개 적은 상태에서 FermiNet을 초월하고 PsiFormer과 동등한 정확도를 달성하여 높은 파rameter 효율성을 입증한다.

Figure 2: Illustration of Moon. 1) We initialize electron embeddings by aggregating their local neighborhood of electrons. 2) Nuclei aggregate electron embeddings via message passing. 3) Nuclei embeddings are iteratively updated. 4) Nuclei embeddings are structurally diffused towards the electrons v