[논문 리뷰] A Self-Attention Ansatz for Ab-initio Quantum Chemistry

Ingrid von Glehn, James S. Spencer외 1명|arXiv (Cornell University)|2022. 11. 24.

Machine Learning in Materials Science인용 수 30

한 줄 요약

Psiformer는 자가 주의 기반 파동함수 Ansatz로서 ab-initio 양자 화학에서 기저상태 에너지를 개선하며, 특히 큰 분자에서 FermiNet 및 FermiNet+SchNet에 비해 향상된다.

ABSTRACT

We present a novel neural network architecture using self-attention, the Wavefunction Transformer (Psiformer), which can be used as an approximation (or Ansatz) for solving the many-electron Schrödinger equation, the fundamental equation for quantum chemistry and material science. This equation can be solved from first principles, requiring no external training data. In recent years, deep neural networks like the FermiNet and PauliNet have been used to significantly improve the accuracy of these first-principle calculations, but they lack an attention-like mechanism for gating interactions between electrons. Here we show that the Psiformer can be used as a drop-in replacement for these other neural networks, often dramatically improving the accuracy of the calculations. On larger molecules especially, the ground state energy can be improved by dozens of kcal/mol, a qualitative leap over previous methods. This demonstrates that self-attention networks can learn complex quantum mechanical correlations between electrons, and are a promising route to reaching unprecedented accuracy in chemical calculations on larger systems.

연구 동기 및 목표

• 자 external data 없이 다전자 슈뢰딩거 방정식을 풀기 위한 자가 주의 기반 신경망 Ansatz를 개발한다.
• 분자 시스템의 기저상태 에너지 계산 정확도를 기존 신경망 Ansatz와 비교하여 향상시킨다.
• 자가 주의 기반 구조를 통해 더 큰 분자까지의 확장성을 입증한다.
• Jastrow 요인 및 엔벨로프 함수와 같은 물리적 제약을 도입한다.

제안 방법

• Psi_theta를 exp(J_theta) 곱하기 determinant 합으로 정의하며, J_theta는 Jastrow 유사 사전 인자이고 Phi_theta^k는 self-attention 네트워크를 통해 학습된다.
• 입력 피처에 전자-핵 간 상호작용을 모델링하기 위해 스핀을 입력으로 연결한 단일 스트림 다중 헤드 self-attention 아키텍처를 사용한다.
• 큰 간격 기하를 처리하기 위한 입력 재스케일링과 지수 엔벨로프를 통한 경계 조건 강제를 적용한다; 잔차 연결과 tanh 비선형성을 적용한다.
• Rayleigh 몬테카를로 샘플링과 변분적 최적화를 사용하는 변분 양자 몬테카를로로 학습하며, Kronecker-팩터화된 근사 곡률(KFAC)로 최적화한다.
• 소형 및 대형 분자 벤치마크(예: 벤젠 포함)에서 Psiformer를 FermiNet 및 FermiNet+SchNet과 비교하여 정확도 향상을 평가한다.

실험 결과

연구 질문

• RQ1 자가 주의 기반 신경망이 ab-initio 양자 화학에서 다전자 슈뢰딩거 방정식에 대해 경쟁력 있거나 우수한 Ansatz를 제공할 수 있는가?
• RQ2 Psiformer가 기존의 신경망 Ansatz에 비해 특히 큰 분자 시스템에서 기저상태 에너지 정확도를 향상시키는가?
• RQ3 스핀 도입, Jastrow를 통한 cusp 처리, 주의 기반 전자-핵 결합이 시스템 크기에 따라 학습과 안정성에 어떤 영향을 미치는가?
• RQ4 자체 주의 아키텍처가 이전 방법에 비해 수십 ~ 수십 개의 전자를 가진 분자까지 확장 가능성은 어느 정도인가?

주요 결과

• Psiformer는 FermiNet 및 FermiNet+SchNet에 비해 에너지를 크게 개선하며, 특히 더 큰 시스템에서 그렇다.
• 소형 분자에서 Psiformer의 대형 모델은 ozone와 같은 일부 시스템에서 CCSD(T)/CBS에 약 1 kcal/mol 이내로 도달한다.
• 42–74 전자 분자(벤젠에서 CCl4까지)에서 Psiformer는 FermiNet보다 우수하며 일부 발표된 DMC 에너지를 능가하고 벤젠과 CCl4에서 주목할 만한 이점을 보인다.
• Psiformer 에너지는 시스템 크기가 커짐에 따라 향상되며, 대형 분자에서 일반적인 기준선 대비 최대 약 ~75 kcal/mol (100 mHa) 개선을 보인다.
• LayerNorm 및 다른 학습 선택은 최적화에 영향을 주며, Psiformer는 일반적으로 구성을 바꿔도 우수한 정확도를 유지한다.