[论文解读] Bridging the Gap between Deep Learning and Frustrated Quantum Spin System for Extreme-scale Simulations on New Generation of Sunway Supercomputer

Mingfan Li, Junshi Chen+8 more|arXiv (Cornell University)|Aug 31, 2021

Quantum many-body systems参考文献 50被引用 23

一句话总结

本论文提出了一种基于卷积神经网络（CNN）的变分波函数，并结合迁移学习与二阶优化（随机重配置，Stochastic Reconfiguration），在申威超算上模拟了二维受挫自旋-1/2 J1–J2海森堡模型。通过利用超过3000万个异构核心及通过ScaLapack实现的分布式线性代数，该方法实现了最先进的精度，成功模拟了此前难以处理的24×24量子自旋系统，展示了深度学习与高性能计算协同在量子多体模拟中前所未有的规模与效率。

ABSTRACT

Efficient numerical methods are promising tools for delivering unique insights into the fascinating properties of physics, such as the highly frustrated quantum many-body systems. However, the computational complexity of obtaining the wave functions for accurately describing the quantum states increases exponentially with respect to particle number. Here we present a novel convolutional neural network (CNN) for simulating the two-dimensional highly frustrated spin-$1/2$ $J_1-J_2$ Heisenberg model, meanwhile the simulation is performed at an extreme scale system with low cost and high scalability. By ingenious employment of transfer learning and CNN's translational invariance, we successfully investigate the quantum system with the lattice size up to $24 imes24$, within 30 million cores of the new generation of sunway supercomputer. The final achievement demonstrates the effectiveness of CNN-based representation of quantum-state and brings the state-of-the-art record up to a brand-new level from both aspects of remarkable accuracy and unprecedented scales.

研究动机与目标

• 为克服二维受挫自旋系统中量子多体态表示的指数级增长问题。
• 实现自旋-1/2 J1–J2海森堡模型在18×18极限之外的极端规模模拟。
• 将深度学习与高性能计算（HPC）相结合，实现可扩展、低成本且高精度的量子态模拟。
• 通过异构并行与分布式线性代数，解决蒙特卡洛采样与二阶优化（SR方法）中的计算瓶颈。
• 在前所未有的系统规模下，展示基于CNN的波函数表示方法的可行性与优越性。

提出的方法

• 采用具有平移不变性的卷积神经网络（CNN）表示量子多体波函数，利用局部连接与参数共享实现高效计算。
• 在多个晶格尺寸（如10×10至24×24）之间应用迁移学习，实现知识迁移并降低大规模系统的训练成本。
• 采用随机重配置（Stochastic Reconfiguration, SR）方法作为二阶优化技术，通过利用梯度中统计相关性，提升收敛性与精度。
• 通过swDNN库加速CNN的推理与训练，将计算密集型操作卸载至申威处理器上的异构核心。
• 重构并行蒙特卡洛采样，使每个进程支持多个独立链，提升吞吐量并减少通信开销。
• 使用分布式ScaLapack执行SR方法所需的大型稠密线性代数运算，实现对超过10万个变分参数的协方差矩阵的高效处理。

实验结果

研究问题

• RQ1 基于CNN的变分波函数能否在极端规模下，实现对二维受挫量子自旋系统的高精度与高可扩展性模拟？
• RQ2 如何有效利用迁移学习，以降低计算成本并实现从较小系统（如10×10）向更大系统（如24×24）的模拟扩展？
• RQ3 当深度学习与申威等极端规模HPC系统结合时，尤其在通信密集型工作负载下，可实现怎样的性能与可扩展性？
• RQ4 在蒙特卡洛样本有限的条件下，二阶优化（SR）在优化量子多体波函数方面相较于一阶方法有多大优势？
• RQ5 深度学习与HPC的协同作用能否克服强关联体系中量子态表示的指数复杂度？

主要发现

• 该方法成功模拟了24×24正方晶格上的二维自旋-1/2 J1–J2海森堡模型，这是首次使用基于CNN的变分方法研究如此大规模的系统。
• 模拟在申威超算超过3000万个异构核心上完成，490,000个进程的并行效率达到92.5%。
• 通过迁移学习与分布式优化，CNN模型实现了106,529个变分参数，是此前基于CNN模型的30倍之多。
• 地面态能量计算达到最先进的精度，显著优于以往在18×18晶格上的结果，并将搜索空间扩展至2252倍之大。
• 通过使用分布式ScaLapack，高效计算了SR方法中的大型协方差矩阵，克服了二阶优化中内存与性能瓶颈。
• 迁移学习降低了训练成本，并通过从较小预训练模型初始化，实现了24×24系统的稳定收敛，展示了通往更大量子系统模拟的可扩展路径。