Wasserstein Quantum Monte Carlo: A Novel Approach for Solving the Quantum Many-Body Schrödinger Equation
Authors: Kirill Neklyudov
Jannes Nys
Luca Thiede
Juan Felipe Carrasquilla Alvarez
qiang liu
Max Welling
Alireza Makhzani
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Quantum Monte Carlo
Schrödinger equation
Wasserstein Fisher-Rao gradient flow
TL;DR: We minimize the energy functional of quantum systems by following a projected Wasserstein gradient flow in the space of distributions, which results in faster convergence compared to the conventional Quantum Variational Monte Carlo.

Abstract: Solving the quantum many-body Schrödinger equation is a fundamental and challenging problem in the fields of quantum physics, quantum chemistry, and material sciences. One of the common computational approaches to this problem is Quantum Variational Monte Carlo (QVMC), in which ground-state solutions are obtained by minimizing the energy of the system within a restricted family of parameterized wave functions. Deep learning methods partially address the limitations of traditional QVMC by representing a rich family of wave functions in terms of neural networks. However, the optimization objective in QVMC remains notoriously hard to minimize and requires second-order optimization methods such as natural gradient. In this paper, we first reformulate energy functional minimization in the space of Born distributions corresponding to particle-permutation (anti-)symmetric wave functions, rather than the space of wave functions. We then interpret QVMC as the Fisher--Rao gradient flow in this distributional space, followed by a projection step onto the variational manifold. This perspective provides us with a principled framework to derive new QMC algorithms, by endowing the distributional space with better metrics, and following the projected gradient flow induced by those metrics. More specifically, we propose "Wasserstein Quantum Monte Carlo" (WQMC), which uses the gradient flow induced by the Wasserstein metric, rather than the Fisher--Rao metric, and corresponds to *transporting* the probability mass, rather than *teleporting* it. We demonstrate empirically that the dynamics of WQMC results in faster convergence to the ground state of molecular systems.

Wasserstein Quantum Monte Carlo:
A Novel Approach for Solving the Quantum Many-Body Schrödinger Equation
Kirill Neklyudov
Vector Institute
Jannes Nys
Luca Thiede
Institute of Physics & Center for Quantum Science and Engineering École Polytechnique Fédérale de Lausanne (EPFL)
Vector Institute
Univeristy of Toronto
Juan Felipe Carrasquilla
Qiang Liu UT Austin
Max Welling Microsoft Research
Alireza Makhzani
Vector Institute Univeristy of Waterloo
Vector Institute
A14Science
Univeristy of Toronto
Abstract
Solving the quantum many-body Schrödinger equation is a f'undamental and chal- lenging problem in the fields of quantum physics, quantum chemistry, and material sciences. One of the common computational approaches to this problem is Quan- tum Variational Monte Carlo (QVMC), in which ground-state solutions are obtained by minimizing the energy of the system within restricted family of parameterized wave functions. Deep learning methods partially address the limitations of tradi- tional QVMC by representing rich family of wave functions in terms of neural networks. However, the optimization objective in QVMC remains notoriously hard to minimize and requires second-order optimization methods such as natural gradi- ent. In this paper, we first reformulate energy functional minimization in the space of Born distributions corresponding functions, rather than the space wave functions. We then interpret QVMC as the Fisher-Rao Rao gradient flow in this distributional space, followed by projection step onto the variational manifold. This perspective provides us with a princi- pled framework to derive QMC algorithms, by endowing the distributional space with better metrics, and following the projected gradient flow induced by those metries. More specifically, we propose "Wasserstein Quantum Monte Carlo" (WQMC), which uses the gradient flow induced by the Wasserstein metric, rather than the Fisher-Rao metric, and corresponds transporting the probability mass, rather than teleporting it. We demonstrate empirically that the dynamics of WQMC results in faster convergence the ground state of molecular systems.
perrilee-eermutation (atti))symmetri wave
Introduction
Access the wave function of quantum many-body system allows us study strongly correlated quantum matter, starting from the fundamental building blocks. For example, the solution of the time- independent electronic Schrödinger equation provides all the chemical properties of given atomic state, which have numerous applications in chemistry and materials design. However, obtaining the exact wave function is fundamentally challenging, with complexity scaling exponentially with
37th Conference Neural Information Processing Systems (NeurIPS 2023). Correspondence to: <k. necluddvgggma.. com>,

---

the number of degrees of freedom. Various computational techniques have been developed in the past, including compression techniques based on Tensor Networks (White, 1992). and stochastic approaches such as Quantum Monte Carlo (QMC) (Ceperley et al., 1977). Quantum Variational Monte Carlo (QVMC) (McMillan, 1965; Ceperley et al., 1977) a well-known subclass of the latter that can, in principle, be used to estimate the lowest-energy state (i.e. ground state) of quantum many-body system. The method operates by parameterizing the trial wave function and minimizing the energy of the many-body system w.r.t. the model parameters.
The choice of parametric family of the trial wave function is a crucial component of the QVMC framework. Naturally, deep neural networks, being family of universal approximators, have demonstrated promising results for quantum systems with discrete (Carleo & Troyer, 2017; Choo et al., 2020; Hibat- Allah et al., 2020), as well as continuous degrees of freedom (Pfau et al., 2020; Hermann al., 2020; Pescia et al., 2022; Gnech et al., 2022; von Glehn et al., 2022). However, the optimization process is challenging, especially for rich parametric families of the trial wave functions. This requires the use of advanced optimization techniques that take into account the geometry of the parametric manifold. The most common technique used in QVMC is referred to as 'Stochastic Reconfiguration' (SR) (Sorella, 1998), and can be seen as the quantum version of Natural Gradient Descent (Stokes et al., 2020). While for large neural networks with up to millions of parameters, efficient and scalable implementations of SR are available (Vicentini et al., 2022), it is also possible to use approximate methods such as K-FAC (Martens & Grosse, 2015; Pfau et al., 2020). Higher order optimization techniques are considered to be essential to obtain the necessary optimization performance to accurately estimate ground states of quantum many-body Hamiltonians (see e.g. (Pescia et al., 2023; Pfau et al., 2020)). Therefore, studies of the optimization procedure are an important direction for further development of the QVMC approach.
In this paper, we consider the energy minimization dynamics as a gradient flow on the non-parametric manifold of distributions. First, an example of the proposed methodology, we demonstrate that the imaginary-time Schrödinger equation can be described as the gradient flow under the Fisher-Rao metric on the conpaaraetttic manifold. Then, the QVMC algorithm can be seen as a projection of this gradient flow onto a parametric manifold (see Section for details). Second, the gradient flow perspective gives us an additional degree of freedom in the algorithm. Namely, metric under which we define the gradient flow. Thus, we propose and study a different energy- minimizing objective function, which we derive gradient flow under the Wasserstein metric (or Wasserstein Fisher-Rao Rao metric) (Chizat et al., 2018; Kondratyev et al., 2016).
choose the
In practice, we demonstrate that incorporating the Wasserstein metric into the optimization procedure allows for faster convergence to the ground state. Namely, we demonstrate up to 10 times faster convergence of the variance of the local energy for chemical systems. Intuitively, incorporating the Wasserstein metric regularizes the density evolution by forbidding or regularizing non-local probability mass "teleportation" (as done by Fisher Rao metric). This might facilitate faster mixing of the MCMC rumning along with the density updates.
2 Background
2.1 Quantum variational Monte Carlo
Consider a quantum many-body system subject to the Hamiltonian operator, which we will assume to be of the following form,
1 - VV. H =
-
where x a given many-body configuration and V is the potential operator. The dime-dependent Schrödinger equation determines the wave function (xx t) of the quantum system
d (xx (x,
(2)
dt
As is often the case, we will target the stationary solutions, for which we focus on solving the time- inmeiindependent Schrödinger equation
He(()) =v)())
where the energy of the state sb The ground state of a quantum system is obtained by solving the time-independent Schrödinger equation, by targeting the eigenstate &e of the above Hamiltonian
2

---

with the lowest cigenvalue (energy) E. Hereby, we must restrict the Hilbert space to wave functions that are antisymmetric under particle permutations in the case of fermionic particles, and symmetric for bosons. The latter takes into account the indistinguishability of the particles. Given the Born density ((x:) = (x)), the energy of given quantum state can rewritten a functional form,
El pp (a) [F (x)), ((oc(2) H u
(x))
We will focus on the case where the Hamiltonian operator is Hermitian and time-reversal symmetric. In this case, its eigenfunctions and eigenvalues are real, and the energy can be recast into a functional of the Born probability density (see also Pfau et al. (2020), where the expressions are given in terms of loglehh 2+))
Ely] = Bg((s) [x)c) (x=
1 v log (
$ I log (())
(5)
under the strong condition that q(x) is the Born probability density derived from an (anti-)symmetric wave function: q(x) (x)) The latter will always be tacitly assumed from hereon.
The Rayleigh- Ritz principle guarantees that the E[g] is lower-bounded by the true ground-state energy of the system, i.e. Elg] 2 Eo, if the corresponding wave function & is valid state of the corresponding Hilbert space. Quantum Variational Monte Carlo (QVMC) targets ground states by parametrizing the wavefunction (b(r;, ()) and by minimizing (lg( The solution the minimiza- tion problem 0g = arg ming (0) is obtained by gradient- methods using the following expression for the gradient w.r.t. parameters O
VoElg(0)] BEgt..0)
(Ehuc(ar: 0) Egg(,,0) [E
() 0)] Vo log q(r, 0)
(6)
In sum, the above leads to an iterative procedure in which Monte Carlo sampling is used generate configurations from the current trial state q(2, f) = (22 0), which allows computing the corre- sponding energy and its parameter gradients, and to update the model accordingly. In practice, the parametric model specifies the density q(x, 0) only up to a normalization constant, i.e., it outputs q(x, 0) x q(z, 9). However, the gradient w.r.t. 0 does not depend on the normalization constant; hence, throughout the paper, we refer to the model as the normalized density ((x,,)) 0).
2.2 Gradient flows under the Wasserstein Fisher-Rao metric
In the previous section, we introduced QVMC in terms of Born probability functions and formulated the problem as the minimization of a functional of probability functions constrained tional/parametric manifold. The latter is a more common problem often tackled in machine leamning and by forging connections between both fields, we will be able to derive an alternative to QVMC.
varia-
In Euclidean space, we can minimize a function f Rdd - IR. by following the de T2 = -V (xz() which can be viewed as the continuous version of standard gradient descent. Similarly, we can minimize a functional in the space of probability distributions (or in general any Riemannian manifold), by following an ODE on this manifold. However the notion of gradient on a manifold is more complicated, and relies on the Riemannian metric that the manifold is endowed with. Different Riemannian metrics result in different gradient flows, and consequently different optimization dynamics. For a thorough analysis of gradient flows, we refer the reader to Ambrosio et al. (2005).
Gradient Flows
ODE d
Wasserstein Fisher-Rao gradient flows Consider the space of distributions P2 with finite second moment. This space can be endowed with a Wasserstein Fisher-Rao metric with the corresponding distance. In particular, the Wasserstein Fisher- Rao (WFR) distance (Chizat et al., 2018) is defined by extending the Benamou & Brenier (2000) dynamical optimal transport formulation by a term involving the norm of the growth rate It. and by accounting for the growth term in the modified continuity equation. Namely, the distance between probability densities po and pı is defined as
WFRx(po.P1) (1). = inf / 0e.gh.9
Bya(z)) [[l((=))1 + {ge(2)3] dt, subj. to
(7)
(x) Ve (z(2)(()(( + (x) and (z) = po(z). qi (x) = p1 p1(x),
0t
3

---

where (x(x)) is the vector field defining the probability flow, 9t (x()) is the growth term controlling the creation and annihilation of the probability mass, and the coefficient balancing the transportation and teleportation costs. Note that by setting one of the terms to zero we get 2- 2-Wasserstein distance (ge (ar) 0) and Fisher-Rao distance (x)( = 0). Section 3, we also consider the general case of Wasserstein distance, where c is convex COSt function on the tangent space.
Given functional this manifold, F(g] P2 IR, can define the gradient flow of the function F under any Riemannian metric including the Wasserstein metric, the Fisher-Rao metric, or the Wasserstein Fisher- Rao metric. For example, the gradient flow that minimizes the functional Flg under the Wasserstein Fisher- Rao metric is given by the following PDE (which is shown with detailed derivations Appendix A)
0q: (x) - (1)(z) Øt
8F dqe (x))
(oF(g)) 8qc.
[Flaa) öqc
V
(x))
(x)
(v)
9e (e(),,
continuity equation
growth term
(8)
where 5F(g]/5u is the first-variation of of F with respect the L2 metric. The physical explanation of the terms in Eq. (8) is as follows. The continuity equation defines the change of the density when the samples x ge follow the vector field v2(2) -v25F(q.) ōq. The second term of the PDE defines the creation and annihilation of probability mass, and is proportional to the
growth field gr (z)) <Flui (x) Eg.(s) 6Flae (y) Note that E. [9o] = O, and so while mass can dy
be "teleported", the total mass (or probability) will remain constant. The mechanisms can be considered independently by defining the evolution terms under the 2- Wasserstein and Risher-Rao metrics respectively, i.e.
dq: Ot
Vx (a(())( -V 8F(ae (x)
(x)
V
2- Wasserstein Gradient Flow,
(9)
8q:
0q:
8F(m
[5F(u)]
(x:
(x:)
Exy.(a)
(v)
qt (c(
Fisher-Rao Gradient Flow.
(10)
Öt
5qt
89e
now becomes evident that the stationary condition for all the considered PDEs is
FF
FF(ge] (a)
(x)
= +
constant.
(11)
8q:
5qB
In Appendix A, we provide derivations illustrating that Eqs. (8) (10) correspond to the gradient flow under the Wasserstein Fisher-Rao, Wasserstein, and Fisher-Rao metrics, respectively, and hence they all minimize F[g]. For detailed analysis, we refer the reader to Kondratyev al. (2016); Liero et al. (2016).
3 Methodology
In Section 3. we first demonstrate that the imaginary-time evolution of the Schrödinger equation can be viewed as gradient flow under the Fisher-Rao metrie. Afterwards, in Section 3.2, we discuss how density evolution can be projected the parametric variational family and show that doing so for the Fisher-Rao gradient flow yields the QVMC algorithm. Taking this perspective, we propose the Wasserstein Quantum Monte Carlo by considering Wasserstein (and Wasserstein Fisher-Ran gradient flows, followed by the projection onto the parametric manifold (see Section 3.3).
3.1 Imaginary-Time evolution as the gradient flow under the Fisher-Rao metric
The ground state of a quantum system can theory obtained by imaginary-time evolving any valid quantum state vi (with a non-vanishing overlap with the true ground state) to infinite times. The state evolved according to the imaginary-time Schrödinger equation, which defines the energy- minimizing time evolution of the wavefunction eh, and is expressed the following PDE (which is

---

Fisher- Rao Gradient Flow
4
At
Wasserstein Gradient Flow
f
16 1 10pp
VFn F
A/E3
Vw2
A0wa
Anwa
aay min Dadlgras as (eead))
9t q0
Parametric Manitold of
Figure I: W(FR)QMC: A graphical illustration the gradient flow according to the Wasserstein and Fisher-Rao metrics, and the corresponding projection onto the variational manifold g(xz, 0).
the Wick-rotated version Eq. (2), see e.g. (McArdle et al., 2019; Yuan et al., 2019)),
@v)(z) - ((e)e())))
(H
(12)
where again (x)) = (z). The last term proportional to the energy E [1 comes from enforcing normalization (contrary to real-time evolution, imaginary time evolution non-unitary).
Theorem 3.1. Eg. (12) defines the gradient of the energy functional E[g] under the Fisher- Rao metric.
Proof Sketch. The energy functional E[g] has the following derivative
8Eid (x) Vix)
1 v log ((x)
| V log [x))( = (x))
(13)
óq
4
g
Thus, the gradient flow under the Fisher- metric (see Eq. (10))
dq. (x)
(x) g (Eioe qc(x))
(14)
Ot
which is equivalent (up multiplicative constant) to the imaginary-time Schrödinger Equation in Eq. (12) as shown in the complete proof in Appendix B.
We believe that this result be derived following the derivations from Stokes al. (2020), but not introducing the manifold of parametric distributions. However, considering the evolution of the density the non-parametric manifold first helps derive method and relating to QVMC. In the following subsection, we discuss how to project this non-parametric evolution to a parametric manifold.
3.2 Following the gradient flow by a parametric model
By choosing a metric the distributional space and following the energy-mingt merryyyiiiiizzingg. gradient flows, we can design various algorithms for estimating the ground state wave function. Indeed, in principle, by propagating the samples or the density according any gradient flow (e.g., Eqs. (8) (10)), we can eventually reach the ground state. However, these dynamics are defined on the non-parametric and infinite-dimensional manifold of distributions, which do not allow tractable computation of log densities, and thus tractable evolution. Therefore, we project these dynamics onto the parametric manifold of our variational family, and follow the projected gradient flows instead, which tractable.
Suppose the current density on the parametric manifold is (x = (zz, 0) (see Figure 1). We first evolve this density using (non-parametric) gradient flow method (e.g., Eqs. (8) to (10)) for time At, which will take qe (z) off the parametric manifold to (z)) We then have to update current trial model (2 f) to match qi+ (+)) enabling us to propagate the density further. In order to do so, define the optimal update of parameters A0" as the minimizer of the Kullback-Leibler divergence between (+ and the distributions on the parametric manifold, i.e.
A0" arg min DRI. (g..(ge4del (e<))=l)) A0)).
(15)
A0
,
5

---

In practice, we evaluate the parameters update using the following proposition. Proposition 3.2. For (x))) q 8), the gradient of reverse KL--divergecce can be approximated as
/|00 At
- -VoBm ae loo (1x, 0)] = qr a++2eV Vo log ((2, 0) dx ~
V log a(x, 0) dr
01
qi Valogy(z dr At ʃ
8
Ōt 9: Vo log g 0) dæ
(16)
Using this approximation, the optimal update from Eq. (15) of parameters becomes
A6, - -
/
A arg min 29 ||
Ot (z) log (xx 0) dz
x
Ot (x) log (z, dr.
(17)
where ( denotes the inner product, and should not be confused with the bra-ket notation.
The following Corollary states that QVMC can be viewed as the projected gradient flow of the energy functional with respect to the Fisher-Rao metric.
Corollary 3.3. Consider the Fisher-Rao gradient flow (or imaginary time evolution, which is equivalent, as shown in Theorem 3.1). Then, the parameters update (Eq. (17)) matches the gradient of the conventional QVMC loss, i.e.
(Enc Exg,( (2c, 01) Vo logg(z;;0)
o (gg((e))
(18)
This perspective lays the foundation for deriving our WQMC method in Section 3.3, by following the Wasserstein WFR gradient flows rather than Fisher-Ras gradient flows.
Natural Gradient Preconditioning In order to update the parametric model g(z, 0). instead of following the update in Eq. (17), we can exploit the information geometry of the statistical manifold of qxx, 0), and define the update using the Fisher information matrix Fo
/
2
/ -
2 oe
A arg min
A0,
log (xx 0) dx ar (x) Va
Fo
9 log Vogg dæ,
A#
0t
|اممد
log (zz 0) log (ax T)T
2 00
Fg = E (a.00)
(19)
This update is analogous to the natural gradient update (Amari, 1998). Note that the choice Fisher information the metric the statistical manifold for preconditioning the gradient and updating 0 is independent of the choice of metric on the non-pan non-parametric Wasserstein manifold (e.g., Wasserstein or approximation of the natural
(asher-Rao) for evolving qe (z). In practice, gradient (K-FAC) (Martens & Grosse, 2015).
3.3 Wasserstein quantum Monte Carlo
In the previous sections, we formulated imaginary-time evolution governed by the Schrödinger equation as the energy-minimizing gradient flow under the Fisher-Rao metric. Furthermore, demonstrated that projecting the evolved density to the parametric manifold at every iteration corresponds to the QVMC algorithm.
Naturally, we can consider another metric on the (non nonn-paamettrc) space of distributions, which results in different gradient flow and corresponds to different algorithm. Namely, we propose to consider the gradient descent in 2-Wasserstein space as the energy- mneggy-miiimizing density evolution, as introduced in Section 2.2.
Theorem 3.4. The energy-minimizing 2- Wasseesteii gradient flow is defined by the continuity equation
(q.(z) Vz (q((x) E
(20)
Øt
Proof. In Theorem 3.1, we show that 8F(g) = Fanc Plugging this into the Wasserstein gradient flow defined in Eq. (9), yields the result in Eq. (20).

---

C- Wasserstein Metric This result can be further generalized to the c- c-Wasserstein metric with any convex cost function e: Rd E on the tangent space. The Wasserstein distance between po and is defined as follows
/ (c(c) dt, subj. to
We(po, PI) = inf
(21)
(q((z) Vn (()()))(())), and go po, pt
(22)
Øt
Proposition 3.5. The energy-minimizing c- Wasserstein gradient flow is defined by the following equation
@q((x)
0L = -V
(23)
where (-) is the convex conjugate function of cl-), and Ve (y) is its gradient at y.
Proof. See Appendix D.
Theorem 3.4 can be viewed as a special case of Proposition 3.5 where <(-) = 111111 Introducing a different e than I norm translates to a non-linear transformation of the gradient , Floc In Appendix we demonstrate how choose e such that it corresponds the coordinate-wise application of tanh to the gradient, which we use in practice.
Finally, using Proposition 3.5 in Eq. (17), we get the expression for the parameter update, i.e.
/ (a())Vo (a(t)V
A <
Va log (og 0) dr.
(24)
Similar to the discussion of the previous section for QVMC, we can precondition the gradient with the Fisher Information Matrix, exploiting the geometry of the parametric manifold.
In Algorithm I, we provide pseudocode for the proposed algorithms. The procedure follows closely QVMC but introduces a different objective. When using gradients both from Eqs. (18) and (24), follow the gradient flow under the Wasserstein Fisher-Rao metric with the coefficient A. For + oc, the cost of mass teleportation becomes infinite and we use only the gradient from Eq. (24), which corresponds
the gradient flow under the c- e-Wassesttein metric (we refer to this algorithm as WQMC). For X + O, the cost of mass teleportation becomes negligible compared to the transportation cost and the resulting algorithm becomes QVMC, which uses the gradient from Eq. (18). In practice, we consider the extreme cases (A - 0, oc) and the mixed case = 1.
Algorithm W(FR QMC
Require:
samples (x( J.v ~ ((=o(x)
Require:
potential function (x)
while not converged do (x((())))) V(x(0) 2 log (1(() 0) V/ log ((((((0))''
(see Eq. 5)
V, (1)) = stop_gradient V, Flon (x()))
A0 E
N V
Ve (- (-V,,, Bio (1))) )), V Vogg(((((),))
(see Eq. 24)
Fire (ar (0) = stop poop_gradient Pare (1)
A - ΣN (x)
Σx Bac(ool))) Vo log logg(a(()))) ((x())))
(see Eq. 18)
8' optimizer( FF,
(sce Eq. 19)
update (9) by sampling from q(E, 0¹) via MCMC
end while return model q(x, 0*), samples ( ~ q(x, 0*)

---

Be
L12 (6)
H1D (10)
CAVIC
OYMEC INFRCAME
CAPECE
MITAIONE
Ce4
Relative
1e4
OPMEC
OPPIC
CTAME
Energy variance (14
BROMC
163
1e2
1n/2
10:3
162
WAT POOVE
RORROM
WAFROOME
Be.1
Gradient
36.1
56-1
36.1
3050 100000
2030
6202
0202
0000 E000 2000
5220
11202
33320 20000
Figure 2: Optimization results for different chemical systems (every column corresponds to given molecule). The number of electrons is given in the brackets next to systems" names. Throughout the optimization, we monitor three values: the mean value of the local energy (lower better), the variance of the local energy, and the median value of the gradient norm of the local energy. In the first row plots, we average (removing 5% of outliers from both sides) the energy over 1000 iterations and report the relative error to the actual ground-state energy: (E Eo) En. In the second row, report standard deviation averaged over 1000 iterations (removing 5% of outliers from both sides), In the third row, report the median gradient norm averaged over 1000 iterations (removing 5% of outliers from both sides). See the descriptions of methods in the text.
Experiments
For the empirical study of the proposed method, we consider Born- Caan-Oppenhemme approximation of chemical systems. Within this approximation, the wave function of the electrons in molecule be studied separately from the wave function of the atomic nuclei. Namely, we consider the following Hamiltonian
----V² i<j
Z1 X
Σ
Z,Z, X X, IJJ
(25)
Izy
where z; the coordinates of electrons, X1, Zy are the coordinates and charges of nuclei. The first kinetic term contains derivatives with respect to the electron positions x. Indeed, the positions the function vxxx), which
nuclei are given and fixed, and target the ground state of the electronic is an explicit function of the electron positions only. Solving the electronic Schrödinger equation notoriously difficult task, and is topic of intense research in quantum chemistry and material sciences.
'Code reproducing experiments is available at eithub.com//eeeeeee wwtp/////looo///ggo...... The method is also implemented in the FermiNet library: http://co//ogggldddeei/////ii/ee

---

Method name
QVMC
WQMC
W(FR)QMC QVMC
WQMC
W(FR)QMC
Molecule
Be (4)
B (5)
Relative energy error Energy variance (Ha²)
1.50e-5 5.53e-3
3.79e-6
1.04e-6
9.01e-6
9.58e-6
2.69e-6
1.08e-3
1.07e-3
1.96e-2
1.84e-2
1.19c-2
Molecule
Lig (6)
H1o (10)
Relative energy error
4.43e-5 7.21e-3
3.7le-5
3.66e-5 5.59e-4
4.24e-4
3.90c-4
3.88e-4
Energy variance (Ha")
5.76e-4
1.89e-3
2.30e-4
2.45e-4
Table 1: Energy and variances estimates for all systems after 10k iterations (20k for H1o).
Since electrons are indistinguishable fermions, we restrict the Hilbert space to states vi that are antisymmetric under electron permutations (see Section 2.1). This can be achieved by incorporating Slater determinants into the deep neural network, which parametrizes the wave function (z 6), as proposed in various recent works (Luo & Clark, 2019; Hermann et al., 2020; Pfau et al., 2020; Hermann et al., 2022). The density is then given by the Born rule g(x,,0) = (x, 0) For all our experiments, we follow (von Glehn et al., 2022) and use the "psiformer" architecture together with preconditioning the gradients via K-FAC (Martens & Grosse, 2015).
In our method, we apply several tricks which stabilize the optimization and improve convergence speed. Firstly, we have observed that applying a tanh non-linearity coordinate- wise to the gradient V. (x) significantly improves convergence speed. This corresponds to a different cost function in the Wasserstein metric, as we discuss in Proposition 3.5 and Appendix D. Also, we remove samples from the batch whose norm V log q(<, 0)ll significantly exceeds the median value. Namely, we estimate the deviation from the norm as Ex(2,,9]11 [|| Vz q(z, 0) median(I) V. log ((x, O)IDI and remove samples whose norm exceeds five deviations from the median. When including the gradient from Eq. (18), we clip the local energy values as proposed in (von Glehn et al., 2022). i.c. by estimating the median value and clipping to five deviations from the median, where the deviation is estimated in the same way as for the norm of the gradient.
We consider different chemical systems and compare against QVMC as a baseline. We run our novel method with the same architecture and hyperparameters as the baseline QVMC-based approach in (von Glehn et al., 2022). For the chemical systems, we consider Be, and B atoms, the Li molecule and the hydrogen chain Hio from (Hermann et al., 2020). The exact values of energies for Be, B, Li2 are taken from (Pfau et al., 2020), the exact value of the energy for H is from (Hermann et al., 2020). All the hyperparameters and architectural details are provided in the supplementary material.
In Figure 2. we demonstrate the convergence plots for the baseline (QVMC) and the proposed methods (WQMC and W(FR)QM, see Algorithm 1). For all the considered systems, both WQMC and W(FR)QMC yield more precise estimations of the ground state energy (the first row of Figure 2). To assess convergence, we also monitor the variance of the local energy and the gradient norm of the local energy. As we discuss in Eq. (11), both metrics must vanish at the ground state. More fundamentally, the variance of the local energy can be shown to vanish for eigenstates of the Hamiltonian (Wu et al., 2023), referred to as the zero-variance property. First, we point out that obtaining the ground states of the considered molecules with QVMC is challenging, and even with powerful deep-learning architecture, discrepancies remain with the ground state. Since we use existing statc-of-the-as architectures as backbone, our results are also limited by the limitations of the latter (Gao& & Günnemann, 2023). Developing novel architectures is out of the scope of this work.
However, in Figure 2, we clearly observe that both WQMC and W(F yield significantly faster convergence of the aforementioned metrics compared to QVMC. In particular, for the larger molecules Lig and H10, we observe that we consistently obtain lower energies within 10k steps (20k for HдN) with a more stable convergence. For the smaller molecules we observe that QVMC obtains lower energies in the first few iterations, but its convergence slows down significantly, after which our approach steadily yields improved energies below QVMC. Overall, our experiments demonstrates that taking into account the Wasserstein metric allows for faster convergence to accurate approximations of the ground state. See the final metries in Table 1.

---

5 Limitations
Since our method requires an extra gradient evaluation compared to QVMC, we include the runtime of all the algorithms in Figure 3, Appendix E. Namely, for a proper comparison, instead of reporting the metrics per iteration, we report them per wall time in seconds. Note that all the claims of the experimental section still hold in terms of wall time. All the models were benchmarked on four A40 GPUs. Third-order derivatives can be efficiently computed using modern deep learning frameworks such as JAX (Bradbury et al., 2018), which we used to implement our method.
Potentially, one alleviate the extra cost of the iteration by coming up with more efficient Monte Carlo schemes or other updates of the samples. Indeed, in our experiments, we observed that the proposed method requires fewer MCMC steps (not included in the paper). Moreover, the evaluated gradient of the local energy as the proposal vector field for updating the samples. This would allow decrease the number of MCMC steps at the low cost of additional hyprparraeeter tuning.
can
6 Discussion and conclusion
Conclusion In the current paper, we propose a novel approach to solving the quantum many-body Schrödinger equation, by incorporating the Wasserstein metric on the space of Born distributions. Compared to the Fisher-Rao metrie, which allows for probability mass "teleportation" the Wasser- stein metric constrains the evolution of the density to local changes under the locally-optimal transportation plan, i.c., following fluid dynamics. This property is favorable when the evolution of the parametric model is accompanied by the evolution of samples (performed by MCMC algorithm). Indeed, by forbidding or regularizing non-local mass "teleportation" the density change, one prevents the appearance of distant modes in the density model, which would require longer MCMC runs for proper mixing of the samples.
In practice, we demonstrate that following the gradient flow under the Wasserstein (or Wasserstein Fisher-Rao) metric results better convergence to the ground state wave function. This is expected to be due to our proposed loss, which takes into account the gradient of the local energy and achieves its minimum when the norm of the gradient vanishes, therefore explicitly minimizing the norm of the local energy gradient.
We believe that our new theoretical framework for solving the inneeindeeendent Schrödinger equation for time-reversal symmetric Hamiltonians based optimal transport will open new avenues to develop improved numerical methods for quantum chemistry and physics.
Connection to Energy- -ased and Score- Based Generative Models The developed ideas of this paper, i.c., projecting gradient flows under different metrics onto a parametric family, can be extended to generative modeling by swapping the energy functional with the KL-divergence. More precisely, as we show in Appendix C. using the KL-divergence as our objective functional, the Fisher-Rao gradient flow yields energy-based training scheme, while the 2- 2-sssserttein gradient flow corresponds to the mathe-ng,,hing, which is used for training diffusion generative models.
7 Acknowledgement
The authors thank Rob Brekelmans and anonymous reviewers for helpful discussions and feedback. J.N. was supported by Microsoft Research. J.C. acknowledges support from the Natural Sciences and Engineering Research Council (NSERC), the Shared Hierarchical Academic Research Computing Network (SHARCNET), Compute Canada, and the Canadian Institute for Advanced Research (CIFAR) AI Chairs program. A.M. acknowledges support from the Canada CIFAR AI Chairs program. Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute www www.vectorinstittueee ai ai/#partners.
10

---

References
Amari, S.-I. Natural gradient works efficiently in learning. Neural computation, 10(2):5512776 1998.
Ambrosio, L.,, Gigli, N., and Savaré, G. Gradient flows:
metric spaces and the space of probability measures. Springer Science & Bussness Media, 2005.
Benamou, I.-D. and Brenier, Y. A computational fluid mechanics solution to the konte-Kantorovich mass transfer problem. Numerische Mathematik, 84(3):375-393, 2000.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J.. Leary, C., Maclaurin, D., Necula. G., Paszke, A., VanderPlas, J., "Vanderman-Milne S., and Zhang, Q. JAX: composable transformations of Python+NumPy programs, 2018. URL http: /tt:://gittt github. com/ mom//goole//// j/jae
Carleo, G. and Troyer, M. Solving the quantum many-body problem with artificial neural networks. Science, 355(6325):602-666) 2017.
Ceperley, D., Chester, G. V., and Kalos, M. H. Monte Carlo simulation of a many-fermion study. Physical Review B, 16(7):3081, 1977.
Chizat, L., Peyré. G., Schmitzer, B., and Vialard, F.-. An interpolating distance between optimal transport and Fisher-Rao metrics. Foundations of Computational Mathematics, 18(1):1-44, 2018.
Choo, K., Mezzacapo, A., and Carleo, G. Fermionic neural-network states for ab-initio electronic structure. Nature communications, 11(1):2368, 2020.
Du, Y. and Mordatch, Implicit generation and generalization in energy-based models. arXiv preprint arXiv: 1XX3:/9033.000 006898 2019.
Gao, N. and Günnemann, Generalizing neural wave functions. arXiv preprint arXiv:2302.0416. 2023.
Gnech, A., Adams, C., Brawand, N.,
Carleo, G., Lovato, A., and Rocco, N. Nuclei with up nucleons with artificial neural network wave functions. Few-Body Systems, 63,));7, 2022.
Hermann, J., Schätzle, Z., and Noé, F. Deep-neural-networf solution of the electronic Schrödinger equation. Nature Chemistry, 12110):891-877, 2020.
Hermann, J.. Spencer, J., Choo, K., Mezzacapo, A., Foulkes, W., Pfau, D., Carleo, G., and Noé. F. Ab- initio quantum chemistry with neural-network wavefunctions. arXiv preprint arXiv:2208.125000 125900 2022.
Hibat-Allah, M., Ganahl, M., Hayward, L. E., Melko, R. G., and Carrasquilla, J. Recurrent neural network wave functions. Physical Review Research, 2(2):023358, 2020.
Hyvärinen, A. and Dayan, P. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
Kondratyev, S., Monsaingeon, L., and Vorotnikov, D. A new optimal transport distance on the space of finite Radon measures. Advances in Differential Equations, 111 2016.
Liero, M., Mielke, A., and Savaré, G. Optimal transport in competition with reaction: The Hellinger- Kantorovich distance and geodesic curves. SIAM Journal on Mathematical Analysis, 48(4): 2869-2911,
Luo, D. and Clark, B. K. Backflow transformations via neural networks for quantum many-body wave functions. Physical review letters, 122(22):226401, 2019.
Martens, J. and Grosse, R. Optimizing neural networks with kromecker-factored approximate curvature. In International conference machine learning, pp.
PMLR, 2015.
McArdle, S., Jones, T., Endo, S., Li, Y., Benjamin, S. C., and Yuan, X. Variational ansatz-based quantum simulation of imaginary time evolution. npj Quantum Information, 5(1):75, 2019. 2019.

---

McMillan, W. L. Ground state of liquid he 4. Physical Review, 138(2A):A442, 1965.
Neal, R. M. Annealed importance sampling. Statistics and computing, 11:125-139,
Pescia, G., Han, J., Lovato, A., Lu, J., and Carleo, G. Neural-network quantum states for periodic systems in continuous space. Physical Review Research, 4(2):023138, 2022.
Pescia, G., Nys, J., Kim, J., Lovato, A., and Carleo, G. Message-passing neural quantum states for the homogeneous electron gas, 2023.
Pfau, D., Spencer, S., Matthews, A. G., and Foulkes, W. M. C. Ab initio solution of the many-
electron Schrödinger equation with deep neural networks. Physical Review Research, 2(3):033429, 2020.
Song. Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13450, 2020.
Sorella, S. Green function Monte Carlo with stochastic reconfiguration. Phys. Rev. Lett., 80:4558- 4561, May 1998. doi: 10.1 SasPPhyyRevLLeet.00555555 URL https///: aps org/doi/100 1103/PhysRevLett.c 80.80004 4558.
Stokes, J., Izaac, J., Killoran, N., and Carlco, G. Quantum natural gradient. Quantum, 4:269, 2020.
Vicentini, F., Hofmann, D., Szabó, A., Wu, D., Roth, C., Giuliani, C., Pescia, G., Nys. J., Vargas- Calderón, V., Astrakhantsev, N., et al. Netket 3: Machine learning toolbox for many-body quantum systems. SciPost Physics Codebases, pp. 007, 2022.
von Glehn, I., Spencer, J. S., and Pfau, D. A self-attention ansatz for ab-initio quantum chemistry. arXiv preprint 13X1::2211.13672, 2022.
White, S. R. Density matrix formulation for quantum renormalization groups. Physical review letters, 69(19):2863, 1992.
Wu, D., Rossi, R., Vicentini, F,, Astrakhantsev, N., Becca, F., Cao, X., Carrasquilla, Ferrari, F., Georges, A., Hibat-Allah, M., et al. Variational benchmarks for quantum many-body problems. arXiv preprint arXiv:2302.0491 2023.
Xie, J., Lu, Y., 7hu, S.-C., and Wu, Y. theory of generative convnet. International Conference on Machine Learning, pp. 2635-2644. PMLR, 2016.
Yuan, X., Endo, S., Zhao, Q., Li, Y., and Benjamin, S. C. Theory of variational quantum simulation. Quantum, 3:191, 2019.
12

---

A Gradient flows under Wasserstein Fisher-Rao metric
We, first, remind the concept of the functional derivative. The change of the functional Flyi : P2 IR. along the direction h can be expressed as
Flg h] Flgl + hF] + o(IiI)), dF(h] = / dz h(x)
FF
(x)
(26)
ôq
derivative
Consider change of the density in time, the change of the functional can be defined through the differential as
qt + At 0q:
dqu ((Il/)))
.qq
-f dx
(q)(z)
8F(ad
Flad + At dF
+
dF
(x).
01
al
dgt
(27)
In partieular, have
d Flad
Oge:
= =f
)y)(2)) 8 (x))
5F((1)
dF
dr
(28)
dt
Øt
5q
A.1
Minimizing movement scheme for 2-Wasserstein distance and Kullback-Leibler divergence
Gradient flow under W₂ Consider the following minimizing movement scheme (MMS)
1
inf F(o] Fla)
At 2 W2(a,'),
(29)
where the change of the density is restricted to the continuity equation, i.e.,
dq. ØL v (z)))) and d'(x) ge(x) A (()(1) + o(At).
(30)
Using the static formulation of W2 distance, we have
w2(4.0') = f dr ((>)]<= (x)) At2 / dr (x)))
(31)
where T'(x)) is the optimal transportation plan, and o* oxx) is the corresponding optimal gradient field. Thus, we can rewrite the MMS problem as
inf Flal At f dx V. (a(x))(x))
FF(ad] (x) - Flal + Al / dz g(>))lu(>)11
(32)
8q:
2
inf / de ((()) ((xx), V 8F(a (xx)
ʃ dr ((<))l=(>>)1
(33)
5q.
inf do / ddo)(z)))))))((())) g(z)) u(x) + V FF(gi]
(x)
(34)
òq:
From the last optimization problem, have
8 F
v(x) V
(x).
(35)
dge.
Gradient flow under KL Consider the following minimizing movement scheme (MMS)
inf Flg'] - Fla + Al DKL. DKI.('''ll (g'll4)
(36)
where the change of the density is restricted to the following weighting scheme
dq:
Øt ge (x)) hence
(37)
d(xx) = qe (x) + At (x)))((()) + o(At), and
(38)
At2
log g' (x) = log (z) + At (x( 2 gi (z) + o(At)).
(39)
13

---

The KL-divergence is then
PxI(g'lla))= / de ( (tt (i(())
At
2 )i(()) + ((A
(40)
Ar2
2 / ddz ((()))(())
- At dr (()((())(()))
(41)
+ A12 / do (0)))) (x) o(A13)
(42)
A/
/ dr (()))))))))) dAr)).
(A)))
(43)
2
In the last equation we are using
9:90 dr
(44)
Thus, we can rewrite the MMS problem as
inf Flad + At / do gc())c((x)
8Flad
At / dr ((()))
(x) (x)- F[ael
//
(45)
8ge.
inf / dx (1)))))()))
8Fl1d
(x:) Bgg(())
[5F(qt]
(2))
-
1
/ dr (x)))
(46)
det
dq:
2
inf f dx (((x)) (c(x)+
(
(FF(qc) (x) Bg((v)
[81
(vr)
1)))
1
(47)
8q:
5qc
From the last optimization problem, we have
gc(x)) [8F(9d) (x) (E(()))
[FF(9
(y)
(48)
5qc
5q.
Note, however, that (-1-( is not the same as the Fisher-Rao metric. The derivations here demonstrate that the Fisher-Ras gradient flow can be derived as the MMS scheme with the KL- divergence.
A.2 Minimizing movement scheme for the Wasserstein Fisher-Rao metric
Consider the Wasserstein Fisher- distance
/'' Bu,cas) [10((=0)1)+ + (0x(2)) di, subj. to
(FRx = inf Vh.
(49)
0q((x) V. (z)(z)((()) + (x)))q(()), go(x) po(z), (ı(x) p))x).
(50)
The minimizing movement scheme (MMS) for this distance
1 WF
inf F Flgl +
(51)
At2
where the change of the density is given by the continuity equation with the growth term
dq: - (q((=)v((=)) + (x()))(()).
(52)
Ōt
For close enough y and g'. ('') can be estimated via the metric derivative |/4]² that is defined as
(FRA((e.4++:))) 9++2:)
WFRx (
1n412 lim
(53)
Ar-0
At
hence,
('')(g.9')2 21:114122 At2 dz g(x)) 110 (x))1 Ag (x)3
(54)
14

---

Thus, the MMS problem can be written as
inf Flgel At / de V. V..(a (2(())) gests
8F 9il
(x) + / dx (e()))0(x)
8F(a))
(x) Flae]
gq
5qe
2 do ul(x) [|or(w0)))== Agr(z))2
(55)
/ dx ana (loo)( me((r), V 8F(ad (a))
(8F(gc)
inf St
dr (x())))(())
(x) Bee(s))
(1)
dq
59c
8%/
+ - do (1()) (1:((0)))) + )gn(t))
(56)
/
dx /dr ( (1(z)) V
8F(4 (x))
inf Se
8q
+>f de q((x) (((()))
1
( 16FFlg
6Flad
[5F(qu]
(2) 1))
(x) (x)
(57)
de:
From the last optimization problem, we have
8Flae) (x), (=x)= = - ( [F(aa) (x: Bae(s)) 5ge
[F(u (v)
(x) -
(58)
óge
8qu
Note that different values of À result in different gradient flows. For instance, considering the limit + oo we have g(x) - 0 and Eq. (52) just becomes W2 gradient flow, which natural since we have infinite penalty for the mass teleportation in our metric. Setting + 0 requires some additional consideration, since then the growth term explodes, and all the mass will be teleported without any cost. Indeed, for X + 0, metric does not penalize for the mass teleportation at all, but our change density (Eq. (52)) is still able to teleport mass, hence, it will be doing so "for free".
A.3 PDEs demonstrating the convergence
Consider the change of the density under the continuity equation with the vector field (x()) and the growth term gr (z)
Dg: V.
((((x))u()))) + (x)))(())
(59)
Øt
Thus, the change the functional F'g]]
d Flad - - ff dr V. ((:())) 8F(ur] (1r) + du ge(x))
8F(g
(x:)
(60)
dt
dqc
5gt
/ dæ (2)) oz(2), V
FF(q(]
(x) + dæ (c(x))e(1)
(x)).
(61)
dqp
From this equation, we can clearly see that vy (2r) and gi (x) derived in the previous section minimize F(g). Indeed, taking
vx(x) = -v [ad (x), gx(x) = - ( [f(gg)) (xx) Eg((s)
6F(a.
(w) w))
(62)
5gt
5g.
5qg
we get
d
5F(a
(5F((a)
5F(ae)
Flgd = - dx g((x))
V=
(x) X
dx
ge(x))
(x)
Eyn(s)
(4)
df
dqp
8q.
59c
(63)
A f dx (x)
1
(8F(ge]
8Fl(g)]
[F(qg)
(x))
Egg.
(v)
Ege(=)
(z)
< 0. (64)
89:
=0
15

---

Note that the considered growth term preserves the normalization of the density, i.e.,
[ dæ OL dqc (z) - / dæ ((()) = - f dæ (x(x)
(a)=
(F qil
8qt (œ)
(v)
[F1le ògr
(w)
- 0.
(65)
Thus, functional Flg decreases when the density Qi evolves according the PDE
0t. dqe (1) = - ((n(t))
8Fl
Fl
6F (g
V
(x))
(x)
E (v)
(v)
(e(),,
dqe
ōqe
the continuity equation
growth
(66)
and reaches its stationary point when
V
8Figd oqe
(ar)
0, i.e.,
5F1gel 69e
(1) constant.
Note, that in the same way, we can consider the continuity equation and the growth term separately, which defines the gradient flows under 2- Wasserstein and Risher-Rao metrics respectively. The corresponding PDEs are
0q: (z) -V (a(())
8F de: (xx) (1)))
V
(67)
01
0q:
(1) -(91
(8 [qr]
[F(ge] 8qg
(0)) Dalc))
(x)
Egge(s)
(68)
Øt
dqt
B Imaginary-time Schrödinger equation as the gradient flow under Fisher-Rao Metric
Theorem. Eq. (12) defines the gradient flow of the energy functional Eigl under the Fisher-Rao metric.
Proof. First, we derive the functional derivative of the energy functional Elg]. We denote the differential of the functional F(g) along the direction h as
/ dx h 81 8g
dF dF(g)(h)
(69)
where [ is the functional derivative. Consider the energy functional
Elgl - dr
V
logg
V log q|1
(70)
4
The functional derivative of this functional is as follows
db(g)] 0F(g + 8. h))
de
==0
f du h V
V2 logg
| loggl ||2
dx
[ 22 h 1 vv
log
v.
For the last term, we do integration by parts and get
/ dr q
h 1 VV loglo V Ag
f dx Vaq, (o
h ) dr V.q.
V 0
0.
4
(71)
Thus, we have
dE(g)[h] / dx v v2 log
- v log /
1
(72)
8
6Eigl/69
16

---

and we see that the derivative coincides with the local energy, i.e.
5 (x) (x) (ax)
1 V logg V logg(()))
(73)
4
Using the results from Section 2.2, the energy-minimizing gradient flow under Fisher-Rao metric is
(a)()) (x) Elgel (x))..
(74)
Öt
Second, we derive the PDE for the time-evolution of the density under the imaginary-time Schrödinger equation.
V2 the (V F(
(75)
Öt
0
2
ii V2rr 2(V E(
(76)
Øt
Dg: di
V2dr 2(V Elggl)ge
(77)
(78)
Using the identity
V V. Vv z log (+1) Vv 2 V log (+)) + V log |b]
(79)
V V log 9) + 2 2 log 7-
1 u/ V log 0/1 2 v logg:
1
(80)
4
we have
0q: - 2 V
1 V2 log I1 log qall 2= EE ạr
(81)
Øt
4
0y - 2 Flac(x) - Fla)) plei)((()) ((()))
(82)
Öt
which is equivalent to Eq. (74).
C Analogies to generative modeling and variational inference literature
To draw connection generative models and variational inference literature, consider the reverse and forward KL-divergences as an objective to minimize (instead of the energy) and derive projected gradient flows for them.
C.1 Generative modeling: reverse KI. divergence functional
Consider the reverse KI. divergence objective, and its first variation
Flud]
(KI. (:)
(83)
6F
p
(84)
dg. 4t
where p(z) is the data distribution given empirically.
Fisher-Rao Gradient Flow
Using equations Eq. (10), the PDE that defines the Fisher- Rao gradient
flow is
dq. 01
(x) = p((x)
p(v) (v)
Eg (x)
(1((x) p(x) gr(x)
(85)
2: (x)
Thus we have
9: exp(-t)) + (1 eex(xx))))
(86)
where qo is the initial distribution. This results in distributional path constructed by the linear average (mixture) of the end-point distributions. Note that changing the mixture parameter corresponds to teleporting the particles.
17

---

Projected Fisher-Rao Gradient Flow update corresponding the projected Fisher Rao gradient flow using Proposition 3.2 and Eq. (17),
Having the functional minimizing PDEs, we can find the
A x
qt
V log ((2, da:
(87)
Eq:
p
qr
V log (2x, 0) dr
(88)
9r
VeDp(=))
log q(,, 0)
(89)
This corresponds to the standard maximum likelihood objective that is used for training of energy- based models (Xie et al., 2016; Du & Mordatch, 2019).
Wasserstein Gradient Flow gradient flow is
Using the equation Eq. (9), the PDE that defines the Wasserstein
0y: (x) --)=-Vx' (a(e) Ot
p(x) g)(x))))
(90)
Projected Wasserstein Gradient Flow Having the functional minimizing PDEs, we can find the update corresponding the projected Wasserstein gradient flow using Proposition 3.2 and Eq. (17).
oF ōqu"
A
Vo x
V
Va log g( 0)
dr
(91)
V
90
V x
V= log ((,, 01) dx
(92)
qt
-1
Va
91
V Vn log 11=
(x) 0) dx + Vo
ai p
Vu. V, loga(s, 0) dr (93)
qi
= Vo / pl Vo logp. Ve log (2z 8)) dx Vo
V log q. v log logg(1, 8) dx (94)
V p(x)
V Vogp(()) V, log (22 8)
(95)
Note that similar to Proposition 3.2, (x()) g 0) the density equal to the model density but detached from the parameters 6. This objective corresponds to the score-matching objective (Hyvärinen & Dayan, 2005), which widely used in the diffusson-based generative models (Song et al., 2020).
C.2 Variational inference: forward KL divergence functional
Consider the forward KI. divergence objective, and its first variation
Flgi = DxL p)
(96)
6.F + log 4t
(97)
8gc
where p(z) is the target distribution for variational inference.
Fisher-Rao Gradient Flow
Using equations Eq. (10), the PDE that defines the Fisher-Rao gradient
flow is
grad
(98)
01
8.F
8F
E gt
(99)
8q
8qu
-1
log
+ Eyg
+ log
(100)
p
log
+ DKL (1p)
(101)
p
18

---

Thus we have
0
log qt
log log q: + DKL llp)
(102)
0t
= log at
exp( -t) log + (1 exp(-t) logs - log 2t
(103)
expl 1) p (1-exp(-()
2, = /
exp(-t) (1--)))p-ppppp
(1-oxp(-1))
qt
Zt 40
dz 90
(104)
where is the initial distribution, and Z, is the partition function. This results in distributional path constructed by the geometric average of the end-point distributions. Note that changing the geometric parameter corresponds to teleporting particles. These distributional paths are commonly used in Annealed Importance Sampling (Neal, 2001).
Projected Fisher-Rao Gradient Flow Having the functional minimizing PDEs, find the update corresponding to the projected Fisher-Rao gradient flow using Proposition 3.2 and Eq. (17),
85
x
V log ()) da
(105)
dqe
ạt
4t
1 + log
Vo log ((x; dr
(106)
f log at
Voq(z, 8) dr
(107)
p
/ log qe Vog(a, 0) dx + Vo / g(x,, 0) lo
(108)
/ q((, log (2. 0) dr + Vo / q(x, 0) pog
Vo
(109)
V / (1(z, 0) log g(=,0) 0) dx
(110)
P = vo kk. (((,, 0) p)
(111)
This corresponds to the standard variational inference objective.
Wasserstein Gradient Flow gradient flow is
Using the equation Eq. (9), the PDE that defines the Wasserstein
0i 4:
F
(112)
8F
V
4t V
(113)
5qp
1 + log 9t
Va
V.
(114)
p
Vp (q. V (log q logp))
(115)
- (g: V, log p) + V (g. V log
(116)
-.. (q. V log + V
(117)
This is the Fokker-Planck equation, which characterize the movement of the particles according to Langevin dynamics:
dx. V, logpdt V2WW.
This is a common approach for sampling from the unnormalized density p(x) via Markov Chain Monte Carlo algorithms.
19

---

Projected Wasserstein Gradient Flow Having the functional minimizing PDEs, we can find the update corresponding the projected Wasserstein gradient flow using Proposition 3.2 and Eq. (17),
8F
x
V
90
V log g(2, dr
(118)
ōgr
V
01
1 + log
at
V. logg(o; 0) dr
(119)
Vo
91
x
log
qt
V log g(2, log
0) dr
(120)
glt 0)
V
v log (x, V. log
(121)
which resembles the score-matching objective (Hyvärinen & Dayan, 2005) for variational inference.
D C- C-Wasserstein gradient flow
Wasserstein distance with the convex cost function c: R.d - R is defined
/ T(x, y)c(x y) dr dy,
We(po. P1)
inf xEI(po.P))
(122)
where ['(Po, pi) is the set of all possible couplings of po and p1. The dynamic formulation of this distance is the following
(
We(Po.P1) inf oe
(=())((((()))) dt, subj.
(123)
(q((1) V (z)()))((z)), and 2o Pa. ạ1 P1-
(124)
dt
Proposition. The energy-minin meergy-mnnmmizing c- Wasserstein gradient flow defined by the following PDE
(q((z) - (aa)
(125)
Ot
where (+) is the convex conjugate function of c(-).
Proof. The movement minimizing scheme for Wel:s.) ) is the following optimization problem
inf Flg'] F(g)
At We(q.4').
(126)
Assuming that the density changes according to the continuity equation ('(x) = (x) AtV (ge (x)) + o(A(), and At is small enough so that vx(x) defines the optimal transportation plan, we have
5Flge inf Fiqi At lv. (c.(((()))(())) "
(x) dr F(gdd + At
A Ar (((
(127)
8qc
/ (((=)) cl(x), v
8F(y (x)
Atinf
dr ((((()(((((()))
(128)
ōgt
/ (((e)) (aler((ta(())) ((va(x))) (1((z), -
8Flul
Atinf
(x) dr
(129)
0t
ōg:
At aale)o (-V
8F oF[gel dqi
(x)) dx
(130)
and the infimum is achieved at
2 (x) = Vc (-c(-V.
8F(9e] ōge
(x)
(131)
which gives the formula for the vector field. Using the energy gradient from Theorem 3.1 Erse: we get the result.
6E
(x) 5g.
=
20

---

Proposition. Coordinate-wise application tanh to the vector field, i.e.
Qq((z)
-v (g((x) tanh(- V Eloc (x)))
(132)
ôt
corresponds to gradient descent with c-Wasserstein distance, where Rdd RR is the following cost function
c(x) = Σ9 ( (z; 1) log(2; + +) + (1 x;) log(1 - 2:i) d log 2.
(133)
Proof. Consider c¹ (a(x) = E, (xge + exp(-x;)) It corresponds to applying hyperbolic langent non-linearity coordinate-wel wise to the vector field field, i...
(x) exp(x;) exp(-a;) tanh(r;).
(134)
exp(x;) + exp(-x,)
The corresponding cost function c(x) is the following
clz)
sup(a, u) c² (v)
(135)
sup
Σ
F:Yi log(exp(14) exp((-ui)
(136)
1
1 x
sup
(Ir: + 1)9; - log((xp((22)))) + +1)
1/Br = log
(137)
2
1
d
1
1 + Zi
1 + 3z +11 1)
Σ
(z) + 1) log
llog
(138)
1- Zz
1- Tz
d
Σ 2
1
(r; + +ogg0 + +1) + (1
log([ m11) dlog2
(139)
21

---

E Additional experimental results
COMEC
CVMIC MITRION:
CUMICC VIFAION:
WAFRINN
WITRICOS
163
WONC
Relative
203
2+.5
CONIC
OPPEC
GUANC
GAVC WORROON:
$el1
IHFAIOVC
Energy variance (Ma7)
20-1
16:
36-
COMEC
DXME WOME
260
DVVC 192MC
345
WOMC
IMOIC
Gradient
100
262
26-1
06-2
361
2000
5200
1300 5000 4200
0200
Seee 1rkce 20000 T0000
Figure 3: Optimization results for different chemical systems (every column corresponds to a given molecule). The number of electrons is given in the brackets next to systems' names. Throughout the optimization, we monitor three values: the mean value of the local energy (lower is better), the variance of the local energy, and the median value of the gradient norm of the local energy. In the first row of plots, average (removing 5% of outliers from both sides) the energy over 1000 iterations and report the relative error the actual ground-state energy: (E Fo)/ Feg- the second row, we report standard deviation averaged over 1000 iterations (removing 5% of outliers from both sides). In the third row, we report the median gradient norm averaged over 1000 iterations (removing 5% of
outliers from both sides). All the metrics are plotted versus the wall time of computations performed on four Nvidia A40 GPUs


Reviews:
[{"review_id": "N14McG5sTB", "reviewer": "Reviewer_wiqn", "summary": "The paper shows that quantum variational Monte Carlo can be written as energy minimizing gradient flow under Fisher-Rao metric, and proposes an alternative to traditional quantum Monte Carlo by following Wasserstein gradient flow. Experiments with Be, B, Li2, and H10 chain show lower energies achieved by including Wasserstein gradient flow. ", "questions": "Why does W(FR)QMC seem to be better than WQMC alone for most of the experiments? What was the hyper-parameter controlling the W/FR tradeoff, and how sensitive is this hyper-parameter?\n\nWere other c- cost functions tried for the c-Wasserstein gradient flow? why does the coordinate-application of tanh work well? \n\nstop_gradient in algorithm box was referring to ?", "limitations": "Do authors expect including WQMC to converge better for all physical systems? For ex, Does the strong/weak correlation of the physics affect the method convergence relative to QMC? \n\nCould authors include the additional computation complexity (theoretically, empirically) required for the Wasserstein gradient ? ", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "Paper proposes a novel optimization approach to quantum Monte Carlo, an important problem for computational physics/chemistry. Presentation and writing is very clear, mathematically rigorous. The results show better convergence and better optima reached with proposed method.", "weaknesses": "only 4 physical systems and one NN architecture were tested. \nPaper seems to state that WQMC is better for all cases; if so, more experiments would make this claim stronger.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback and time spent. We are glad that you find the presentation of our ideas to be clear and rigorous. In what follows we answer the questions asked.\n\n> Why does W(FR)QMC seem to be better than WQMC alone for most of the experiments? What was the hyper-parameter controlling the W/FR tradeoff, and how sensitive is this hyper-parameter?  \n\nWe observe that the Fisher-Rao metric demonstrates faster convergence in the beginning and the Wasserstein metric converges faster later into the optimization. Our intuition is that the usage of both metrics at the same time allows us adaptively decide which distance to the ground state is larger at the moment and minimize it.  \nIndeed, one can introduce a hyperparameter controlling the tradeoff between metrics. In our experiments, we haven’t investigated the influence of this hyperparameter, and simply take the sum of two metrics for W(FR)QMC. We leave an extensive study of this parameter for future work.\n\n> Were other c- cost functions tried for the c-Wasserstein gradient flow? why does the coordinate-application of tanh work well?\n\nIn our experiments, we concluded that for stable training it’s important that the c-cost function clips high values of the vector field. Our intuition is that it is related to the instabilities created by cusps close to nuclei, which are common to all variational Monte Carlo methods. We also experimented with the sign non-linearity, i.e. applying it coordinate-wise to the local energy gradient. Although it exhibits faster convergence at the beginning of the optimization, we find that it demonstrates higher variance later on.\n\n> stop_gradient in algorithm box was referring to?\n\nFor better presentation, and to provide some intuition on the implementation, by stop_gradient we denote the operation that detaches the node from the computational graph. That is, we evaluate the value of the local energy gradient, but then we only use its value, not backpropagating w.r.t. inputs or parameters so it’s not affected by other gradient operators. We will clarify this in the next revision of the paper, thank you!\n\n> Do authors expect including WQMC to converge better for all physical systems? For ex, Does the strong/weak correlation of the physics affect the method convergence relative to QMC?\n\nAlthough our approach is agnostic to the functional (and, correspondingly, Hamiltonian), we expect our method to improve over traditional QMC as the effect of correlations increases. In the limit where the system becomes non-interacting, a single Slater determinant can represent the ground state and hence consists of single-particle orbitals only. More importantly, however, the cases we expect our methodology to work best in, are where the target Born distribution is multi-modal, e.g. due to multiple atoms in molecular systems. Due to the latter, our significantly improved results for H10 are not unexpected.\n\n> Could authors include the additional computation complexity (theoretically, empirically) required for the Wasserstein gradient?\n\nAs requested by most of the reviewers, we include the runtime of all the algorithms in the plots (see the new plots in the pdf attached to the general response). Namely, for a proper comparison, instead of reporting the metrics per iteration we report them per wall time in seconds. Note that all the claims in the paper still hold in terms of wall time. We will include these plots in the next version of the paper. All the models were benchmarked on four A40 GPUs."}, {"review_id": "ZFz9E1k2cm", "reviewer": "Reviewer_f58M", "summary": "The paper proposes a new way to compute gradients for parameters in quantum Monte Carlo. The authors propose to update the network parameters by following the Wasserstein Fishier-Rao gradient flow, which is composed of a continuity equation and a  growth term. The authors show that when only the latter term is considered, the WFR gradient flow is equivalent the conventional energy loss. The authors then propose to leverage the first term and propose the Wasserstein quantum Monte Carlo. With some tricks the proposed optimization method leads to a faster convergence. Experiments are conducted on small systems with atoms.", "questions": "- Since some tricks are used in experiments. How does these tricks influence the performance? Are those tricks also used for baseline methods?\n-  Which would be the recommended method? The WQMC or the W(FR)QMC? What would be the trade-off?", "limitations": "None", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is very well written.\n- The proposed method is shown to include the conventional method.\n- The experimental results are improved.", "weaknesses": "- It's not explicitly mentioned in the paper. But I think the extra gradient computation may cause more computation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback and time spent. We are glad you found our paper well-written. In what follows we address the concerns raised and answer the questions.\n\n> It's not explicitly mentioned in the paper. But I think the extra gradient computation may cause more computation.\n\nIndeed, the proposed method uses extra gradient computation. We will clarify it in the paper. In the general response, we added the plots comparing the methods in terms of runtime instead of iterations to take this extra computation into account. Note that all the claims in the paper still hold in terms of wall time. We will include these plots in the next version of the paper.\n\n> Since some tricks are used in experiments. How does these tricks influence the performance? Are those tricks also used for baseline methods?\n\nFiltering out outliers and clipping the gradients are standard practices in variational Monte Carlo methods for molecules. For the baseline, we precisely follow [1] where different clippings were used to stabilize the optimization procedure. For the proposed method, we design analogs for the new objective and use the same tricks as proposed in [1] where it is possible.\n\n> Which would be the recommended method? The WQMC or the W(FR)QMC? What would be the trade-off?\n\nWe observe that the Fisher-Rao metric demonstrates faster convergence in the beginning and the Wasserstein metric converges faster later into the optimization. Our intuition is that the usage of both metrics at the same time allows us adaptively decide which distance to the ground state is larger at the moment and minimize it.   \nIndeed, one can introduce a hyperparameter controlling the tradeoff between metrics. In our experiments, we haven’t investigated the influence of this hyperparameter, and simply take the sum of two metrics for W(FR)QMC. We leave an extensive study of this parameter for future work.\n\n[1] von Glehn, Ingrid, James S. Spencer, and David Pfau. \"A Self-Attention Ansatz for Ab-initio Quantum Chemistry.\" The Eleventh International Conference on Learning Representations. 2022."}, {"review_id": "v26BPL4ma6", "reviewer": "Reviewer_n7Uw", "summary": "The authors propose a novel approach named Wasserstein Quantum Monte Carlo, which uses the gradient flow induced by the Wasserstein and Wasserstein Fisher-Rao metrics and improves the convergence rate of Quantum Variational Monte Carlo. The numerical results show that following the gradient flow under Wasserstein and Wasserstein Fisher-Rao metrics results in better convergence to the ground state wave function.", "questions": "The numerical results show that the proposed method outperforms the QVMC. But it seems not that good compared to variational quantum eigensolvers (such as UCCSD or adaptVQE) on the ground-state energy estimation problems. Could you illustrate the difference between the proposed method and those VQE methods?", "limitations": "Wrong track to submit this paper to NeurIPS.", "rating": 5, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper is technically sound with detailed derivations.\n2. Seeking new approaches to better solve quantum chemistry and quantum physics problems is very important to the quantum machine learning society.", "weaknesses": "I'm not familiar with this subject but I strongly believe that this paper belongs to a Physics conference or journal instead of NeurIPS.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time spent and valuable feedback.\n\n> I'm not familiar with this subject but I strongly believe that this paper belongs to a Physics conference or journal instead of NeurIPS\n\nWe respectfully disagree that this paper does not belong to NeurIPS. The NeurIPS community has a rich tradition of ML applications to other fields of science including natural sciences, e.g. numerous AI4Science workshops evidence for this (e.g., see NeurIPS 2022 \"Machine Learning and the Physical Sciences\" where many papers including the opening keynote were about quantum variational Monte Carlo). Moreover, different approaches to this particular problem were recently published at the top ML conferences including NeurIPS [1,2,3]. While in this paper, we introduced our approach in the context of energy functional minimization for quantum systems, it is quite general-purposed: it can be applied to minimize any functional directly in the space of distributions, making it relevant to the wider audience within the NeurIPS community. For example, instead of using the quantum energy functional one might consider the KL divergence functional $\\text{KL}(p \\Vert q)$ between the data distribution $p$ and the model distribution $q$. Then the Fisher--Rao metric corresponds to the energy-based model training and the Wasserstein metric corresponds to score-matching. We will draw these connections and expand on this discussion in the next revision of the paper.\n\n> Could you illustrate the difference between the proposed method and those VQE methods?\n\nThe considered Variational Monte Carlo methods (both the baseline and the proposed method) are classical computation approaches to the quantum chemistry problem. On the other side, (adapt)VQE is an algorithm for quantum devices, and cannot be efficiently run until efficient quantum devices are available. Hence, we provide no comparison. If compared to CCSD, which is the classical version of UCCSD, the approaches considered in the paper are known to be much more scalable (see Fig. 1 in [4]).\n\n[1] Gerard, Leon, et al. \"Gold-standard solutions to the Schrödinger equation using deep learning: How much physics do we need?.\" Advances in Neural Information Processing Systems 35 (2022): 10282-10294.  \n[2] Gao, Nicholas, and Stephan Günnemann. \"Generalizing neural wave functions.\", International Conference on Machine Learning 2023.  \n[3] von Glehn, Ingrid, James S. Spencer, and David Pfau. \"A Self-Attention Ansatz for Ab-initio Quantum Chemistry.\" The Eleventh International Conference on Learning Representations. 2022.  \n[4] Hermann, Jan, et al. \"Ab-initio quantum chemistry with neural-network wavefunctions.\" arXiv preprint arXiv:2208.12590 (2022)."}, {"review_id": "gyKzVrVcyf", "reviewer": "Reviewer_Nk4c", "summary": "The authors show that the optimization objective typically used in variational Monte Carlo (VMC) solutions of the Schrödinger equation can be interpreted as the Fisher-Rao gradient flow in the space of Born distributions. Based on this insight, they suggest to substitute the Fisher-Rao metric with Wasserstein metrics, which has the effect that during optimization, probability mass is \"transported gradually\" rather than \"teleported\" in space. Since the VMC optimization process typically relies on sampling \"walkers\" from the probability distribution, which evolve over time according to a Markov process, a gradual \"transport\" of the probability mass promises improved convergence. The authors empirically demonstrate that their proposed method leads to lower ground state energies (and variance) for several small systems.", "questions": "- Have the authors considered applying their method also to larger test systems (e.g. benzene), or even just small molecular test systems (say ethanol, methane, ...)? Currently the method is only tested for two atoms, the lithium dimer, and the H10 chain. I encourage the authors to extend their empirical evaluation to more complicated examples (if they can afford it).\n\n- The plots in Figure 2 show the number of iterations, but they do not show the associated runtime costs. Is a single iteration of WQMC comparable in cost to a single iteration of QVMC?\n\n- A single iteration of VMC consists of several individual steps including the parameter update followed by sampling walker positions for the updated probability density. The sampling of walker positions typically uses a certain number of steps in a Markov process. Considering the \"more well-behaved\" updates of the probability density with WQMC compared to QVMC, is it possible to reduce the required number of sampling steps for the walker positions? If the authors have not considered this, I encourage them to try. Perhaps it is possible to decrease the runtime costs of a single iteration this way?", "limitations": "The authors do not discuss limitations of their work, nor its potential for negative societal impact. While the latter is acceptable in my opinion (this work is about an algorithmic method to solve a fundamental problem in quantum chemistry and a discussion of negative societal impacts would probably appear contrived), a discussion of limitations should be added in a revised version.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper introduces a very simple (although non-trivial!) idea and is clearly presented. The proposed optimization methods show clear practical advantages and the underlying theory may lead to even more sophisticated optimization methods for VMC solutions of the Schrödinger equation in the future.", "weaknesses": "The empirical evaluation of the proposed method is limited to relatively few and small systems. While the results are very promising, it is therefore unclear how the proposed method compares to the standard approach when applied to more complicated, larger systems (e.g. benzene). This being said, I find this limitation acceptable considering the large computational cost of running VMC calculations for larger systems (which the authors may not be able to afford).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful and positive feedback. We are glad that you appreciated both the simplicity and non-trivial nature of our idea. We also believe that our underlying theory has many potential implications for developing even better VMC solutions. In what follows we address the raised concerns and answer your questions.\n\n> The plots in Figure 2 show the number of iterations, but they do not show the associated runtime costs. Is a single iteration of WQMC comparable in cost to a single iteration of QVMC?\n\nWe include the runtime of all the algorithms in the plots (see attachment in the general response). Namely, for a proper comparison, instead of reporting the metrics per iteration we report them per seconds elapsed. We will include these plots in the next version of the paper. All the models were benchmarked on four A40 GPUs.\n\n> Considering the \"more well-behaved\" updates of the probability density with WQMC compared to QVMC, is it possible to reduce the required number of sampling steps for the walker positions? If the authors have not considered this, I encourage them to try. Perhaps it is possible to decrease the runtime costs of a single iteration this way?\n\nIn our experiments, we observed that the proposed method indeed can be trained with fewer MCMC steps (not included in the paper). Moreover, one can use the evaluated vector field (gradient of the local energy) as the first proposal update to the samples. We leave the detailed study of this degree of freedom beyond the scope of the current work, instead, we apply minimal changes to the optimization procedure to highlight the effect of the newly introduced metric.\n\n> I encourage the authors to extend their empirical evaluation to more complicated examples (if they can afford it).\n\nThe systems considered in the experiments are not “simple”. Indeed, a single run for the H10 system takes more than 33 GPU hours on Nvidia A40 GPU, and we can see that there is already a gap between the baseline accuracy and the chemical accuracy. In this work, we focus on extending the Variational Monte Carlo algorithms family, and illustrating that the proposed extensions might be indeed beneficial for the performance. Scaling the proposed method to bigger molecules requires a more comprehensive empirical study and substantial engineering infrastructure, which are beyond the scope of this project.\n\n> Discussion of limitations should be added in a revised version\n\nAs suggested by the reviewer, we will include a limitation section in the final paper. The main downside of the proposed method is the additional evaluation of the derivative, which results in longer iterations. Another issue, common to all Variational Monte Carlo methods, is the singularities around nuclei, which create instabilities during training. In the general response, we discuss how these limitations can be addressed in future work.\n"}, {"review_id": "p8u9U1wWkg", "reviewer": "Reviewer_Qbaw", "summary": "The paper provides a framework to derive a new loss for Variational Monte Carlo methods. They introduce Wasserstein Quantum Monte Carlo, which uses a gradient flow based on the Wasserstein metric. The method is empirically tested for fermionic systems, e.g. a Hydrogen chain and compared against another Quantum Monte Carlo approach based on the Fisher-Rao metric. Faster convergence in terms of number of epochs is observed. ", "questions": "1.\tIn algorithm 1 you stop the gradient for the gradient of the local energies. It is not clear to me from eq. 24, why this is the case?", "limitations": "All limitations are addressed.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The paper derives with a theoretical framework a new loss function for Quantum Monte Carlo calculations. \nThey underline their findings with experiments and improve upon state-of-the-art variational Monte Carlo calculations with respect to the number of epochs on rather small molecular systems.\nA lot of recent improvements in this field came mainly from architectural changes (e.g. PsiFormer), it is great to see also improvements on the loss function and optimization side.\n", "weaknesses": "The derived loss function is novel and improves with respect to the number of epochs current state-of-the-art calculation, I have the following concerns:\n\n1.\tThe lack of immediate relevance to the ML community, due to the rather small experimental section and the newly derived loss fits rather to the Quantum Chemistry community. \n2.\tWhat is the QVMC baseline in section 4, is it PsiFormer? What kind of settings did you use for the PsiFormer architecture? The authors of PsiFormer proposed two variants, a small and a larger one. \n3.\tThe loss WQMC loss seems to converge faster in terms of epochs but based on the Algorithm 1 my understanding is that you have to perform third derivatives which is computationally quite expensive. Could the authors add to Fig. 2 the timings per epoch for each method and the type of GPUs used to perform the experiments? \n4.\tIn case it is possible it would be also interesting to see a broader experimental section rather than these “simple” molecules to better understand the improvements gained through the newly defined loss. For example, the potential energy surface of the Nitrogen dimer is a difficult system (especially considering the bond breaking geometry), which was also analyzed by the authors from FermiNet. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback and the time spent. We are glad that you appreciate our developments in the design of new optimization procedures for variational Monte Carlo methods. In what follows we address the concerns raised and answer the questions.\n\n> The lack of immediate relevance to the ML community, due to the rather small experimental section and the newly derived loss fits rather to the Quantum Chemistry community.  \n\nWe respectfully disagree. The NeurIPS community has a rich tradition of ML applications to other fields of science including natural sciences, e.g. numerous AI4Science workshops evidence for this (e.g., see NeurIPS 2022 \"Machine Learning and the Physical Sciences\" where many papers including the opening keynote were about quantum variational Monte Carlo). Moreover, different approaches to this particular problem were recently published at the top ML conferences including NeurIPS [1,2,3]. While in this paper, we introduced our approach in the context of energy functional minimization for quantum systems, it is quite general-purposed: it can be applied to minimize any functional directly in the space of distributions, making it relevant to the wider audience within the NeurIPS community. For example, instead of using the quantum energy functional one might consider the KL divergence functional $\\text{KL}(p \\\\| q)$ between the data distribution $p$ and the model distribution $q$. Then the Fisher--Rao metric corresponds to the energy-based model training and the Wasserstein metric corresponds to score-matching. We will draw these connections and expand on this discussion in the next revision of the paper.\n\n> What is the QVMC baseline in section 4, is it PsiFormer? What kind of settings did you use for the PsiFormer architecture? The authors of PsiFormer proposed two variants, a small and a larger one.\n\nFor the baseline, we precisely follow [3]. Thus, as considered in [3] for the systems of similar size, we use small Psiformer architecture, the same hyperparameters, and the same techniques to stabilize optimization.\n\n> The loss WQMC loss seems to converge faster in terms of epochs but based on the Algorithm 1 my understanding is that you have to perform third derivatives which is computationally quite expensive. Could the authors add to Fig. 2 the timings per epoch for each method and the type of GPUs used to perform the experiments?\n\nAs suggested by the reviewer, for a proper comparison, we updated Fig. 2 with wall-clock time as the x-axis and uploaded the new figure in our global response. The type of GPUs used to run experiments is also mentioned in the global response. In short, all the claims of our paper still hold when considering wall-clock time. Third-order derivatives can be efficiently computed using modern deep learning frameworks such as JAX, which we used to implement our method. We will include these updates in the final paper.\n\n> In case it is possible it would be also interesting to see a broader experimental section rather than these “simple” molecules to better understand the improvements gained through the newly defined loss. For example, the potential energy surface of the Nitrogen dimer is a difficult system (especially considering the bond breaking geometry), which was also analyzed by the authors from FermiNet.\n\nThe systems considered in the experiments are not “simple”. Indeed, a single run for the H10 system takes more than 33 GPU hours on Nvidia A40 GPU, and we can see that there is already a gap between the baseline accuracy and the chemical accuracy. In this work, we focus on extending the Variational Monte Carlo algorithms family, and illustrating that the proposed extensions might be indeed beneficial for the performance. Scaling the proposed method to bigger molecules requires a more comprehensive empirical study and substantial engineering infrastructure, which are beyond the scope of this project.\n\n> In algorithm 1 you stop the gradient for the gradient of the local energies. It is not clear to me from eq. 24, why this is the case?\n\nThank you for bringing this up! The local energy in (24) is evaluated at $q_t(x) = q(x,\\theta)$ but does not depend on parameters $\\theta$. Indeed equation (24) is obtained by using (23) for the density time derivative in (17) and integration by parts. Note that the gradient w.r.t. $\\theta$ is applied only to the log density in (17) and to the gradient of log density after the integration by parts. We will further clarify this in the next revision of the paper.\n\n[1] Gerard, Leon, et al. \"Gold-standard solutions to the Schrödinger equation using deep learning: How much physics do we need?.\" Advances in Neural Information Processing Systems 35 (2022): 10282-10294.  \n[2] Gao, Nicholas, and Stephan Günnemann. \"Generalizing neural wave functions.\", International Conference on Machine Learning 2023.  \n[3] von Glehn, Ingrid, James S. Spencer, and David Pfau. \"A Self-Attention Ansatz for Ab-initio Quantum Chemistry.\" The Eleventh International Conference on Learning Representations. 2022.\n"}]

Author rebuttal:
We thank all reviewers for their time spent and valuable feedback on the paper. In what follows we would like to address the common concerns raised by most of the reviewers.

**Runtime.** As requested by most of the reviewers, we include the runtime of all the algorithms in the plots (see the new plots in the attached pdf to this response). Namely, for a proper comparison, instead of reporting the metrics per iteration we report them per wall time in seconds. Note that all the claims in the paper still hold in terms of wall time. We will include these plots in the next version of the paper. All the models were benchmarked on four A40 GPUs.

**Limitations.** As suggested by reviewer Nk4c, we will include a limitation section in the final paper. The main downside of the proposed method is the additional evaluation of the derivative, which results in longer iterations. Another issue, common to all Variational Monte Carlo methods, is the singularities around nuclei, which create instabilities during training.

**Future work.** We believe that the aforementioned limitations can be addressed in future work with an extensive empirical study. 
1. One direction is to come up with more efficient Monte Carlo schemes (updates of the samples). Indeed, in our experiments, we observed that the proposed method indeed requires fewer MCMC steps (not included in the paper). Moreover, one can use the evaluated vector field (gradient of the local energy) as the first proposal update to the samples. This would require additional tuning of the hyperparameters (such as the proposal step) but would allow to decrease the runtime.
2. Another direction is the usage of the effective core potentials for larger atoms, which help alleviate instabilities created by cusps around nuclei (see, e.g. [1]). 
3. Finally, one can study alternative metrics by choosing the tradeoff between Fisher-Rao and Wasserstein metrics. In our experiments, we observe that such a tradeoff exists since the joint metric W(FR)QMC (without tuned coefficient) stably outperforms both metrics used separately.

**Scaling experiments.** The systems considered in the experiments are not “simple”. Indeed, a single run for the H10 system takes more than 33 GPU hours on Nvidia A40 GPU, and we can see that there is already a gap between the baseline accuracy and the chemical accuracy. In this work, we focus on extending the Variational Monte Carlo algorithms family, and illustrating that the proposed extensions might be indeed beneficial for the performance. Scaling the proposed method to bigger molecules requires a more comprehensive empirical study and substantial engineering infrastructure, which are beyond the scope of this project.

[1] Li, Xiang, et al. "Fermionic neural network with effective core potential." Physical Review Research 4.1 (2022): 013021.

Meta-review: This paper provides new insights into quantum variational Monte Carlo (QVMC) by interpreting QVMC as a different objective, leading to alternative QMC algorithms based on Wasserstein gradient flow. Overall, reviewers came to a consensus about accepting this paper and appreciated the clear presentation and the empirical results presented with the proposed approach. Finally, reviewers noted that in addition to the method being a interesting contribution on its own, the theoretical insights from this paper may also lead to the development of novel VMC algorithms.
Decision: Accept (spotlight)