Foundation Neural-Network Quantum States for Molecular Potential Energy Surfaces in Second Quantization

Lizhong Fu; Jianan Wei; Wenguan Wang; Honghui Shang

physics.chem-phSep 29, 2026

Foundation Neural-Network Quantum States for Molecular Potential Energy Surfaces in Second Quantization

Authors: Lizhong Fu, Jianan Wei, Wenguan Wang, Honghui Shang

Organizations: State Key Laboratory of Precision and Intelligent Chemistry, University of Science and Technology of China

Abstract

Second-quantized neural-network quantum states have achieved accurate molecular energies, but extending them across molecular geometries requires a shared representation of the geometry-dependent wavefunction coefficients. We introduce geometry-conditioned foundation neural-network quantum states for molecular electronic structure in second quantization. A single autoregressive model learns a family of ground states from sparse anchor geometries and provides wavefunctions at untrained geometries without further optimization. Orbital alignment matches orbital identities and transports their phases, establishing an aligned orbital basis across geometries. Frozen energies reach chemical accuracy at every untrained query geometry for N $_2$ , CO, and H $_4$ . On additional molecular paths, the energy-trained wavefunctions yield dipoles, quadrupoles, and natural occupations without property labels. Across three paired N $_2$ training seeds, orbital alignment lowers the mean absolute energy error over all untrained query geometries from 34-37 mHa to 0.049-0.085 mHa. At approximately 1 mHa mean absolute error, frozen evaluation reduces the per-geometry cost by $986\times$ relative to independent optimization, yielding an estimated $25.8\times$ end-to-end GPU-cost reduction on a 161-point N $_2$ grid.

Figures & tables

Figure 1: Geometry-conditioned FNQSs. (a) Starting from sparse anchor geometries, restricted Hartree–Fock (RHF) orbitals are matched and phase-aligned to construct consistent occupation bases and molecular Hamiltonians. Joint variational Monte Carlo (VMC) training produces a shared wavefunction model. At new geometries, the frozen model yields energies and the one-particle reduced density matrix (1-RDM). (b) Occupation prefixes pass through causal Transformer blocks, with geometry and orbital descriptors supplied through feature-wise linear modulation (FiLM), to produce normalized occupation probabilities and hence the wavefunction amplitude. A separate multilayer perceptron (MLP) maps the full occupation string and geometry to a phase. Amplitude and phase combine into the wavefunction.

Figure 2: Potential energy surfaces. (a) N 2 , (b) CO, and (c) H 4 . Each panel combines NQS, RHF, and coupled-cluster singles and doubles (CCSD) relative energies (upper) and NQS signed errors against FCI (lower), with the 1.6 mHa band below. Energies are relative to each method’s path origin. CCSD uses a restricted reference. Squares mark anchors, dots queries, and bars one sampling standard error.

Figure 3: Orbital alignment. Alignment enables transfer between N 2 training anchors. The upper plot tracks aligned RHF orbital energies; the degenerate pair 5/6 is drawn as one curve. Arrows label changes in the canonical indices of the indicated continuous branches. Both plots share the N–N distance axis; vertical bands mark the two orbital-ordering exchanges (0.92–0.94 and 1.18–1.20 Å). The lower plot compares absolute energy errors over three paired seeds: lines show means and colored ribbons the minimum–maximum range. Squares mark anchors, dots queries, and the dotted line 1.6 mHa. Architecture, optimization, and sampling budgets are matched within each pair.

Figure 4: Electronic properties. (a) NH 3 dipole, (b) CO 2 quadrupole, and (c) ethylene natural occupations 6/7. Each panel combines the observable (upper) and its NQS-minus-FCI error (lower). Upper curves show NQS estimates connected by solid lines, RHF (dashed), and CCSD (dash-dotted). Circles mark queries, squares anchors, and bars one sampling standard error.

Figure 5: Computational cost. The N 2 workflows use fixed training budgets with MAEs near 1 mHa. (a) Average cost per additional geometry, including optimization where applicable and energy evaluation. The warm-start range spans the measured means at 41- and 161-point spacings. (b) Estimated total costs, including initial training, on logarithmic axes. Dots and vertical guides mark frozen evaluation on the 41- and 161-point grids. Independent and shared-adaptation costs are estimated from eight geometries selected by stratified sampling; warm-start dashed and solid curves use the 41- and 161-point spacings, respectively. Each curve holds its measured average per-geometry cost fixed as the query count varies.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Study	Steps	Geometries/update	Distinct configurations
N 2	20,000	5	6,000–50,000
CO	20,000	5	600–2,400
H 4 /cc-pVDZ	20,000	4	2,000–4,000
NH 3	20,000	5	6,000–20,000
Ethylene	20,000	7	6,000–20,000
CO 2	20,000	5	6,000–20,000

Appendix

Table 1: Training budgets and per-geometry distinct-configuration targets.

System	Basis	$e^{-}$ /orbitals	Coordinate	Anchors
N 2	STO-3G	14/10	bond length	5
CO	STO-3G	14/10	bond length	5
H 4 rectangle	cc-pVDZ	4/20	rectangle angle	4
NH 3	6-31G, frozen core	8/14	umbrella height	5
CO 2	STO-3G, frozen core	16/12	C–O distance	5
C 2 H 4	STO-3G, frozen core	12/12	torsion angle	7

Appendix

Table 2: Molecular systems. Electron/orbital counts refer to the retained space.

System	Molecular path	Anchors / queries	Anchor MAE	Query MAE	Query max.
N 2	bond stretch	5 / 36	0.036	0.049	0.095
CO	bond stretch	5 / 36	0.197	0.253	0.585
H 4	rectangle	4 / 21	0.088	0.076	0.120

Appendix

Table 3: Frozen molecular predictions, in mHa relative to FCI. Anchor and query metrics use the same final checkpoint.

System	Molecular path	Anchors / queries	Anchor MAE	Query MAE	Query max.
C 2 H 4	torsion	7 / 28	0.919	0.946	1.233
NH 3	inversion	5 / 20	0.872	1.036	1.456
CO 2	symmetric stretch	5 / 20	0.591	0.535	0.895

Appendix

Table 4: Energy errors for the three molecular-property benchmarks in Fig. 4 , in mHa relative to FCI. Anchor and query metrics use the same final checkpoints as the property comparisons; queries include all non-anchor grid points.

Figure 6: Frozen energy predictions. (a) NH 3 inversion, (b) CO 2 symmetric stretching, and (c) ethylene torsion. Each panel combines frozen NQS energies relative to each path origin (upper) and the corresponding signed complete-Hamiltonian errors against FCI (lower), with the

\pm 1.6

mHa band. Hollow squares mark anchors; dots mark queries. Error bars show one sampling standard error. The corresponding property panels appear in Fig. 4 .

Observable	RHF	CCSD	NQS
NH 3 dipole (D)	0.06057	0.00575	0.0058
CO 2 quadrupole ( $ea_{0}^{2}$ )	1.07714	0.06971	0.0403
C 2 H 4 occupations	0.27589	0.00485	0.0009

Appendix

Table 5: Query property MAEs against FCI. All methods use identical geometries, retained spaces, and observable conventions.

Figure 7: Conditioning ablations. All variants use the aligned CO orbital basis. Open circles show query MAEs for three seeds and filled diamonds their means. The geometry condition

\mathbf{c}

and orbital descriptors

\mathbf{d}

enter either the input embeddings or post-block FiLM adapters. The backbone, phase network, training budget, and evaluation geometries are matched.

Model	Seed 11	Seed 23	Seed 37
N 2 unaligned	36.806	33.609	35.027
N 2 aligned	0.049	0.085	0.053
CO $\mathbf{c}$ input	0.292	0.414	0.575
CO $\mathbf{c}$ FiLM	0.328	0.333	0.247
CO $\mathbf{c}+\mathbf{d}$ input	0.364	0.438	0.443
CO $\mathbf{c}+\mathbf{d}$ FiLM	0.277	0.194	0.382

Appendix

Table 6: Per-seed component ablations (MAE in mHa). N 2 uses all 36 query geometries; CO uses the eight query geometries. All runs use 20,000 updates with post-block FiLM where applicable; N 2 uses a fixed learning rate of

10^{-3}

and CO uses cosine decay. Each cell reports the error of a fixed final checkpoint.

Workflow	Updates	MAE (mHa)	41 points (GPU-h)	161 points (GPU-h)
Independent	4,000	1.027	10.785	42.351
Shared + fine-tune	2,000 + 200/point	1.064	1.682	4.162
Sequential warm-start	3,000 + 100/next point	1.173 / 0.988	0.664	1.948
Joint + frozen	4,000	0.793	1.612	1.644

Appendix

Table 7: Fixed-budget N 2 comparison. MAEs use the common eight geometries. Costs include training and final energy evaluation; independent and shared costs are estimated from the sampled geometries. The two warm-start MAEs correspond to the 41- and 161-point grids, respectively.

$R$ (Å)	Independent	Shared + fine-tune	Warm (41)	Warm (161)	Joint + frozen
0.60	0.284	0.458	0.365	0.365	0.327
0.76	0.663	0.473	0.559	0.458	0.527
0.84	1.054	0.601	0.705	0.566	0.573
0.96	1.097	0.728	0.898	0.787	0.723
1.00	0.709	0.790	0.998	0.872	0.788
1.18	1.868	1.477	1.666	1.392	1.088

Appendix

Table 8: Absolute full-Hamiltonian energy errors (mHa) at the eight common N 2 comparison geometries. Warm-start columns correspond to the two complete sequential paths.

Neural Quantum States

Neural quantum states use neural networks to represent many-body wavefunctions and enable variational simulations beyond exact diagonalization.

See all 25 papers

One More Time: Revisiting Neural Quantum States from a Reinforcement Learning Perspective
Jul 2, 2026Juan Agustín Duque, Sergio García Heredia, Vinicius Hernandes +4Neural Quantum States Second-Order Optimization
Mila Quebec AI Institute, Université de Montréal, Montréal, Canada · Applied Quantum Algorithms ⟨aQaL⟩, LIACS & LION, Leiden University, Leiden, Netherlands · QuTech and Kavli Institute of Nanoscience, Delft University of Technology, Delft, Netherlands +1
Excited Pfaffians: Generalized Neural Wave Functions Across Structure and State
Mar 15, 2026Nicholas Gao, Till Grutschus, Frank Noé +1Variational Monte Carlo Neural Quantum States
Department of Computer Science & Munich Data Science Institute, Technical University of Munich · CuspAI · Free University of Berlin +2
NucleoNet: A Nuclear-Aware Determinant Ansatz for Multi-Geometry Neural Variational Monte Carlo
Oct 9, 2026Qiao Gao, Yang Huang, Enze Hou +3Variational Monte Carlo Computational Quantum Chemistry
School of Artificial Intelligence and Data Science, University of Science and Technology of China, Hefei 230026, China · University of Science and Technology of China, Hefei 230026, China · Suzhou Institute for Advanced Research, University of Science and Technology of China, Suzhou 215213, China +5

Jul 2, 2026J/K move · Enter open · S save

One More Time: Revisiting Neural Quantum States from a Reinforcement Learning Perspective

Juan Agustín Duque, Sergio García Heredia, Vinicius Hernandes, Eliška Greplová+3

Mila Quebec AI Institute, Université de Montréal, Montréal, Canada · Applied Quantum Algorithms ⟨aQaL⟩, LIACS & LION, Leiden University, Leiden, Netherlands · QuTech and Kavli Institute of Nanoscience, Delft University of Technology, Delft, Netherlands+1

Neural Quantum States Second-Order Optimization Policy Gradient Methods Trust-Region Optimization Quantum Many-Body Physics

Open paper AskPDF arXiv

Neural quantum states (NQS) provide a flexible and scalable framework for approximating quantum many-body wavefunctions. Among NQS parameterizations, autoregressive models are especially attractive because they enable exact, independent sampling from the Born distribution, avoiding the autocorrelation and mixing issues of Markov chain methods. Yet their optimization remains comparatively underexplored: Adam is a scalable method but ignores function space geometry, while stochastic reconfiguration is principled but costly and numerically fragile in large models. To address this gap, we show that variational energy minimization can be viewed as an advantage policy-gradient problem over the Born distribution, motivating trust-region optimization for NQS training. We introduce Proximal Wavefunction Optimization (PWO), a principled trust-region algorithm that clips probability-ratio changes in the amplitude channel and phase increments in the phase channel. PWO avoids explicit matrix inversion, reuses samples across multiple updates, and combines the scalability of first-order optimization with theoretical guarantees. Across Ising and frustrated

J_1

J_2

one- and two-dimensional spin systems, PWO improves stability and wall-clock convergence over Adam, minSR, and SPRING. Finally, we fine-tune a

1.5

B-parameter RWKV-7 model, demonstrating NQS optimization at a scale over three orders of magnitude beyond prior work.

Abstract

Neural Quantum States

Explore similar work

One More Time: Revisiting Neural Quantum States from a Reinforcement Learning Perspective

Excited Pfaffians: Generalized Neural Wave Functions Across Structure and State

NucleoNet: A Nuclear-Aware Determinant Ansatz for Multi-Geometry Neural Variational Monte Carlo