Excited Pfaffians: Generalized Neural Wave Functions Across Structure and State

Nicholas Gao; Till Grutschus; Frank Noé; Stephan Günnemann

cs.LGMar 15, 2026

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 ·

Abstract

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(N_s^4)$ -scaling natural excited states while training $>200\times$ 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

\mathcal{O}({N_{\mathrm{s}}}^{0.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 $\Phi_{\text{Pf}}$ and state-specific selector matrices $A_{s}$ that combine via ${\mathrm{Pf}}(\Phi A_{s}\Phi^{T})$ . (b) The standard Slater determinant approach requires separate neural network passes and learned orbitals $\Phi_{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 ( $\tau$ ) and pyramidalization ( $\phi$ ). Top: rotation about the C–C bond. Bottom: pyramidalization showing a state crossing $\phi\approx$99\text{\,}\mathrm{\SIUnitSymbolDegree}$$ (inset). — Figure 7: Potential energy surfaces of ethylene along torsion ( $\tau$ ) and pyramidalization ( $\phi$ ). Top: rotation about the C–C bond. Bottom: pyramidalization showing a state crossing $\phi\approx$ 99\text{,}\mathrm{\SIUnitSymbolDegree}$$ (inset).

Figure 8: Spatial overlap measured through Bhattacharyya coefficients $\langle|{\Psi}_{s}|||{\Psi}_{t}|\rangle$ 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

\mathrm{ESS}\cdot{N_{\mathrm{s}}}/{N_{\mathrm{b}}}

of the MSIS overlap estimator across the second-row atoms (blue,

{N_{\mathrm{s}}}=10

) and 12 molecules (orange,

{N_{\mathrm{s}}}=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.

$N_{\mathrm{mol}}$	1	2	4	8	16	32	64
${N_{\mathrm{s}}}$
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

{N_{\mathrm{s}}}

and molecule counts

N_{\mathrm{mol}}

. The total batch size is kept constant at

4096

. Times are given in

\mathrm{m}\mathrm{s}

Figure 12: Convergence of the total energy on 33-state Be for the standard Excited Pfaffian (

{N_{\mathrm{k}}}=16

), a single-determinant Excited Pfaffian (

{N_{\mathrm{k}}}=1

), and the naive multi-state architecture at

{N_{\mathrm{k}}}=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

{N_{\mathrm{k}}}=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 15: Spatial overlap of atomic excited states measured as Bhattacharyya coefficients $\langle|{\Psi}_{s}|||{\Psi}_{t}|\rangle$ .

Figure 16: Spatial overlap of 33 states of beryllium measured as Bhattacharyya coefficients $\langle|{\Psi}_{s}|||{\Psi}_{t}|\rangle$ .

Figure 17: Per-step training time as a function of the number of states

{N_{\mathrm{s}}}

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

{N_{\mathrm{s}}}

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\times$ fewer	$15$ – $60\times$ fewer	$4$ – $40\times$ fewer	$2$ – $30\times$ 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	$\sum$	$\frac{\Sigma}{\Sigma_{\mathrm{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\times$
Pfau et al. (2024)	KFAC	5.8	7.1	10.0	12.2	15.2	17.1	22.8	26.2	116.4	$7\times$
Szabo et al. (2024)	Adam	3.3	3.8	4.3	5.8	7.0	7.6	9.3	10.1	51.2	$51\times$
Szabo et al. (2024)	KFAC	11.4	12.4	13.2	13.7	15.4	16.4	18.2	18.8	119.5	$7\times$
This Work	Adam	-	-	-	-	-	-	-	-	1.0	$1\times$
This Work	SPRING	-	-	-	-	-	-	-	-	16.8	$1\times$

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	$\sum$	$\frac{\Sigma}{\Sigma_{\mathrm{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\times$
(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\times$
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\times$
(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\times$
This Work	Adam	-	-	-	-	-	-	-	-	-	-	-	-	3.9	$1\times$
This Work	SPRING	-	-	-	-	-	-	-	-	-	-	-	-	14.0	$1\times$

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	$\frac{\mathrm{total}}{\mathrm{our\,total}}$	Step time	Total 23 structures	$\frac{\mathrm{total}}{\mathrm{our\,total}}$
Pfau et al. (2024)	Adam	49.0	490.4	$208\times$	3.9	89.2	$24\times$
Pfau et al. (2024)	KFAC	49.9	499.2	$46\times$	4.2	97.5	$9\times$
Szabo et al. (2024)	Adam	14.0	140.0	$59\times$	5.5	125.6	$33\times$
Szabo et al. (2024)	KFAC	24.8	248.0	$23\times$	7.4	170.2	$16\times$
This Work	Adam	-	2.4	$1\times$	-	3.8	$1\times$

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.995\text{\,}\mathrm{\text{Å}}

. “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)	$\Delta E$ (Ha)
$\mathrm{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)
$\mathrm{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)	$\Delta 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	$X^{1}\Sigma_{g}^{+}$	$A^{1}\Pi_{u+}$	$A^{1}\Pi_{u-}$	$B^{\prime 1}\Sigma_{g}^{+}$	$B^{1}\Delta_{g}^{(2)}$	$B^{1}\Delta_{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)	$\Delta E$ (Ha)	$\Delta E_{\mathrm{exp}}$ (Ha)	$\Delta\Delta E_{\mathrm{exp}}$ (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

\Delta E_{\mathrm{exp}}

( Sansonetti, 2003 ) and errors

\Delta\Delta E_{\mathrm{exp}}

are also given. Energies are given in atomic units.

Variational Monte Carlo

See all 11 papers

Foundation Neural-Network Quantum States for Molecular Potential Energy Surfaces in Second Quantization
Sep 29, 2026Lizhong Fu, Jianan Wei, Wenguan Wang +1Neural Quantum States Computational Quantum Chemistry
State Key Laboratory of Precision and Intelligent Chemistry, University of Science and Technology of China
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
Parallel Scan Recurrent Neural Quantum States for Scalable Variational Monte Carlo
May 13, 2026Ejaaz Merali, Mohamed Hibat-Allah, Mohammad Kohandel +2Neural Quantum States Variational Monte Carlo
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

Sep 29, 2026J/K move · Enter open · S save

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

Lizhong Fu, Jianan Wei, Wenguan Wang, Honghui Shang

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

Neural Quantum States Computational Quantum Chemistry Quantum Chemistry

Open paper AskPDF arXiv

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.

Abstract

Variational Monte Carlo

Explore similar work

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

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

Parallel Scan Recurrent Neural Quantum States for Scalable Variational Monte Carlo