Variational Monte Carlo for Ab Initio Electronic Structure and Quantum Chemistry with Neural Quantum States, Natural-Gradient Optimization, and Selective State-Space Models

Title
Variational 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(Nelog⁡Ne)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.