NucleoNet: A Nuclear-Aware Determinant Ansatz for Multi-Geometry Neural Variational Monte Carlo
Authors: Qiao Gao, Yang Huang, Enze Hou, Yuzhi Liu, Jingrun Chen, Han Wang
Organizations: 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 · Science, Mathematics and Technology Cluster, Singapore University of Technology and Design, 8 Somapah Road, 487372 Singapore · AI for Science Institute, Beijing 100080, China · School of Mathematical Sciences, University of Science and Technology of China, Hefei 230026, China · National Laboratory of Computational Physics, Institute of Applied Physics and Computational Mathematics, Beijing 100088, China · HEDPS, CAPT, College of Engineering, Peking University, Beijing, China
Neural-network variational Monte Carlo provides an accurate route to solving the electronic Schrödinger equation, but many neural wavefunctions are optimized for a single nuclear geometry, requiring separate calculations to map out a potential energy surface. We introduce NucleoNet, a nuclear-aware determinant ansatz in which electrons and nuclei form a unified particle set propagated through a shared electron--nucleus backbone that jointly evolves one-particle and two-particle representations. Nuclei carry learned representations that evolve alongside the electronic representations and condition orbital construction, while antisymmetry is imposed only by electronic determinants. On equilibrium small-molecule benchmarks, NucleoNet achieves accuracy comparable to established neural-wavefunction baselines with substantially fewer parameters. For each molecular system, a single model optimized jointly across geometries describes the N2 dissociation curve, Li displacement in LiCN, and the ethane and formamide minimum-energy paths for bond dissociation. Sparse-geometry training generalizes to unseen geometries for N2 and LiCN without further optimization, and N2 ablations show that nuclear streams substantially improve this interpolation. These results show that a compact wavefunction with jointly evolved nuclear representations can describe the geometries of a molecular system with a single set of parameters.
Figures & tables
Figure 1: Overview of NucleoNet. (a) Wavefunction construction. Electrons and nuclei are represented in a unified particle set and processed by a shared electron–nucleus backbone. The final electronic representations produce the Jastrow factor and orbital matrices, while the final nuclear representations generate envelope parameters. These quantities are combined to form the Slater matrices, determinant expansion, and final wavefunction. (b) Initialization layer. Fixed geometric features and particle descriptors are mapped to the initial one-particle and two-particle representations. (c) Shared electron–nucleus backbone. The initial representations are propagated through L electron–nucleus layers. (d) Electron–nucleus layer. Pair-gated messages and permutation-invariant group summaries enter the one-particle update, while the two-particle representations are updated through a shared nonlinear map. (e) Mean pooling. Representations are averaged separately over nuclei, spin-up electrons, and spin-down electrons and concatenated to form the group summaries. For simplicity, the figure uses ∥ for concatenation, ⋅ for elementwise multiplication, and × for scalar multiplication.
Figure 2: Energy error versus number of trainable parameters for equilibrium small-molecule benchmarks. The lower-left region corresponds to lower errors and fewer parameters. Results for the same molecule use the same marker style and are connected by gray lines as a visual guide. NucleoNet lies in the low-parameter, low-error region relative to FermiNet, Psiformer, and LapNet baselines. Benchmark geometries and reference energies follow the protocols of Refs. [ 23 , 29 , 20 ] .
Figure 3: Energy errors along the N 2 dissociation curve relative to the MLR 4(6,8) reference. A single NucleoNet model is jointly optimized on 32 bond lengths uniformly spaced between 1.6 and 6.4a0 . MAEs reported in the legend are computed over the full bond-length range.
Figure 4: Signed energy differences across the five LiCN geometries as a function of the Li–C distance, relative to independently optimized single-geometry NucleoNet energies. The Orbformer results correspond to the model fine-tuned from the LAC checkpoint; the Orbformer and CCSD(T)/CBS energies are taken from Ref. [ 3 ] . The black zero line denotes the single-geometry NucleoNet baseline, and the MSD values in the legend are averaged over the five geometries.
Figure 5: Multi-geometry dissociation benchmarks for ethane and formamide. Panels (a) and (b) show representative molecular geometries along the corresponding minimum-energy paths (MEPs). Panels (c) and (e) show mean-centered relative energies, while panels (d) and (f) show signed energy differences relative to the reference calculations reported in Ref. [ 3 ] . The MSD values are averaged over all geometries along each path. NucleoNet uses a single set of shared parameters for all geometries of each molecular system.
Figure 6: Generalization from sparse geometry sets. Blue triangles mark the training geometries, blue curves show the model optimized on the training geometries and evaluated at all geometries, and red curves show full-geometry joint optimization. The upper-right values give the corresponding benchmark metric (MAE for N 2 and MSD for LiCN), using the same references and units as in the preceding sections. For N 2 , six uniformly spaced bond lengths are used for optimization; for LiCN, only the two endpoint geometries are used.
Figure 7: Energy errors along the N 2 dissociation curve for NucleoNet and its ablation variants. Triangles mark the six training geometries, and the legend reports the MAE on the evaluation grid. The inset shows the full error range, including the large error of NucleoNet-FE-EB near the short-bond region.
Model
Envelope
Propagated streams
NH 3 ΔE
N 2 ΔE
N 2 training/evaluation ΔMAE
NucleoNet
Conditioned
Electrons+nuclei
0.00
0.00
0.00 / 0.00
NucleoNet-FE
Fixed
Electrons+nuclei
+ 0.01
+ 0.11
+ 0.21 / + 0.23
NucleoNet-FE-EB
Fixed
Electrons only
+ 0.05
+ 0.22
+ 0.92 / + 38.86
Table 1: Diagnostic ablation results on NH 3 single-geometry, N 2 single-geometry, and N 2 multi-geometry benchmarks. Positive values indicate worse performance. The N 2 training MAE is computed on the six training geometries, and the evaluation MAE on the grid of Sec. III.6 . All relative errors are reported in m Eh .
Appendix figures & tables14 assets
Supplementary material from the paper’s appendix.
Appendix
Category
Hyperparameter
Value
SPRING
Learning-rate schedule
(1+10−4t)−1
Initial learning rate
5×10−2
Initial damping λ0
0.001
History coefficient μ
0.99
Norm constraint
0.001
Pretraining
Optimizer
Adam
Appendix
Table S1: Default optimizer, pretraining, and VMC hyperparameters used in the reported calculations.
Category
Hyperparameter
Value
Small molecules
One-particle width h1
64
Two-particle width h2
16
Multi-geometry systems
One-particle width h1
128
Two-particle width h2
32
Shared architecture
Number of electron–nucleus layers
4
Number of determinants K
16
Appendix
Table S2: Architecture hyperparameters used in the reported calculations.
System
Psiformer(small)
Psiformer(large)
FermiNet(small)
FermiNet(large)
LapNet
NucleoNet
LiH
Energy
-8.070528(5)
-8.070536(4)
-8.07050(1)
-8.070515(8)
-8.070523(4)
-8.070544(2)
# parameters
1.59M
6.37M
0.68M
2.61M
1.88M
0.11M
Li 2
Energy
-14.99486(1)
-14.99485(2)
-14.99480(2)
-14.99484(2)
-14.99485(1)
-14.994913(4)
# parameters
1.60M
6.40M
0.70M
2.68M
1.89M
0.11M
NH 3
Energy
-56.56367(2)
-56.56381(2)
-56.56347(4)
-56.56387(3)
-56.56359(2)
-56.56389(2)
# parameters
1.62M
6.47M
0.74M
2.82M
1.93M
0.12M
Appendix
Table S3: Total energies and parameter counts for eight small molecular systems. Energies are reported in Hartree. FermiNet and Psiformer results are taken from [ 29 ] , and LapNet results are taken from [ 20 ] . Energies are reported in Hartree, and values in parentheses denote the statistical uncertainty in the last reported digit. Bold values indicate the lowest reported energy for each system.
System
Number of geometries
Effective batch size
N 2 dissociation
32
128×32
LiCN displacement
5
820×5
Ethane dissociation
20
128×20
Formamide dissociation
20
128×20
Appendix
Table S4: Batch sizes used in the multi-geometry experiments.
Bond length L ( a0 )
Energy ( Eh )
Bond length L ( a0 )
Energy ( Eh )
1.6000
−109.22020(8)
4.0774
−109.19612(6)
1.7548
−109.42544(7)
4.2323
−109.19126(6)
1.9097
−109.51676(7)
4.3871
−109.18787(6)
2.0645
−109.54171(7)
4.5419
−109.18552(6)
2.2194
−109.52889(6)
4.6968
−109.18367(6)
2.3742
−109.49621(6)
4.8516
−109.18246(5)
Appendix
Table S5: N 2 energies on 32 bond lengths uniformly spaced between 1.6 and 6.4a0 . Energies are reported in Hartree, and values in parentheses denote the statistical uncertainty in the last reported digit.
Geometry ID
Li–C distance ( a0 )
Single-geometry NucleoNet ( Eh )
Multi-geometry NucleoNet ( Eh )
1
3.590479
−100.38046(3)
−100.38045(3)
2
3.779452
−100.37951(3)
−100.37945(3)
3
3.968425
−100.37624(3)
−100.37625(3)
4
4.346370
−100.36622(3)
−100.36614(4)
5
4.724315
−100.35389(3)
−100.35385(4)
Appendix
Table S6: Li–C distances and NucleoNet energies for the five LiCN configurations. The C and N nuclei remain fixed across all configurations, while the Li position varies along the molecular axis. The single-geometry energies are obtained from independent optimization at each geometry and serve as the baseline used in the main text. The multi-geometry energies are obtained from a single NucleoNet model optimized jointly across all five configurations. Distances are reported in a0 , energies in Hartree, and values in parentheses denote the statistical uncertainty in the last reported digit.
Image ID
Ethane ( Eh )
Formamide ( Eh )
Image ID
Ethane ( Eh )
Formamide ( Eh )
0
−79.82448(8)
−169.8992(2)
10
−79.64326(7)
−169.8623(2)
1
−79.82254(7)
−169.8995(2)
11
−79.64816(7)
−169.8478(2)
2
−79.81594(7)
−169.8957(2)
12
−79.64788(7)
−169.7359(2)
3
−79.77018(8)
−169.8853(2)
13
−79.64827(7)
−169.6903(2)
4
−79.70948(7)
−169.8650(2)
14
−79.64937(7)
−169.6905(2)
5
−79.65531(7)
−169.7275(2)
15
−79.65006(7)
−169.6891(2)
Appendix
Table S7: NucleoNet energies for the ethane and formamide dissociation benchmarks. Image IDs identify the 20 geometries along each dissociation path. Energies are reported in Hartree, and values in parentheses denote the statistical uncertainty in the last reported digit.
System
Optimization geometries
Effective batch size
N 2
6
684×6
LiCN
2
2048×2
Appendix
Table S8: Batch sizes used in the sparse-geometry experiments.
Bond length L ( a0 )
Energy ( Eh )
1.6000
−109.22046(5)
2.4800
−109.46837(5)
3.3600
−109.25831(5)
4.2400
−109.19088(5)
5.1200
−109.18093(4)
6.0000
−109.17897(4)
Appendix
Table S9: N 2 optimization geometries and corresponding energies for the sparse-geometry experiment. Energies are reported in Hartree, and values in parentheses denote the statistical uncertainty in the last reported digit.
Bond length L ( a0 )
Energy ( Eh )
Bond length L ( a0 )
Energy ( Eh )
1.6000
−109.22032(7)
4.0774
−109.19600(6)
1.7548
−109.42535(7)
4.2323
−109.19123(7)
1.9097
−109.51670(7)
4.3871
−109.18768(6)
2.0645
−109.54159(8)
4.5419
−109.18527(6)
2.2194
−109.52875(7)
4.6968
−109.18351(7)
2.3742
−109.49613(7)
4.8516
−109.18239(6)
Appendix
Table S10: N 2 energies on a 32-point bond-length grid uniformly spaced between 1.6 and 6.4a0 for the sparse-geometry experiment. Energies are reported in Hartree, and values in parentheses denote the statistical uncertainty in the last reported digit.
Geometry ID
Role
Sparse-geometry energy ( Eh )
1
Optimization
−100.38049(3)
2
Evaluation
−100.37944(4)
3
Evaluation
−100.37626(4)
4
Evaluation
−100.36622(3)
5
Optimization
−100.35382(3)
Appendix
Table S11: LiCN energies evaluated with the model optimized on the two endpoint geometries. The model is optimized only on the two endpoint geometries and evaluated on all five configurations without further wavefunction optimization. Energies are reported in Hartree, and values in parentheses denote the statistical uncertainty in the last reported digit.
System
NucleoNet
NucleoNet-FE
NucleoNet-FE-EB
NH 3
-56.56389(2)
-56.56388(2)
-56.56384(2)
N 2
-109.54195(3)
-109.54184(3)
-109.54173(4)
Appendix
Table S12: Absolute variational energies for the NH 3 and N 2 single-geometry ablation benchmarks. Energies are reported in Hartree, and values in parentheses denote the statistical uncertainty in the last reported digit.
Bond length L ( a0 )
NucleoNet
NucleoNet-FE
NucleoNet-FE-EB
1.6000
−109.22046(5)
−109.22033(7)
−109.21963(9)
2.4800
−109.46837(5)
−109.46817(7)
−109.46775(9)
3.3600
−109.25831(5)
−109.25818(8)
−109.25729(9)
4.2400
−109.19088(5)
−109.19091(7)
−109.19006(9)
5.1200
−109.18093(4)
−109.18032(6)
−109.17995(8)
6.0000
−109.17897(4)
−109.17877(5)
−109.1777(1)
Appendix
Table S13: Absolute N 2 energies at the six geometries used to optimize the sparse-geometry ablation models. Energies are reported in Hartree, and values in parentheses denote the statistical uncertainty in the last reported digit.
Bond length L ( a0 )
NucleoNet
NucleoNet-FE
NucleoNet-FE-EB
1.6000
−109.22032(7)
−109.22020(7)
−109.2197(1)
1.7548
−109.42535(7)
−109.42536(7)
−109.4011(5)
1.9097
−109.51670(7)
−109.51666(7)
−108.997(3)
2.0645
−109.54159(8)
−109.54156(7)
−108.896(4)
2.2194
−109.52875(7)
−109.52881(7)
−109.5015(8)
2.3742
−109.49613(7)
−109.49617(7)
−109.4950(2)
Appendix
Table S14: Absolute N 2 energies on the evaluation grid for the sparse-geometry ablation models. Energies are reported in Hartree, and values in parentheses denote the statistical uncertainty in the last reported digit.
Neural-network wave functions in Variational Monte Carlo (VMC) have achieved great success in accurately representing both ground and excited states. However, achieving sufficient numerical accuracy in state overlaps requires increasing the number of Monte Carlo samples, and consequently the computational cost, with the number of states. We present a nearly constant sample-size approach, Multi-State Importance Sampling (MSIS), that leverages samples from all states to estimate pairwise overlap. To efficiently evaluate all states for all samples, we introduce Excited Pfaffians. Inspired by Hartree-Fock, this architecture represents many states within a single neural network. Excited Pfaffians also serve as generalized wave functions, allowing a single model to represent multi-state potential energy surfaces. On the carbon dimer, we match the O(Ns4)-scaling natural excited states while training >200× faster and modeling 50% more states. Our favorable scaling enables us to be the first to use neural networks to find all distinct energy levels of the beryllium atom. Finally, we demonstrate that a single wave function can represent excited states across various molecules.
Nicholas Gao, Till Grutschus, Frank Noé +1
Department of Computer Science & Munich Data Science Institute, Technical University of Munich · CuspAI · Free University of Berlin +2
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 N2, CO, and H4. On additional molecular paths, the energy-trained wavefunctions yield dipoles, quadrupoles, and natural occupations without property labels. Across three paired N2 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× relative to independent optimization, yielding an estimated 25.8× end-to-end GPU-cost reduction on a 161-point N2 grid.
Lizhong Fu, Jianan Wei, Wenguan Wang +1
State Key Laboratory of Precision and Intelligent Chemistry, University of Science and Technology of China
We introduce Large Electron Model, a single neural network model that produces variational wavefunctions of interacting electrons over the entire Hamiltonian parameter manifold. Our model employs the Fermi Sets architecture, a universal representation of many-body fermionic wavefunctions, which is further conditioned on Hamiltonian parameter and particle number. For interacting electrons in a two-dimensional harmonic potential, a single trained model accurately predicts the ground state wavefunction while generalizing across unseen coupling strengths and particle-number sectors, producing both accurate real-space charge densities and ground state energies, even up to 50 particles. Our results establish a foundation model method for material discovery that is grounded in the variational principle, while accurately treating strong electron correlation beyond the capacity of density functional theory.
Timothy Zaklama, Max Geier, Liang Fu
Department of Physics, Massachusetts Institute of Technology, Cambridge, Massachusetts 02139, USA