🔬 physics

Fully GPU-accelerated semistochastic Epstein-Nesbet perturbation theory

This paper presents a fully GPU-accelerated implementation of distributed semistochastic second-order Epstein-Nesbet perturbation theory that fuses key computational steps and utilizes CountSketch arrays to drastically reduce memory and communication costs, achieving speedups of up to 37.53 times over CPU-based methods while maintaining high parallel efficiency for large-scale quantum chemistry applications.

Original author: Lizhong Fu, Honghui Shang, Jinlong Yang

Published 2026-11-11

📖 4 min read☕ Coffee break read

CC BY 4.0

Original author: Lizhong Fu, Honghui Shang, Jinlong Yang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the author. For technical accuracy, refer to the original paper. Read full disclaimer

In the microscopic world of atoms and molecules, scientists seek to understand how electrons move and interact to hold matter together. To do this, they build mathematical models of the "wave function," a description that captures the complex, probabilistic nature of these tiny particles. For decades, the most accurate way to solve these models has been to list every possible arrangement of electrons, a method known as full configuration interaction. However, as molecules grow larger, the number of possible arrangements explodes so rapidly that even the world's most powerful supercomputers cannot list them all. To cope, researchers use a strategy called selected configuration interaction, which keeps only the most important arrangements and discards the rest. While this makes the problem manageable, it leaves a gap: the discarded arrangements still influence the energy of the molecule. To get a truly precise answer, scientists must estimate the effect of these missing pieces. This estimation, known as second-order Epstein–Nesbet perturbation theory, is a crucial final step that turns an approximate answer into a highly accurate one. Yet, performing this calculation for large molecules has traditionally been a slow, memory-hungry process that struggles to keep pace with modern computing hardware.

A team of researchers at the University of Science and Technology of China has now solved this bottleneck by creating a new way to run these calculations entirely on graphics processing units, or GPUs. These are the powerful chips originally designed for video games but now widely used in scientific computing because they can perform millions of simple calculations simultaneously. The researchers, led by Honghui Shang and Jinlong Yang, developed a software system that fuses three difficult steps—generating potential electron arrangements, filtering out the insignificant ones, and calculating their energy contributions—into a single, streamlined operation. Instead of storing a massive, unwieldy list of every possible electron arrangement that could exist outside the main calculation, their method uses a clever statistical trick. It compresses the information about these millions of discarded arrangements into a small, fixed-size memory structure. This allows the computer to sum up the contributions of all these missing pieces without ever needing to write them down individually or shuffle huge amounts of data between different processors.

The results of this new approach are striking. When tested on a single modern GPU, the system completed these complex energy corrections between eight and thirty-seven times faster than previous methods running on sixty-four traditional computer processors, all while maintaining the same level of statistical precision. The researchers demonstrated that their method scales efficiently, meaning it continues to work well even when spread across multiple GPUs working together. In tests involving a large iron-oxo complex and a model of an iron-porphyrin molecule, the system handled reference wave functions containing up to 200 million electron arrangements. For the iron-porphyrin molecule, the calculation took roughly 85 seconds on a single GPU, a task that would have been prohibitively slow or impossible with older techniques. The method also proved effective when applied to wave functions generated by artificial neural networks, a newer approach to modeling electrons that does not rely on traditional diagonalization. In these cases, the new tool successfully refined the neural network's predictions, bringing the calculated energies down to a level of accuracy known as "chemical accuracy," where the error is small enough to be useful for predicting real-world chemical behavior.

The significance of this work lies not just in the speed, but in how it changes what is possible. By removing the memory barriers that previously forced scientists to simplify their models or wait days for results, this GPU-accelerated tool allows researchers to tackle much larger and more complex chemical systems. The team applied their method to study the delicate balance of energy states in iron-based molecules, which are central to biological processes and industrial catalysis. Their calculations provided clearer insights into whether these molecules prefer a triplet or a quintet state, resolving ambiguities left by earlier, less precise studies. Furthermore, the ability to apply these corrections to neural network-based wave functions suggests a new path forward: using artificial intelligence to generate a rough guess of a molecule's behavior and then using this fast, precise tool to polish it into a definitive answer. The researchers have made their software available to the community, opening the door for others to explore the electronic structures of complex materials with a speed and efficiency that was previously out of reach.

Technical Summary: Fully GPU-Accelerated Semistochastic Epstein–Nesbet Perturbation Theory

Problem Statement
Second-order Epstein–Nesbet perturbation theory (EN-PT2) is a critical tool for refining energies obtained from Selected Configuration Interaction (SCI) methods by estimating correlation energy from determinants outside the reference space. However, applying EN-PT2 to large reference wave functions on GPU clusters faces three primary bottlenecks:

Excitation Generation: Existing implementations often generate coupled determinants on CPUs, creating a bottleneck due to indirect table accesses and variable-length outputs that require complex buffer management.
Device Memory: Storing the full set of external determinants and their partial contributions exceeds the memory capacity of current GPUs, as the external space can be orders of magnitude larger than the reference space.
Communication: In distributed settings, contributions to the same external determinant from different GPUs must be aggregated before squaring. This typically requires irregular all-to-all communication of explicit determinant records, the volume of which scales poorly with the number of generated records.

Methodology
The authors present a fully GPU-accelerated, distributed implementation of semistochastic EN-PT2 designed to overcome these limitations through three core innovations:

Fused GPU Kernel: The authors implement a single fused kernel that performs excitation generation, screening, and PT2 weighting simultaneously on the GPU. This kernel processes reference determinants (or weighted samples), generates single and double excitations from precomputed tables, applies screening thresholds ( $\epsilon_c$ and $\epsilon_f$ ), and immediately accumulates contributions. This eliminates the need to store explicit lists of external determinants or coordinate parallel writes to variable-length buffers.
CountSketch Accumulation: To avoid storing the massive external determinant set, the method accumulates contributions directly into fixed-size random-phase CountSketch arrays. Each external determinant is mapped to a specific "bucket" via a hash function and assigned a random phase ( $\pm 1, \pm i$ ). Contributions to the same determinant share the same bucket and phase, allowing them to sum coherently. Contributions to different determinants in the same bucket introduce cross-terms, but their expectation value vanishes due to independent phases. This allows the method to estimate the PT2 energy without storing the full external set.
Bounded Collective Communication: Because the CountSketch arrays are fixed-size and use identical hashing and phase schemes across all GPUs, the arrays can be summed across the cluster using a standard AllReduce operation (via NCCL). This replaces the expensive, irregular all-to-all exchange of explicit determinant records with a bounded collective summation of small, fixed-size arrays.

The implementation follows a semistochastic strategy, combining a deterministic control (traversing a subset of large coefficients) with a sampled correction (stochastically sampling the remainder) to reach a target statistical standard error (SE).

Key Contributions

Algorithmic Fusion: The integration of excitation generation, screening, and accumulation into a single GPU traversal, removing CPU-GPU data transfer overheads for intermediate steps.
Memory Efficiency: The use of CountSketch reduces peak memory usage from scaling with the number of external determinants (gigabytes to terabytes) to a fixed size determined by the number of buckets (megabytes), enabling calculations on devices with limited memory.
Scalable Communication: The replacement of record-based all-to-all communication with AllReduce on fixed-size arrays, drastically reducing communication volume and enabling strong scaling.
General Applicability: The code accepts explicit determinants and coefficients from various variational methods, including traditional SCI and Neural Network Quantum States (NNQS), making it a versatile post-processing tool for the GPU era.

Results
The paper validates the method through numerical accuracy, speedup benchmarks, and large-scale applications:

Accuracy: Comparisons with the CPU-based Arrow code across 10 molecules show that the GPU implementation reproduces PT2-corrected energies with high fidelity. For a Fe-porphyrin system, the empirical standard deviation of 64 independent runs matched the reported statistical standard error, confirming the reliability of the CountSketch estimator.
Speedup: On a single NVIDIA A100 GPU, the implementation is 8.21 to 37.53 times faster than Arrow running on 64 CPU cores for systems requiring the same statistical standard error (10 $\mu E_h$ ).
Parallel Efficiency: Strong-scaling tests on up to 16 GPUs (across two nodes) achieved 93.5% parallel efficiency for a fixed workload, demonstrating excellent scalability.
Memory Reduction: For a reference with 2 million determinants, the sketch-based accumulation reduced peak tensor memory per GPU by a factor of approximately 35 compared to exact accumulation (0.325 GiB vs. 11.35 GiB).
Large-Scale Applications:
- Iron-Oxo Complex: Applied to a (60e, 187o) active space with 150 million determinants, refining the triplet-quintet splitting and confirming the quintet preference without the reversal suggested by smaller calculations.
- Iron-Porphyrin: Handled references with 200 million determinants per state, refining the triplet-quintet gap to 5.93 m $E_h$ (extrapolated), consistent with other high-level methods.
- Neural Network Quantum States (NNQS): Successfully applied to autoregressive NNQS wave functions for H $_2$ O, NH $_3$ , and HCN. The PT2 correction reduced signed errors to chemical accuracy (within 1.6 m $E_h$ ) even when the variational errors were significantly larger, demonstrating the utility of perturbative post-processing for neural network states.

Significance
The paper claims that this work establishes a practical and efficient route for applying perturbative corrections to large reference wave functions in the GPU era. By fusing generation and accumulation and utilizing CountSketch, the method removes the memory and communication barriers that previously hindered GPU-accelerated EN-PT2. The authors emphasize that this allows researchers to directly apply high-accuracy perturbative corrections on the same GPU clusters used to generate the reference wave functions, whether those references come from traditional SCI methods or emerging neural-network-based approaches. The results suggest that this approach can significantly reduce the variational training required for NNQS to reach chemical accuracy.

Technical Summary: Fully GPU-Accelerated Semistochastic Epstein–Nesbet Perturbation Theory

More like this