Equivariant Continuous Normalizing Flows with Offline Sampling for Fermionic Ground State Estimation

Research Approach

This paper addresses the computational bottleneck of Markov chain Monte Carlo (MCMC) in fermionic Variational Monte Carlo (VMC) by replacing the standard neural quantum state approach with a Continuous Normalizing Flow (CNF). Instead of learning the entire wavefunction from scratch, the authors use a fixed, antisymmetric base wavefunction (such as a Slater determinant) and apply a learned, permutation-equivariant neural ODE to transform it. This transformation captures complex many-body correlations that the base wavefunction misses while preserving the necessary antisymmetry of the fermionic state.

Key Innovations

To make this framework practical for large-scale systems, the authors introduce three primary technical contributions:

• Novel Architectures: They propose three permutation-equivariant vector field architectures—Pairwise Deep Sets (PDS), FermiNet Vector Fields (FVF), and Pairwise Deep Sets Gradient (PDSG)—which allow the model to learn interparticle correlations with varying degrees of computational complexity.
• Augmented Dynamics: By co-evolving derivative quantities as part of the ODE state, the authors eliminate the need to differentiate through the ODE trajectory during kinetic energy calculation. This significantly reduces memory usage and wall-clock time.
• Offline Sampling: Because the flow is a diffeomorphism, samples can be generated by pushing pre-cached base samples through the forward ODE. This decouples sampling from parameter optimization, allowing for embarrassingly parallel training across multiple GPUs.

Findings and Scalability

The framework demonstrates strong performance on systems of harmonically trapped spinless electrons, achieving ground-state energies that surpass CISD reference values for systems up to 35 electrons. Furthermore, the authors demonstrate near-ideal strong scaling on up to 128 NVIDIA A100 GPUs, showing that the method is well-suited for large-scale high-performance computing environments.