Efficient optimization of neural network backflow for ab-initio quantum chemistry

Abstract: The ground state of second-quantized quantum chemistry Hamiltonians is key to determining molecular properties. Neural quantum states (NQS) offer flexible and expressive wavefunction ansatze for this task but face two main challenges: highly peaked ground-state wavefunctions hinder efficient sampling, and local energy evaluations scale quartically with system size, incurring significant computational costs. In this work, we overcome these challenges by introducing a suite of algorithmic enhancements, which includes efficient periodic compact subspace construction, truncated local energy evaluations, improved stochastic sampling, and physics-informed modifications. Applying these techniques to the neural network backflow (NNBF) ansatz, we demonstrate significant gains in both accuracy and scalability. Our enhanced method surpasses traditional quantum chemistry methods like CCSD and CCSD(T), outperforms other NQS approaches, and achieves competitive energies with state-of-the-art ab initio techniques such as HCI, ASCI, FCIQMC, and DMRG. A series of ablation and comparative studies quantifies the contribution of each enhancement to the observed improvements in accuracy and efficiency. Furthermore, we investigate the representational capacity of the ansatz, finding that its performance correlates with the inverse participation ratio (IPR), with more delocalized states being more challenging to approximate.