Neural Pfaffians: Solving Many Many-Electron Schrödinger Equations

We thank the reviewer for their invaluable feedback and hope to address their concerns. Firstly, we would like to highlight the broad range of new experimental evidence we present in the general comment. The following details how the experiments relate to the reviewer's concerns.

Ablation studies
To better isolate the contribution of our Neural Pfaffians, we performed several new ablation studies on the TinyMol dataset. In particular, we also trained Globe (+ Moon) [1] on the TinyMol dataset and combined our Neural Pfaffians with FermiNet [2] and PsiFormer [3].

Globe shares the same embedding technique as our Neural Pfaffian and, thus, only differs in how to enforce the fermionic antisymmetry. The results of Globe are depicted in Figure 1 of the general response. There, we can see that while Globe starts similar to our Neural Pfaffians, it cannot reach similar accuracies and converges to significantly higher, i.e., worse, energies.

Since the Neural Pfaffian is not restricted to Moon as an embedding technique, we perform an ablation study where we replace Moon with FemriNet and PsiFormer, respectively. The convergence plots are depicted in Figure 2 of the general response. There, we can see that our Neural Pfaffian outperforms TAO and Globe, independent of the choice of embedding model. However, consistent with [1], Moon performs better in generalized wave functions.

Additionally, we would like to point the reviewer to the ablation study in Appendix F, where we replace the Pfaffians with AGPs $\Psi=\det(\Phi_\uparrow\Phi_\downarrow^T)$, a special case of the Pfaffian that is faster to evaluate. The rest of the network stays the same. There, we find that AGPs are significantly faster to compute, though they cannot match the accuracy of our Neural Pfaffians.

Comparison to CCSD(T)
Our NeurPf does not match the CCSD(T) energies in Figure 5 of the paper, as convergence typically requires 100k-200k steps [2,3]. We only trained for 32k steps to match the setup from [4]. To better illustrate final convergence, we added an evaluation to Table 1 of the general response, where we train a NeurPf for 128k steps. Our NeurPf comes within chemical accuracy ($\leq 1.6mE_h$) or surpasses CCSD(T) even on the larger dataset.

Odd numbers of electrons
We are happy to present an alternative solution to deal with odd numbers of electrons in the general comment and would appreciate the reviewer's opinion on this. In short, instead of appending a learnable vector to the orbital matrix, we pad the orbitals $\Phi$ in both dimensions with an identity block to obtain $\hat{\Phi}=\begin{pmatrix}\Phi&0\\\\0&1\end{pmatrix}$. Additionally, we also pad the antisymmetrizer $A$ to $\hat{A}=\begin{pmatrix}A&1\\\\-1&0\end{pmatrix}$ such that one obtains $\text{Pf}(\hat{\Phi} \hat{A}\hat{\Phi}^T)\propto\det\Phi$ if $\Phi$ is square.

We repeat the experiment with this new formulation on the second-row elements experiment in Figure 6 of the general comment. We find little difference. Going forward, we will adopt the new padding technique as it requires no additional parameters.

Construction of skew-symmetric matrix
This is a great point raised by the reviewer. We agree with the reviewer that there are various ways of parametrizing skew-symmetric matrices. We decided to go with $\Phi A\Phi^T$ as a general construction. For instance, $A=\begin{pmatrix}0&I\\\\ -I&0\end{pmatrix}$ and $\Phi=(\Phi_1 \hspace{1em} \Phi_2)$ corresponds to the reviewer's suggestion $\Phi A\Phi^T=\Phi_1\Phi_2^T - \Phi_2\Phi_1^T$. We experimentally verify the advantage of having $A$ being fully learnable in Figure 3 of the general response, where we compare our approach to a fixed non-learnable $A$.

The results indicate that having $A$ as learnable grants an accuracy benefit in the later stages of training. During the development of our method, we experimented with several other parametrizations, e.g., $A$ being learnable block-diagonal or other fixed forms like $A=\text{diag}\left(\begin{pmatrix}0&I\\\\-I&0\end{pmatrix}, ...\right)$. Still, all of these resulted in very similar training trajectories to the one depicted in Figure 3 of the general response. We also experimented with parametrizing the skew-symmetric matrix for the Pfaffian directly via pair-orbitals $\text{Pf}(A)$ with $A_{ij}=\phi(h_i, h_j) - \phi(h_j, h_i)$ but found this to be numerically unstable for molecular systems, especially with heavier atoms. To better communicate this, we will add this experiment among all the other new experimental evidence to the paper.

Final remarks
We hope to have addressed the reviewer's concerns and were able to isolate our contribution better with our additional experimental evidence. We are happy to discuss further concerns and look forward to an engaging discussion.

[1] Gao et al. "Generalizing Neural Wave Functions"
[2] Pfau et al. "Ab-Initio Solution of the Many-Electron Schrödinger Equation with Deep Neural Networks"
[3] von Glehn et al. "A Self-Attention Ansatz for Ab-initio Quantum Chemistry"
[4] Scherbela et al. "Towards a transferable fermionic neural wavefunction for molecules"

We are delighted by the reviewer's positive feedback and want to address the remaining concerns. Firstly, we would like to highlight the broad range of new experimental evidence we present in the general comment. The following details how the experiments relate to the reviewer's concerns.

Time convergence plot
For our new ablation studies with NeurPf (+ Moon/PsiFormer/FermiNet) and Globe (+ Moon), we added convergence plots regarding compute time in Figure 10 of the general response.

One can see that despite the additional computational overhead through the Pfaffian, our NeurPf is approximately as fast as Globe with the same embedding. This comes mainly from the fewer operations, as we do not require Globe's extra message-passing steps from atoms to orbitals.

Total energy
We agree with the reviewer that the total energy of all elements in the training set is an imperfect measure. For the ablation study, we decided to plot the total energy as this represents the optimization objective and is less noisy than individual molecular energies. Further, all energies are within the same order of magnitude ($-78.5E_h$ to $-114.49E_h$).

Nonetheless, to better communicate the error per molecule, we added Table 1 in the general response to attribute the error on a per-molecule basis. Generally, models that perform well on one of the molecules also perform well on the others. We also would like to highlight the results in Figure 8 of the manuscript, where we break down the error per molecule with error bars indicating the different structures for finer details. Figure 8 shows that our Neural Pfaffian results in more consistent relative energies than TAO.

Pfaffian vs. determinant
Great question; we will extend the Appendix with the following discussion to clarify this. While it is a simple result to show that a Pfaffian can generalize Slater determinants, i.e., a Pfaffian can represent every Slater determinant, it is non-obvious that Pfaffians are naturally better in convergence or expressiveness. Empirical evidence in classical QMC [1] finds little to no improvements in molecular systems. However, in non-molecular systems, Pfaffians had greatly improved accuracy [2].

In our new ablation study in Figure 3 of the general response, we replace the learnable $A$ in $\text{Pf}(\Phi A\Phi^T)$ with a fixed one and investigate the impact on convergence. Here, we find that the parametrization is an essential factor in the accuracy of our Neural Pfaffians.

In summary, it is unclear whether Pfaffians are generally better suited for modeling molecular systems. However, as we demonstrate in this work, they can achieve identical accuracy and are well-suited for generalized wave functions.

Envelopes
We appreciate the keen eye of the reviewer; the classical envelopes are indeed faster than our memory-efficient envelopes. While our envelopes and Pfau et al. [3] aim to reduce memory requirements, they require more operations. In particular, the full envelopes require $O(N_bN_dN_nN_e^2)$ operations. In contrast, our memory efficient envelopes require $O(N_bN_dN_nN_\frac{\text{env}}{\text{atom}}N_e^2)$ (for $N_o=N_e$) operations where $N_b$ is the batch size, $N_d$ is the number of determinants, $N_n$ the number of nuclei, $N_e$ the number of electrons and $N_\frac{\text{env}}{\text{atom}}$ is the number of envelopes per atom in our memory-efficient envelopes. Our envelopes primarily reduce the memory from $O(N_bN_dN_nN_e^2)$ for the full envelopes to $O(N_bN_dN_nN_\frac{\text{env}}{\text{atom}}N_e)$ where $N_\frac{\text{env}}{\text{atom}}\ll N_e$.

We most likely attribute the empirical performance to the increased number of wave function parameters. While the $\sigma$ tensor is reduced from $N_d \times N_n \times N_e$ to $N_d \times N_n \times N_\frac{\text{env}}{\text{atom}}$, the $\pi$ tensor is enlarged from $N_d \times N_n \times N_e$ to $N_d \times N_n \times N_\frac{\text{env}}{\text{atom}} \times N_e$. For instance, for $N_e=20, N_n=4, N_d=16, N_\frac{\text{env}}{\text{atom}}=8$, we get the following parameter counts:

We will update Appendix A to better reflect memory and compute requirements in the context of the full envelopes.

Hydrogen chain results
We agree with the reviewer that the H2 case is the simplest of all structures. However, it is arguably the most distinct from the other structures as the chain has no "middle" elements. We hypothesize that due to this, it has the lowest accuracy as no fine-tuning has been performed in this experiment.

Odd numbers of electrons
We are happy to present an alternative solution to deal with odd numbers of electrons in the general comment and would appreciate the reviewer's opinion on this. In short, instead of appending a learnable vector to the orbital matrix, we pad the orbitals $\Phi$ in both dimensions with an identity block to obtain $\hat{\Phi}=\begin{pmatrix}\Phi&0\\\\0&1\end{pmatrix}$. Additionally, we also pad the antisymmetrizer $A$ to $\hat{A}=\begin{pmatrix}A&1\\\\-1&0\end{pmatrix}$ such that one obtains $\text{Pf}(\hat{\Phi} \hat{A}\hat{\Phi}^T)\propto\det\Phi$ if $\Phi$ is square.

We repeat the experiment with this new formulation on the second-row elements experiment in Figure 6 of the general comment. We find little difference. Going forward, we will adopt the new padding technique as it requires no additional parameters.

Final remarks
We will make sure to correct any typos in our manuscript. We hope to have adequately addressed the reviewer's concerns and questions. We welcome any additional feedback from the reviewer and eagerly await their response.

[1] Bajdich et al. "Pfaffian pairing and backflow wavefunctions for electronic structure quantum Monte Carlo methods"
[2] Kim et al. "Neural-network quantum states for ultra-cold Fermi gases"
[3] Pfau et al. "Natural Quantum Monte Carlo Computation of Excited States"

We thank the reviewer for their positive feedback on our manuscript and would like to take the opportunity to address the few concerns that were raised. Firstly, we would like to highlight the broad range of new experimental evidence in the general comment. The following details how the experiments relate to the reviewer's concerns.

Comparison to Globe
To include more baselines in our empirical evaluation, we trained Globe [1] on both TinyMol datasets and plotted the convergence in Figure 1 of the general response. While Globe is initially close to our NeurPf, it converges to higher, i.e., worse energies.

Comparison to PESNet
We want to stress that PESNet [2] can only perform generalization across different geometric configurations, while Neural Pfaffians tackle the more complex problem of generalization across arbitrary molecular compounds. Nonetheless, we are happy to add comparisons of PESNet on the N2 energy surface. We also added FermiNet [3] from [4] to cover a broader range of methods. However, it should be noted that Globe (Ethene) and our (Ethene) have been trained on a more challenging augmented dataset, while PESNet is only optimized on the N2 energy surface directly. For FermiNet, each structure is optimized independently.

The results in Figure 8 of the general response demonstrate NeurPf's high accuracy on energy surfaces.

Comparison to CCSD(T) CBS
In Figure 5 of the manuscript, we replicated the setting from [5] and, thus, only trained for 32k steps. However, neural wave functions typically require between 100k and 200k steps to converge [3,6]. Therefore, we extend the training of our Neural Pfaffian to 128k steps and compute energies for the converged model. The results are displayed in Table 1 of the general response. There, our long-trained NeurPf significantly outperforms CCSD(T) CBS on 3 of the 4 large molecules while being within chemical accuracy ($\leq 1.6mE_h$) on the last one.

Extended training on second-row elements
As the reviewer suggested, we trained our Neural Pfaffian for 200k steps on the second-row elements in Figure 6 of the general response. These results strongly suggest that a single Pfaffian can learn the ionization potentials with higher accuracy with additional training.

Final remarks
We hope to have answered the reviewer's questions and look forward to an engaging discussion. We appreciate any further feedback and questions from the reviewer.

[1] Gao et al. "Generalizing Neural Wave Functions"
[2] Gao et al. "Ab-Initio Potential Energy Surfaces by Pairing GNNs with Neural Wave Functions"
[3] Pfau et al. "Ab-Initio Solution of the Many-Electron Schrödinger Equation with Deep Neural Networks"
[4] Fu et al. "Variance extrapolation method for neural-network variational Monte Carlo"
[5] Scherbela et al. "Towards a transferable fermionic neural wavefunction for molecules"
[6] von Glehn et al. "A Self-Attention Ansatz for Ab-initio Quantum Chemistry"

We thank the reviewer for their invaluable feedback and suggestions. We hope to address their concerns. Firstly, we would like to highlight the broad range of new experimental evidence we present in the general comment. The following details how the experiments relate to the reviewer's comments.

TinyMol baselines
In addition to the results from [1], we also trained Globe (+ Moon) from [4] on the TinyMol datasets and plot the convergence in Figure 1 of the general response. While starting similarly to our NeurPf, Globe does not converge to the same energies but saturates at energy errors at around $9.4mE_h$ and $47.8mE_h$ on the small and large datasets, respectively.

Unfortunately, we cannot provide FermiNet or PsiFormer energies for all structures as these require $\approx$ 6k A100 GPU hours each.

Embedding
We agree with the reviewer's suggestion to ablate the embedding method of our NeurPfs. To demonstrate that they also perform well with different embedding methods, we train them with FermiNet [2], PsiFormer [3], and Moon [4] on both TinyMol test sets. We omitted LapNet due to its similarity with PsiFormer but happily add it to the next version of the paper.

The results are depicted in Figure 2 of the general response. NeurPfs outperform Globe and TAO independently of the choice of embedding network. Though consistent with the results from [4], Moon performs best in generalized wave functions.

Transferability
We agree that transferability is an exciting aspect of generalized neural wave functions but would also like to stress that it has not been the focus of this work as it typically cannot achieve chemical accuracy. Nonetheless, we are happy to present additional experiments. Like Scherbela et al. [1], we first train our NeurPf on the TinyMol training set and then transfer it to the unseen molecules in the test sets.

The results are depicted in Figure 4 of the general response. These suggest that even after pretraining, both methods still require significant fine-tuning to lower errors, but only NeurPf can reach chemical accuracy.

In addition to the TinyMol results, we would like to point the reviewer to the hydrogen chain experiment in Appendix E, where we extrapolate to larger hydrogen chains without finetuning.

Joint vs. separate training
We compare separately optimized wave functions to our NeurPf trained on the 30 and 40 structures, respectively. Since training separate wave functions for all 70 molecules in TinyMol exceeds our computational resources ($\approx$ 6k A100 hours), we picked one structure for each of the 7 molecules. We trained a separate NeurPf for each to estimate the convergence.

The results are shown in Figure 5 of the general response. At a fixed cost, our generalized wave function generally offers higher accuracy than separately optimizing wave functions. However, we also find that the additional degrees of freedom (higher ratio of parameters/molecule) and specialized optimization offer better final accuracies for separate optimization.

Computational efficiency
We measure the compute time of the Pfaffian operation in Figure 8 of the general response. Our Pfaffian is five times slower than the determinant. This is primarily due to optimized CUDA kernels for the latter. Note that here, we only measure the Pfaffian and determinant, not the rest of the network.

We benchmark the effect of the batch composition on the time per training step in Figure 9 of the general response for different batches composed of two molecules. There, we see a small overhead for small structures, while for large $N_e$, the time per batch converges to the geometric mean of the individual structures.

When training systems of different sizes, we optimize with various techniques. We work with flattened representations for the embedding network. For the Pfaffian operation (or determinant), we switch to sequential processing for each molecule in a batch (but batch different conformer). This also allows us to use different $N_o$ for each molecule. We want to highlight that this maintains a high level of parallelism thanks to the batch of electronic configurations per molecule.

For a comparison to APG, $\det (\Phi_\uparrow\Phi_\downarrow^T)$, we would like to point the reviewer to our ablation study in Appendix F. AGPs are a special case of Pfaffians. There, the AGP wave function is faster and reaches 32k steps in 70h compared to our Pfaffian's 95h. However, the AGP cannot match the accuracy of our NeurPf.

Number of orbitals
When going to large systems, the number of orbitals must increase with the number of electrons. As described in Section 4.4 and further detailed in Appendix C.3, we accomplish this by predicting a set of orbitals per atom:

[...] we grow the number of orbitals $N_o$ with the system size by defining $N_\text{orb/nuc}$ orbitals per nucleus, as depicted in Fig. 2.

Thus, a system with twice the number of atoms (assuming they are the same atoms) has twice the number of orbitals, while the generalized wave function still has the same number of parameters. The computational scaling of neural network wave functions (incl. FermiNet/PsiFormer/...) to hundreds of electrons remains an issue for future work. Still, NeurPf remains well-defined independent of the system size, thanks to the orbitals growing with system size.

Final remarks
Again, we thank the reviewer for their detailed assessment and suggestions for improving our manuscript. We hope to have addressed their concerns and look forward to an engaging discussion period. We appreciate any further feedback or questions from the reviewer.

[1] Scherbela et al. "Towards a transferable fermionic neural wavefunction for molecules"
[2] Pfau et al. "Ab-Initio Solution of the Many-Electron Schrödinger Equation with Deep Neural Networks"
[3] von Glehn et al. "A Self-Attention Ansatz for Ab-initio Quantum Chemistry"
[4] Gao et al. "Generalizing Neural Wave Functions"