Excited Pfaffians: Generalized Neural Wave Functions Across Structure and State
Authors: Nicholas Gao, Till Grutschus, Frank Noé, Stephan Günnemann
Organizations: Department of Computer Science & Munich Data Science Institute, Technical University of Munich · CuspAI · Free University of Berlin · Rice University · Microsoft Research AI4Science
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.
Figures & tables
Figure 1: Computational scaling with the number of states. While previous works scale superlinearly, our method exhibits near-constant scaling O(Ns0.14) .
Figure 2: Illustration of single-state importance sampling ( Szabó et al., 2024 ) and our Multi-State Importance Sampling (MSIS).
Figure 3: Overview of wave function architectures for excited states. (a) Our Excited Pfaffian architecture shares a single backbone network across all states, producing shared orbitals ΦPf and state-specific selector matrices As that combine via Pf(ΦAsΦT) . (b) The standard Slater determinant approach requires separate neural network passes and learned orbitals Φs for each state.
Figure 4: Second-row atoms’ (left) and molecules’ (right) excitation energies. We trained an Excited Pfaffian for each group and compare to individual models ( Pfau et al., 2024 ; Szabó et al., 2024 ) . References from NIST ( Sansonetti, 2003 ) or QUEST ( Véril et al., 2021 ) .
Figure 5: Potential energy surfaces for the lowest 12 states of C 2 . Results compared against SHCI. singlet, triplet.
Figure 6: 32 excited states of Be.
Figure 7: Potential energy surfaces of ethylene along torsion ( τ ) and pyramidalization ( ϕ ). Top: rotation about the C–C bond. Bottom: pyramidalization showing a state crossing ϕ≈99\text{,}\mathrm{\SIUnitSymbolDegree}$$ (inset).
Figure 8: Spatial overlap measured through Bhattacharyya coefficients ⟨∣Ψs∣∣∣Ψt∣⟩ for all molecules in fig. 4 .
Figure 9: Potential energy surface of C 2 near a state crossing.
Figure 10: Ablation study on MSIS for the 12-state C 2 PES and the 33-state Be. The x-axis shows the number of VMC steps and the y-axis shows the total energy (left) and overlap (right).
Appendix figures & tables20 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 11: Normalized effective sample size ESS⋅Ns/Nb of the MSIS overlap estimator across the second-row atoms (blue, Ns=10 ) and 12 molecules (orange, Ns=5 ). Left: per system, averaged over states with error bars indicating the standard deviation. Middle: as a function of the number of valence electrons. Right: as a function of the number of atoms. The dashed red line at 1 marks the effective sample size of single-state importance sampling.
Nmol
1
2
4
8
16
32
64
Ns
1
0.040
0.082
0.080
0.081
0.085
0.084
0.085
2
0.790
0.675
0.617
0.583
0.539
0.555
0.518
4
0.819
0.649
0.614
0.587
0.579
0.587
0.567
8
0.848
0.693
0.684
0.670
0.664
0.671
0.731
16
0.996
0.815
0.812
0.829
0.900
0.885
0.894
Appendix
Table 1: Step-time contribution of the iterative bridge sampling algorithm measured on an Nvidia H100 for state counts Ns and molecule counts Nmol . The total batch size is kept constant at 4096 . Times are given in ms .
Figure 12: Convergence of the total energy on 33-state Be for the standard Excited Pfaffian ( Nk=16 ), a single-determinant Excited Pfaffian ( Nk=1 ), and the naive multi-state architecture at Nk=1 (the largest that fits in memory). Left: per training step. Right: per wall-clock hour. Energies are reported relative to the converged energy of the standard Nk=16 Excited Pfaffian.
Figure 13: Comparison of C 2 potential energy surfaces across methods. Each row shows singlet (left) and triplet (right) states for Excited Pfaffian (top), the method of Schätzle et al. (2025) (middle), and NES ( Pfau et al., 2024 ) (bottom). For Excited Pfaffian and the method of Schätzle et al. (2025) , we give the raw energies before correcting state swaps and averaging degeneracies. Raw energies for NES were not available. SHCI reference curves are shown in red/orange; method results are in blue. All three methods capture the qualitative PES structure, but the method of Schätzle et al. (2025) shows increased deviation from the SHCI baseline ( Holmes et al., 2017 ) at stretched geometries.
Figure 14: Potential energy surfaces of ethylene along torsion ( τ ) and pyramidalization ( ϕ ). The Excited Pfaffian PES for both slices is computed in a single neural wave function. Left: rotation about the C–C bond. Right: pyramidalization showing a state crossing ϕ≈99\text{,}\mathrm{\SIUnitSymbolDegree}$$ (inset).
Figure 15: Spatial overlap of atomic excited states measured as Bhattacharyya coefficients ⟨∣Ψs∣∣∣Ψt∣⟩ .
Figure 16: Spatial overlap of 33 states of beryllium measured as Bhattacharyya coefficients ⟨∣Ψs∣∣∣Ψt∣⟩ .
Figure 17: Per-step training time as a function of the number of states Ns on Be (left; 1, 10, 33 states) and C 2 (right; 1, 8, 12 states). Top: time normalized to the single-state-per-structure cost. Bottom: absolute time per step. Excited Pfaffian (blue) and NES ( Pfau et al., 2024 ) (orange) are compared with Adam (solid) and their native second-order preconditioner (dashed; SPRING for ours, KFAC for NES). For C 2 , filled markers show the per-structure cost. Hollow markers show the cumulative cost across the 10 dissociation geometries (one shared model for Excited Pfaffian, one model per geometry for NES). Excited Pfaffian maintains near-flat scaling, while NES grows roughly quadratically with Ns .
Method
Atoms
Molecules
C 2 PES
Ethylene PES
Pfau et al. (2024)
64B
48B
32B
27.6B
Szabo et al. (2024)
16B
12B
8B
9.2B
Schätzle et al. (2025)
-
-
3.2B
1.6B
Ours
800M
800M
800M
800M
Sample reduction range
20 – 80× fewer
15 – 60× fewer
4 – 40× fewer
2 – 30× fewer
Appendix
Table 2: Overview of total sample counts used in training in different experiments across works and improvement range for this work.
Method
Optimizer
Li
Be
B
C
N
O
F
Ne
∑
ΣourΣ
Pfau et al. (2024)
Adam
5.4
6.7
9.7
11.8
14.7
16.6
22.3
25.6
112.7
113×
KFAC
5.8
7.1
10.0
12.2
15.2
17.1
22.8
26.2
116.4
7×
Szabo et al. (2024)
Adam
3.3
3.8
4.3
5.8
7.0
7.6
9.3
10.1
51.2
51×
KFAC
11.4
12.4
13.2
13.7
15.4
16.4
18.2
18.8
119.5
7×
This Work
Adam
-
-
-
-
-
-
-
-
1.0
1×
SPRING
-
-
-
-
-
-
-
-
16.8
1×
Appendix
Table 3: Measured training step time in s for 10 states of atomic systems. Measurements for Adam and KFAC are taken for 1000 walkers and up-scaled to 4000 walkers. Measurements for SPRING are directly taken for 4000 walkers.
Method
Optimizer
BH
HCl
H 2 O
H 2 S
BF
CO
C 2 H 4
CH 2 O
CH 2 S
HNO
HCF
H 2 CSi
∑
ΣourΣ
Pfau et al.
Adam
3.5
4.9
7.3
4.8
12.3
12.3
14.8
14.7
12.4
14.6
14.6
9.4
125.4
32×
(2024)
KFAC
3.9
5.2
7.7
5.2
12.8
12.8
15.4
15.4
13.0
15.2
15.2
9.9
131.7
9×
Szabo et al.
Adam
3.4
15.2
5.6
15.2
9.2
9.4
10.6
10.4
24.8
10.4
11.2
21.2
146.6
37×
(2024)
KFAC
7.4
19.6
9.4
19.6
13.2
13.2
15.2
15.0
29.8
15.0
15.0
25.6
198.0
14×
This Work
Adam
-
-
-
-
-
-
-
-
-
-
-
-
3.9
1×
SPRING
-
-
-
-
-
-
-
-
-
-
-
-
14.0
1×
Appendix
Table 4: Measured training step time in s for 5 states of molecular systems. Measurements for Adam and KFAC are taken for 1000 walkers and up-scaled to 4000 walkers. Measurements for SPRING are directly taken for 4000 walkers.
C 2
Ethylene
Method
Optimizer
Step time
Total 10 structures
ourtotaltotal
Step time
Total 23 structures
ourtotaltotal
Pfau et al. (2024)
Adam
49.0
490.4
208×
3.9
89.2
24×
KFAC
49.9
499.2
46×
4.2
97.5
9×
Szabo et al. (2024)
Adam
14.0
140.0
59×
5.5
125.6
33×
KFAC
24.8
248.0
23×
7.4
170.2
16×
This Work
Adam
-
2.4
1×
-
3.8
1×
Appendix
Table 5: Measured training step time in s for the PES reported in the main body. Measurements for Adam and KFAC are taken for 1000 walkers and up-scaled to 4000 walkers. Measurements for SPRING are directly taken for 4000 walkers.
Method
S1
S2
S3
S4
S5
S6
S7
S8
S9
S10
S11
Ours
0.4294
2.0161
2.0438
3.0738
3.1192
3.2308
3.2352
4.3443
4.7957
5.3375
5.3634
NES (paper)
0.4293
2.0369
2.0369
3.1051
3.1051
3.2752
3.2752
—
—
—
—
NES (4.5k)
1.3817
3.5633
3.7173
4.4938
4.7929
5.1994
6.6696
12.7497
14.1597
14.6219
15.0104
Appendix
Table 6: Excitation energies (eV) for C 2 at a bond distance of 0.995A˚ . “NES (paper)” lists the values reported in Pfau et al. (2024) , who only report energies averaged over degenerate states (entries repeated within each degeneracy here). “NES (4.5k)” is NES retrained from scratch at the per-geometry compute budget equivalent to our 12-state, 10-geometry C 2 PES.
Category
Parameter
Value
Wave Function Architecture
Embedding network
Moon
Hidden dimension
256
Number of layers
4
Envelopes per nucleus
8
Orbitals per nucleus and determinant
8
Appendix
Table 7: Hyperparameters used in all experiments.
Experiment
A100 GPU-hours
Atomic excitations ( fig. 4 , left)
640
Molecules ( fig. 4 , right)
792
Carbon dimer PES ( fig. 5 )
600
Ethylene PES ( fig. 7 )
640
33 states of Be ( fig. 6 )
231
Appendix
Table 8: Total compute times. Nvidia H100 GPU-hours correspond to 2 Nvidia A100 GPU-hours.
Molecule
Energy (Ha)
ΔE (Ha)
BH
-25.28897(8)
-
-25.24029(8)
0.04868(12)
-25.24021(8)
0.04876(12)
-25.18341(9)
0.10555(12)
-25.18336(9)
0.10561(12)
HCl
-15.5959(4)
-
Appendix
Table 9: Energies and excitation energies of the molecules in fig. 4 (right) obtained from the jointly trained Excited Pfaffian in atomic units.
Atom
Energy (Ha)
ΔE (Ha)
Li
-7.47807(3)
-
-7.41013(4)
0.06794(5)
-7.41005(4)
0.06802(5)
-7.40994(4)
0.06813(5)
-7.35424(3)
0.12382(4)
-7.33542(9)
0.14264(9)
Appendix
Table 10: Energies and excitation energies of the atoms in fig. 4 (left) obtained from the jointly trained Excited Pfaffian in atomic units.
Bond length
X1Σg+
A1Πu+
A1Πu−
B′1Σg+
B1Δg(2)
B1Δg(1)
0.9952
-75.77445(12)
-75.66149(12)
-75.65982(12)
-75.59821(14)
-75.57830(12)
-75.57735(12)
1.1204
-75.89580(11)
-75.82222(11)
-75.82083(11)
-75.76439(11)
-75.76767(11)
-75.76991(11)
1.1823
-75.91752(10)
-75.85855(10)
-75.85629(10)
-75.80813(11)
-75.81718(10)
-75.82005(10)
1.2431
-75.92326(10)
-75.87808(10)
-75.87552(10)
-75.83539(10)
-75.84666(10)
-75.84959(10)
1.3064
-75.91801(10)
-75.88474(10)
-75.88113(10)
-75.84960(10)
-75.86225(10)
-75.86555(10)
1.3685
-75.90679(10)
-75.88355(10)
-75.88037(10)
-75.85471(10)
-75.86803(10)
-75.87119(10)
Appendix
Table 11: Absolute energies of the carbon dimer upon dissociation ( fig. 5 ) obtained from a jointly trained Excited Pfaffian. We report raw energies as tracked by our generalized wave function, before averaging over degenerate states and correcting for state swaps. Energies are in Hartree, bond lengths in Angstrom.
Table 12: Absolute energies of the two lowest singlet states of ethylene ( fig. 7 ) under torsion and pyramidalization. A single Excited Pfaffian is trained jointly across both PES. Energies are in Hartree, angles in degrees.
State
Energy (Ha)
ΔE (Ha)
ΔEexp (Ha)
ΔΔEexp (Ha)
0
-14.66719(3)
-
-
-
1
-14.56714(6)
0.10005(7)
0.100149
0.000100
2
-14.56707(3)
0.10012(4)
0.100149
0.000028
3
-14.56706(3)
0.10013(4)
0.100149
0.000020
4
-14.47333(3)
0.19386(4)
0.193942
0.000083
5
-14.47327(3)
0.19392(4)
0.193942
0.000020
Appendix
Table 13: Absolute and excitation energies of the 33 states of beryllium ( fig. 6 ). Reference energies from NIST ΔEexp ( Sansonetti, 2003 ) and errors ΔΔEexp are also given. Energies are given in atomic units.
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
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.
Qiao Gao, Yang Huang, Enze Hou +3
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
Neural-network quantum states have emerged as a powerful variational framework for quantum many-body systems, with recent progress often driven by massively parallel architectures such as transformers. Recurrent neural network quantum states, however, are frequently regarded as intrinsically sequential and therefore less scalable. Here we revisit this view by showing that modern recurrent architectures can support fast, accurate, and computationally accessible neural quantum state simulations. Using autoregressive recurrent wave functions together with recent advances in parallelizable recurrence, we develop variational ansätze, called parallel scan recurrent neural quantum states (PSR-NQS), which can be trained efficiently within variational Monte Carlo in one and two spatial dimensions. We demonstrate accurate benchmark results and show that, with iterative retraining, our approach reaches two-dimensional spin lattices as large as 52×52 while remaining in agreement with available quantum Monte Carlo data. Our results establish recurrent architectures as a practical and promising route toward scalable neural quantum state simulations with modest computational resources.
Ejaaz Merali, Mohamed Hibat-Allah, Mohammad Kohandel +2
Department of Physics and Astronomy, University of California, Davis, California 95616, USA · Department of Applied Mathematics, University of Waterloo, Waterloo, ON N2L 3G1, Canada · Vector Institute, Toronto, Ontario, M5G 0C6, Canada +1