[Paper Review] 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 systems50 references23 citations

TL;DR

This paper presents a scalable convolutional neural network (CNN)-based variational ansatz combined with transfer learning and second-order optimization (Stochastic Reconfiguration) to simulate the 2D frustrated spin-1/2 J1–J2 Heisenberg model on the Sunway supercomputer. By leveraging over 30 million heterogeneous cores and distributed linear algebra via ScaLapack, the method achieves state-of-the-art accuracy and simulates a 24×24 quantum spin system—previously intractable—demonstrating unprecedented scale and efficiency in quantum many-body simulations using deep learning and high-performance computing synergy.

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.

Motivation & Objective

• To overcome the exponential scaling of quantum many-body state representation in 2D frustrated spin systems.
• To enable extreme-scale simulations of the spin-1/2 J1–J2 Heisenberg model beyond the 18×18 limit.
• To integrate deep learning with high-performance computing (HPC) for scalable, low-cost, and accurate quantum state simulations.
• To address computational bottlenecks in Monte Carlo sampling and second-order optimization (SR method) via heterogeneous parallelism and distributed linear algebra.
• To demonstrate the feasibility and superiority of CNN-based wave function representations at unprecedented system sizes.

Proposed method

• A convolutional neural network (CNN) with translational invariance is used to represent the quantum many-body wave function, exploiting local connectivity and parameter sharing for efficient computation.
• Transfer learning is applied across multiple lattice sizes (e.g., from 10×10 to 24×24), enabling knowledge transfer and reducing training cost for larger systems.
• The Stochastic Reconfiguration (SR) method is employed as a second-order optimization technique to improve convergence and accuracy by leveraging statistical correlations in gradients.
• The CNN inference and training are accelerated via the swDNN library, offloading compute-intensive operations to heterogeneous cores on the Sunway processor.
• Parallel Monte Carlo sampling is restructured to allow multiple independent chains per process, improving throughput and reducing communication overhead.
• Distributed ScaLapack is used to perform large-scale dense linear algebra operations required by the SR method, enabling efficient handling of the covariance matrix for over 100k variational parameters.

Experimental results

Research questions

• RQ1 Can a CNN-based variational ansatz achieve high accuracy and scalability in simulating 2D frustrated quantum spin systems at extreme scales?
• RQ2 How can transfer learning be effectively leveraged to reduce computational cost and enable simulation of larger lattices (e.g., 24×24) from smaller ones?
• RQ3 What performance and scalability can be achieved when combining deep learning with extreme-scale HPC systems like Sunway, particularly under communication-intensive workloads?
• RQ4 To what extent does second-order optimization (SR) outperform first-order methods in optimizing quantum many-body wave functions with limited Monte Carlo samples?
• RQ5 Can the synergy between deep learning and HPC overcome the exponential complexity of quantum state representation in strongly correlated systems?

Key findings

• The method successfully simulated the 2D spin-1/2 J1–J2 Heisenberg model on a 24×24 square lattice, marking the first time such a large system has been studied using a CNN-based variational approach.
• The simulation was performed on over 30 million heterogeneous cores of the Sunway supercomputer, achieving a parallel efficiency of 92.5% across 490,000 processes.
• The CNN model achieved 106,529 variational parameters—30 times larger than previous CNN-based models—by leveraging transfer learning and distributed optimization.
• The ground state energy was computed with state-of-the-art accuracy, significantly outperforming prior results on 18×18 lattices and extending the record to a 2252-fold larger search space.
• The use of distributed ScaLapack enabled efficient computation of the large covariance matrix in the SR method, overcoming memory and performance bottlenecks in second-order optimization.
• Transfer learning reduced training cost and enabled stable convergence for the 24×24 system by initializing from smaller, pre-trained models, demonstrating a scalable path to larger quantum systems.