Excited Pfaffians: Generalized Neural Wave Functions Across Structure and State

Motivation

• Known: Variational Monte Carlo(VMC) 기반 neural-network wave function은 ground state와 excited state 모두에서 매우 높은 정확도를 달성해왔으며, ground state의 경우 generalized wave function을 통해 여러 분자 구조에 대한 학습 비용을 amortize하는 방법들이 이미 존재한다.
• Gap: 그러나 excited state에 대해서는 이러한 amortization 기법이 존재하지 않으며, 기존 방법들은 state 수가 늘어날수록 (1) pairwise overlap ⟨Ψs|Ψt⟩ 추정에 필요한 Monte Carlo 샘플 수 증가와 (2) 각 state를 별도 wave function 또는 지나치게 큰 네트워크로 표현해야 하는 문제로 인해 계산량이 superlinear~quartic하게 증가한다.
• Why: Excited state의 정확한 계산은 약물 설계, 광전지 등 광화학적 과정과 분광학적 성질을 다루는 응용에 필수적이지만, 현재는 계산 하나에 수일이 소요되고 state 수에 따라 비용이 급격히 증가하여 실용적 적용이 제한되므로, state 수에 대한 스케일링 문제를 해결하는 것은 이 분야의 확장성 있는 실용적 응용에 중요하다.
• Approach: Hartree-Fock 이론에서 영감을 받아 여러 state가 대부분의 파라미터를 공유하고 경량의 state별 selector만 다르게 갖는 Excited Pfaffians 구조를 제안하고, 모든 state의 샘플을 통합하여 overlap을 추정하는 MSIS를 함께 사용하여 state 수에 거의 무관한 계산 스케일링을 달성한다.

Achievement

1. 근-상수 스케일링(near-constant scaling) 달성: 기존 방법들이 state 수에 대해 superlinear(O(Ns^1.11))~quartic(O(Ns^3.07))하게 스케일링되는 것과 달리, 본 방법은 O(Ns^0.14)의 거의 상수에 가까운 스케일링을 보인다(Fig. 1).
2. carbon dimer(C2)에서의 효율성 및 정확도: O(Ns^4)-스케일링의 natural excited states(NES) 방법과 동등한 정확도를 유지하면서 200배 이상 빠른 학습과 40배 적은 샘플로 8개 대신 12개(50% 더 많은) state를 모델링했다.
3. beryllium 원자의 전체 에너지 준위 최초 계산: 유리한 스케일링 덕분에 neural network를 이용해 beryllium 원자의 모든 구별되는 에너지 준위를 최초로 계산했다.
4. 분자 간 일반화된(generalized) excited state wave function 최초 시연: 단일 wave function 모델이 여러 분자에 걸친 excited state를 표현할 수 있음을 최초로 입증했다.
5. state crossing에서의 ground state 성능 개선: Excited Pfaffian이 ground-state 전용 방법 대비 state crossing 지점에서 ground state 성능을 향상시킴을 확인했다.

How

• Multi-State Importance Sampling(MSIS): 기존의 단일 state 기반 importance sampling과 달리 모든 state에서 얻은 샘플을 pool하여 pairwise overlap ⟨Ψs|Ψt⟩를 추정하며, 이를 통해 고정 batch size Nb에서 추정량의 평균 분산을 O(Ns/Nb)에서 O(1/Nb)로 감소시킨다.
• Excited Pfaffians 아키텍처: Hartree-Fock에서 서로 다른 orbital selector Πs를 사용해 직교하는 여러 wave function ψHF,s를 구성하는 방식에서 영감을 받아, Neural Pfaffian(Pf(ΦAΦᵀ)) 구조에 state별 경량 selector를 결합함으로써 하나의 신경망으로 여러 state를 동시에 표현한다.
• penalty 기반 excited state 목표함수: 기존의 orthogonality penalty term(Σ ω_st|⟨Ψs|Ψt⟩|²)을 사용한 최적화에 MSIS로 얻은 overlap 추정치를 결합하여 학습 안정성을 확보한다.
• generalized wave function 프레임워크와의 통합: 기존 Neural Pfaffian의 generalized wave function 프레임워크(구조·화학 조성에 대한 일반화)와 Excited Pfaffians의 state 일반화를 결합하여, 구조·화학 조성·전자 상태 세 가지 축 모두에 걸쳐 일반화하는 최초의 단일 신경망 모델을 구현한다.
• 실험 검증: carbon dimer의 potential energy surface(PES), beryllium 원자의 에너지 준위, 여러 second-row atoms 및 분자의 여기 에너지에 대해 정확도와 계산 비용을 기존 NES(Pfau et al., 2024), Szabo et al.(2024) 방법과 비교한다.

Originality

• Hartree-Fock의 orbital selector 개념을 neural excited-state wave function 아키텍처에 적용하여, 여러 state가 파라미터를 공유하되 경량 selector로 구분되는 최초의 구조(Excited Pfaffians)를 제안했다.
• 모든 state의 샘플을 통합해 pairwise overlap을 추정하는 MSIS를 도입함으로써, 기존 single-state importance sampling의 O(Ns/Nb) 분산 문제를 O(1/Nb)로 근본적으로 개선했다.
• 구조(structure), 화학 조성(chemical composition), 전자 상태(electronic state)라는 세 축 모두에 대해 동시에 일반화하는 최초의 neural wave function 프레임워크를 제시했다.

Limitation & Further Study

• beryllium 원자와 carbon dimer 등 비교적 작은 시스템에 대한 검증에 그쳐, 더 큰 분자계나 복잡한 전자 구조로의 확장성은 추가 검증이 필요하다.
• 논문 저자들도 언급하듯 미지의(unseen) 분자 구조로의 transfer 성능은 본 연구의 범위를 벗어나며, 이는 실제 스크리닝 응용에서 중요한 한계로 남는다.
• ωst 추정치의 과대추정(overestimation) 문제와 관련된 하이퍼파라미터 민감도, 그리고 state 수가 매우 커질 때(near-constant이지만 완전한 상수는 아님) 장기적 확장성에 대한 정량적 분석이 제한적이다.
• 후속 연구로 더 다양한 분자 계열 및 더 많은 state 수에 대한 체계적 스케일링 검증, 그리고 미지 구조에 대한 일반화 성능 향상이 필요하다.

Evaluation

총평: Excited state neural wave function의 고질적인 계산 스케일링 문제를 아키텍처와 샘플링 기법 양쪽에서 동시에 해결한 독창적이고 실질적인 기여를 담은 논문으로, beryllium 전체 에너지 준위 및 분자 간 일반화된 excited-state wave function이라는 최초의 성과를 제시하지만, 더 큰 시스템과 미지 구조로의 확장성 검증은 후속 과제로 남아 있다.