The introduction of Neural Quantum States (NQS) has recently given a new twist to variational Monte Carlo (VMC). The ability to systematically reduce the bias of the wave function ansatz renders the approach widely applicable. However, performant implementations are crucial to reach the numerical state of the art. Here, we present a Python codebase that supports arbitrary NQS architectures and model Hamiltonians. Additionally leveraging automatic differentiation, just-in-time compilation to accelerators, and distributed computing, it is designed to facilitate the composition of efficient NQS algorithms.
The introduction of Neural Quantum States (NQS) has recently given a new twist to variational Monte Carlo (VMC). The ability to systematically reduce the bias of the wave function ansatz renders the approach widely applicable. However, performant implementations are crucial to reach the numerical state of the art. Here, we present a Python codebase that supports arbitrary NQS architectures and model Hamiltonians. Additionally leveraging automatic differentiation, just-in-time compilation to accelerators, and distributed computing, it is designed to facilitate the composition of efficient NQS algorithms.
The introduction of Neural Quantum States (NQS) has recently given a new twist to variational Monte Carlo (VMC). The ability to systematically reduce the bias of the wave function ansatz renders the approach widely applicable. However, performant implementations are crucial to reach the numerical state of the art. Here, we present a Python codebase that supports arbitrary NQS architectures and model Hamiltonians. Additionally leveraging automatic differentiation, just-in-time compilation to accelerators, and distributed computing, it is designed to facilitate the composition of efficient NQS algorithms.
When combined with highly expressive ansatz functions such as neural quantum states, variational Monte Carlo (VMC) constitutes a versatile numerical approach to tackle the quantum many-body problem in and out of equilibrium. However, its traditional formulation exhibits a subtle estimation bias leading to inaccuracies, which can be particularly detrimental when addressing real time dynamics. In this work, we investigate two avenues to circumvent said estimation bias. First, we propose an unbiased variant of time-dependent VMC using self-normalized importance sampling with respect to a cutoff-based deformation of the Born distribution. We demonstrate the feasibility and accuracy of the approach in pathological and generic cases of quench dynamics. Furthermore, we explore an alternative sampling strategy based on active learning via the tensor cross interpolation (TCI). While we find that our choice of tensor network architecture lacks the required low rank property, the proposed TCI-based algorithm complements the conventional importance sampling paradigm, providing an alternative perspective that may be further explored in future work.
We present jQMC, a Python-based computational package for {\it ab initio} Quantum Monte Carlo (QMC) simulations, designed for modern GPU-accelerated computing environments. jQMC implements two well-established QMC algorithms: Variational Monte Carlo (VMC) and the lattice-regularized variant of Diffusion Monte Carlo (LRDMC). The employed wave function is a Jastrow factor combined with the antisymmetrized geminal power with spin-singlet and spin-triplet pairings, which contains the single Slater determinant as its special lowest-rank case. The wave function can be initialized from external Hartree-Fock/Density Functional Theory calculations through the TREX-IO library (a common wave-function format across electronic-structure packages) and optimized by stochastic reconfiguration and linear-method energy minimization. One of the prominent features of jQMC is its use of JAX, which enables automatic differentiation for wave function optimization and atomic force calculations, and allows the main QMC algorithms to be Just-In-Time (JIT) compiled and portable across CPU and GPU. jQMC is vectorized over walkers at the top level of the QMC algorithms, providing efficient intra-GPU~(CPU) vectorization. The multi-GPU~(CPU) parallelization is also supported through MPI and JAX sharding. To assess the practical performance of this implementation, we benchmarked jQMC performance on NVIDIA GPUs (A100 and H100) and analyzed CUDA kernels. For the test cases analyzed here, with system sizes up to 160 electrons, the current version of jQMC is faster than TurboRVB, a Fortran90 code implementing the same algorithms and wave functions, once jQMC is run on GPUs. In terms of wall-time, the gain can reach an order of magnitude for VMC, while it is more moderate for LRDMC.
Neural Network Variational Monte Carlo (NNVMC) has emerged as a promising paradigm for solving quantum many-body problems by combining variational Monte Carlo with expressive neural-network wave-function ansätze. Although NNVMC can achieve competitive accuracy with favorable asymptotic scaling, practical deployment remains limited by high runtime and memory cost on modern graphics processing units (GPUs). Compared with language and vision workloads, NNVMC execution is shaped by physics-specific stages, including Markov-Chain Monte Carlo sampling, wave-function construction, and derivative/Laplacian evaluation, which produce heterogeneous kernel behavior and nontrivial bottlenecks. This paper provides a workload-oriented survey and empirical GPU characterization of four representative ansätze: PauliNet, FermiNet, Psiformer, and Orbformer. Using a unified profiling protocol, we analyze model-level runtime and memory trends and kernel-level behavior through family breakdown, arithmetic intensity, roofline positioning, and hardware utilization counters. The results show that end-to-end performance is often constrained by low-intensity elementwise and data-movement kernels, while the compute/memory balance varies substantially across ansätze and stages. Based on these findings, we discuss algorithm--hardware co-design implications for scalable NNVMC systems, including phase-aware scheduling, memory-centric optimization, and heterogeneous acceleration.
TitleVariational Monte Carlo for Ab Initio Electronic Structure with Neural Quantum States and Selective State-Space Models Overview This repository provides a research implementation of Variational Monte Carlo (VMC) for ab initio electronic-structure calculations using Neural Quantum States (NQS). The work focuses on stable natural-gradient optimization, physically motivated wavefunction constraints, and the exploration of alternative architectural inductive biases for modeling electronic correlation. The project is released as an exploratory research artifact. All results, limitations, and open questions are stated explicitly. The work is motivated by a central question in modern neural quantum chemistry:can alternative sequence-based architectures reduce the quadratic scaling bottleneck of electron–electron interaction modeling while preserving physical fidelity? Scientific Scope The framework targets: Atomic systems (H through Ne) Small molecules (e.g., H₂, LiH, H₂O) Exploratory extensions, including: Periodic boundary conditions Excited-state considerations Berry-phase estimation Entanglement diagnostics Time-dependent VMC via McLachlan’s variational principle At present, numerical validation is limited to a single-electron benchmark (Hydrogen) and is explicitly presented as a sanity check of the physics engine and optimization loop, not as evidence of improved expressivity for correlated fermions. Core Methodology The implementation follows the standard Variational Monte Carlo formalism, in which a parameterized wavefunctionψ is optimized by minimizing the energy expectation value Optimization is performed using Stochastic Reconfiguration (SR), a natural-gradient method that accounts for the geometry of the variational manifold and significantly improves stability compared to first-order optimizers. Key numerical features include: Log-domain arithmetic for wavefunction evaluation Stable determinant computation using slogdet Hutchinson trace estimation for Laplacian evaluation Explicit enforcement of Kato cusp conditions via analytic Jastrow factors Multi-determinant antisymmetric ansatz to satisfy fermionic exchange symmetry SSM-Backflow Hypothesis The primary research contribution explored in this repository is the SSM-Backflow hypothesis. Most contemporary NQS models (e.g., FermiNet, PauliNet) model electron–electron interactions using explicit quadratic aggregation, leading to O(Ne2)O(N_e^2)O(Ne2) computational complexity. This work investigates an alternative inductive bias: For each electron, neighboring electrons are sorted by inter-electron distance The resulting sequence is processed using a selective state-space model (SSM) inspired by Mamba architectures Correlation information is accumulated sequentially through a learned recurrence Aggregated messages are used to parameterize backflow transformations or orbital functions This approach yields a theoretical complexity of O(NelogNe)O(N_e \log N_e)O(NelogNe) per layer (dominated by sorting) and is motivated by the physical observation that electronic correlation typically decays approximately exponentially with distance. Importantly, this is treated strictly as a hypothesis. The manuscript explicitly notes that: SSMs are not permutation-invariant by construction Sequential ordering may introduce inductive biases misaligned with fermionic symmetry Empirical validation on multi-electron systems is required before drawing conclusions Results and Current Validation Status At present, the repository reports results for atomic hydrogen (Ne = 1): Exact energy: −0.500000 Ha Achieved energy: −0.5001 Ha Error: 0.12 mHa (well below chemical accuracy) Local-energy variance: 0.0454 Optimizer: Stochastic Reconfiguration Convergence: ~100 optimization steps Because hydrogen has no electron–electron correlation, this result is interpreted solely as a sanity check of: Hamiltonian implementation Local-energy estimator Natural-gradient optimization loop Numerical stability mechanisms The manuscript clearly states that multi-electron benchmarks (e.g., He, Li, or H₂) are required to meaningfully evaluate the SSM-Backflow hypothesis and are planned for future versions. Stability and Diagnostics The implementation includes multiple practical safeguards designed to stabilize VMC optimization: Dynamic energy-based divergence detection Variance-based walker reset mechanisms Controlled respawning of walkers near nuclear centers In addition, an exploratory commutator-based diagnostic (informally termed the “Noether Inverse”) is included as a monitoring signal during training. This quantity is explicitly not claimed to represent a new conserved symmetry and is treated purely as an empirical observation requiring further theoretical and numerical investigation. Intended Use and Limitations This repository is released as: A research implementation A transparent record of an ongoing investigation A foundation for future validation and benchmarking It is not presented as a production-ready electronic-structure solver, nor as evidence of state-of-the-art performance. All limitations are explicitly documented in the manuscript. Licensing and Reuse The work is shared to encourage: Reproducibility Methodological discussion Extension and validation by the community If you use or build upon this work, please cite the Zenodo DOI associated with this release.
Quantum many-body problems are some of the most challenging problems in science and are central to demystifying some exotic quantum phenomena, e.g., high-temperature superconductors. The combination of neural networks (NN) for representing quantum states, coupled with the Variational Monte Carlo (VMC) algorithm, has been shown to be a promising method for solving such problems. However, the run-time of this approach scales quadratically with the number of simulated particles, constraining the practically usable NN to - in machine learning terms - minuscule sizes (<10M parameters). Considering the many breakthroughs brought by extreme NN in the +1B parameters scale to other domains, lifting this constraint could significantly expand the set of quantum systems we can accurately simulate on classical computers, both in size and complexity. We propose a NN architecture called Vector-Quantized Neural Quantum States (VQ-NQS) that utilizes vector-quantization techniques to leverage redundancies in the local-energy calculations of the VMC algorithm - the source of the quadratic scaling. In our preliminary experiments, we demonstrate VQ-NQS ability to reproduce the ground state of the 2D Heisenberg model across various system sizes, while reporting a significant reduction of about ×10 in the number of FLOPs in the local-energy calculation.