MPS-VMC：一种高性能的矩阵积态变分蒙特卡洛求解器，用于从头算量子化学计算

《Computer Physics Communications》：MPS-VMC: A high-performance matrix product state variational Monte Carlo Solver for ab initio quantum chemistry

量子化学计算中，矩阵产品态（MPS）与变分蒙特卡洛（VMC）结合提出MPS-VMC方法，创新采用四叉树批处理自回归采样和GEMM优化算法，显著提升大系统（76自旋轨道子）的收敛精度与计算效率。

　　

徐志倩|吴阳军|傅立中|郭楚|尚红辉

中国科学技术大学精密与智能化学国家重点实验室，合肥，中国

摘要

量子蒙特卡洛（Quantum Monte Carlo, VMC）为解决量子化学中的复杂量子多体问题提供了一个强大的变分框架。然而，其实际应用常常受到局部能量计算和采样效率的重大瓶颈限制。为了克服这些限制，我们提出了MPS-VMC，这是一种将矩阵积态（Matrix Product States, MPS）集成到变分蒙特卡洛框架中的新型量子化学求解器。利用MPS的固有特性，我们引入了一种基于四叉树的高性能批量自回归深度优先采样算法。此外，我们还提出了一种基于样本的高性能批量推理算法，该算法将通用矩阵-向量（General Matrix-Vector, GEMV）操作合并为通用矩阵-矩阵（General Matrix-Matrix, GEMM）操作，从而显著提高了推理效率。实验结果表明，与其他基于VMC的方法相比，MPS-VMC在具有多达76个自旋轨道的分子系统中表现出更优的收敛精度。

引言

量子多体系统对于理解微观层面的物理化学现象至关重要。在从头算量子化学中，主要挑战在于求解时间无关的多体薛定谔方程，该方程描述了决定分子结构和性质的核与电子之间的复杂相互作用。这个问题的计算复杂性随粒子数呈指数级增长，使得仅对最简单的系统才能得到解析解[1]。为了应对这一挑战，计算方法通常使用有限基函数对问题进行离散化，将多体量子态表示为反对称Slater行列式的组合。然而，行列式空间的指数级增长严重限制了高精度方法的适用性。全配置相互作用（Full Configuration Interaction, FCI）虽然能够在给定基组内提供精确解，但仅限于小分子。尽管像耦合簇（Coupled Cluster, CC）[2]这样的传统近似方法提供了实用的替代方案，但它们在强关联系统中会遇到困难。最近的进展试图通过创新的重构来克服这些限制。例如，坐标下降FCI（Coordinate Descent FCI）[3]、[4]、[5]方法将FCI计算重新表述为优化问题，使得在更大的行列式空间中进行近乎精确的计算成为可能。尽管取得了这些进展，CDFCI仍然面临根本的可扩展性限制，因为它最终必须遍历所有Slater行列式的空间。

神经网络量子态（Neural Network Quantum States, NNQS）作为一种强大的变分方法，已成为分子基态的候选者[6]，为传统量子化学方法提供了有前景的替代方案。尽管具有潜力，但NNQS在应用于较大分子系统时仍面临关键的可扩展性挑战。计算瓶颈主要来自两个方面：局部能量计算，其成本随哈密顿项数量的增长而增加，导致内存需求过高；以及采样效率低下，传统的马尔可夫链蒙特卡洛（Markov Chain Monte Carlo, MCMC）方法在高维空间中表现不佳。最近的进展试图通过创新架构来解决这些限制。自回归波函数模型（Autoregressive Wavefunction Models）在克服采样瓶颈方面显示出特别的前景[7]、[8]。一些方法同时解决了这两个挑战——一些利用多GPU并行化的自回归NNQS[9]，其他方法则结合自回归解码器和近似局部能量计算以及GPU加速[10]。然而，这些方法可能在收敛速度和精度方面存在局限性。

与这些发展并行的是，密度矩阵重整化群（Density Matrix Renormalization Group, DMRG）及其变分方法矩阵积态（Matrix Product State, MPS）[11]、[12]、[13]、[14]、[15]已成为解决量子多体问题的强大工具[16]、[17]。DMRG在从头算量子化学中取得了显著成功[18]、[19]、[20]、[21]。并行化的DMRG实现[22]、[23]可以处理高达CAS(114e, 73o)的完整活性空间。对于如此庞大的活性空间，DMRG的计算复杂性和内存开销变得过大，可扩展性和并行效率受到显著限制。相比之下，基于VMC的方法提供了更好的可扩展性和高度可并行化的计算方案，当MPO键维数较大时（这在量子化学应用中很常见），这一点在我们的复杂性分析（第2.1节）中有详细说明。NNQS框架虽然表现出优异的可扩展性，但往往缺乏物理直觉。已经有人尝试在物理模型系统中优化MPS方法与VMC方法的结合[24]。然而，这种方法尚未在分子系统中得到探索，在分子系统中，VMC过程中所需的局部能量计算需要不同于具有有限相互作用项的模型系统的处理方式，而且哈密顿MPO键维数的扩展性也对传统DMRG算法构成了挑战。在这项工作中，我们旨在结合DMRG和NNQS方法的优点，扩展变分蒙特卡洛（VMC）方法在量子化学问题中的适用范围，并利用神经网络领域广泛使用的高性能GPU实现来进一步提高该方法的速度。本文将MPS作为VMC框架内的自回归方法引入，并开发了利用MPS固有特性的高性能采样和推理方法。这些进步提高了VMC的准确性和可扩展性，从而建立了MPS-VMC量子化学求解器。

本文的结构如下：第2节提供了关于变分蒙特卡洛算法、矩阵积态以及自回归采样概念的基本背景。第3节详细介绍了MPS-VMC的核心方法，包括提出的基于四叉树的高性能批量自回归深度优先采样算法和高效的基于样本的批量推理算法，该算法将GEMV操作合并为GEMM操作。第4节展示了我们的实验结果，证明了MPS-VMC在分子系统中的优越收敛精度，并分析了我们的采样和推理策略的性能。最后，第5节总结了我们的发现，并讨论了MPS-VMC在大规模量子化学计算中的意义和未来发展方向。

代码片段

用于矩阵积态的变分蒙特卡洛算法

一个量子化学问题涉及找到包含N_e个电子的分子系统的基态。在玻恩-奥本海默（Born-Oppenheimer）近似下，分子哈密顿量可以表示为