中国团队利用新一代神威超级计算机实现高精度量子模拟

光子盒研究院出品

近日，由中国科学技术大学、国家海洋科学与技术试点实验室（青岛）、北京大学数学科学学院、无锡国家超级计算中心和中国海洋大学组成的联合团队在arXiv网站发布了《通过可扩展的深度学习方法解决2^1296指数级复杂的量子多体模拟》论文[1]。

研究表明，一个基于深度学习的模拟协议可以在希尔伯特空间实现最先进的精度，具体表现为自旋系统达2^1296，费米子系统达3^144，并在新一代神威超级计算机上使用HPC-AI（高性能计算-人工智能）混合框架。这项研究的应用具有高度的可扩展性：最高可达4000万个异构核心，测得的弱扩展效率为94%、强扩展效率为72%。

当前一代神威超级计算机

这项工作的完成为在前所未有的晶格规模上以极高的精度模拟自旋模型和费米子模型打开了大门。

强相关的量子多体物理学是凝聚态物理学中最迷人的研究领域之一。在量子多体物理的模型中，有两个经典模型：1）海森堡自旋模型，其中海森堡模型本质是一个自旋模型，它描述了格点上两个自旋量子的相互作用；和2）费米-哈伯德模型，它描述了量子在格点上的运动。

解决这些问题将大大加深我们对自然界基本规律的理解，并指导我们发现新的物理现象和新的量子材料，可能在能源、信息等方面有很大的应用潜力。

尽管量子力学的基本规律是已知的，但解决强相关的量子多粒子问题仍然是极具挑战性的：因为解的希尔伯特空间随着问题大小呈指数级增长。此外，微扰方法虽然在模拟弱相关物质系统方面取得了巨大的成功，但对于强相关系统来说，将完全失败。揭示量子多体系统迷人的物理本质主要依靠非微扰数值方法，如精确对角化（ED）、量子蒙特卡罗方法（QMC）和密度矩阵重整化群（DMRG）。然而，这些方法都有严重的局限性：例如，ED的计算成本随着系统大小呈指数增长，因此ED方法的系统大小被限制在50个位置以下；QMC对于费米和受挫系统有众所周知的符号问题；而DMRG仅限于一维或准一维系统，对于高维系统效果不佳。为了保持相同的精度，DMRG的计算成本随着晶格的宽度（L）呈指数级增长。到目前为止，DMRG模拟的宽度被限制在𝐿≈12。

最近，所谓的张量网络方法，如投影纠缠对状态（PEPS）方法已经被开发出来，它可以模拟费米和受挫系统。一些学者已经开发了大规模的PEPS代码PEPS++，并在神威超级计算机上实现了它。下图演示了在24×24晶格上求解𝐽1-𝐽2模型，边界条件是开放的，连接维度𝐷=16。然而，周期性边界条件（PBC）通常具有比OBC更低的边界效应，更适合于模拟量子多体系统。由于对连接维度𝑂(𝐷18)的高扩展性，目前还没有关于PEPS对连接维度𝐷≥6的PBC进行模拟的报道，这对于捕捉量子多体模型的基本物理现象来说仍然太小了。

由于希尔伯特空间随着粒子数增加而指数增长（组合爆炸），量子多体问题的高精度模拟是对于经典计算机极富挑战性的问题。近几年发展起来的深度学习算法为模拟量子多体提供了新的有效的计算工具。

(a)𝐽1-𝐽2模型和(b)𝑡-𝐽模型，(c)神经网络的自我学习优化程序过程

对目前解决自旋模型和费米子模型的最先进的神经网络方法进行比较。未受挫的海森堡模型已经被RBM解出了相当高的精度。对于𝐽1-𝐽2模型，使用了三种方法。1）单一卷积神经网络（CNN）；2）深度CNN；3）Gutzwiller投影波函数和RBM。由于Gutzwiller投影波函数的计算复杂度很高，所以研究的晶格大小被限制在𝐿=18。深度CNN的方法具有较低的计算复杂度，然而能量的精度取决于深度CNN的表示能力。有了适当的深度CNN结构，并将参数数增加到10^5，能量的精度就会达到最先进的水平。在这项工作中，参数进一步增加到4×10^5，研究的晶格大小增加到𝐿=36。对于费米子模型，由于斯莱特行列式的计算复杂度较高，研究的晶格大小受到限制。在这项工作中，使用无行列式的深度CNN波函数，在方形晶格上求解了𝑡-𝐽模型，宽度最大为𝐿=12。

与传统的深度学习任务（如分类）相比，通过神经网络解决量子多体问题存在一些重大挑战——因为人们必须在指数级大的希尔伯特空间中获得极高的精确基态：

1）神经网络的泛化能力应该足够高，以代表指数级大希尔伯特空间中的量子态；

2）网络参数的双精度是必须的；

3）地面能量应该是能量的全局最小值。基于一阶梯度的优化器如Adam、SGD等并不高效，因为它们很容易陷入局部最小值。在这里，研究团队使用了类似于二阶自然梯度的方法，如随机重构法；

4）在指数级大的希尔伯特空间中寻找解决方案需要极其大量的“马尔可夫链蒙特卡罗算法”（MCMC）样本，这也是精确的“超级分辨率”（SR）优化所需要的。

为了满足这一系列的严格要求，研究人员为自旋模型和费米子模型分别开发了两种深度CNN结构：分别表示为CNN1和CNN2。

CNN1的结构由(a)(b)描述，其中(a)是(b)的深层结构的构建块；CNN2的结构由(c)(d)描述，其中(c)是(d)的深层结构的构建块；(e)CNN1的迁移学习过程；(f)CNN2的初始状态选择过程。

本工作中使用的深度CNN结构与其他用于量子多体问题的CNN结构相比有几个重要的区别：1）文中的深度CNN中的非线性是由“最大池化操作”（maxpooling）引起的，而不是传统的激活函数。最大池化操作在卷积滤波器中拾取最重要的自由度，这类似于重整化群理论中的粗粒度过程；2）波函数系数是由神经元的乘积产生的，这与基于RBM的结构中的指数函数不同。

sw26010pro处理器是sw26010处理器的升级版，用于神威超级计算机。sw26010pro由六个核心组（CG）组成，每个核心组都连接到一个环形网络；每个CG包含一个管理过程元素（MPE，控制核心）和一个由64个计算过程元素（CPE，计算核心）组成的集群，排列成8×8的网状结构。

sw26010pro结构图

本研究使用的神经网络涉及大量的内存密集型OP（即PBC_padding，Conv1d）。因此，为物理学提出的特殊模型在很大程度上受到计算能力和内存访问带宽之间的限制。因此研究人员在新的神威异构处理器上对这些运算器进行了优化，特别是对几个内存密集型的运算器。

多核加速对内存密集型操作的影响。

与MPE基线版本相比，swHMAE实现了巨大的性能改进。Tile的多核加速率达到34.45，Slice达到39.34，PBC达到32，GradPBC达到52.77。GradPBC算子获得了更好的加速，这源于它的计算特性。它不是一个纯粹的内存访问运算器，而是包括一定量的累积计算，这增加了CPE的利用率。有了这些高度优化的核，热点能量计算的整体速度比原来的MPE版本取得了相当大的性能提高。多核加速版本对有106529和421953个参数的CNN计算模型获得了90倍和130倍的速度提升。

对于整个应用测试，研究人员记录了一个完整的训练迭代的总执行时间：MCMC阶段包括基于CNN的随机游走和重要性采样，SR阶段包括建立协方差矩阵和矩阵反演。强扩展（strong scaling）分析和弱扩展（weak scaling）分析结果如下：

详细的强扩展（a）和弱扩展（b）结果。两个阶段（MCMC、SR）有四个类别（随机行走、重要性采样、矩阵乘法和反转）的时间剖析。

MCMC可以完全并行化，并表现出了很好的可扩展性；SR优化有两个变量：矩阵乘法（gemm）和矩阵因式分解（potrf）。

gemm的计算随着行和列的大小而增加，而行和列是总马尔科夫链和参数的数量；同时，potrf的计算只与参数数有关，与系统规模或马尔科夫链数量完全无关。对于具有106529个参数和小批量的模型，potrf的计算量远远高于gemm的计算量，因此主导了SR的执行时间；对于并行的potrf，SR时间随着进程的增加而减少。然而，gemm随着批处理量的增加而增加，它逐渐主导了大部分的执行时间，如下图所示：

不同进程数下gemm（potrf）操作的时间。左图的矩阵大小为106529，右图的矩阵大小为421953。

参与并行搜索的马尔科夫链的数量越多，找到的初始状态的质量就越高。由于希尔伯特空间很大，对于大晶格来说，找到高质量初始状态的难度明显高于小晶格，这个困难可以通过增加链数来克服。对于小晶格量子系统，初始态的质量对能量收敛的影响相对较小。在选定的初始状态下，经过100个SR步骤，12×12晶格的能量从-0.3517下降到-0.4523；16×16晶格从-0.3483下降到-0.4599，24×24晶格从-0.2900下降到-0.4173：

不同数量的马尔科夫链为𝑡-𝐽模型找到的初始状态的平均能量。

𝐽1-𝐽2模型和𝑡-𝐽模型的优化结果如下图所示：

将CNN1得到的不同晶格宽度𝐿上的𝐽1-𝐽2模型的基态能量与Fugaku超级计算机成果进行比较。

通过CNN2与PBC，以及PEPS++与OBC得到的𝑡-𝐽模型的基态能量比较。

与富岳超级计算机上结果进行比较，𝐿=10时，本工作中的能量为-0.497468，比已有报告的-0.497629高3.2×10^-4；对于𝐿=18，本工作中得到的基态能量为-0.496500，比已有报告的-0.496275低4.5×10^-4。由于CNN1出色的可扩展性，使得团队能够在明显大于先前工作的规模上研究该模型。

与通过PEPS计算的OBC的结果进行比较。当𝐿=8时，CNN2获得的能量为-0.60965，比已有报告的能量高3×10^-3；当𝐿=12时，CNN2获得的能量为-0.63265，比已有报告的数值高3×10^-3。然而，不同边界条件的能量不能直接比较，这里显示的结果是证明了CNN2对费米子系统的有效性。

本篇研究展示了𝐽1-𝐽2模型和𝑡-𝐽模型相较已有的研究结果具有精度优势。

因此，研究人员得以对一大类重要的物理模型进行模拟：如量子自旋液体、高温超导、超固体、重费米子、分数量子霍尔效应等，在高精度下对物理学作出有把握的结论。

除此之外，本篇研究还介绍了用深度学习方法解决量子多体问题。

这个框架可能适用于其他问题，例如经典超级计算机上的量子电路模拟：目前张量网络方法在超级计算机上模拟量子电路的效率高达10×10的晶格，基于可扩展的神经网络的量子态表示方法可能能够对更大数量的量子比特进行量子电路模拟。

参考文献：

[1]https://arxiv.org/pdf/2204.07816.pdf