PaperlayerView exact revision on arXiv

Source-linked AI summary

Autoregressive neural-network wavefunctions for ab initio quantum chemistry

Thomas D. Barrett, Aleksei Malyshev, A. I. Lvovsky

arXiv:2109.12606v2physics.chem-phcs.LGphysics.comp-phquant-ph

TL;DR

Full configuration interaction becomes intractable as molecular determinant spaces grow, while existing neural-network quantum states face sampling challenges. This paper uses an autoregressive wavefunction with molecular physical priors and efficient sampling, obtaining FCI-quality solutions on systems with up to 30 spin orbitals. It concludes that further scaling is limited by variational Monte Carlo computation rather than the wavefunction ansatz.

  • Problem

    Full configuration interaction rapidly becomes intractable as the number of molecular Slater determinants grows exponentially, and applying neural-network quantum states to electronic structure faces sampling challenges.

    Source: §I · pp. 1–2

  • Method

    An autoregressive neural network parameterizes molecular wavefunctions, embeds physical priors, and samples with cost scaling in unique configurations.

    Source: §I · p. 2, B. Autoregressive quantum states · p. 3, §IV · p. 8

  • Results

    The ARN approaches FCI-quality solutions on systems with up to 30 spin orbitals, including settings where canonical approximation methods fail.

    Source: §IV · pp. 7–8

  • Takeaways & Limitations

    For larger molecules, evaluating local energies makes VMC the dominant computational expense and the remaining scaling barrier beyond the neural-network ansatz.

    Source: §IV · p. 8

    The fundamental accuracy of a second-quantized approach is limited by the choice of basis set.

    Source: §IV · p. 8

Abstract

from arXiv · show

In recent years, neural network quantum states (NNQS) have emerged as powerful tools for the study of quantum many-body systems. Electronic structure calculations are one such canonical many-body problem that have attracted significant research efforts spanning multiple decades, whilst only recently being attempted with NNQS. However, the complex non-local interactions and high sample complexity are significant challenges that call for bespoke solutions. Here, we parameterise the electronic wavefunction with a novel autoregressive neural network (ARN) that permits highly efficient and scalable sampling, whilst also embedding physical priors reflecting the structure of molecular systems without sacrificing expressibility. This allows us to perform electronic structure calculations on molecules with up to 30 spin-orbitals -- at least an order of magnitude more Slater determinants than previous applications of conventional NNQS -- and we find that our ansatz can outperform the de-facto gold-standard coupled cluster methods even in the presence of strong quantum correlations. With a highly expressive neural network for which sampling is no longer a computational bottleneck, we conclude that the barriers to further scaling are not associated with the wavefunction ansatz itself, but rather are inherent to any variational Monte Carlo approach.

I. INTRODUCTION

Molecular electronic structure is difficult because the many-body problem and its determinant space scale beyond practical exact treatment. The paper introduces an autoregressive neural-network wavefunction with efficient sampling and molecular physical priors to address limitations of existing methods.

I. INTRODUCTION

  • The number of Slater determinants grows exponentially, making full configuration interaction rapidly intractable for larger systems.The paper frames this scaling challenge alongside the many-body electronic problem's NP-hardness [2].Source: §I · p. 1
  • Coupled cluster accesses arbitrarily excited determinants but cannot represent arbitrary superpositions and may fail under strong static correlations.These methods trade expressibility for reduced complexity while often performing well on typical molecular wavefunctions.Source: §I · p. 2
  • Neural networks, specifically RBMs, can parameterise many-body wavefunctions and capture non-trivial correlations within exponentially large encoded Hilbert spaces, but RBMs rely on inherently inefficient sampling such as MCMC and do not reflect molecular physical knowledge.Finding alternative neural-network architectures that overcome these shortcomings remains an important item on the agenda.Source: §I · p. 2
  • The proposed autoregressive ansatz embeds conserved quantities and invariances without sacrificing expressibility, with sampling cost scaling in unique configurations rather than batch size.It approaches FCI solutions on systems at least 25 times larger than earlier conventional NNQS applications.Source: §I · p. 2

II. NEURAL NETWORK WAVEFUNCTION

The wavefunction is optimized to the molecular ground state using variational Monte Carlo: sample determinants, estimate their local energies and energy gradient, then update the ansatz by backpropagation.

II. NEURAL NETWORK WAVEFUNCTION

  • The Jordan–Wigner encoding represents the molecular state as interacting qubits, with Hamiltonian terms expressed using Pauli operators.The paper reserves “spin” for the electron spin-state to distinguish it from qubits.Source: A. Variational Monte-Carlo optimization · p. 2
  • VMC optimizes the parameterized wavefunction toward the ground state by minimizing its energy expectation value [24].The ansatz enforces normalization, and the local energy is defined for each Slater determinant.Source: A. Variational Monte-Carlo optimization · pp. 2–3
  • The mean local energy over determinants sampled from |ψθ(xk)|^2 gives an unbiased energy estimate and can also estimate the parameter gradient.Standard backpropagation then updates the wavefunction parameters.Source: A. Variational Monte-Carlo optimization · p. 3

B. Autoregressive quantum states

The ARN factorizes the wavefunction distribution into sequential conditional distributions, while treating amplitudes autoregressively and predicting phases separately. Its inference and sampling procedures exploit parallel evaluation and aggregate repeated configurations.

B. Autoregressive quantum states

  • The ARN factorizes a joint distribution into sequential conditional distributions, modeling each variable using those that precede it.The network models the absolute wavefunction amplitude; a separate feed-forward network predicts the phase from the full configuration string.Source: B. Autoregressive quantum states · p. 3
  • The network groups each spatial orbital's two spin projections into a unit with four possible occupation configurations.Normalization of the conditional distributions yields a normalized overall wavefunction.Source: B. Autoregressive quantum states · p. 3

C. Inference and efficient sampling

  • Inference evaluates a configuration's log-amplitude in one forward pass, with all subnetworks evaluated in parallel.Each subnetwork receives the partial configuration and outputs normalized log-amplitudes for its orbital's four configurations.Source: C. Inference and efficient sampling · p. 3
  • Sampling generates configuration strings from |ψ|^2 and reports unique strings with their occurrence counts rather than retaining repeated samples.At each orbital, counts are sampled multinomially; only nonzero-count strings continue, preventing exponential dataset growth through the orbitals.Source: C. Inference and efficient sampling · p. 3

end

The batched sampler tracks unique partial configurations and their occurrence counts, avoiding redundant sample storage. Its computational cost scales with unique configurations, and independent subnetworks can process different inputs in parallel.

end

  • The sampling dataset retains unique partial strings, their occurrence counts, and log-amplitudes as orbitals are processed sequentially.Counts are sampled for each orbital configuration; strings with zero occurrences are not carried forward, preventing exponential growth of the dataset.Source: end · p. 4
  • Different inputs to a given subnetwork can be processed in parallel, and generating a sample with its wavefunction coefficient costs one full-network inference pass.Together with unique-sample scaling, these properties distinguish the procedure from standard MCMC sampling approaches.Source: end · p. 4
  • Sampling cost scales with the number of unique configurations rather than the total batch size.This is useful when the Hartree–Fock state dominates and batches contain orders of magnitude fewer unique states than samples.Source: end · p. 4

D. Encoding physical constraints

Only a small fraction of qubit configurations represent physically viable Slater determinants. The ARN incorporates known physical constraints in its architecture rather than having to learn them during training.

D. Encoding physical constraints

  • Only configurations with the molecule’s electron number and compatible physical symmetries represent viable Slater determinants.A naive NNQS may assign nonzero probability across all 2^NO basis configurations.Source: D. Encoding physical constraints · p. 4
  • Encoding the constraints a priori focuses the network on a physically meaningful output subset and can streamline training and reduce computational costs.Source: D. Encoding physical constraints · p. 5
  • The ARN combines a trainable MLP with hard-coded preprocessing and postprocessing to enforce physical priors without sacrificing expressibility.Unlike RBMs, this approach does not rely on modified sampling methods to generate only physically viable states.Source: D. Encoding physical constraints · p. 5

1. Conservation of electron number and multiplicity

The ARN uses known electron number and spin to restrict its configurations. A mask assigns zero probability to partial configurations that violate those constraints.

1. Conservation of electron number and multiplicity

  • Because electron number Ne and multiplicity 2S + 1 are known, the method restricts determinants to |MS| ≤ S and, by spin symmetry, chooses MS = S.The energy does not depend on MS for a given total spin S.Source: §1 · p. 5
  • Each subnetwork’s output is masked to give zero probability to partial configurations that fail the conservation requirements.Source: §1 · p. 5, Figure 2 · p. 5

2. Electron spin-flip symmetry

The paper reports that NAQS energies approach or match FCI results across the tested molecules, including cases where conventional methods do not. Its ARN also encodes spin-flip symmetry, subject to restrictions on the molecular setting.

2. Electron spin-flip symmetry

  • The ARN enforces equal absolute amplitudes for spin-flipped determinants by applying the symmetry to each conditional subnetwork, including input and output transformations.The transformation canonicalizes spin-symmetric inputs; when an input is flipped, a corresponding flip is applied to the output.Source: §2 · pp. 5–6
  • This spin-flip symmetry applies only for MS = 0 molecules, which under the paper’s MS = S choice means closed-shell molecules with S = 0.It also requires a basis whose spatial orbitals do not depend on electron spin.Source: §2 · p. 6

III. RESULTS

A. Obtained molecular ground states

  • NAQS energies approach or match FCI for all tested molecules with up to 20 electrons and 28 spin-orbitals, and outperform CCSD and CISD on Li2O with over 40 million physically valid determinants.Source: A. Obtained molecular ground states · p. 6
  • Compared with the prior RBM-based quantum-chemistry ansatz, ARN scales to systems with more than an order of magnitude more determinants without reported performance loss.The paper reports better results than Ref. on every molecule larger than H2.Source: A. Obtained molecular ground states · p. 6
  • For LiCl and Li2O, training sampled only about 28% and 5% of physically viable determinants with masking, yet the network learned important configurations for accurate wavefunctions.Without masking, the reported sampled fractions were about 20% and 1%, respectively.Source: A. Obtained molecular ground states · pp. 6–7
  • Across N2 bond separations, NAQS nearly matches FCI and outperforms the compared approximations, including where coupled-cluster methods fail under strong correlations.The figure shows the best and mean NAQS results across five seeds at each separation.Source: Figure 3 · p. 6, A. Obtained molecular ground states · p. 6, Table I · p. 7
  • Diagonalizing the subspace of 13,761 determinants sampled in Li2O yielded −87.8911 Ha, below the ARN’s −87.8909 Ha, suggesting a possible hybrid with canonical methods.Those determinants came from 1012 samples at the final step, without masking.Source: A. Obtained molecular ground states · p. 7

B. Ablation studies

The ARN reaches FCI-quality solutions on systems with up to 30 spin-orbitals, while the discussion identifies physical priors and efficient sampling as key design features. Ablations show that the priors improve average performance, especially on larger systems, while local-energy evaluation remains a scaling bottleneck.

B. Ablation studies

  • The ablations remove the constraint to physically viable determinants or the wavefunction's spin-flip invariance; open-shell CH2 and O2 achieved FCI-level performance without the latter.The effect of removing either prior is subtle overall and becomes visible mainly on larger systems, where the configuration space grows exponentially.Source: B. Ablation studies · p. 7
  • Across multiple seeds, encoding physical priors significantly improves average performance; the best ablation runs can reflect a fortuitous seed.Figure 4 compares the Standard model with No mask and No spin sym. across six of the largest molecules.Source: B. Ablation studies · p. 7, Figure 4 · p. 8

IV. DISCUSSION

  • The ARN approaches FCI-quality solutions on systems with up to 30 spin-orbitals, even where canonical approximation methods fail.The authors attribute this success to encoding physical priors and sampling that scales with unique configurations rather than total batch size.Source: §IV · pp. 7–8
  • Local-energy evaluation becomes the dominant computational expense for large molecules, leaving VMC computation—not the neural network—as the remaining scaling barrier.The other major optimization cost is sampling configurations; its cost depends on network forward passes and the number of unique samples.Source: §IV · p. 8
  • Second-quantized accuracy is limited by the chosen basis set; flexible-basis first-quantized alternatives report unconstrained accuracy at greater computational complexity.The paper notes that scaling to larger basis sets poses challenges equivalent to scaling to larger molecules.Source: §IV · p. 8

Network architecture

The architecture uses multilayer perceptrons for conditional amplitudes and phases, with spin symmetry built into the subnetworks. It orders orbitals by decreasing energy, while noting that finding an optimal ordering remains open.

Network architecture

  • The conditional-amplitude MLP has one 64-neuron hidden layer, while the phase network has two 512-neuron hidden layers; both use ReLU activations.The authors describe this architecture as substantially overparameterized and say Hamiltonian matrix-element calculations, not network passes, are the computational bottleneck.Source: Network architecture · p. 9
  • Orbitals are ordered by decreasing energy, with the highest-energy orbital first; optimizing this ordering is left for future research.The autoregressive decomposition permits different orders of conditional wavefunctions.Source: Network architecture · p. 9
  • Subnetworks receive spin-up and spin-down orbital inputs and are designed so spin-symmetric inputs produce correspondingly symmetric outputs.The procedure makes inputs invariant to global spin flips and reconstructs the desired symmetry at the output.Source: Network architecture · p. 9

Training details

Training uses VMC gradients with a variance-reducing baseline and Adam, while adapting batch size to maintain a target number of unique valid samples. The calculations use STO-3G molecular geometries from PubChem.

Training details

  • Subtracting a baseline reduces gradient variance without biasing the sample approximation, and improves training stability.The baseline is described as a standard machine-learning technique for improving gradient estimation.Source: Training details · p. 9
  • Adam training starts at 5 × 10^-3 and drops to 5 × 10^-4 halfway through; batch size adjusts to keep 10,000–100,000 unique valid samples where possible.The initial batch size is 10^6, adjusted by an order of magnitude as needed, with an upper limit of 10^12.Source: Training details · p. 9
  • Reported energies use the best result across five seeds trained for 10,000 steps, except Li2O, which uses one seed because of its greater computational cost.The authors note that random parameter initialization affects final performance.Source: Training details · p. 9

Molecular geometries and calculations

I. FURTHER DETAILS ON RESULTS FROM MAIN TEXT

NAQS outperforms the RBM on every compared molecule larger than H2. The supplementary results also describe multi-seed ablations and the computational cost of converting the Jordan–Wigner Hamiltonian into determinant couplings.

Full comparison of NAQS and RBMs

II. ANALYSIS OF WALL-CLOCK TIMINGS

  • Timing analyses were performed on GeForce RTX 2080 Ti GPUs and 32 shared Intel Xeon E5-2620 v4 processors.Source: §II · p. 11

Pre-calculating coupling terms.

  • Hamiltonian conversion scales with the number of Pauli strings and physically valid configurations, takes minutes for LiCl, and was not done in full for Li2O.Li2O couplings were instead computed on demand and cached during optimisation.Source: Pre-calculating coupling terms · p. 11, Table III · p. 12

Average performance of ablations

  • The ablation results report average ground-state energies across five seeds, with standard deviations and physically valid determinant counts provided for each molecule.The table notes rare failure seeds for the No mask configuration on O2 and missing spin-symmetry entries for CH2 and O2.Source: Average performance of ablations · p. 11, Table II · p. 12

Timing analysis

Timing depends strongly on the number of unique configurations, which changes during optimisation. Sampling, local-energy evaluation, and parameter updates contribute differently across molecules and training stages.

Sampling

  • Batch size is tuned to keep 10^4–10^5 unique samples as long as possible, with limits of 10^12 total and 10^6 unique samples; LiCl is omitted at the largest total batch because of GPU memory.The figure separates stage timings by total batch size and by unique-sample count.Source: Timing analysis · p. 12, Figure 1 · p. 13
  • Multinomial sampling takes (24.3 ± 0.8) ms for 10^12 samples across 10^5 unique states on a MacBook, while numerical sampling contributes 1.5%–2.6% of measured network-sampling time.The remaining contribution is network inference.Source: Timing analysis · pp. 12–13

Evaluating local energy and loss

  • Local-energy and loss evaluation scales with the number of unique samples and becomes rate-determining for larger molecules; on-demand Hamiltonian couplings add cost for Li2O.Source: Timing analysis · p. 13

Update the wavefunction

  • Updating NAQS parameters by backpropagation scales with sampling cost but is strictly cheaper.Source: Timing analysis · p. 13

Overall timings

  • On Li2O, an optimisation step fell from 34 s to 2 s as the batch grew from 100 k samples across 95 k unique configurations to 10^12 samples across 9 k unique configurations.The unique-sample count is highly non-stationary during training.Source: Timing analysis · p. 13
  • The reported implementation leaves speed improvements available, and its fixed training hyperparameters were not tuned for fastest performance.The authors cite sparse-to-dense conversions and compiled auto-grad implementation as possible improvement paths.Source: Timing analysis · p. 13

Related papers

Loading 2109.12606v2…