NNQS-Transformer: an Efficient and Scalable Neural Network Quantum States Approach for Ab initio Quantum Chemistry

TL;DR

This paper introduces NNQS-Transformer, a scalable neural network quantum state method utilizing transformer architecture and advanced parallelization techniques, significantly improving efficiency and accuracy for ab initio quantum chemistry calculations.

Contribution

The paper presents a novel transformer-based neural network quantum state approach with optimized parallel algorithms for scalable and accurate ab initio electronic structure computations.

Findings

Achieves superior accuracy over existing methods.

Demonstrates strong scalability for large molecular systems.

Reduces sampling cost with parallel batch strategies.

Abstract

Neural network quantum state (NNQS) has emerged as a promising candidate for quantum many-body problems, but its practical applications are often hindered by the high cost of sampling and local energy calculation. We develop a high-performance NNQS method for \textit{ab initio} electronic structure calculations. The major innovations include: (1) A transformer based architecture as the quantum wave function ansatz; (2) A data-centric parallelization scheme for the variational Monte Carlo (VMC) algorithm which preserves data locality and well adapts for different computing architectures; (3) A parallel batch sampling strategy which reduces the sampling cost and achieves good load balance; (4) A parallel local energy evaluation scheme which is both memory and computationally efficient; (5) Study of real chemical systems demonstrates both the superior accuracy of our method compared to state-of-the-art and the strong and weak scalability for large molecular systems with up to $120$ spin orbitals.