A Score-Based Model for Learning Neural Wavefunctions
Authors: Xuan Zhang
Shenglong Xu
Shuiwang Ji
Venue: Submitted to ICLR 2023 (2023)
Keywords: Neural wavefunction
Quantum Monte Carlo
Score-based method

Abstract: Quantum Monte Carlo coupled with neural network wavefunctions has shown success in finding the ground state of quantum many-body systems. The existing optimization approaches compute the energy by sampling local energy from an explicit probability distribution given by the wavefunction. In this work, we provide a new optimization framework for obtaining properties of quantum many-body ground state using score-based neural networks. This new framework does not require explicit probability distribution and performs the sampling via Langevin dynamics. Our method is based on the key observation that the local energy is directly related to the score, defined as the gradient of the logarithmic wavefunction. Inspired by the score matching and the diffusion Monte Carlo methods, we derive a weighted score matching objective, which guides our score-based models to correctly converge to the ground state. We first validate our approach with experiments on quantum harmonic traps, and further results show that it can accurately learn the ground states of atomic systems. By implicitly modeling the high-dimensional data distribution, our work paves the way toward a more efficient representation of quantum systems.

Under review as a conference paper at ICLR 2023
A SCORE-BASED MODEL FOR LEARNING NEURAL WAVEFUNCTIONS
Anonymous authors Paper under double-blind review
ABSTRACT
Quantum Monte Carlo coupled with neural network wavefunctions has shown success in computing ground states of quantum many-body systems. Existing op- timization approaches compute the energy by sampling local energy from an ex- plicit probability distribution given by the wavefunction. In this work, we provide a new optimization framework for obtaining properties of quantum many-body ground states using score-based neural networks. Our new framework does not require explicit probability distribution and performs the sampling via Langevin dynamics. Our method is based on the key observation that the local energy is directly related to scores, defined as the gradient of the logarithmic wavefunc- tion. Inspired by the score matching and diffusion Monte Carlo methods, we derive a weighted score matching objective to guide our score-based models to converge correctly ground states. We first evaluate our approach with exper- iments on quantum harmonic traps. and results show that it can accurately learn ground states of atomic systems. By implicitly modeling high-dimensional data distributions, our work paves the way toward more efficient representation of quantum systems.
INTRODUCTION
Understanding the properties of quantum systems lies at the core of many scientific disciplines, such as condensed matter physics, material science, and quantum chemistry. A quantum system is char- acterized by its ground state wavefunction, formally obtained by solving the Schrödinger equation. However, directly solving the Schrödinger equation for quantum systems with many particles is im- practical due the exponentially large Hilbert space. Owning to its strong dimension reduction capabilities, deep learning methods have been used as a strong candidate to approximately solve the Schrödinger equation and extract properties of quantum systems with the desired accuracy. For ex- ample. under the supervised learning setting, deep learning methods have been successfully applied to predict the quantum properties of molecular systems based on training data generated from den- sity functional theory (DFT) calculation (Schütt et al., 2017; Gasteiger et al. 2020 Liu et al. 2022 Wang et al.) (022). However, supervised methods rely on expensive computational simulations to generate large amount of training data, and the accuracy of these methods is fundamentally limited by the data quality. Furthermore, DFT calculations involve various approximations and are not guar- anteed to reach true ground states. A common scheme for approximately solving the Schrödinger equation is the variational principle, which optimizes a trial wavefunction to reach the ground state by minimizing its energy as much as possible via quantum Monte Carlo (QMC). Such a method is called variational Monte Carlo, whose accuracy relies on the expressive power of the trial wave- function. Recently, deep learning methods coupled with variational Monte Carlo have unleashed the potential of both methods (Carleo & Troyer] 2017 Hermann et al.) 2022). Powered by the efficient sampling and optimization framework of quantum Monte Carlo and the universal approximation capability of deep neural networks, neural wavefunctions can accurately model quantum states, and dramatic improvements have been achieved (Pfau et al. 2020: Hermann et al. 2020).
Modeling wavefunction is conceptually similar to modeling probability density. Existing meth- ods model the wavefunction explicitly by training a neural network to directly output the wavefunc- tion values. However, numerous examples in machine learning have shown that implicitly modeling data distributions provides better representations (Kingma & Welling (014:) Goodfellow et al. 2014 Ho al. 2020). As direct reference, score-based methods have demonstrated their strong

---

Under review as a conference paper at ICLR 2023
cess in generative modeling (Song & Ermon 2019 Song al. 2020 A score is defined as the gradient of the log probability. For example, realistic images can be generated from random noise by following dynamics defined by scores. In this paper, show that the quantum wavefunction can be represented by score models and be optimized within the QMC framework.
Our motivation relate score-based models with QMC is based an interesting connection be- tween energy computations and score-based formulations. In QMC, the energy of a system is av- eraged over local energy of plausible quantum states. Our observation is that local energy only involves gradients of the logarithmic wavefunction, which we define as the score of the wavefunc- tion. As result, to minimize energy. the score must be explicitly computed. On the other hand, the actual wavefunction values are only used for sampling and optimization. To this end, we propose a new optimization framework for QMC where sampling and optimization are also achieved by using score functions alone, eliminating the need to explicitly compute the wavefunction value. In our proposed score-based framework, the sampling done via Langevin dynamics and optimization done through a new loss function inspired by diffusion Monte Carlo.
Our score-based method enables the possibility of performing QMC computation with only score functions, which is infeasible in existing optimization frameworks. A direct benefit is that, by pre- dicting gradients, we avoid the need to recompute it from the wavefunction. Moreover, score func- tions can be interpreted as the force of quantum systems, by implicitly modeling distributions with score functions, the dynamics of quantum systems could be hetter captured. Our experimental re- sults show that with our score-based optimization framework, ground states of quantum systems can be accurately learned.
2 BACKGROUND AND RELATED WORK
2.1 QUANTUM MANY-BODY WAVEFUNCTION IN CONTINUOUS SPACE
We use R.xx to denote the coordinates of N particles in d-dimension. The quantum state of system is defined by its wavefunction ti IR Nxd + R. By definition va is normalized (s, [2=)
1) and (22 gives the probability density of observing x. Any wavefunction 1/ can be expressed as linear combination of eigenfunctions via. which are solutions to the time-indepen indeeeendn Schrödinger equation H (ax) = E, tha (z:), where H an linear operator known as the Hamiltonian, and E,, is a scalar giving the energy of the n-th eigen state. The Hamiltonian is defined as
Hv(()) - 1 E V(()))++(<))(e))))
(1)
where the index i runs over all of the N dimensions in the summation. The first term in the Hamiltonian takes the sum of the second-order partial derivatives of the wavefunction and is related to the kinetic energy of the system. The second term in the Hamiltonian multiplies the wavefunction by a scalar value and is related to the potential energy of the system. The kinetic term is intrinsic to the Schrödinger equation and always takes the same form, whereas the potential function V : N + IR varies for different physics problems. Note that although wavefunction can be complex-valued in general, we can let v be real-valued because H is real.
Our objective to find the ground state v/o which is the eigen state associated with the lowest energy Eg. In our notation, the coordinates æ can be either viewed as N d-dimensional vectors or as a flattened N dimensional vector. In the rest of this paper, depending on the context, we may use bold x; to denote the i-th particle or use the regular font x; to denote the i-th scalar component of the flattened vector.
2.2 VARIATIONAL MONTE CARLO
The variational Monte Carlo (VMC) method uses a parameterized function v/ Rxd IR. (called the Ansatz) to model a wavefunction, where 0 denotes the parameters to be optimized. The normal- ization of vhe is not required. The energy expectation of the computed as:
J (v(x)) ? Avo(w) da
(o(x)) (o (())= I ( (x))d
ve(z)
E Ez 0),
(2)
[ the (x) 2dx
2

---

Under review as a conference paper at ICLR 2023
where EL (2; 0) (mr)
Va(a) is called the local energy of æ. The expectation is evaluated numerically
on sparse samples. Markov Chain Monte Carlo (MCMC) sampling is employed to drive sample
distribution converge to the target density v v/a
We can approach ground state wavefunctions by minimizing the energy expectation with gradient descent. The unbiased gradient of (()) w.r.t parameters is given by:
Val(#) 2 / fo [( (Er((:: 0) Bm // Se Ez (xxx 011 Vo log [o
(3)
where expectations are evaluated by average over samples. The derivation of this loss makes use of the fact that / is Hermitian (Ceperley et al. 1977). A detailed derivation can be found in Appendix Li et al. (202T). D. The optimized energy expectation value gives the approximated ground state energy, and the eigenvalue formulation ensures the estimation be variational, that is, the energy expectation defined by the Ansatz is always above the true ground state energy. Other properties of quantum systems can also be estimated by taking expectations over corresponding operators.
2.3 RELATED WORK
Neural quantum state based on the restrictive Boltzmann machine was initially studied for simple quantum spin models lattices (Carleo & Troyer 2017). More complicated neural network archi- tectures, such as convolutional neural network and graph neural network are then extended to more complicated spin systems to capture frustration due to the lattice structure and next nearest neigh- boring interaction (Kochkov et [022]] Fu et al. (022). Similar ideas are also studied for ab-initio simulation in quantum chemistry (Han et al. 20199 9 study properties of small molecules. Recently, FermiNet Pfau et al. (020) and PauliNet Hermann et al. (2020) greatly improve the accuracy of neu- ral wavefunction and apply larger molecular systems. Follow-up works contribute to improved performance (Gerard et al. (022). joint training for multiple geometries (Gao & Günnemann 2021
Scherbela et al. [2022] Gao & Günnemann 2022) or solving for excited state (Entwistle al.] 2022))
All of the existing methods use neural networks to explicitly model the wavefunction, while in this
work we propose to implicitly model the quantum state using the score function.
Diffusion Monte Carlo (DMC) Toulouse et al. 2016 Ceperley 2004) employs score-based diffu- sion to improve upon VMC. In VMC, the quality of approximations is limited by the capacity of Ansatz. DMC additionally assigns weight for each sample in a way that the weighted distribution gives a better approximation of ground state. In DMC's formulation, each sample is defined by a walker, and each walker is assigned weight. In aaportance-samplee DMC, trial function रoy is used. At each iteration, the walkers randomly diffuse by following score of the trial function and the weights updated according to the imaginary time evolution that the repeated diffusing and weighting procedure projects out ground states from the trial wavefunction. At convergence, the ground state energy is estimated by the weighted average of local energies. Ceperley & Alder (1880) applies DMC to electonic gas. Ren et al. (2022) and Wilson et al. 2021 apply fixed-node DMC starting from the optimized FermiNet Ansatz.
Score-based methods have shown great success in generative modeling (Song & Ermon 2019: Song
al. 2022)). The effectiveness of implicit density modeling has also been demonstrated by other successful generative models, such as VAEs (Kingma & Welling, 2014), GANs (Goodfellow et al.) 2014 Brock al. 2018), and diffusion models (Sohl-Dickstein al. 2015 Ho et al.) 2020).
3
METHOD
For the Hamiltonian A defined in Equation the local energy can be expressed in terms of the logarithmic wavefunction as:
1 Σ
82 log wo(x:) + +
2 We (
(z; 0) = -
+ V(x).
(4)
dxf
Ox,
One can easily prove this expression by using the fact that 2oo (x)l O2s
v, 02. If we take closer look at this expression, we can notice that the local energy depends only the gradient of
3

---

Under review as a conference paper at ICLR 2023
MCMC
vo(x) E E
VMC
v
Energy Loss
X E IR Nxd
Langevin Dynamics
(x
Rxd
Diffusiion VMC (Proposed)
Weighted
Sq
Score Matching Loss
E Bxd
Figure Comparison between the pipeline of VMC and our proposed DiffVMC. In VMC, the wavefunction is explicitly modeled with a neural network, who directly outputs the wavefunction value. The sampling is carried out with MCMC and the optimization achieved by minimizing the energy loss. In the proposed DC we instead model the score of the wavefunction. The samples are generated via Langevin dynamics and the proposed weighted score matching objective is employed to optimize the score neural network.
the logarithmic wavefunction. In fact, if we define the function so Rxd R Nxd such that Se (z); log We(a)) we can rewrite the local energy as:
dx.
Er, = - (Ir(Vz.o(x))) + [[0())) + V(x).
(5)
The function sø is called the quantuum force or the drift in quantum Monte Carlo. Here follow the convention from the machine learning community and call the score. This direct connection between local energies and scores motivates us ask the question: using only sø? As a direct benefit, this can avoid recomputing the first-order derivatives in evaluating the local energy and consequently provide more succinet representation. It is true that will lose access to the unnormalized wavefunction. However, in practice, our primary goal is to estimate observables, such as kinetic energy density, which be estimated from sample distributions.
represent quantum states
To this end, our score-based optimization framework follows the gradient descent formulation in VMC, and design new loss function inspired by DMC. As we mix the flavors of these two methods, we call our new optimization framework the Diffusion Variational Monte Carlo (Dif- FVMC). The pipelines of VMC and DiffVMC are shown in Figure
Please note that the original score Hyvärinen & Dayan 2005 is defined in terms of probability densities, so we have Soriginal (x) of the score in our definition. By abuse of notation we define our score in terms of wavefunctions. We do this for two reasons. First, wavefunctions are more natural to deal with in quantum mechanics. Second, as we will show later, in addition to moving samples locally, we also use sx(x) to describe distributions, which shares the same physical meaning as the original score definition.
logp(z) = V log (x))
2 log (x)),, which twice
3.1 SCORE-BASED NEURAL WAVEFUNCTION ANSATZ
We use a parameterized score function to implicitly model the wavefunction. Formally, for sys- tems with particles in d dimension, its score function 8ø R xd - Nxd maps a set of input coordinates set of output scores, which are vectors having the same dimensions inputs. In- tuitively, after training, the output score should tell the particles the direction toward regions with higher density.
In our case, we are dealing with indistinguishable particles in quantum mechanics. So ex- changing two particles will not change the probability density: p(...,,,..... ......)? vi( ((...: Bi.............................. Di....)²
As result, the score, which is the gradient of logarithmic wavefunction,
4

---

Under review as a conference paper at ICLR 2023
is permutation equivariant. Formally,
ne log (ib(..,,, B.,...., -. = Vaz log (1 Bj....,,, "i....,,,i...................... Fi....)]
(6)
Essentially the equivariance means that two input particles exchange their positions, their cor- responding output score vectors will also exchange positions. The proof is straightforward (Ap-
pendix [B. We parameterize the score function with neural network. The equivariance can be
easily achieved by considering the input coordinates as set (Zaheer et al. 2017
Qi ct al. 2017).
3.2 SAMPLING VIA LANGEVIN DYNAMICS
To generate samples that follow the distribution implicitly defined by the score, we use Langevin dynamics, which is similar in the score-based generative models (Song & Hrmon 2019). Given the samples at the current time step Fe, the coordinates for the new samples are computed as:
L++1 = + Vac + (2:e)
where a is an hyperparameter defining the step size, and ~ (V,0, I Nd) is random vector sampled from the standard multivariate Gaussian distribution. The process is the same as in DMC and can be understood as first doing random diffusion, then drifting by following scores. We can prove that when a is small, the distribution converges to the distribution defined by the score function.
In Langevin dynamics, similar to the Metropolis-Hasting rejection step in MCMC, accept/reject procedure can be employed to alleviate the finite time error. Although omitted in the context of generative modeling (Song& Ermon (2019). this standard step in DMC. The original rejection step computes the ratio
exp
x x'
so
the (x')²
Pace
2
(8)
exp
2nx
sa())))))) (x))
After each Langevin dynamics move, we decide whether accept or reject the move based on the P Pace- Concretly, we first sample random number uniformly between 0 and 1. The move is accepted Pace is larger than the random number and is rejected otherwise. However, the expression of Pace involves the wavefunction values the (a)) and voo(""), which are generally not available in our score- based framework. To still be able to use the rejection step, we propose an approximated estimation of Pace which involves only the score function. This is achieved by approximating log [so[==))
[a]]
log (2') log (oo using the average gradient (Vz log (z)) + V log (=a()))) then we can show that the ratio can be approximated only in terms of the score as (derivation in Appendix C
Pace Ro exp
[==)())))²))
(9)
3.3 NEURAL WAVEFUNCTION OPTIMIZATION
The energy loss in VMC (Equation 3) depends explicitly on wavefunctions. The expression of the energy loss only in terms of scores is unknown. So we need to find new loss to optimize scores towards ground states. We motivate our new loss from the imaginary time evolution in DMC.
The imaginary time evolution operator e - projects out ground states when T + co. At short T,
by Taylor expansion, e - (xx) x vxx) - - (2) (1x He(a)) (=)
(x))
Thus, evolving (x)) in imaginary time for a short time 7 can be approximated as:
<l(x) + v))))>))))))>)))))))))))))))))>) - (xx).
(10)
a/ is closer to the ground state than v because higher energy eigenstates decays exponentially faster than the ground state. Therefore, for our score-based model, we can minimize energies by letting so(x) approach the evolved score Vz log [x)). We achieve this via score matching.
Score matching (Hyvärinen Dayan 2005) provides way to make so converge to the true score of the sample distribution. Assume at current step the samples follow 2/2 ʃ şb2 and sø(x) = Vz log (.)), we can make 8 (x)))) converge to Vn log (/(()) by minimizing the implicit score matching (ISM) objective:
ISM(O) 2 m n (1r(V=s0(z)) +
( []=](=))
(11)
5

---

Under review as a conference paper at ICLR 2023
The problem is that we do not have samples following v) v/? We can solve this by transforming the ISM into a weighted version based on current samples:
2 " (x)
tr V (o(o))))) + [|s0())) dx
ISM(0)
(12)
x )d
2 2 20 (a)) (x)
(r(Vaoo(=)) +
[|s0()))]])) ( dx f (
(13)
1/2 (x))d
[ (1/2(())
2E 20 / exp( 1 E2 (x)) (r((V=s(xx)) + ||s@ (x))l dx C,
(14)
where C dm)
(an)a exp(-2r (c)))))) (=)))(a)
( do
F RRawwp//// / I w2 exp(- (21 F.r. (x:))
a constant
independent of 0. In practice, using small T makes the loss too small, we thus replace 27 with hyperparameter 3. We also subtract the sample mean of local energies to improve numerical stabil- ity. Put everything together, given a batch of M samples, we define the weighted score matching (WSM) objective as:
1 Ex((x) - (1r(Va.0(0))) + [|s0(a))1))) + V(x)
(15)
M
E ai
(2;) = (xx)) - (EL) = (x,)
Eg (x;)
(16)
M
M
exp (-3 (axi))
WSM(O) 2
tr V so(x)) ) + [0:))))))
(17)
M 1= exp (-3E diar (a::)
Σ
softmax (-3E amn (x:)
[tr(V (0:)))))) + (x::))
(18)
Docs not differentiate w.E.L.
Differentatee w.r.t. 0
where softmax(a); Σ, exp(a,) Note that the weighting terms are treated constant when differentiating the loss w.r.t. the parameters O. Only the ISM terms require computing gradient during back-propagatton. This weighting scheme is similar to the attention mechanism in machine learning where the attention scores are based on the local energies and ,3 plays a similar role to the temperature variable in the Gumbel softmax (Jang et al., 2016). As interesting observation, the ISM (Equation II is -4 times the kinetic part in local energy. So we only need to compute one of them and reuse the computation.
exp(as)
We can prove that WSM is unbiased by using the zero-variance property of ground states. At ground (e(
(z) = (z)) Hence E (x) = va(=) Eo. Consequently, the WSM states we have (v)
weights will be identical at ground states and equal to due to softmax. As a result, the WSM loss reduces to the regular ISM loss at the ground state. Moreover, thanks to Langevin dynamies, every state is a local minimum of the ISM loss because the estimated score is the same with the true score of the sample distribution. Hence the ground state local minimum of the WSM loss.
Also due to the zero-variance property, when close to ground states, the weights should be close 1. We can create more imbalanced weights by further normalizing Eas by dividing its standard deviation. We thus have the Scaled- WSM defined as:
M
softmax (-8 Eanr(as)) std (Earr)
Cealed-WSMIp =
(r((vso((z)) + (=
(19)
As in VMC, we update according to the loss for one step, then we run Langevin dynamics for several steps to make samples follow the updated score. In our experiments, we use B ~ 1. So the short time approximation from imaginary time evolution does not hold strictly. However, since each time we only update parameters for small step size and that the target wavefunction s is constantly updated, the approximation may still be valid to some extend. But this should be more as compare
an intuition than an proof. To have further intuition for the convergence behavior, the WSM loss to the energy gradient 3 In fact, in both cases, we try to increase the probability density for regions with smaller local energies and decrease the probability density for regions with

---

Under review as a conference paper at ICLR 2023
larger local energies. However, by far do not have rigorous proof for the convergence property. Nevertheless, in our experiments, systems converge to ground states consistently with the proposed WSM loss.
3.4 THE DIFFUSION VARIATIONAL MONTE CARLO ALGORITIMM
The overall procedure is similar to VMC. We first perform several steps of Langevin dynamics to equilibrate the sampling. Then compute the weighted score matching loss and update the network parameters via gradient descent. The algorithm is summarized in Algorithm
Algorithm Diffusion Variational Monte Carlo (DiffVMC)
Input: Randomly initialized sample æ, score network 80, step size a, number of iterations T, number of Langevin dynamics step Ni
Output: New sample æ; optimized score network 89
for t 1 T do
for i = 1 to N₂ la do
x + Vae + (s)( (O.
D Langevin dynamics (Section 3.2
P = exp ( (1]s ((²))
Sample ~ (0,
if ncc then
x x²' end if end for loss Sealed
D Accept the move
b Weighted score matching objective (Section 3.3)
Update A via gradient descent to minimize the loss.
end for
EXPERIMENTS
We first show that our score-based method correctly finds the ground state by solving the harmonic trap. Then we showcase the applicability of our method on simple atomic systems, i.e. interacting fermions with coulomb potential. We will release our code after the review process.
4.1 WAVEFUNCTION ANSATZ FOR BOSONS AND FERMIONS
Since the unmormalized density (-) is invariant under exchange of particle positions (Section 3.1 there are two cases for wavefunction values. The first case is that the wavefunction does not change its sign, iee, , 2; ......) = bosons. The second case is that the wavefunction change its sign, i.e., v'( ( T. ... .j....) -v...., ....... mentally change ground states. The fermion ground state must consider the antisymmetric constraint and has higher ground state energy than the boson ground state.
Dt.... such particles are called
......., Ti,...), ), such particles are called fermions. These different symmetries will funda-
Bosons. We use a feed-forward neural network (Figure Baj composed of multi-layer perceptrons (MLPs) which satisfies the permutation equivariance (Zaheer et al. 2017 Qi et al. 2017). Input features are coordinates and distances the center. Each particle is simultaneously transformed into two feature vectors with two different MLPs. We perform average pooling for the first feature vector over all particles and obtain a global feature. The global feature is then concatenated with the second feature vector. A final MI.P is used to predict the score of each particle. All MLPs are shared among different particles.
The antisymmetric property of fermions is very challenging to model (Ceperley) 1991). Different from bosons, due to the change sign, the fermion wavefunction has both positive regions and negative regions. The region where the wavefunction equals zero is called the nodal surface (or the node). 1-d, the nodal surface exactly the set where two particles coincide. However, for higher dimension the nodal surface be arbitrary. The fermion sign structure gives rise some intrinsic difficulties to parametrize the score function, The score function diverges inversely propor- tionally the distance away from the nodal surface (Umrigar et al. 1993). Since neural networks
Fermions.
7

---

Under review as a conference paper at ICLR 2023
Enoti (EL)
Elorggy (FR))
200
100
600
800
1003
1000
2000
5000
6000
5000
0000
7000
5000
Iteration
Iteration
(a) Bosons
(b) Fermions
Figure 2: Train energies for bosons and fermions 2d quantum harmonic trap. N is the number of particles. All runs converge correctly to ground states, despite slight fluctuations for fermions.
essentially models continuous functions, using feed-forward neural networks to model the score function is prohibitive. In fact, even the networks manage to approximately model the discontinuity, in order to compute the local energy also requires the derivatives of the score to be accurate, which is more difficult.
Nevertheless, are still able to illustrate how score-based framework works with fermions. We can do so by modeling the score as the gradient an antisymmetric wavefunction Ansatz. We use the score computed from FermiNet (Pfau et al. 2020), where Slater determinants ensure the antisymmetry. We call Ferminet (Figure 3b Although by doing so the wavefunc- tion value is practically computed, fermion setting. Currently this stands as the only feasible option to overcome fermion sign problem. More generally we can model scores for fermions as the sum of the gradient of antisymmetric wavefunction and symmetric wavefunction:
employed to
goal is evaluate the method in this more challenging
Feemion
x Vu (x:) (x, x.
(20)
In our setting, fo, is FermiNet and we can model ge, using a feed-forward network (FFN). In our experiment, FFN uses FermiNet encoder to get feature vector for each particle, which is then mapped to score via linear layer. The output of FermiNet score and FFN score are summed
together to get the final score. We call it Ve.FeeriNNe+FFF (Figure 3c density vanishes boundary.
We do - ensure
4.2 2D QUANTUM HARMONIC TRAP
The Schrödinger equation with harmonic potential
VOHO (x) - -
(21)
describes quantum harmonic oscillator, of the most famous quantum systems that is solv- able. The energy levels and eigenstates are known analytically. In particular, for n bosons in the dimensional harmonic potential, the ground state energy and wavefunction are E ad (1b) 20d73 exp(- ||= Therefore, the quantum harmonic oscillator provides an analytical benchmark for our numerical method. Furthermore, despite its simplicity, the quantum harmonic oscillator is of important experimental relevance and can realized in cold atom systems by more com-
trapping atoms using lasers (Dalfovo al. 1999). Ground states for fermions plex due to the antisymmetry constraint. The ground state energies are E, = En= e; where Slater determinants of Hermite
1, 2, 3, 3, and the ground state wavefunctions polynomials. V FermiiiNe+FFN used here. We conduct simulation with various number parti- cles. The results shown in Figure All runs converge correctly to ground states.
8

---

Under review as a conference paper at ICLR 2023
4.3 ATOMIC SYSTEMS
For atoms, we work in the Bom-Oppenheimer where we assume the atoms are fixed in space and only the electrons allowed move. Hence, inputs are the coordinates of the electrons. The potential due to the Coulomb interactions is:
(x)
(22)
Atomn
2
23
}>1
where Z is the atom charge. As mentioned above, compute the score by taking the gradient of FermiNet (Pfau et al. 2020). Following Lin et al. (2021), four atoms are tested, including Boron, Carbon, Nitrogen and Oxygen. To show the effect of step size, We trained and tested the score network with a=le-3 and a=le-2. We and optimize for 200k iterations for all systems. We use 8=2 for Oxygen and B=1 for all other atoms. We do Langevin dynamics for 20 steps between two parameter update. The norm of the score is clipped at 20 to increase numerical stability. The complete hyperparameters are Appendix D
VeremiNeet VereeniNet a=le-3
VeFermiNetFFFN a=le-3
FermiNet
HF
Reference
a=le-2
B 224.64893(22)
22..65233(3)
22..63537(46)
-24.65370(3)
22.533316
-24.63391
C 37.84414(13)
33.83445(5)
-37.83736(43)
-37.84471(5)
-37.6938
-37.8450
N -4.58177((9)
-54.53580(12)
-54.58221(28)
-54.58882(6)
-54.4047
-54.5892
0 -7..06587(63) -74.94880(709) -75.06410(53)
-75.06655(7)
-74.8192
-75.0673
Table 1: Experimental results on Atoms in Hartree (Ex). Underline denotes our results within
chemical accuracy. References taken from Lin et al. (2021).
The results are summarized on Table The chemical accuracy is defined as 1.594 mE, (Pfau al. 2020). Except for the Nitrogen atom, all atoms can enter the chemical accuracy under one setting, although the effect of step size is mixed. In many cases a higher level of variance is observed. The reason may be due to the difficulty in optimizing the gradient network. Moreover, FermiNet works with the second order optimizer KFAC (Martens & Grosse 2015), which is adapted for wavefunction outputs and may not be suitable for our socre-based optimization. Nevertheless, demonstrated the correct convergence behaviour for the challenging electronic potential. The results for N are improved with FFN. For other atoms the results this setting we show that in terms of network architectures we can go beyond the FermiNet score.
framework has
improved but are comparable. With
5 COMPARISON TO VMC AND DMC
Compared to DMC, DifIVMC closer to VMC's paradigm where Ansatz is optimized by estimating loss based on samples following Ansatz distribution. With DiffVMC are able directly model scores, which are more fundamental to the optimization problem and could be more expressive than modeling wavefunctions. On the other hand, DiffVMC and DMC both use scores update samples. DMC starts with an optimized trial wavefunction Ansatz. Energy minimization is achieved by constantly adjusting weights of walkers. However, the walkers are always guided by the trial function's score and the parameters of the trial Ansatz cannot be updated DMC. In contrast, Di is able to update the guiding score function directly. Nevertheless, there is no conflict between MC and DMC. DMC walkers can be guided with a score optimized with DiffVMC and further project toward ground state through weighting.
CONCLUSION
Inspired by the connection between the score-based formulation of local energy, explore the possibility to implicitly model the quantum wavefunction. With the weighted score matching objec- tive, the proposed DC enables the possibility to optimize the score network toward the ground state. Experiments show that our proposed method can accurately find the ground state for both bosons and fermions.

---

Under review as a conference paper at ICLR 2023
REFERENCES
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
//github. ctttuu..oo///ogg/e joagg////ao////oo//////a
Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale gan training for high fidelity natural image synthesis. arXiv preprint arXiv:1809.1109 2018.
Giuseppe Carleo and Matthias Troyer. Solving the quantum many-body problem with artificial neural networks. Science, 355(6325))602-60)) 2017.
David Ceperley, Geoffrey V Chester, and Malvin Kalos. Monte carlo simulation of a many- fermion study. Physical Review B, 16(7):3081, 1977.
David M Ceperley. Fermion nodes. Journal of statistical physics, 63(5):1237-1267, 1991.
David M Ceperley. Introduction quantum monte carlo methods applied to the electron gas. In The Electron Liquid Paradigm in Condensed Matter Physics, pp. 3-42. IOS Press, 2004.
David M Ceperley and Berni J Alder. Ground state of the electron gas by stochastic method. Physical review letters, 45(7)5566, 1980.
Franco Dalfovo, Stefano Giorgini, Lev Pitaevskii, and Sandro Stringari. Theory of bose-einstein condensation in trapped gases. Reviews of modern physics, 71(3):463, 1999.
Mike Entwistle, Zeno Schätzle, Paolo A Erdman, Jan Hermann, and Frank Noé. Electronic excited states in deep variational monte carlo. arXiv preprint (-Xiv:2203,09477 2022.
Cong Fu, Xuan Zhang, Huixin Zhang, Hongyi Ling, Shenglong Xu, and Shuiwang Ji. Lattice convolutional networks for learning ground states of quantum many-body systems. arXiv preprint arXiv:2206.0733 2022.
Nicholas Gao and Stephan Günnemann. Ab-initio potential energy surfaces by pairing gnns with neural wave functions. In International Conference on Learning Representations, 2021.
Nicholas Gao and Stephan Günneemann Sampling-feee free inference for ab-initio potential energy sur- face networks. arXiv preprint arXiv:2205. 14962, 2022.
Johannes Gasteiger, Janek GroB, and Stephan Günnemann. Directional message passing for molec- ular graphs. In International Conference on Learning Representations (ICLR), 2020.
Leon Gerard, Michael Scherbela, Philipp Marquetand, and Philipp Grohs. Gold-standard solutions to the schrödinger equation using deep learning: How much physics do we need? arXiv preprint 0:Xiv:2205.009438, 2022.
Ian J Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. 2014.
Jiequn Han, Linfeng Zhang, and E Weinan. Solving many-electron schrödinger equation using deep neural networks. Journal of Computational Physics, 399: 2019.
Jan Hermann, Zeno Schätzle, and Frank Noé. Deep-neural-networkk solution of the electronic schrödinger equation. Nature Chemistry, 12(10):9118877, 2020.
Jan Hermann, James Spencer, Kenny Choo, Antonio Mezzacapo, WMC Foulkes, David Pfau, Giuseppe Carleo, and Frank Noé. Ab-initio quantum chemistry with neural-network wavefunc- tions. arXiv preprint arXiv: 220. 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
Aapo Hyvärinen and Peter Dayan, Estimation of non non-normalized statistical models by score match- ing, Journal of Machine Learning Research, 6(4), 2005.
10

---

Under review as a conference paper at ICLR 2023
David Pfau James S. Spencer and FermiNet Contributors. FermiNet, 2020.
URL http://
github. cot.uuu..comm/eemminndn.. ferpaapmiind/eeminnee
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizztioon with gumbel-softmax, arXiv
preprint arXiv:1611.00 1.0.111101100 .0..01144 144, 2016.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. International Conference on Learning Representations, 2014.
Dmitrii Kochkov, Tobias Pfaff, Alvaro Sanchez-Gonzalez, Peter Battaglia, and Bryan K Clark. Learning ground states of quantum hamiltonians with graph networks. arXiv preprint 10.06390 2021.
Jeffmmi Lin, Gil Goldshlager, and Lin Lin. Explicitly antisymmetrizee neural network layers for variational monte carlo simulation. arXiv preprint arXiv:21 2021.
Yi Liu, Limei Wang, Meng Liu, Yuchao Lin, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical message passing for 3d molecular graphs. In International Conference Learning Representa- tions, 2022. URL /
net wwt.:e/oorrmllddiiieeeee.eo.
James Martens and Roger Grosse. Optimizing neural networks with Lronecker-factoree approximate curvature. In International conference on machine learning, pp. 2408-2417. PMLR, 2015.
Tianyu Pang, Kun Xu, Chongxuan Li, Yang Song, Stefano Ermon, and Jun Zhu. Efficient learn- score matching. Advances in Neural Information
ing of generative models via Processing Systems, 33:19175- 2020.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary Veti,o, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
D. Pfau, J.S. Spencer, A.G. de G. Matthews, and W.M.C. Foulkes. Ab-initio solution of the many-clectron schrödinger equation with deep neural networks. Phys. Rev. Research, 2:033429, 2020. doi: 10. Rew.eehh.. 1103/PPR htte///yeeeeeeeeaechh.o
(Rthttt::///iii was.kkanns............m...m....a org/ /oi////
033429
Charles R Qi, Hao Su, Kaichun Mo, and Leonidas Guibas. Pointnet: Deep learning on point sets for 3d classification and segmentation. In Proceedings the IEEE conference compuler vision and pattern recognition, pp. 652-660, 2017.
Weiluo Ren, Weizhong Fu, and Chen. Towards the ground state molecules via diffusion monte carlo neural networks. arXiv preprint arXiv:2204. 13903, 2022.
Michael Scherbela, Rafael Reisenhofer, Leon Gerard, Philipp Marquetand, and Philipp Grohs. Solv- ing the electronic schrödinger equation for multiple nuclear geometries with weight-sharing deep neural networks. Nature Computational Science, 2(5):331-341 2022.
Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Rober Müller, Schnet: continuous-fin convolutional neural network for modeling quantum interactions. Advances in neural information processing systems, 30, 2017.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan. and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learn- ing. pp. 2256-2265 PMLR, 2015.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 32, 2019.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna- tional Conference on Learning Representations, 2020.
Julien Toulouse, Roland Assaraf, and Cyrus J Umrigar. Introduction to the variational and diffusion monte carlo methods. Advances in Quantum Chemistry, volume 73, pp. 285-314, Elsevier, 2016.
11

---

Under review as a conference paper at ICLR 2023
CJ Umrigar, MP Nightingale, and KJ Runge. A diffusion monte carlo algorithm with very small time-step errors. The Journal of chemical physics. 99(4):2865-288) 2890, 1993.
Limei Wang, Yi Liu, Yuchao Lin, Haoran Liu, and Shuiwang Ji. ComENet: Towards complete and efficient message passing for 3D molecular graphs. In The 36th Annual Conference on Neural Information Processing Systems, 2022.
Max Wilson, Nicholas Gao, Filip Wudarski, Eleanor Rieffel, and Norm M Tubman. Simulations of state-of-the-art fermionic neural network wave functions with diffusion monte carlo. arXiv preprint arXiv:2103.125700 2021.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander Smola. Deep sets, Advances in neural information processing systems, 30, 2017.
A COMPUTATIONAL COMPLEXITY
Compared to VMC, DiffVMC avoids recomputing gradient of wavefunctions during energy evalu- ation, but the loss is more complicated. So the overall computational cost should be similar with VMC. In our experiments with FermiNet, DiffVMC is slower than VMC because need to back- propagate through the gradient network. However, this should be mitigated with more adapted net- work architectures. In general, the major computation bottleneck of QMC is evaluating tr( tr(V (x)) which cannot be paralleled efficiently. This is a common problem for all QMC methods. However, by noticing the connection between kinetic energy and score matching, we might be able accel- erate this computation via more efficient score matching techniques (Song & Ermon
2019
Song
al. 2020 Pang et al. 2020).
B PROOF FOR EQUIVARIANCE
We show it for two 1-d particles:
log 11( + Ar, x2)l
log (1((,1 #2)
lim 0 log (11,
(23)
Ax-+0
Ar
log (1222, + Ar) log (12,
lim
0g log (22,
(24)
Ar
Ar
C DERIVATION OF THE APPROXIMATED DETAILED BALANCING
The original detailed balancing in Langevin Monte Carlo is
x' so(x )a) 12) th (x')² (a(x)a?? 20(x)²
P Pace (x' (x)
exp
2x
(25)
exp
x
x
In our score-based framework, generally do not have access to wavefunction value. So our goal here is to get rid of the explicit dependencies on (x') and the (x). We can do so by assuming the score is constant between æ and x'. We first transform into the log domain:
log Pace
(||æ x' 8o(x")al12 ||l x so(x)al)3) + 2(log (wo')') - log (() (xx) 2x
(26)
where we use A and B to denote the two components and log Pase = A+ B.
We first simplify A. By expanding the norms as | x' s0(x (x 'aa 2
1 x x112 2(x
x'. Sp (=())a + (x))) and | - x so(a)a|12 = || x
2(xx'
x, 30 (0
[|s(())) ||a,, we obtain:
12

---

Under review as a conference paper at ICLR 2023
A = (a x', 80 + s0(x)) + a
(11=0(=))]1 (=0()=)
(27)
To simplify B, we use the approximation:
log (2)((')) - log [10((=)]
=
2 (V log ((o
V log (x))) x' -x)
(28)
2 (00(2') + s0())) x)
(29)
Such approximation is commonly employed in finite difference methods. The approximation should be valid since a and are close space. With this approximation, we have:
= 2(log [='((=') log (() = -- x', ss((o + so(x:)
(30)
Finally,
log Pace = + B a [|| ((0 (13)
(31)
2
D COMPUTATIONAL SETTINGS
We implement DifIVMC in Pytorch (Paszke et al.j 2017) for experiments with Bosons and in
Jax (Bradbury et al. 2018) for experiments with Fermions. All gradient computation is done with automatic differentiation from the respective libraries. For experiments with FermiNet, we adapt based on the official implementation of FermiNet (James S. Spencer & Contributors 2020
Hidden dim
32
Network
# MLP layers
Optimizer
Adam
Batch size
256
Learning rate Optimization steps
5e-4
Optimization
2000
Clip local energy WSM B
std
Langevin dynamics steps Metroplis Hasting rejection Step size Score norm clip
20
Approximate (Equation
Langevin dynamies
0.01
Table 2: Hyperparameters for Bosons
13

---

Under review as a conference paper at ICLR 2023
# determinants
16 (for atoms) or I (for QHO)
Single layer hidden dim Double layer hidden dim # embedding layers
256
FermiNet
32
Single layer hidden dim Double layer hidden dim # embedding layers
64
Side network (FFN)
8
4
Pertraining steps
500 (for B, C) or 1000 (for N, O)
Batch size
256 (for B, C) or 512 (for N, O)
Optimizer
Kfac
Optimization
Kfac damping
0.001
Kfac norm constraint
0.001
Initial learning rate
5e-2
Learning rate decay
Imi (1 + 1/10000)
Clip local energy
5 x std
WSM
1 or 2
Langevin dynamics steps Metroplis Hasting rejection Step size a Optimization steps Score norm clip
20
Approximate (Equation 9
Langevin dynamics
0.001 or 0.01 200,000
20
Table 3: Hyperparameters for Fermions
[x, ||x||]] E (x(d+1)
X E Vxd
x E
IR Rxd
MLP 1
Mean pool
MLP 2
Broadcast
FermiNet Encoding
Concat
FermiNet
FermiNet
MLP 3
I Nxd
Linear
V
Linear
SiLU
V
MLP
Linear
SiLU
5
IR Nx
Linear
E Nxdd
(a) Network for Bosons
FerFenniNet
V FermiiNeeLFNN
Figure 3: Network architectures. For simplicity, - is omitted in V FemmiNeetFFNN
14


Reviews:
[{"review_id": "xw90xbxC3W", "reviewer": "Reviewer_PhHJ", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work proposes a score-based method for learning neural networks to model wavefunctions minimize quantum potentials. The work adapts an objective function from the VMC model that minimizes local energy under current model samples. The main innovation is the proposal to model only the gradients of the log wave function instead of the wave function itself, which is motivated by an insight into the functional form of the local energy. DMC weighting is introduced to overcome optimization difficulties that arise when the wavefunction is not directly available. The proposed method is evaluated on a boson potential and a challenging fermion potential, and experimental results show a good match to physical observations.", "summary_of_the_review": "I enjoyed reading this paper and appreciate the efforts to develop deep learning methods that are useful in computational quantum physics. I have little knowledge of quantum physics so I can only comment on ML aspects of the paper. Replacing the explicit wave function with a score network is technically relevant and mirrors similar developments of EBM and score models in deep generative modeling. The integration of VMC and DMC is interesting but would be stronger with a more principled formulation if possible. I cannot provide any context on the quality or significance of the experimental results. A major experimental weakness is that the work uses a network with explicit wavefunction for the central experiment even though the motivation of the work is to avoid explicitly modeling the wave function. It seems there is currently no suitable architecture for the proposed method. Overall, I recommend the paper, but I am on the fence.", "strength_and_weaknesses": "STRENGTHS:\n\n1. The work is based on an interesting insight into the functional form of the wavefunction that appears in the local energy loss. Replacing the gradient of an explicit wavefunction with a function that directly estimates the score function is a compelling proposal that is aligned with similar developments in generative modeling.\n2. The hybridization of DMC and VMC is an interesting solution to the optimization difficulties that arise when replacing the explicit wave function with the score function. DMC provides a way to move towards states with low potential, while VMC provides a more principled probabilistic proposal for DMC based on the score matching objective.\n3. The work explores an exciting area at the intersection of deep learning and computational physics.\n\nWEAKNESSES:\n\n3. The Ferminet network used for experiments already provides an estimate for $\\psi_0$, so the motivation for score-based modeling is not as clear. What is the motivation or benefit of learning the gradient of a scalar-valued net, instead of learning a network that outputs the estimated gradient directly as typically done in score-based modeling? The fermion restrictions are a difficult challenge with limited options available and the proposed method might be useful for future architectures, but at the present it seems roundabout to use score-based modeling for a scalar valued net.\n4. The experimental results appear somewhat limited, where only a Gaussian-like density and fermion density are investigated. However, I am not familiar with the quantum physics literature and cannot provide perspective on the scope or significance of the results. The fermion results do not match the performance of Ferminet, which might limit the appeal of the proposed method especially since the same architecture is used.\n5. While interesting, the formulation of the WSM objective might be somewhat ad-hoc. Is there a more principled way to integrate DMC and VMC beyond an intuitive connection?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "oMXTr1U_B2O", "reviewer": "Reviewer_FLe5", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a method for solving the quantum many body problem using neural networks as variational ansatz. In contrast to more standard approaches which model the wave function directly, authors use a score-based model (hence modeling the gradient of the underlying probability distribution associated with the wavefunction). Authors draw on techniques from score based modeling and diffusion Monte Carlo to formulate a schema for training such models. They demonstrate that the proposed approach works well for the bosonic quantum harmonic trap. Authors also achieve good results for fermionic atomic systems by using gradients of a FermiNet as a model for the score function.", "summary_of_the_review": "I enjoyed reading the paper and I think the direction taken by the authors is both promising and interesting. Results and ideas are presented clearly and concisely. Authors honestly point out challenges associated with nodal structure in fermionic systems.\n\nWhile I’m very excited about the field and the research direction considered in this paper, I’m a bit concerned that the proposed method addresses problems that are already “solved” by classical techniques such as diffusion Monte Carlo (DMC). In most cases the resort to parametric variational Monte Carlo (VMC) methods is motivated by the sign problem which results in poor estimates in DMC techniques. This is however an exciting direction and would be of great interest to a more narrow audience in computational physics (and a potentially major milestone if such method could be effectively generalized to fermionic systems). Hence I’m leaning towards the score of (5), but could be convinced otherwise.\n", "strength_and_weaknesses": "Strengths:\n* The paper studies an underrepresented approach to modeling quantum states based on gradients of the logarithm of the wave function amplitude.\n* Paper explains methodology clearly and accurately\n\nWeaknesses:\n* The proposed method seems to be most suited for problems without the “sign problem”. In this setting diffusion monte carlo and similar methods already provide very good results and then at the very least should be considered as a baseline.\n* Some approximations made to set up the optimization loop are not fully justified (or I might have missed the argument, e.g. eqs. 14->15;)\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "eOHl1P-aew1", "reviewer": "Reviewer_Lxrt", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors propose Diffusion Variational Monte Carlo (DVMC) for simulating many-body quantum systems based on the neural score function. Contrast to a conventional Variational Monte Carlo (VMC), the proposed method requires estimating the score function only, which is a gradient of the logarithm of the wavefunction, rather than evaluating the wavefunction itself. To this end, the authors formulate a new loss function, named weight score matching (WSM, which is analogous to the score matching used for energy-based models (EBMs)), that can learn a ground state of the target wavefunction with the score function only. By combinding such a loss function with the Lagevin dynamics, the proposed DVMC framework samples and learns the targeted ground state simultaneuosly. The authors validate their approach for two quantum systems.", "summary_of_the_review": "Although the paper is interesting, I think there is still much room for improvement. I think the paper is not ready for publication in its current form.", "strength_and_weaknesses": "Strength:\n\n1. The motivation of the propsed framework, i.e., learning and sampling ground states by using the score function only, is clear. First, MCMC step of VMC learning can be removed. It takes a similar advantage as by-passing expensive MCMC likelihood estimation in the score matching method of EBMs. Second, as the authors mention, one can omit the necessity of computing the score function from the wavefunction.\n\n2. The authors introduce a permutation-equivariant architecture [1] to guarantee the indistinguishability of quantum mechanical particles. To me, it is a novel application of [1] (although not sure because I am not an expert of machine-learning-based quantum simulations).\n\nWeakness:\n1. The authors derive the local energy of QMC, and conclude its kinetic term is equal to the implicit score matching (ISM) [2] (up to a constant). However, it is not true considering a coefficient of the norm of the score function $||s_{\\theta}(\\mathbf{x})||^2$. Note that the kinetic term of the local energy of QMC is given by $\\text{tr}(\\nabla_{\\mathbf{x}}s_{\\theta}(\\mathbf{x})) + ||s_{\\theta}(\\mathbf{x})||^2$ while the original ISM is $\\text{tr}(\\nabla_{\\mathbf{x}}s_{\\theta}(\\mathbf{x})) + 0.5||s_{\\theta}(\\mathbf{x})||^2$. Thus, it is not entirely clear the connection between the local energy and ISM. \n\n2. It seems that the major advantage of the proposed DVMC is its learning and sampling efficiency. However, there is no benchmark on the computational efficiency of DVMC.\n\n3. The authors introduce four assumptions (in page 5) to eliminating the wavefunction-dependency and distributional shift of the porposed framework. However, they do not dicuss whether the introduced assumptions are reasonable (acceptable) or not.\n\n4. For fermion systems, the proposed method requires computing the wavefunction. Therefore, the main claim of this paper is partially applicable (only for boson systems). \n\n5. Experimental results are weak. The harmonic trap example is very simple and can be solved analytically. Atomic systems (B/C/N/O) are more intersting, but they are fermion systems, which doe not support the main claim of this paper as I mentioned in 5.\n\n6. The authors calim \"numerous examples in machine learning have shown that implicitly modeling the data distribution provides better representation\" and \"properly chosen latent embedding can better describe the struvture of the data space\" with refering VAE and GAN literatures in page 1 of the paper. Because the used algorithm (score matching) is not based on the latent representation and the used data (quantum state) seems to be not an embedded manifold, such claims are not entirely clear to me.\n\nSome minor issues:\n1. In (3), it seems that $\\mathbf{x} \\sim \\psi_{\\theta}^{2}(\\mathbf{x})$ should be $\\mathbf{x} \\sim \\psi_{\\theta}^{2}(\\mathbf{x}) / \\int \\psi_{\\theta}^{2}(\\mathbf{x}) d\\mathbf{x}$.\n2. In (7), the ISM should be $\\mathbb{E} \\[ \\text{tr}(\\nabla_{\\mathbf{x}}s_{\\theta}(\\mathbf{x})) + 0.5||s_{\\theta}(\\mathbf{x})||^2 \\]$.\n3. In page 6, the authors mention \"The equivariance can be easily achieved by considering the input coordinates as a set\". Please elarborate it more detaily or add a reference (might be [1]?)\n4. The potential $V(\\mathbf{x})$ is not included in (18 - 19).\n\n** *\n[1] Zaheer, M., Kottur, S., Ravanbakhsh, S., Poczos, B., Salakhutdinov, R. R., & Smola, A. J. (2017). Deep sets. Advances in neural information processing systems, 30.\n\n[2] Hyvärinen, A., & Dayan, P. (2005). Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4).\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "JTDNQiYcDDM", "reviewer": "Reviewer_afAv", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Consider the problem of finding the ground-state of the Hamiltonian $H = -\\frac{1}{2} \\nabla^2 + V$, i.e. its lowest eigenfunction $\\psi$. This corresponds to minimising the loss function $\\mathcal{L}(\\theta) = \\mathbb{E}_{|\\psi|^2}\\left[ E(x,\\psi) \\right]$ where $E(x, \\psi) = \\frac{H \\psi(x)}{\\psi(x)}$ and the wave-function $\\psi$ is parametrised with a neural-network with weight $\\theta \\in \\Theta$. We have that\n\n$$ \\nabla \\mathcal{L}(\\theta) = \\mathbb{E}_{|\\psi|^2}\\left[  \\left(E(x,\\psi)  -  \\mathcal{L}(\\theta) \\right)  \\nabla \\log |\\psi(x)|  \\right].$$\n\nInstead of parametrising the wave function, the articles proposes to parametrize (half-of) the score $s(x) = \\nabla \\log |\\psi(x)|$ of the probability density proportional to $|\\psi^2|$. In order to learn the score of the ground-state, the proposed approach maintains at iteration $t \\geq 0$ an approximation of the score function $s^{(t)} = \\nabla \\log |\\psi^{(t)}|$ as well as a population of $N$ particles $[x^t_i]_{i=1}^N$ approximately distributed according to $|\\psi|^2$.\n\n1. at iteration $t \\geq 0$, for a current approximation of the score $s^{(t)}$ and population $[x^t_i]_{i=1}^N$, use (approximate) Langevin dynamics with $s^{(t)}$ and finally compute the usual Diffusion-Monte-Carlo (DMC) weights. This procedure gives a slightly better approximation of the ground state expressed as a weighted population of particles (i.e. it is the working principle of DMC)\n2. Instead of resampling with these weights as a standard DMC method would do, use (weighted) score matching to get an approximation go this weighted population of particles. Since the weighted population of particles was a better approximation of the ground state, this results in a better estimate of the score of the ground state.\n3. iterate\n\n", "summary_of_the_review": "In many settings, modelling the score of a distribution has advantages when compared to modelling the distribution itself (eg. generative modelling). The text proposes to investigate this claim in the context of finding the ground state of the Schrodinger equation.\n\nAfter reading the text, it is not entirely clear to me whether modelling the score has advantages over modelling the wave-function itself. Still,  the text seems to demonstrate that this approach is viable and deserves more numerical investigations.", "strength_and_weaknesses": "**Disclaimer:** I am not an expert in ML methods for computing the ground state of quantum systems. I am **not** on top of the current literature.\n\n**Strength:** the paper is accessible and reads reasonably well. I believe that the fundamental motivation for the proposed methodology is the claim that modelling the score of the ground state distribution is somehow easier than modelling the distribution itself (or equivalently, parametrising the score is easier that parametrising the wave function with a neural net). I am finding it interesting that the authors are trying to study this claim. I do not know the literature well enough to comment on the novelty of the approach.\n\n**Weakness:**\n1. parametrising the score (instead of the wave function) means that it is not possible to use an exact Metropolis-Hastings (MH) accept/reject. The authors choose to completely ignore the problematic quantities within the MH and I am finding this step a bit hard to swallow. More comments and simulations would be welcome.\n2. All the quantities in Equation (4) can be expressed in term of the score: the function $E$ is a function of the score, and the gradient term is exactly the score. Why not discretize this quantity to get an approximation of $\\nabla \\mathcal{L}(\\theta)$ and update the score? Doing this, we do not even need to compute the DMC weights, no?\n3. In order to substantiate the claim that modelling the score is better than directly modelling the wave function, the authors should run the method that consists in evolving a population of $N$ particles with the Langevin-dynamics with **exact** accept-reject, with the same computation of the DMC weights, and update of wave function through Equation (4).\n4. It is a bit disappointing that the authors are still relying (as they readily admit it) on the parametrisation of the wave-function (i.e. FermiNet parametrisation), which seems to defeat the whole purpose of the whole study. \n\nAgain, I would like to re-iterate that I am not an expert in this domain (and may have mis-understood some aspects of the proposed method) and I am looking forward to hearing the authors comments.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject