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.
翻译:神经网络量子态(NNQS)已成为解决量子多体问题的有前景候选方法,但其实际应用常受限于高成本的采样和局域能量计算。我们提出了一种用于从头算电子结构计算的高性能NNQS方法。主要创新包括:(1) 基于Transformer架构的量子波函数拟设;(2) 一种数据中心化的变分蒙特卡洛(VMC)算法并行化方案,该方案保持数据局部性并能良好适应不同计算架构;(3) 一种并行批次采样策略,可降低采样成本并实现良好的负载均衡;(4) 一种兼顾内存和计算效率的并行局域能量评估方案;(5) 对实际化学体系的研究表明,我们的方法在精度上优于现有最先进技术,且对于包含多达120个自旋轨道的大型分子系统展现出优异的强扩展性和弱扩展性。