Variational Monte Carlo on a Budget -- Fine-tuning pre-trained Neural Wavefunctions

Summary

The authors propose a pre-trained deep learning wavefunction ansatz for variational Monte Carlo (VMC) solutions of the Schrödinger equation. The ansatz is trained on a chemically diverse set of (different conformations of) molecules. The zero-shot accuracy of the proposed method is a significant improvement over the previous state-of-the-art of pre-trained wavefunction ansätze, but additional fine-tuning steps are still necessary to obtain accurate relative energies. Due to the relatively simple model architecture, the achievable accuracy (in the limit of unlimited fine-tuning steps) is lower compared to more complex model architectures like PsiFormer, which however, need to be trained from scratch for every system.

Strengths

The paper adresses the important problem of reducing the computational overhead of VMC methods using a deep learning wavefunction ansatz. The description of the proposed method is clear and well-written. Further, the authors include their code for reproducibility and the method is tested on a large number of systems.

Weaknesses

While the paper compares results with previous methods, these comparisons consider only metrics like the number of optimization steps, but not the raw computational cost in terms of CPU/GPU time, which is arguably more important in practice. It is difficult to assess how "practical" the presented ansatz is compared to more traditional quantum chemistry methods.

Questions

- Looking at the accuracy vs. number of heavy atoms curves shown in Figure 2, I wonder how "practical" the current method is compared to running coupled cluster calculations. For example, how expensive is the calculation of a system with six heavy atoms at the CC-4Z level compared to running 4k fine-tuning steps with the presented method? Are we at a stage already where VMC methods are comparable in cost, or even cheaper than conventional quantum chemistry methods? - The authors show in Figure 3 and discuss in the text that the energies obtained with their method are not invariant w.r.t. global rotations. It is correct that the wavefunction is not invariant w.r.t. rotations, but it should be equivariant. The authors mention that they experimented with fully equivariant architectures, but observed worse results compared to their (non-equivariant) method. I assume this is because some ground state wavefunctions require "symmetry-breaking", e.g. the ground state of a carbon atom is not spherically symmetric. Still, once a "symmetry-broken" wavefunction is obtained, it should be equivariant under rotations. Have the authors already looked into the distinction between equivariance/symmetry-breaking? In my opinion, any insights regarding this problem are valuable and should probably be discussed in more detail (in the appendix/supplement). - It is well known that the energy variance in VMC methods allows to assess how close the found solution is to the true ground state wavefunction, which would have a variance of 0. The authors show many plots with absolute/relative energies, but how do energy variances compare between different methods? Related to this: How large are energy variances when comparing the zero-shot to the few-shot setting? - I believe there is a small mistake on line 122: Based on the description in the following lines ($l=0,\dots,L$ and $m=-l,\dots,l$), it should be $x^{\mathrm{nuc}} \in \mathbb{R}^{N_\mathrm{nuc}\times(L+1)^2\times N_{\mathrm{channels}}}$ (instead of $x^{\mathrm{nuc}} \in \mathbb{R}^{N_\mathrm{nuc}\times(2L+1)^2\times N_{\mathrm{channels}}}$). Can you please check this?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors fairly discuss limitations of their work. While its potential for negative societal impact is not discussed, I think this is acceptable considering that this work is about the solution to a fundamental problem in quantum chemistry and a discussion of negative societal impacts would probably appear contrived.

Summary

This paper concentrates on applying deep learning techniques to electronic structure computations. Unlike the strategy adopted in [1], the authors introduce a pretraining process to enhance the few-shot or zero-shot performance of Deep Learning Variational Monte Carlo (DL-VMC) methods. As their initial step, they design a novel architecture for the conventional model. Drawing inspiration from [2], the paper proposes a learnable orbital model for adjusting to molecular structures with a disparate number of electrons. The orbital model employs a neural network to predict the Fock and overlap matrices in the initial stages, and subsequently, it resolves a generalized eigenvalue problem to derive the orbital coefficients. Additionally, the network within the orbital model predicts the energy and selects new configurations for training. With the improved architecture, the authors pretrain the entire model using a dataset containing 98 distinct molecules. The pretrained model is benchmarked under various metrics, including absolute energy and relative energy. Furthermore, the authors conduct extensive experiments, comparing their zero-shot performance with Density Functional Theory (DFT) methods on the QM7 dataset. The empirical results demonstrate the effectiveness of their methodologies. [1] David Pfau et al., Ab Initio Solution of the Many-Electron Schrödinger Equation with Deep Neural Networks [2] Nicholas Gao and Stephan Günnemann, Generalizing Neural Wave Functions

Strengths

1. Improved architecture. Compared to the [1, the authors employ a refined architecture to extract embedding features from electrons and nuclei. Rather than using orbital coefficients derived from the Hartree-Fock method, they directly predict the Fock matrix and overlap matrix, subsequently solving the generalized eigenvalue problem. This strategy eliminates the reliance on the Hartree-Fock method, an inaccurate solution for electronic structure. Ablation studies reveal that this enhanced structure significantly improves absolute energy performance in both zero-shot and few-shot settings. 2. Enhanced Sampling Methods. Despite the insignificance of the MCMC burn-in step in the standard FermiNet-like training process, it can be time-consuming in few-shot/zero-shot settings. To counteract this issue, the authors refine the initialization of MCMC walkers, achieving equilibrium within a shorter MCMC burn-in stage. Beyond this initialization enhancement, they also identify and address the burn-in requirements during the pretraining stage by adjusting the movement of electrons relative to the shifts in nuclei position. Both these strategies improve the convergence of MCMC walkers, leading to superior training performance. 3. Larger Pretraining Dataset. In addition to the technical enhancements, the authors also construct a more substantial dataset for pretraining, a development that could prove beneficial for future research in this area. 4. Extensive Empirical Study. The authors execute extensive experiments to showcase the effectiveness of their methods, including comparisons with the Density Functional Theory (DFT) method on a large-scale dataset. [1] Michael Scherbela, Leon Gerard, and Philipp Grohs, Towards a Foundation Model for Neural Network Wavefunctions.

Weaknesses

1. The relative energy experiments may not be convincing. Unlike the absolute energy, if the relative energy is not converged, the discrepancy between CCSD(T) and the proposed method could magnify with additional finetuning steps. Although model with 4k finetuning steps achieves relative energy nearly identical result as CCSD(T), the reviewer suggests the authors to supply the relative energy result of model with an increased number of finetuning steps (e.g., 8k and 16k) to demonstrate energy convergence. 2. The comparison with Density Functional Theory (DFT) methods should also encompass speed comparison. Typically, Deep Learning Variational Monte Carlo (DL-VMC) methods should benchmark their energy results against post-HF methods and Quantum Monte Carlo (QMC) results, given that the total computation cost of these methods is comparable. Despite this paper's focus on zero-shot performance, the computational cost of evaluation of DL-VMC remains high. Therefore, the reviewer suggests that the authors provide speed comparisons to elucidate the effectiveness of their method.

Questions

Questions about experiments have been shown in the weakness. Questions about methods: 1. Why do the authors use the Hessian matrix of predicted energy to guide the sampling of configuration? 2. How to train the network to predict $E_{phis}$? 3. Why did the authors choose to solve the generalized eigenvalue problem to produce orbital coefficients? How is the gradient computed through this block?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

Compared to other ab-initio methods, the proposed method requires pretraining on a large molecule dataset. Therefore, its performance highly depends on the similarity between pretraining and evaluation datasets, which limits the application of the method.

Summary

The paper proposes a meta-learning model for neural wavefunctions. The neural model is first pre-trained by solving the variational optimization problem on many molecules. To adapt the model for a new problem (such as a new molecule), relatively few fine-tuning steps are needed to arrive at an accurate solution. The method consists of an orbital model based on PhisNet (with SE-3 equivariance), an improved electron embedding using message passing on the nuclear embeddings from PhisNet , and a Slater determinant to ensure antisymmetric property of the wavefunciton. The model is trained on a dataset of 98 small molecules with up to 4 heavy atoms. The resulting "base model" shows stronger zero-shot transfer across different number of heavy atoms than recent meta-learning wavefunction models. With a small amount of fine-tuning (4k steps), the model achieves very accurate energies, outperforming CCSD(T)-4Z on molecules up to 5 heavy atoms. It also achieves excellent accuracy on challenging systems like conformer energies, dissociation curves, and relative energies of transition states.

Strengths

- Use of the improved architecture embeddings based on PhisNet that preserves SE-3 equivariance. - The sampling of initial electron positions to run MCMC for VMC are improved by using an exponential distribution from the nuclei. - For sampling geometries of the molecules, the covariance of the Gaussian is rectified by inverse Hessian of the energy, yielding better geometries that make more sense physically.

Weaknesses

- Relative energies, which depend on small energy differences, are not very accurate without fine-tuning on that specific system. The transferability is not as good for relative energies. - The overall end-to-end training is built upon previous works.

Questions

- During sampling of molecular geometries, are they non-physical geometries discarded? - Are the electron embeddings SE-3 equivariant? (I assume the MPNN used does not preserve equivariance.)

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

N.A.

Summary

This paper tackles the problem of solving quantum Monte Carlo optimization on multiple molecules. A PhisNet model is used to encode geometric information from nuclei and other GNN models are then used to embed electrons. Additionally, the authors propose a method to generate more diverse nuclear geometries by a specially designed distortion method. Experiments show that the method can achieve good accuracy when trained smaller molecules and test on larger molecules. After some fine-tuning, the method can achieve improved results.

Strengths

- The authors provide detailed description of the architecture. - The paper uses diverse datasets to evaluate their method.

Weaknesses

- The improvements of the proposed method mainly comes from an improved architecture. To my understanding, the novelty seems to be limited. - The proposed method does not yield improved results when more optimization steps are used. - The presentation of the training and evaluation steps is not clear. Specifically, how is pre-training and fine-tuning done.

Questions

- Among the many experimented datasets, which one(s) most decisively show the benefit of the proposed method?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

None

Summary

The paper proposes a neural network wave function model which generalizes to different molecules and geometries. The neural network includes electron model, orbital model, both with message passing, Slater determinants, and an energy estimate for geometry sampling. The network is pretrained on 98 molecules and results show lower energies with one-shot and fine-tuning, on new molecules, compared to other generalized wavefunctions and some computational chemistry approaches. The paper also includes an ablation study and prediction of absolute energies for 250 molecules.

Strengths

The idea is original since there are very few methods of generalizable neural networks for wavefunctions. The normal mode of sampling geometries from an inverse hessian of energy seems to be original. The experiments and figures are high quality. The results are very good, it seems to be the first generalizable, ab initio quality wavefunction, that can be trained with quantum Monte Carlo. Generalizability can be computationally efficient if not every optimization needs to be started from scratch. The problem of solving Schrodinger equation is an important and significant one.

Weaknesses

I'm not so familiar with CC-2Z, -3Z etc methods the paper compares against. Not sure how the proposed method compares with something like Slater-Jastrow variational Monte Carlo (which is quite fast and generalizable for an abinitio method).

Questions

From Fig. 3c, it seems that the errors are somewhat dependent on the rotation angle. Is that actually happening? How much of the improvement attributed to "improved architecture" may be due to Phisnet?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

4 excellent

Limitations

The authors are upfront about the limitations: potential limited expressiveness with more optimization iterations, prediction is worse for increasing extrapolation and out-of-distribution predictions.