Expectation-Maximization-Based Optimization of Neural Quantum States for Ab Initio Quantum Chemistry

Accurate solutions for ab initio quantum chemistry are critical for understanding chemistry on the atomic scale. Recently, optimization methods for neural quantum states (NQSs) based on deep neural networks have shown great potential for improving computational accuracy over traditional methods, but they inevitably introduce the challenge of substantial computational costs. To address this challenge, this study proposed expectation-maximization-based optimization of NQS (EMO-NQS), a novel method developed to accelerate optimizations through cost amortization. Specifically, the EMO-NQS method reformulates time-intensive operations as the E-step, amortizing their costs through replays of efficient network-updating operations in the M-step to reduce the overall computational load. Furthermore, an active learning (AL)-based sampling method was introduced to enhance the robustness of EMO-NQS. The performance of EMO-NQS was validated by optimizing the NQS of the equilibrium structures for various molecules. The results demonstrated that EMO-NQS achieved optimization speeds at least twice as fast as those of the vanilla NQS method without compromising accuracy. Furthermore, the consistent acceleration of optimization and the robustness of accuracy facilitated by the AL-based EMO-NQS method were demonstrated by an additional optimization of the construction of the potential energy surface of the N₂ molecule.