Title: Wasserstein Quantum Monte Carlo: A Novel Approach for Solving the Quantum Many-Body Schrödinger Equation
Venue: NeurIPS 2023
Accepted: yes
Keywords: Quantum Monte Carlo

Schrödinger equation

Wasserstein Fisher-Rao gradient flow
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

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.

1 Introduction

Access to the wave function of a quantum many-body system allows us to study strongly correlated quantum matter, starting from the fundamental building blocks. For example, the solution of the timeindependent electronic Schrödinger equation provides all the chemical properties of a given atomic state, which have numerous applications in chemistry and materials design. However, obtaining the exact wave function is fundamentally challenging, with a complexity scaling exponentially with

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) is a well-known subclass of the latter that can, in principle, be used to estimate the lowest-energy state (i.e. ground state) of a 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 a 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 et 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, as 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 non-parametric manifold. Then, the QVMC algorithm can be seen as a projection of this gradient flow onto a parametric manifold (see Section 3 for details). Second, the gradient flow perspective gives us an additional degree of freedom in the algorithm. Namely, we can choose the metric under which we define the gradient flow. Thus, we propose and study a different energyminimizing objective function, which we derive as a gradient flow under the Wasserstein metric (or Wasserstein Fisher-Rao metric) (Chizat et al., 2018; Kondratyev et al., 2016).

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 running 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,

where x a given many-body configuration and V is the potential operator. The time-dependent Schrödinger equation determines the wave function ψ ( x, t ) of the quantum system

As is often the case, we will target the stationary solutions, for which we focus on solving the time-independent Schrödinger equation

where E is the energy of the state ψ . The ground state of a quantum system is obtained by solving the time-independent Schrödinger equation, by targeting the eigenstate ψ of the above Hamiltonian

with the lowest eigenvalue (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 q ( x ) = | ψ ( x ) | 2 , the energy of a given quantum state can be rewritten in a functional form,

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 log | ψ | )

under the strong condition that q ( x ) is the Born probability density derived from an (anti-)symmetric wave function: q ( x ) = ψ 2 ( x ) . The latter will always be tacitly assumed from hereon.

The Rayleigh-Ritz principle guarantees that the E [ q ] is lower-bounded by the true ground-state energy of the system, i.e. E [ q ] ≥ E 0 , if the corresponding wave function ψ is a valid state of the corresponding Hilbert space. Quantum Variational Monte Carlo (QVMC) targets ground states by parametrizing the wavefunction ψ ( x, θ ) and by minimizing E [ q ( θ )] . The solution to the minimization problem θ 0 = arg min θ E [ q ( θ )] is obtained by gradient-based methods using the following expression for the gradient w.r.t. parameters θ

In sum, the above leads to an iterative procedure in which Monte Carlo sampling is used to generate configurations from the current trial state q ( x, θ ) = ψ 2 ( x, θ ) , which allows computing the corresponding energy and its parameter gradients, and to update the model accordingly. In practice, the parametric model specifies the density q ( x, θ ) only up to a normalization constant, i.e., it outputs ˜ q ( x, θ ) ∝ q ( x, θ ) . However, the gradient w.r.t. θ does not depend on the normalization constant; hence, throughout the paper, we refer to the model as the normalized density q ( x, θ ) .

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 to a variational/parametric manifold. The latter is a more common problem often tackled in machine learning, and by forging connections between both fields, we will be able to derive an alternative to QVMC.

Gradient Flows In Euclidean space, we can minimize a function f : R d → R by following the ODE d d t x t = -∇ x f ( x t ) , 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 a 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).

Wasserstein Fisher-Rao gradient flows Consider the space of distributions P 2 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 g t , and by accounting for the growth term in the modified continuity equation. Namely, the distance between probability densities p 0 and p 1 is defined as

where v t ( x ) is the vector field defining the probability flow, g t ( x ) is the growth term controlling the creation and annihilation of the probability mass, and λ is the coefficient balancing the transportation and teleportation costs. Note that by setting one of the terms to zero we get 2-Wasserstein distance ( g t ( x ) ≡ 0 ) and Fisher-Rao distance ( v t ( x ) ≡ 0 ). In Section 3, we also consider the general case of c -Wasserstein distance, where c is a convex cost function on the tangent space.

Given a functional on this manifold, F [ q ] : P 2 → R , we 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 F [ q ] under the Wasserstein Fisher-Rao metric is given by the following PDE (which is shown with detailed derivations in Appendix A)

where δF [ q ] /δq is the first-variation of of F with respect to the L 2 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 ∼ q t ( x ) follow the vector field v t ( x ) = -∇ x δF [ q t ] /δq t . The second term of the PDE defines the creation and annihilation of probability mass, and is proportional to the growth field g t ( x ) = δF [ q t ] δq t ( x ) -E q t ( y ) [ δF [ q t ] δq t ( y ) ] . Note that E q t [ g t ] = 0 , and so while mass can be 'teleported', the total mass (or probability) will remain constant. The two mechanisms can be considered independently by defining the evolution terms under the 2-Wasserstein and Fisher-Rao metrics respectively, i.e.

It now becomes evident that the stationary condition for all the considered PDEs is

In Appendix A, we provide derivations illustrating that Eqs. (8) to (10) correspond to the gradient flow under the Wasserstein Fisher-Rao, Wasserstein, and Fisher-Rao metrics, respectively, and hence they all minimize F [ q ] . For detailed analysis, we refer the reader to Kondratyev et al. (2016); Liero et al. (2016).

3 Methodology

In Section 3.1, we first demonstrate that the imaginary-time evolution of the Schrödinger equation can be viewed as a gradient flow under the Fisher-Rao metric. Afterwards, in Section 3.2, we discuss how a density evolution can be projected to 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-Rao) 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 in theory be obtained by imaginary-time evolving any valid quantum state ψ (with a non-vanishing overlap with the true ground state) to infinite times. The state is evolved according to the imaginary-time Schrödinger equation, which defines the energyminimizing time evolution of the wavefunction ψ t , and is expressed as the following PDE (which is

the Wick-rotated version of Eq. (2), see e.g. (McArdle et al., 2019; Yuan et al., 2019)),

where again q t ( x ) = ψ 2 t ( x ) . The last term proportional to the energy E [ ψ t ] comes from enforcing normalization (contrary to real-time evolution, imaginary time evolution is non-unitary).

Theorem 3.1. Eq. (12) defines the gradient flow of the energy functional E [ q ] under the Fisher-Rao metric.

Proof Sketch. The energy functional E [ q ] has the following derivative

Thus, the gradient flow under the Fisher-Rao metric is (see Eq. (10))

which is equivalent (up to a 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 can be derived following the derivations from Stokes et al. (2020), but not introducing the manifold of parametric distributions. However, considering the evolution of the density on the non-parametric manifold first helps us to derive our method and relating it 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 in the distributional space and following the energy-minimizing 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 to any gradient flow (e.g., Eqs. (8) to (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 is tractable.

Suppose the current density on the parametric manifold is q t ( x ) = q ( x, θ ) (see Figure 1). We first evolve this density using a (non-parametric) gradient flow method (e.g., Eqs. (8) to (10)) for time ∆ t , which will take q t ( x ) off the parametric manifold to q t +∆ t ( x ) . We then have to update current trial model q ( x, θ ) to match q t +∆ t ( x ) enabling us to propagate the density further. In order to do so, we define the optimal update of parameters ∆ θ ∗ as the minimizer of the Kullback-Leibler divergence between q t +∆ t ( x ) and the distributions on the parametric manifold, i.e.

In practice, we evaluate the parameters update using the following proposition.

Proposition 3.2. For q t ( x ) = q ( x, θ ) , the gradient of reverse KL-divergence can be approximated as

Using this approximation, the optimal update from Eq. (15) of parameters becomes

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.

This perspective lays the foundation for deriving our WQMC method in Section 3.3, by following the Wasserstein or WFR gradient flows rather than Fisher-Rao gradient flows.

Natural Gradient Preconditioning In order to update the parametric model q ( x, θ ) , instead of following the update in Eq. (17), we can exploit the information geometry of the statistical manifold of q ( x, θ ) , and define the update using the Fisher information matrix F θ

This update is analogous to the natural gradient update (Amari, 1998). Note that the choice of Fisher information as the metric on the statistical manifold for preconditioning the gradient and updating θ is independent of the choice of metric on the non-parametric Wasserstein manifold (e.g., Wasserstein or Fisher-Rao) for evolving q t ( x ) . In practice, we use Kronecker-factored approximation of the natural 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, we 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-parametric) space of distributions, which results in a different gradient flow and corresponds to a different algorithm. Namely, we propose to consider the gradient descent in 2-Wasserstein space as the energy-minimizing density evolution, as introduced in Section 2.2.

Theorem 3.4. The energy-minimizing 2-Wasserstein gradient flow is defined by the continuity equation

Proof. In Theorem 3.1, we show that δE [ q ] /δq = E loc ( x ) . Plugging this into the 2-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 -Wasserstein metric with any convex cost function c : R d → R on the tangent space. The c -Wasserstein distance between p 0 and p 1 is defined as follows

Proposition 3.5. The energy-minimizing c -Wasserstein gradient flow is defined by the following equation

where c ∗ ( · ) is the convex conjugate function of c ( · ) , and ∇ c ∗ ( y ) is its gradient at y .

Theorem 3.4 can be viewed as a special case of Proposition 3.5 where c ( · ) = 1 2 ∥·∥ 2 . Introducing a different c than L 2 norm translates to a non-linear transformation of the gradient -∇ x E loc ( x ) . In Appendix D, we demonstrate how to choose c such that it corresponds to 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.

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 1, we provide a pseudocode for the proposed algorithms. The procedure follows closely QVMC but introduces a different objective. When using gradients both from Eqs. (18) and (24), we follow the gradient flow under the Wasserstein Fisher-Rao metric with the coefficient λ . For λ →∞ , the cost of mass teleportation becomes infinite and we use only the gradient from Eq. (24), which corresponds to the gradient flow under the c -Wasserstein metric (we refer to this algorithm as WQMC). For λ → 0 , 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 ( λ → 0 , ∞ ) and the mixed case λ = 1 .

1e-3

1e-4

1e-5

1e-6

Relative energy error

Be (4)

W(FR)QMC

WQMC

chemical accuracy

1e-3

1e-4

1e-5

1e-6

B (5)

WVMC

chemical accuracy

1e-3

5e-4

1e-4

2e-5

Li2 (6)

W(FR)QMC

WQMC

chemical accuracy

1e-3

6e-4

3e-4

1e-4

H10 (10)

W(FR)QMC

WQMC

chemical accuracy

5e-1

1e-1

1e-2

1e-3

Energy variance (

Ha

2

)

W(FR)QMC

WQMC

3e0

1e0

1e-1

1e-2

WVMC

1e0

1e-1

1e-2

1e-3

W(FR)QMC

WQMC

4e-2

1e-2

1e-3

2e-4

W(FR)QMC

WQMC

2000

4000

6000

8000

10000

2e0

1e0

3e-1

1e-1

Gradient norm

W(FR)QMC

WQMC

2000

4000

6000

8000

10000

2e-1

5e-1

2e0

5e0

WVMC

2000

4000

6000

8000

10000

1e-1

3e-1

8e-1

2e0

W(FR)QMC

WQMC

2000

6500

11000

15500

20000

4e-2

8e-2

2e-1

5e-1

W(FR)QMC

WQMC

4 Experiments 1

For the empirical study of the proposed method, we consider Born-Oppenheimer approximation of chemical systems. Within this approximation, the wave function of the electrons in a molecule can be studied separately from the wave function of the atomic nuclei. Namely, we consider the following Hamiltonian

where x i are the coordinates of electrons, X I , Z I are the coordinates and charges of nuclei. The first kinetic term contains derivatives with respect to the electron positions x . Indeed, the positions of the nuclei are given and fixed, and we target the ground state of the electronic wave function ψ ( x ) , which is an explicit function of the electron positions only. Solving the electronic Schrödinger equation is a notoriously difficult task, and is a topic of intense research in quantum chemistry and material sciences.

Since electrons are indistinguishable fermions, we restrict the Hilbert space to states ψ 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 ψ ( x, θ ) , 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 q ( x, θ ) = | ψ ( x, θ ) | 2 . 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 ∇ x E loc ( 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 ∥∇ x log q ( x, θ ) ∥ significantly exceeds the median value. Namely, we estimate the deviation from the norm as E q ( x,θ ) [ |∥∇ x log q ( x, θ ) ∥ -median ( ∥∇ x log q ( x, θ ) ∥ ) | ] 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.e. 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 2 molecule and the hydrogen chain H 10 from (Hermann et al., 2020). The exact values of energies for Be, B, Li 2 are taken from (Pfau et al., 2020), the exact value of the energy for H 10 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 state-of-the-art architectures as a 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(FR)QMC yield significantly faster convergence of the aforementioned metrics compared to QVMC. In particular, for the larger molecules Li 2 and H 10 , we observe that we consistently obtain lower energies within 10 k steps ( 20 k for H 10 ) 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 metrics 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 can 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, one can use the evaluated gradient of the local energy as the proposal vector field for updating the samples. This would allow to decrease the number of MCMC steps at the low cost of additional hyperparameter tuning.

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 metric, which allows for probability mass 'teleportation', the Wasserstein metric constrains the evolution of the density to local changes under the locally-optimal transportation plan, i.e., following fluid dynamics. This property is favorable when the evolution of the parametric model is accompanied by the evolution of samples (performed by an MCMC algorithm). Indeed, by forbidding or regularizing non-local mass 'teleportation' in 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 in 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 time-independent Schrödinger equation for time-reversal symmetric Hamiltonians based on optimal transport will open new avenues to develop improved numerical methods for quantum chemistry and physics.

Connection to Energy-Based and Score-Based Generative Models The developed ideas of this paper, i.e., 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-Wasserstein gradient flow corresponds to the score-matching, which is used for training diffusion generative models.

References

Amari, S.-I. Natural gradient works efficiently in learning. Neural computation , 10(2):251-276, 1998.

Ambrosio, L., Gigli, N., and Savaré, G. Gradient flows: in metric spaces and in the space of probability measures . Springer Science & Business Media, 2005.

Benamou, J.-D. and Brenier, Y. A computational fluid mechanics solution to the Monge-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., Wanderman-Milne, S., and Zhang, Q. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax .

Carleo, G. and Troyer, M. Solving the quantum many-body problem with artificial neural networks. Science , 355(6325):602-606, 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.-X. 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, I. Implicit generation and generalization in energy-based models. arXiv preprint arXiv:1903.08689 , 2019.

Gao, N. and Günnemann, S. Generalizing neural wave functions. arXiv preprint arXiv:2302.04168 , 2023.

Gnech, A., Adams, C., Brawand, N., Carleo, G., Lovato, A., and Rocco, N. Nuclei with up to a= 6 nucleons with artificial neural network wave functions. Few-Body Systems , 63(1):7, 2022.

Hermann, J., Schätzle, Z., and Noé, F. Deep-neural-network solution of the electronic Schrödinger equation. Nature Chemistry , 12(10):891-897, 2020.

Hermann, J., Spencer, J., Choo, K., Mezzacapo, A., Foulkes, W., Pfau, D., Carleo, G., and Noé, F. Abinitio quantum chemistry with neural-network wavefunctions. arXiv preprint arXiv:2208.12590 , 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 , 21(11/12):1117-1164, 2016.

Liero, M., Mielke, A., and Savaré, G. Optimal transport in competition with reaction: The HellingerKantorovich distance and geodesic curves. SIAM Journal on Mathematical Analysis , 48(4): 2869-2911, 2016.

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 kronecker-factored approximate curvature. In International conference on machine learning , pp. 2408-2417. 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.

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, 2001.

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, J. S., Matthews, A. G., and Foulkes, W. M. C. Ab initio solution of the manyelectron 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.13456 , 2020.

Sorella, S. Green function Monte Carlo with stochastic reconfiguration. Phys. Rev. Lett. , 80:45584561, May 1998. doi: 10.1103/PhysRevLett.80.4558. URL https://link.aps.org/doi/10. 1103/PhysRevLett.80.4558 .

Stokes, J., Izaac, J., Killoran, N., and Carleo, G. Quantum natural gradient. Quantum , 4:269, 2020.

Vicentini, F., Hofmann, D., Szabó, A., Wu, D., Roth, C., Giuliani, C., Pescia, G., Nys, J., VargasCalderó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 arXiv: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, J., Ferrari, F., Georges, A., Hibat-Allah, M., et al. Variational benchmarks for quantum many-body problems. arXiv preprint arXiv:2302.04919 , 2023.

Xie, J., Lu, Y ., Zhu, S.-C., and Wu, Y . A theory of generative convnet. In 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.

A Gradient flows under Wasserstein Fisher-Rao metric

We, first, remind the concept of the functional derivative. The change of the functional F [ q ] : P 2 → R along the direction h can be expressed as

Consider a change of the density in time, the change of the functional can be defined through the differential as

In particular, we have

A.1 Minimizing movement scheme for 2-Wasserstein distance and Kullback-Leibler divergence

Gradient flow under W 2 Consider the following minimizing movement scheme (MMS)

where the change of the density is restricted to the continuity equation, i.e.,

Using the static formulation of W 2 distance, we have

where T ∗ ( x ) is the optimal transportation plan, and v ∗ ( x ) is the corresponding optimal gradient field. Thus, we can rewrite the MMS problem as

From the last optimization problem, we have

Gradient flow under KL Consider the following minimizing movement scheme (MMS)

where the change of the density is restricted to the following weighting scheme

The KL-divergence is then

In the last equation we are using hence,

Thus, we can rewrite the MMS problem as

From the last optimization problem, we have

Note, however, that D KL ( ·∥· ) is not the same as the Fisher-Rao metric. The derivations here demonstrate that the Fisher-Rao gradient flow can be derived as the MMS scheme with the KLdivergence.

A.2 Minimizing movement scheme for the Wasserstein Fisher-Rao metric

Consider the Wasserstein Fisher-Rao distance

The minimizing movement scheme (MMS) for this distance is

where the change of the density is given by the continuity equation with the growth term

For close enough q and q ′ , WFR λ ( q, q ′ ) 2 can be estimated via the metric derivative | µ ′ t | 2 that is defined as

Thus, the MMS problem can be written as

From the last optimization problem, we have

Note that different values of λ result in different gradient flows. For instance, considering the limit λ → ∞ we have g ( x ) → 0 and Eq. (52) just becomes W 2 gradient flow, which is natural since we have an 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 λ → 0 , our metric does not penalize for the mass teleportation at all, but our change of 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 q t under the continuity equation with the vector field v t ( x ) and the growth term g t ( x )

Thus, the change of the functional F [ q ] is

From this equation, we can clearly see that v t ( x ) and g t ( x ) derived in the previous section minimize F [ q ] . Indeed, taking

we get

Note that the considered growth term preserves the normalization of the density, i.e.,

Thus, our functional F [ q ] decreases when the density q t evolves according to the PDE

and reaches its stationary point when ∥ ∥ ∥ ∇ x δF [ q t ] δq t ( x ) ∥ ∥ ∥ 2 = 0 , i.e., δF [ q t ] δq t ( x ) ≡ 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 Fisher-Rao metrics respectively. The corresponding PDEs are

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 E [ q ] under the Fisher-Rao metric.

Proof. First, we derive the functional derivative of the energy functional E [ q ] . We denote the differential of the functional F ( q ) along the direction h as

where δF [ q ] /δq is the functional derivative.

Consider the energy functional

The functional derivative of this functional is as follows

For the last term, we do integration by parts and get

Thus, we have

and we see that the derivative coincides with the local energy, i.e.,

Using the results from Section 2.2, the energy-minimizing gradient flow under Fisher-Rao metric is

Second, we derive the PDE for the time-evolution of the density q t under the imaginary-time Schrödinger equation.

(78)

Using the identity

which is equivalent to Eq. (74).

C Analogies to generative modeling and variational inference literature

To draw a connection to generative models and variational inference literature, we 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 KL divergence functional

Consider the reverse KL divergence objective, and its first variation

where p ( x ) is the data distribution given empirically.

Fisher-Rao Gradient Flow Using equations Eq. (10), the PDE that defines the Fisher-Rao gradient flow is

Thus we have

where q 0 is the initial distribution. This results in a distributional path constructed by the linear average (mixture) of the end-point distributions. Note that changing the mixture parameter corresponds to teleporting the particles.

we have

Projected Fisher-Rao Gradient Flow Having the functional minimizing PDEs, we can find the update corresponding to the projected Fisher-Rao gradient flow using Proposition 3.2 and Eq. (17),

This corresponds to the standard maximum likelihood objective that is used for training of energybased models (Xie et al., 2016; Du & Mordatch, 2019).

Wasserstein Gradient Flow Using the equation Eq. (9), the PDE that defines the Wasserstein gradient flow is

Projected Wasserstein Gradient Flow Having the functional minimizing PDEs, we can find the update corresponding to the projected Wasserstein gradient flow using Proposition 3.2 and Eq. (17).

Note that similar to Proposition 3.2, we use q t ( x ) = q ( x, θ ) as the density equal to the model density but detached from the parameters θ . This objective corresponds to the score-matching objective (Hyvärinen & Dayan, 2005), which is widely used in the diffusion-based generative models (Song et al., 2020).

C.2 Variational inference: forward KL divergence functional

Consider the forward KL divergence objective, and its first variation

where p ( x ) 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

Thus we have

where q 0 is the initial distribution, and Z t is the partition function. This results in a 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, we can find the update corresponding to the projected Fisher-Rao gradient flow using Proposition 3.2 and Eq. (17),

This corresponds to the standard variational inference objective.

Wasserstein Gradient Flow Using the equation Eq. (9), the PDE that defines the Wasserstein gradient flow is

This is the Fokker-Planck equation, which characterize the movement of the particles according to Langevin dynamics:

This is a common approach for sampling from the unnormalized density p ( x ) via Markov Chain Monte Carlo algorithms.

Projected Wasserstein Gradient Flow Having the functional minimizing PDEs, we can find the update corresponding to the projected Wasserstein gradient flow using Proposition 3.2 and Eq. (17),

which resembles the score-matching objective (Hyvärinen & Dayan, 2005) for variational inference.

D c-Wasserstein gradient flow

c -Wasserstein distance with the convex cost function c : R d → R is defined

where Γ( p 0 , p 1 ) is the set of all possible couplings of p 0 and p 1 . The dynamic formulation of this distance is the following

Proposition. The energy-minimizing c -Wasserstein gradient flow is defined by the following PDE

where c ∗ ( · ) is the convex conjugate function of c ( · ) .

Proof. The movement minimizing scheme for W c ( · , · ) is the following optimization problem

Assuming that the density changes according to the continuity equation q ′ ( x ) = q t ( x ) -∆ t ∇ x · ( q t ( x ) v t ( x )) + o (∆ t ) , and ∆ t is small enough so that v t ( x ) defines the optimal transportation plan, we have

and the infimum is achieved at

which gives the formula for the vector field. Using the energy gradient from Theorem 3.1 δE [ q t ] δq t ( x ) = E loc, we get the result.

Proposition. Coordinate-wise application of tanh to the vector field, i.e.

corresponds to gradient descent with c -Wasserstein distance, where c : R d → R is the following cost function

Proof. Consider c ∗ ( x ) = ∑ i log(exp( x i ) + exp( -x i )) . It corresponds to applying hyperbolic tangent non-linearity coordinate-wise to the vector field field, i.e.,

The corresponding cost function c ( x ) is the following

E Additional experimental results

1e-3

1e-4

1e-5

1e-6

Relative energy error

Be (4)

W(FR)QMC

WQMC

chemical accuracy

1e-3

1e-4

1e-5

1e-6

B (5)

W(FR)QMC

WQMC

chemical accuracy

1e-3

5e-4

1e-4

2e-5

Li2 (6)

W(FR)QMC

WQMC

chemical accuracy

1e-3

6e-4

3e-4

1e-4

H10 (10)

W(FR)QMC

WQMC

chemical accuracy

5e-1

1e-1

1e-2

1e-3

Energy variance (

Ha

2

)

W(FR)QMC

WQMC

3e0

1e0

1e-1

1e-2

W(FR)QMC

WQMC

1e0

1e-1

1e-2

1e-3

W(FR)QMC

WQMC

4e-2

1e-2

1e-3

2e-4

W(FR)QMC

WQMC

1000

2000

3000

4000

wall time (seconds)

2e0

1e0

3e-1

1e-1

Gradient norm

W(FR)QMC

WQMC

1000

2000

3000

4000

5000

wall time (seconds)

2e-1

5e-1

2e0

5e0

W(FR)QMC

WQMC

2000

4000

6000

wall time (seconds)

1e-1

3e-1

8e-1

2e0

W(FR)QMC

WQMC

5000 1000015000200002500030000

wall time (seconds)

4e-2

8e-2

2e-1

5e-1

W(FR)QMC

WQMC

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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\nWEAKNESSES:\nThe 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\nQUESTIONS:\n1.\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?\n\nLIMITATIONS:\nAll limitations are addressed.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-23"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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).\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-28"}

comments: 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.

STRENGTHS:
1. This paper is technically sound with detailed derivations.
2. 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.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

comments: 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.

STRENGTHS:
- The paper is very well written.
- The proposed method is shown to include the conventional method.
- 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.

QUESTIONS:
- Since some tricks are used in experiments. How does these tricks influence the performance? Are those tricks also used for baseline methods?
-  Which would be the recommended method? The WQMC or the W(FR)QMC? What would be the trade-off?

LIMITATIONS:
None
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-16

comments: 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.

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. 
Paper seems to state that WQMC is better for all cases; if so, more experiments would make this claim stronger.

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?

Were other c- cost functions tried for the c-Wasserstein gradient flow? why does the coordinate-application of tanh work well? 

stop_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? 

Could authors include the additional computation complexity (theoretically, empirically) required for the Wasserstein gradient ?
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-22