Gold-standard solutions to the Schrödinger equation using deep learning: How much physics do we need?

Abstract: Finding accurate solutions to the Schrödinger equation is the key unsolved challenge of computational chemistry. Given its importance for the development of new chemical compounds, decades of research have been dedicated to this problem, but due to the large dimensionality even the best available methods do not yet reach the desired accuracy.
Recently the combination of deep learning with Monte Carlo methods has emerged as a promising way to obtain highly accurate energies and moderate scaling of computational cost. In this paper we significantly contribute towards this goal by introducing a novel deep-learning architecture that achieves 40-70% lower energy error at 6x lower computational cost compared to previous approaches. Using our method we establish a new benchmark by calculating the most accurate variational ground state energies ever published for a number of different atoms and molecules.
We systematically break down and measure our improvements, focusing in particular on the effect of increasing physical prior knowledge.
We surprisingly find that increasing the prior knowledge given to the architecture can actually decrease accuracy.

Reviews and author responses:

 Thank you for the detailed and clarifying response.

 Thank you once more, for reviewing our paper and helping us to improve it!

 Thank you for reviewing our paper and your constructive feedback!

Yes, for the systems such as 4th row atoms (K, Fe), and large molecules (e.g. Glycine), there are no published results using FermiNet, so we used the best variational reference energies we could find in the literature. 

 Thanks for the kind words! 

We're glad to hear you enjoyed the paper, and are very grateful for your support and helpful feedback.

 thanks again again for the answers to all questions. it has been a joy to review this paper, and I will strongly champion it in the further discussions.

 Thank the authors for the reply. I understand the full theoretical is out of the scope of this work. And I think the motivations for these modifications that can improve accuracy given by the authors are reasonable. Along with the good performance of the presented method, I will increase my score to 7 and recommend the acceptance of this paper.

 Thanks for the clarification. Adding Table 6 makes the speedup easier to quantify. As for the convergence, I am convinced by the authors' response that the QMC evaluation is accurate and the uncertainty is smaller than the chemical accuracy. I also notice for the systems with energy well below the reference (K, Fe, Glycine), the authors mainly compare with the MRCI method and FermiNet (VMC or DMC) is not used in comparison (Appendix A, Table 1). I have limited knowledge about the non deep learning methods. If the reference energies are indeed state-of-the art, then I think the improvement of deep QMC over classical methods is well demonstrated by the authors. 

With the above considerations, I will increase my score and will recommend acceptance of this paper. To further make the convergence / speedup more evident, the authors could consider showing some plots of training energy versus training iterations. However, this should be optional.

 Thank you for your reply!

I agree that finding out the right amount of prior knowledge to bake into the model is indeed a walk on a tightrope. Providing proper evidence that the physics prior is harmful is refreshing, in particular as there are so many armchair scientists at conferences and/social media who happily share their (usually strong) opinions about the need of “principled physics”, but then don’t provide any experiments. It is much appreciated indeed that the authors instead proceed in the opposite way, providing ample empirical evidence while staying humble, and I do recommend that this paper will make its way into the conference.

 Thank you very much for your review! Regarding your specific question:

**W1**: This paper demonstrates great experimental results, but lacking the necessary theoretical explanation to address the accuracy improvement under reduced computational cost. Section 4 discussed in detail the improvements obtained by each individual change and their combined effects, but the discussion is still restricted within the experimental point of view. Adding theoretical explanation on why the proposed method can obtain accuracy better than the classical as well as existing deep learning method can be beneficial.

**A**: While a full theoretical analysis of deep learning based Variational Monte Carlo is beyond the scope of our work, we explain the motivation and probable cause of our accuracy improvements below. 
We will add those points to section 4 in the final version of the manuscript, when the page limit is increased to 10 pages.

- *Embedding*: Our embedding is strictly more general than both PauliNet and FermiNet and our ansatz can therefore better approximate the true solution and reach lower energies.
    As discussed in Sec. 2.1 we can represent two-particle interactions more expressive than FermiNet and one-particle effects more expressive than PauliNet. 
    In addition we have a dedicated electron-nucleus embedding which further increases expressiveness.
- *Dense determinants*: Dense determinants are a direct generalization of block-diagonal determinants and therefore naturally more expressive. There is empirical evidence [1] that this improves the description of the wavefunction's nodal surface, but we are not aware of a thorough theoretical analysis.
- *Local input features*: As argued in Sec. 2.2. and depicted in Fig. 3, local coordinates capture physically more meaningful inputs than raw cartesian coordinates and can therefore be seen as physics-inspired feature engineering, leading to modest improvements in energy.
    The key improvement however is that enforce the required symmetries, enable transfer-learning of the wavefunction across different molecules or geometries.
- *Envelope initialization*: Initializing the wavefunction parameters closer to their optimal values accelerates optimization, because fewer update steps are needed, and the Monte Carlo sampling provides a better distribution of samples. By using the explicit analytical solution of the one-electron problem, we obtain a well founded estimator for the envelope exponents, which enables initialization of these parameters closer to their optimal values.
- *Hyperparameters*: The two key changes we propose (besides smaller batch size) are lower learning rate and higher gradient norm-constraint (i.e. less gradient clipping).
    This ensure that more steps are taken according to the curvature estimate by the KFAC-optimizer and fewer steps are limited by the clipped gradient norm.

[1]: *Explicitly antisymmetrized neural network layers for variational Monte Carlo simulation*, Lin et al., arxiv.org:2112.03491

 Thank you very much for your review! Regarding your specific points:

**Q1**: Some results in Figure 4 are lower than the reference energy, especially for larger molecules where the differences are significant (>100mHa) compared to the chemical accuracy (~2mHa). To which extent are we certain that the QMC results are correct? Although the variational principe could grantee the variational energy to be upper bound but doesn't it assume the sampling to be accurate in the first place? For example, in the Figure 3 of the FermiNet paper [Pfau et al. 2020], during VMC training, the energy could overshot when the MCMC step size is not large enough.

**A**: We are highly confident in the accuracy of our QMC energies to within 0.1-0.5 mHa, well beyond the 100mHa difference observed for larger molecules.
We have two independent estimates for this error: Firstly, we can estimate the Monte Carlo uncertainty from the variance within the Monte Carlo samples of a single run. Secondly, we can analyze the energy variance between independent optimization runs as we do for the $N_2$ study.
Both estimates yield uncertainties well below the inter-method energy differences.
We ensure that our Monte Carlo samples are indeed independent and unbiased during energy evaluation: We perform 20 Metropolis-Hastings steps between energy sampling to ensure independence, and do not optimize wavefunction parameters during final energy optimization to avoid bias.
Similar to our work [Pfau et al. 2020] also uses a separate evaluation run without optimization to avoid this bias in the final energies.
Lastly, we want to stress that the overshooting in Figure 3 of [Pfau et al. 2020] is depicted on a logarithmic scale and may therefore appear exaggerated on first glance. It is far less than the actual 100mHa by which we outperform classical methods.

&nbsp;

**Q2**: Under which setting the claimed "40-70% lower energy error at 8x lower computational cost" is obtained? As shown in Table 5, the runtime per epoch improves from ~6s to 3.6s. How is the convergence time defined?

**A**: The claims are obtained for the total time of computation under the prerequisite of the same hardware compared to the previously most accurate deep-learning-based method FermiNet. 
The two key factors leading to the speed-up are that we require fewer training iterations and smaller sample sizes.
Initially, we roughly estimated the speed-up by 4x fewer training iterations each being 2x faster due to a 50% smaller batch size. Since our architectural changes slightly slow down each iteration by roughly 20%, we agree that this is an slight overstatement and therefore adjust the mentioned speed-up in the manuscript to be around a factor of 6. 
To increase transparency, we have added in the appendix C of the manuscript a complete table stating for each system the computed speed-ups and accuracy improvements with regard to FermiNet. 
We want to highlight that for certain systems (e.g. Ar, P, S) we reach similar accuracy as FermiNet 10x faster, and that we reach 70-90% higher accuracy with a speed-up of 6x for medium sized molecules (e.g. C$_4$H$_4$, Ne).

&nbsp;

**W1**: Since the paper essentially integrates FermiNet and PauliNet, the novelty of proposed methods are very limited. In particular, the proposed coordinate transform only has sight effects.

**A**: While our work heavily builds on prior work such as FermiNet and PauliNet, we want to emphasize that it goes well beyond those two approaches:

- Our architecture is based on FermiNet and PauliNet, but in particular the embedding is an extension that is strictly more expressive than either predecessor, leading to better results.
- Several proposed improvements, such as the local coordinate transformation and envelope initialization are entirely novel to the best of our knowledge. The local coordinates specifically have low impact on single energies, but enable transferring a wavefunction ansatz to a different geometry or molecule [1,2], as we discuss in Sec. 6.
- Our surprising empirical result that maximizing prior physical knowledge can actually deteriorate accuracy can be of importance not only for the quantum chemistry community, but potentially many other domains.

**W2**: Although the overall improvement and speedup are clear, for some claims it is difficult to find reference in the main paper. E.g., "40-70% lower energy error at 8x lower computational cost" as stated in the abstract. It would be clearer to collect these comparisons in one place.

**A**: Please see our answer to your related question and appendix C in the revised manuscript.

&nbsp;

[1]: *Ab-Initio Potential Energy Surfaces by Pairing GNNs with Neural Wave Functions*, Gao et al., ICLR 2022

[2]: *Solving the electronic Schrödinger equation for multiple nuclear geometries with weight-sharing deep neural networks*, Scherbela et al., Nature Comp. Sci. 2022

 Thank you very much for your thorough review! Regarding your specific questions and potential weaknesses:

**Q1**: The units used in the equation before Eq. (1) should be clarified.

**A**: We use atomic units ($e=\hbar=m_e=a_0=1$) and have clarified this in the revised manuscript.

&nbsp;

**Q2**: The (probably well-known) Rayleigh-Ritz principle mentioned in L57 warrants a citation for the interested reader.

**A**: We added a citation to the original work by Walther Ritz.

&nbsp;

**Q3**: How is $H \psi(r)$ evaluated?

**A**:  The Hamiltonian $H$ applied to the wavefunction $\psi: \mathbb{R}^{3 \times n_\text{el}} \rightarrow \mathbb{R}$ can be evaluated as $H \psi({r}) = -\frac{1}{2} \sum_{i} \nabla^2_{r_i} \psi({r})+ \sum_{i>j} \frac{1}{|{r_i - r_j}|}\psi(r) + \sum_{I>J}\frac{Z_I Z_J}{|{R_I -R_J}|} \psi({r}) - \sum_{i,I} \frac{Z_I}{|{r_i - R_I}|}\psi(r)$ with $H\psi({r}) \in \mathbb{R}$. In particular the kinetic energy term $\nabla^2_{{r}_i}\psi({r})$ is being evaluated using automatic differentiation of our wavefunction model $\psi$.

&nbsp;

**Q4**: In Eq. (2), k seems to be, next to i, an index for an electron. In the previous equations, you use j. Why?

**A**: The index $k$ does not run over electrons, but enumerates different orbitals.
Although the indices $i$, $j$, and $k$ run over the same range of $1\ldots n_\text{el}$, we reserve $i$ and $j$ for electrons and $k$ for orbitals.

&nbsp;

**Q5**: The concrete form of Eq. (1) is not obvious. I assume that $\lambda$ is actually a matrix of shape $n_\text{el} \times n_\text{el}$?

**A**: Yes, we compute the determinant of a matrix of shape $n_\text{el} \times n_\text{el}$. We have update the notation in the equation and now explicitly state this fact in the manuscript.

&nbsp;

**Q6**: Below Eq. (3), you have another k. Is that related to the from before?

**A**: In the original manuscript we had the index $k$ for orbitals and the bold-face vector $\boldsymbol{k}$ denoting embedding vectors of the electron-nucleus stream.
These two concepts are entirely unrelated and we have renamed the embedding from $\boldsymbol{k}$ to $\boldsymbol{v}$ in the revised manuscript to avoid unnecessary confusion.

&nbsp;


**Q7**: I didn't fully understand the challenges you (and the community, I assume) are facing in Section 2.2. How are these requirements different from, let's say, an ML force field.

**A**: One key challenging difference is that the wavefunction $\psi$ has in general lower symmetry than the Hamiltonian, whereas all macroscopic observables of a force field (such as energies and forces) have the same symmetry. For example for a single atom (which has full rotational symmetry) the wavefunction is not in general rotationally invariant. See for example the cited reference [1]  for more detail.

&nbsp;

**Q8**: Do you mention anywhere the total computational cost of each method?

**A**:  In appendix C we report the runtime per optimization epoch for every method used in our ablation study for potassium on a NVIDIA A100 GPU. This includes our proposed method as well as our implementation of FermiNet.
A thorough analysis of the total computational cost for each non-deep-learning-based method mentioned in Fig. 4 is unfortunately out of scope for us, since we did not run these reference calculations ourselves but have taken the results from literature.
Most of our cited conventional reference methods have much higher scaling of computational cost with the number of particles (e.g. $\mathcal{O}(n_\text{el}^7)$ for CCSD(T)) than deep-learning-based VMC.
Additionally, we roughly estimate in appendix C the total computational cost used in our paper with about 50k GPUh. 

&nbsp;

**W1**: Main weakness is the lack of immediate relevance of the model improvements to the broader ML community.

**A**: We believe that our work will be relevant for this audience for two reasons: First, integrating physical- and machine-learning-models is a key active area of research, with many recent publications in venues such as NeurIPS. Examples of such work are, among many others, Physics Informed Neural Networks [2,3], exoplanet detection [4] chemical potential energy surface prediction [5] and a dedicated workshop *Machine Learning and the Physical Sciences* (https://ml4physicalsciences.github.io/2022/) at NeurIPS 2022. Our surprising observation that maximizing prior physical information can deteriorate model performance will be of interest to many members of this broad ML community. 
Secondly, many recent developments in using ML for quantum chemistry [1,6,7,8] have been published in this and similar venues, suggesting that our specific work will be of interest for another substantial audience at this conference.

[1]: Gao ICRL 2022, arxiv:2110.05064

[2]: Krishnapriyan, NeurIPS 2021

[3]: Belbute-Peres, NeurIPS 2021

[4]: Hönes, NeurIPS 2021

[5]: Corzo, NeurIPS 2021

[6]: arxiv 2011.07125

[7]: Schütt, NeurIPS 2017

[8]: Thölke, ICLR 2022


 Thank you very much for your review!
Regarding your specific questions and potential weaknesses:

**Q1**: Would a study of the transition states of the butadiene system (as in the paulinet paper) be instructive?

**A**: There are many interesting benchmark systems one could investigate, one of them being the butadiene transition states. For this work we have purposefully focused most of our attention on ground state equilibrium energies because most accurate benchmark data is available for these geometries and we have limited our investigation of transition states to the challenging $N_2$ system.

**W1**: I found the the description of the method quite unclear, and some notation was not very clearly defined where it appeared in the text

**A**: We further improved notation throughout the manuscript and added relevant citations for clarity.

**W2**: one could argue the paper is "only" a recombination of prior works, but if we take this argument, AlexNet would have needed to be rejected as well.

**A**: While we are humbled by the comparison to AlexNet, we would like to add that our work goes beyond a recombination of prior works: We introduce novel improvements (such as local coordinates and physics-based envelope initialization) and demonstrate the dangers of maximizing prior knowledge built into a model. In particular the latter point could be relevant for applications far beyond quantum chemistry.


 We would like to thank all reviewers for their appreciation of our work and constructive criticism. We are particularly delighted that most reviewers were impressed by our experimental results and that the overall paper was ”very well-written and clear, even for someone with little domain expertise” (Reviewer fBuC). We have answered all questions individually and have additionally updated our manuscript
to incorporate the improvement suggestions. Changes to the original version have been highlighted in red for the reviewers' convenience.

In addition we have incorporated important additional feedback we received during the review process: Several of the reference energies we used as non-variational best-estimate turned out to severely underestimate the true energy. Using more accurate reference calculations reveals that our results are even better for single atoms than originally presented in the manuscript. For example for the F atom our results correspond to 80% lower energy errors, while we had originally only reported 40%. We have updated Fig. 4 accordingly.

 An improved method to approximately solve the Schroedinger equation is described, which combines ideas from the PauliNet and FermiNet papers. A variety of ablation studies are performed. good performance is achieved. Strengths:
- good results
- thorough study and recombination of approaches of prior works, without reinventing the wheel
- ablation studies
- no grandiose claims
- good discussion of limitations

Weaknesses:
- I found the the description of the method quite unclear, and some notation was not very clearly defined where it appeared in the text 
- one could argue the paper is "only" a recombination of prior works, but if we take this argument, AlexNet would have needed to be rejected as well. Would a study of the transition states of the butadiene system (as in the paulinet paper) be instructive?


 yes

 The authors present and analyze a list of improvements that significantly increase accuracy and reduce the computational cost of variational Monte Carlo methods based on deep learning. # Strengths

- Overall, the paper is very well-written and clear, even for someone with little domain expertise.
- The introduction is helpful and allows non-experts to onboard (to the extent one can reasonably expect).
- The key contributions are laid out.
- Differences and similarities to previous methods (especially PauliNet) are exhibited.
- Thoughtful experimental design.
- Impressive experimental results.

# Weaknesses

- A (minor) weakness is the lack of immediate relevance of the model improvements to the broader ML community. - The units used in the equation before Eq. (1) should be clarified.
- The (probably well-known) Rayleigh-Ritz principle mentioned in L57 warrants a citation for the interested reader.
- How is $H \psi(r)$ evaluated?
- In Eq. (2), $k$ seems to be, next to $i$, an index for an electron. In the previous equations, you use $j$. Why?
- The concrete form of Eq. (1) is not obvious. I assume that $\lambda$ is actually a matrix of shape $n_{el} \times n_{el}$?
- Below Eq. (3), you have another $k$. Is that related to the $k$ from before?
- I didn't fully understand the challenges you (and the community, I assume) are facing in Section 2.2. How are these requirements different from, let's say, an ML force field.
- Do you mention anywhere the total computational cost of each method? The authors discuss the limitations of their work.

 This paper focuses on empirical improvements of variational quantum Monte Carlo estimation of molecular ground state energy. In particular, the authors focus on evaluating the various design choices in the FermiNet and the PauliNet, including dense determinant, hyper parameter tuning, the choice of envelope function and the pretraining strategies. On top of these, the authors propose to use a SchNet-like embedding blocks and an input feature transformation to make the input feature to make the feature to be invariant, local and expressive. 

The authors empirically evaluate all these design choices on molecules with up to ~40 electrons and obtain a model which improves upon the existing methods both in terms of accuracy and speed. The authors also conduct ablation studies by removing each design choice one by one. 

From the experiments, although the factors are not totally independent, it seems the most improvements come from the use of dense determinant, hyper parameters tuning, and the SchNet-like embedding. The proposed feature transformation only has minor effects. **Strength:** 

The paper is well written and easy to follow. The experimental results are good and the ablation studies are extensive and well presented. The effects of different design choices can be easily identified. Since PauliNet and FermiNet differ in many aspects while both give high quality quantum Monte Carlo results, it is interesting to compare and clarify the effects of these differences. Making improvements by fusing these two methods is well motivated.

**Weakness:** 
* Since the paper essentially integrates FermiNet and PauliNet, the novelty of proposed methods are very limited. In particular, the proposed coordinate transform only has sight effects. 
* Although the overall improvement and speedup are clear, for some claims it is difficult to find reference in the main paper. E.g., "40-70% lower energy error at 8x lower computational cost" as stated in the abstract. It would be clearer to collect these comparisons in one place. * Some results in Figure 4 are lower than the reference energy, especially for larger molecules where the differences are significant (>100mHa) compared to the chemical accuracy (~2mHa). To which extent are we certain that the QMC results are correct? Although the variational principe could grantee the variational energy to be upper bound but doesn't it assume the sampling to be accurate in the first place? For example, in the Figure 3 of the FermiNet paper [Pfau et al. 2020], during VMC training, the energy could overshot when the MCMC step size is not large enough.

* Under which setting the claimed "40-70% lower energy error at 8x lower computational cost" is obtained? As shown in Table 5, the runtime per epoch improves from ~6s to 3.6s. How is the convergence time defined?

[Pfau et al. 2020] Ab initio solution of the many-electron Schrödinger equation with deep neural networks. Physical Review Research, 2020. 

------
**After rebuttal:**

The authors' response clarifies the concern about the uncertainty and the speedup. Overall I think the proposed method is well motivated and results are convincing. Hence I will increase my score from 5 to 7 and will recommend acceptance of this paper. The authors state the limitations adequately.

 This paper proposes a deep-learning architecture to solve ground state of many-electron systems. The proposed model combines PauliNet-like neural network and envelope function in FermiNet, with additional improvements on embedding, input feature and parameter initialization of previous methods. Experimental results show that the proposed method can reduce errors by 40-70% with 4-8x lower computational cost. This paper also establish a new benchmark of several deep learning based and classical methods over a number of different atoms and molecules. Authors also research into the reason for improvements and find out that including too much physical prior knowledge can deteriorate the accuracy. **Strengths:**
- *originality:* Related works are adequately cited and it is clear to see the difference from these works. This work is a novel combination of PauliNet and FermiNet along with additional improvements.  
- *quality:* This work provides solid experiments to prove the accuracy and efficiency improvements. Also, this work systematically breaks down which changes cause the improvements, and gives insights that including physical prior knowledge can hinder the optimization. 
- *clarity*: This paper is overall well written and easy to follow. The method and results are clearly presented.
- *significance:* This work achieves the best results for the numerical solution of the electronic Schrodinger equation and establishes a new benchmark for current most accurate methods on a number of molecules and atoms. The proposed method is 4-8x faster than FermiNet in terms of optimization.  

**Weaknesses:**
This paper demonstrates great experimental results, but lacking the necessary theoretical explanation to address the accuracy improvement under reduced computational cost. Section 4 discussed in detail the improvements obtained by each individual change and their combined effects, but the discussion is still restricted within the experimental point of view. Adding theoretical explanation on why the proposed method can obtain accuracy better than the classical as well as existing deep learning method can be beneficial. 
 Please see weaknesses. Limitation of this work is well addressed.

Meta-review: There is a clear consensus among the reviewers that this is a quality paper and worthy of acceptance (in fact, this may be the first time I've ever seen 4 reviewers give the exact same score), so I recommend accept.

I do however have one additional comment. I find the current title somewhat unwieldy and wonder if it would be possible for the authors to condense it at all. This is not a critical issue, of course, but one that the authors may want to consider (if the program chairs allow it).

Decision: Accept