Excited Pfaffians: Generalized Neural Wave Functions Across Structure and State

This paper introduces Excited Pfaffians, a generalized neural network architecture combined with Multi-State Importance Sampling, which enables the efficient and accurate representation of multiple excited states and potential energy surfaces with nearly constant computational cost, achieving significant speedups and scalability for systems like the carbon dimer and beryllium atom.

The Gist

Imagine you are trying to predict the weather. You know the basic laws of physics (thermodynamics, fluid dynamics), but the atmosphere is so chaotic and complex that you can't solve the equations perfectly. So, instead, you build a super-smart AI model that learns to guess the weather by looking at millions of past data points.

In the world of chemistry, scientists do something similar. They want to predict how atoms and molecules behave. The "weather" here is the electronic structure of a molecule—how its electrons move and interact. The "laws" are the Schrödinger equation.

For decades, AI has been great at predicting the ground state of a molecule. Think of the ground state as the molecule's "relaxed" mode—like a person sitting comfortably on a couch. It's the most stable, lowest-energy configuration.

But molecules also have excited states. This is like the same person jumping up, dancing, or running a marathon. These states are crucial for things like how plants use sunlight (photosynthesis), how our eyes see color, or how solar panels work. However, calculating these "dancing" states is incredibly hard for computers.

This paper introduces a new method called Excited Pfaffians (with a helper technique called MSIS) that makes calculating these excited states fast, accurate, and scalable. Here is how they did it, explained through simple analogies.

1. The Problem: The "Crowded Room" Bottleneck

Imagine you are in a room with 100 people (representing 100 different excited states of a molecule). You want to know how similar every person is to every other person (this is called calculating "overlaps").

• The Old Way: In previous methods, if you wanted to compare Person A to Person B, you had to ask Person A to stand still and ask Person B to stand still, then compare them. Then you'd do the same for A vs. C, A vs. D, and so on.
  • As you add more people to the room, the number of comparisons explodes. If you have 10 people, it's manageable. If you have 100, the computer has to do billions of comparisons. It gets so slow that scientists usually give up after just a few states.
• The Variance Problem: Also, in these comparisons, sometimes the data is "noisy." It's like trying to hear a whisper in a noisy room. To hear clearly, you have to shout (run more simulations), which takes even more time.

2. The Solution Part 1: Multi-State Importance Sampling (MSIS)

The authors realized they were being inefficient. Instead of asking each person to stand alone, they decided to pool everyone together.

• The Analogy: Imagine you want to compare the height of 100 people. Instead of measuring them one by one in separate rooms, you put them all in one big hall. You take a snapshot of the whole crowd.
• How it helps: By looking at the whole crowd at once, the "noise" cancels out. The computer can estimate the relationship between any two people using the same set of data.
• The Result: This technique (MSIS) means that adding more excited states doesn't make the computer work harder. Whether you calculate 5 states or 50, the time it takes stays almost the same. It turns a "super-linear" nightmare into a "constant" breeze.

3. The Solution Part 2: Excited Pfaffians (The "Swiss Army Knife" Network)

Now, let's talk about the AI model itself.

• The Old Way: Previously, to study 10 different excited states, scientists had to train 10 separate AI models. It was like hiring 10 different chefs to make 10 different versions of the same soup. Each chef needed their own kitchen, their own ingredients, and their own training time.
• The New Way (Excited Pfaffians): The authors built a single, massive "Swiss Army Knife" AI.
  • The Backbone: This AI learns the general rules of how electrons behave (the "soup recipe"). This part is shared by everyone.
  • The Switch: To switch between different excited states (e.g., from "dancing" to "jumping"), the AI just flips a tiny, lightweight switch (a "selector"). It doesn't need to relearn the whole recipe; it just changes the specific mode.
• The Result: One model can now represent dozens of different states simultaneously. It's like having one chef who can instantly switch between making soup, salad, and cake just by turning a dial, rather than hiring a new chef for every dish.

4. Why This Matters: The "Beryllium Breakthrough"

Because of these two tricks, the authors achieved something that was previously impossible:

• The Carbon Dimer: They modeled a molecule with 12 excited states (previous methods could only do about 8) and did it 200 times faster.
• The Beryllium Atom: They calculated all 33 possible excited states of a Beryllium atom. Before this, no one had ever used neural networks to find every single energy level of an atom. It's like finally mapping every single floor of a skyscraper, whereas before we only knew about the first few.
• Generalization: They showed that this single AI model can learn the rules for many different molecules at once. It's like teaching a student chemistry once, and then having them solve problems for water, carbon, and oxygen without needing a new textbook for each.

Summary

Think of this paper as upgrading from a manual transmission car to a self-driving electric vehicle.

• Old methods were like shifting gears manually for every single state, getting slower and slower as the road got steeper (more states).
• Excited Pfaffians use a smart transmission (MSIS) and a single, powerful engine (the shared neural network) that handles any number of states effortlessly.

This breakthrough means scientists can now simulate complex chemical reactions, design better solar cells, and understand light-matter interactions much faster and more accurately than ever before.

Technical Summary

1. Problem Statement

The accurate computation of electronic excited states is critical for applications ranging from drug design to photovoltaics. While Variational Monte Carlo (VMC) with neural network wave functions has achieved unprecedented accuracy for ground states, extending this to excited states faces two major scalability bottlenecks:

1. Sampling Inefficiency: Calculating the overlap $\langle \Psi_s | \Psi_t \rangle$ between $N_s$ states is required to enforce orthogonality. Standard methods require the number of Monte Carlo samples to scale linearly with $N_s$ to maintain low variance in overlap estimates. This leads to a computational cost that scales superlinearly (often $O(N_s^4)$) with the number of states.
2. Representation Inefficiency: Representing multiple excited states typically requires training separate neural networks for each state or using prohibitively large single networks with state-specific parameters. This prevents the amortization of training costs across different molecular structures and states.

Consequently, previous methods are limited to computing only a handful of excited states (typically $<10$) for a single geometry, making the generation of multi-state Potential Energy Surfaces (PES) or generalization across molecules computationally infeasible.

2. Methodology

The authors propose a unified framework combining two novel components: Multi-State Importance Sampling (MSIS) and the Excited Pfaffian architecture.

A. Multi-State Importance Sampling (MSIS)

To address the sampling bottleneck, the authors introduce MSIS, which pools samples from all $N_s$ states to estimate pairwise overlaps.

• Mechanism: Instead of sampling from a single state distribution $\rho_s$, MSIS samples from a mixture distribution $\rho_{mix} = \frac{1}{N_s} \sum \rho_s$.
• Mathematical Advantage:
  • Bounded Variance: The random variable used for estimation, $\frac{\Psi_t(r)\Psi_s(r)}{\rho_{mix}(r)}$, is bounded by $N_s/2$. This eliminates the heavy-tailed outliers that occur near wave function nodes in single-state sampling.
  • Variance Reduction: The variance of the overlap estimator scales as $O(1/N_b)$ (where $N_b$ is the total batch size) rather than $O(N_s/N_b)$. This allows the total sample size to remain constant regardless of the number of states, achieving near-constant computational scaling.
• Normalization: Since wave functions are unnormalized, MSIS employs bridge sampling to estimate the ratios of normalization constants ($Z_s/Z_t$) iteratively.

B. Excited Pfaffian Architecture

To address the representation bottleneck, the authors introduce the Excited Pfaffian, a neural architecture inspired by Hartree-Fock (HF) theory.

• Parameter Sharing: Unlike standard approaches that learn separate orbitals for each state, the Excited Pfaffian shares a single "backbone" network (generating electron embeddings and orbital matrices $\Phi$) across all states.
• State Differentiation: States are distinguished only by lightweight, state-specific orbital selector matrices ($A_s$). The wave function is defined as a linear combination of Pfaffians:
  $$\psi_s(r) = \sum_{k=1}^{N_k} \text{Pf}(\Phi(r)_k A_{sk} \Phi(r)_k^T)$$
• Orthogonality: By varying only the selector $A_s$ (analogous to the orbital selector $\Pi_s$ in HF), the architecture ensures orthogonality between states while drastically reducing the parameter count and memory footprint. This avoids the $O(N_s)$ memory explosion associated with separate orbital weights.

C. Auxiliary Techniques

• Spin-State Snapping: A loss term is added to dynamically snap wave functions to the nearest integer or half-integer spin eigenvalue ($S(S+1)$), preventing convergence to linear combinations of different spin states without requiring prior knowledge of the target spin.
• Generalized Pretraining: A scheme is introduced to propagate Hartree-Fock targets across a Potential Energy Surface (PES) using a Directed Acyclic Graph (DAG) and Procrustes alignment, ensuring smooth state continuity across geometric perturbations.

3. Key Contributions

1. Near-Constant Scaling: The combination of MSIS and Excited Pfaffians reduces the computational scaling with the number of states from superlinear ($O(N_s^3)$ to $O(N_s^4)$) to near-constant ($O(N_s^{0.14})$ empirically).
2. Generalized Multi-State Wave Functions: The first method capable of jointly generalizing across molecular structures, chemical compositions, and electronic states within a single neural network.
3. Stable Convergence: MSIS stabilizes the training of excited states by eliminating the high variance in overlap estimates that previously caused erratic energy optimization.
4. Memory Efficiency: The architecture reduces memory requirements by orders of magnitude compared to naive multi-state extensions, enabling the use of second-order optimizers (like SPRING) for large state counts.

4. Experimental Results

The method was evaluated on atoms, molecules, and dissociation curves:

• Computational Scaling: On a Neon atom, the method achieved near-constant step time as the number of states increased from 1 to 30, contrasting sharply with the superlinear scaling of previous works (Pfau et al., 2024; Szabó et al., 2024).
• Accuracy & Efficiency:
  • Carbon Dimer ($C_2$): The model matched the accuracy of the Natural Excited States (NES) method while using 200$\times$ fewer compute resources and modeling 50% more states (12 states vs. 8).
  • Second-Row Atoms & Molecules: A single model trained on Li–Ne atoms and 12 molecules achieved chemical accuracy ($<1.6$ mHa) for 10 states per atom, using 20–80$\times$ fewer samples than separate training approaches.
• Many-State Computation (Beryllium): Leveraging the favorable scaling, the authors computed all 33 excited states of the Beryllium atom (up to the ionization potential) using neural networks for the first time, with errors $<0.5$ mHa compared to NIST data.
• Potential Energy Surfaces: The method successfully tracked state crossings in the $C_2$ dissociation curve and ethylene torsion/pyramidalization, outperforming previous methods that failed to track states correctly at stretched geometries.
• Ground State Improvement: The authors demonstrated that including excited states in the training objective improves ground-state accuracy near state crossings, a phenomenon where ground-state-only methods often fail due to getting trapped in local minima.

5. Significance

This work represents a paradigm shift in neural quantum chemistry:

• Feasibility of Many-Body Excited States: It makes the computation of high-dimensional excited state manifolds feasible, which was previously restricted by computational cost.
• Foundation Models for Chemistry: By enabling a single model to generalize across structures and states, it paves the way for "foundation models" in quantum chemistry that can predict properties for unseen molecules and complex photochemical processes.
• Applications: The ability to generate accurate, multi-state reference data efficiently opens new avenues for training machine learning force fields, density functionals, and simulating photochemical dynamics and absorption spectra.

In summary, Excited Pfaffians solve the dual problems of sampling variance and parameter redundancy in neural VMC, enabling the first scalable, generalized, and highly accurate computation of excited states across diverse chemical systems.