SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Wasserstein Quantum Monte Carlo:
A Novel Approach for Solving
the Quantum Many-Body Schrödinger Equation
Kirill Neklyudov
Jannes Nys
Institute of Physics & Center for Quantum Science and Engineering École Polytechnique Fédérale de Lausanne (EPFL)
Luca Thiede
Juan Felipe Carrasquilla
Univeristy of Waterloo
Qiang Liu
Alireza Makhzani
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.
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 time-independent 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
Correspondence to: <k.necludov@gmail.com>, <makhzani@vectorinstitute.ai>
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 energy-minimizing 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.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,
$$
H = - \frac {1}{2} \nabla_ {x} ^ {2} + V. \tag {1}
$$
where  $x$  a given many-body configuration and  $V$  is the potential operator. The time-dependent Schrödinger equation determines the wave function  $\psi (x,t)$  of the quantum system
$$
i \frac {\mathrm {d}}{\mathrm {d} t} \psi (x, t) = H \psi (x, t) \tag {2}
$$
As is often the case, we will target the stationary solutions, for which we focus on solving the time-independent Schrödinger equation
$$
H \psi (x) = E \psi (x) \tag {3}
$$
where  $E$  is the energy of the state  $\psi$ . The ground state of a quantum system is obtained by solving the time-independent Schrödinger equation, by targeting the eigenstate  $\psi$  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) = |\psi(x)|^2$ , the energy of a given quantum state can be rewritten in a functional form,
$$
E [ \psi ] = \mathbb {E} _ {q (x)} \left[ E _ {\operatorname {l o c}} (x) \right], \quad E _ {\operatorname {l o c}} (x) := \frac {[ H \psi ] (x)}{\psi (x)} \tag {4}
$$
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 |\psi|$ )
$$
E [ q ] = \mathbb {E} _ {q (x)} [ E _ {\mathrm {l o c}} (x) ], \quad E _ {\mathrm {l o c}} (x) = V (x) - \frac {1}{4} \nabla_ {x} ^ {2} \log q (x) - \frac {1}{8} \| \nabla_ {x} \log q (x) \| ^ {2}, \tag {5}
$$
under the strong condition that  $q(x)$  is the Born probability density derived from an (anti-)symmetric wave function:  $q(x) = \psi^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] \geq E_0$ , if the corresponding wave function  $\psi$  is a valid state of the corresponding Hilbert space. Quantum Variational Monte Carlo (QVMC) targets ground states by parametrizing the wavefunction  $\psi(x,\theta)$  and by minimizing  $E[q(\theta)]$ . The solution to the minimization problem  $\theta_0 = \arg \min_{\theta} E[q(\theta)]$  is obtained by gradient-based methods using the following expression for the gradient w.r.t. parameters  $\theta$
$$
\nabla_ {\theta} E [ q (\theta) ] = \mathbb {E} _ {q (x, \theta)} \left[ \left(E _ {\mathrm {l o c}} (x, \theta) - \mathbb {E} _ {q (x, \theta)} \left[ E _ {\mathrm {l o c}} (x, \theta) \right]\right) \nabla_ {\theta} \log q (x, \theta) \right]. \tag {6}
$$
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,\theta) = \psi^2 (x,\theta)$ , 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,\theta)$  only up to a normalization constant, i.e., it outputs  $\tilde{q} (x,\theta)\propto q(x,\theta)$ . However, the gradient w.r.t.  $\theta$  does not depend on the normalization constant; hence, throughout the paper, we refer to the model as the normalized density  $q(x,\theta)$ .
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: \mathbb{R}^d \to \mathbb{R}$  by following the ODE  $\frac{\mathrm{d}}{\mathrm{d}t} x_t = -\nabla_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  $\mathcal{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
$$
\operatorname {W F R} _ {\lambda} \left(p _ {0}, p _ {1}\right) ^ {2} := \inf  _ {v _ {t}, g _ {t}, q _ {t}} \int_ {0} ^ {1} \mathbb {E} _ {q _ {t} (x)} \left[ \| v _ {t} (x) \| ^ {2} + \lambda g _ {t} (x) ^ {2} \right] \mathrm {d} t, \quad \text {s u b j . t o} \tag {7}
$$
$$
\frac {\partial q _ {t} (x)}{\partial t} = - \nabla_ {x} \cdot \left(q _ {t} (x) v _ {t} (x)\right) + g _ {t} (x) q _ {t} (x), \quad \text {a n d} q _ {0} (x) = p _ {0} (x), \quad q _ {1} (x) = p _ {1} (x),
$$
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  $\lambda$  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)\equiv 0)$  and Fisher-Rao distance  $(v_{t}(x)\equiv 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]: \mathcal{P}_2 \to \mathbb{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)
$$
\frac {\partial q _ {t}}{\partial t} (x) = \underbrace {- \nabla_ {x} \cdot \left(q _ {t} (x) \left(- \nabla_ {x} \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x)\right)\right)} _ {\text {t h e c o n t i n u i t y e q u a t i o n}} - \frac {1}{\lambda} \underbrace {\left(\frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) - \mathbb {E} _ {q _ {t} (y)} \left[ \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (y) \right]\right)} _ {\text {g r o w t h t e r m}} q _ {t} (x), \tag {8}
$$
where  $\delta F[q] / \delta q$  is the first-variation 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\sim q_t(x)$  follow the vector field  $v_{t}(x) = -\nabla_{x}\delta F[q_{t}] / \delta 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) = \frac{\delta F[q_{t}]}{\delta q_{t}} (x) - \mathbb{E}_{q_{t}(y)}\left[\frac{\delta F[q_{t}]}{\delta q_{t}} (y)\right]$ . Note that  $\mathbb{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.
$$
\frac {\partial q _ {t}}{\partial t} (x) = - \nabla_ {x} \cdot \left(q _ {t} (x) \left(- \nabla_ {x} \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x)\right)\right), \quad 2 \text {- W a s s e r s t i n G r a d i e n t F l o w}, \tag {9}
$$
$$
\frac {\partial q _ {t}}{\partial t} (x) = - \left(\frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) - \mathbb {E} _ {q _ {t} (y)} \left[ \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (y) \right]\right) q _ {t} (x), \quad \text {F i s h e r - R a o G r a d i e n t F l o w .} \tag {10}
$$
It now becomes evident that the stationary condition for all the considered PDEs is
$$
\left\| \nabla_ {x} \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) \right\| = 0 \Longleftrightarrow \frac {\delta F [ q _ {t} ]}{\delta q _ {t}} (x) \equiv \text {c o n s t a n t}. \tag {11}
$$
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).
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  $\psi$  (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 energy-minimizing time evolution of the wavefunction  $\psi_t$ , and is expressed as the following PDE (which is
Figure 1: W(FR)QMC: A graphical illustration of the gradient flow according to the Wasserstein and Fisher-Rao metrics, and the corresponding projection onto the variational manifold  $q(x,\theta)$ .
the Wick-rotated version of Eq. (2), see e.g. (McArdle et al., 2019; Yuan et al., 2019)),
$$
\frac {\partial \psi_ {t} (x)}{\partial t} = - (H - E [ \psi_ {t} ]) \psi_ {t} (x), \tag {12}
$$
where again  $q_{t}(x) = \psi_{t}^{2}(x)$ . The last term proportional to the energy  $E[\psi_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  $\mathbb{E}[q]$  under the Fisher-Rao metric.
Proof Sketch. The energy functional  $E[q]$  has the following derivative
$$
\frac {\delta E [ q ]}{\delta q} (x) = V (x) - \frac {1}{4} \nabla_ {x} ^ {2} \log q (x) - \frac {1}{8} \| \nabla_ {x} \log q (x) \| ^ {2} = E _ {\mathrm {l o c}} (x). \tag {13}
$$
Thus, the gradient flow under the Fisher-Rao metric is (see Eq. (10))
$$
\frac {\partial q _ {t} (x)}{\partial t} = - \left(E _ {\mathrm {l o c}} (x) - \mathbb {E} _ {q _ {t} (x)} [ E _ {\mathrm {l o c}} (x) ]\right) q _ {t} (x), \tag {14}
$$
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,\theta)$  (see Figure 1). We first evolve this density using a (non-parametric) gradient flow method (e.g., Eqs. (8) to (10)) for time  $\Delta t$ , which will take  $q_{t}(x)$  off the parametric manifold to  $q_{t + \Delta t}(x)$ . We then have to update current trial model  $q(x,\theta)$  to match  $q_{t + \Delta t}(x)$  enabling us to propagate the density further. In order to do so, we define the optimal update of parameters  $\Delta \theta^{*}$  as the minimizer of the Kullback-Leibler divergence between  $q_{t + \Delta t}(x)$  and the distributions on the parametric manifold, i.e.
$$
\Delta \theta^ {*} = \underset { \begin{array}{c} \Delta \theta \\ \| \Delta \theta \| _ {2} = 1 \end{array} } {\arg \min } D _ {\mathrm {K L}} \left(q _ {t + \Delta t} (x) \| q (x, \theta + \Delta \theta)\right). \tag {15}
$$
In practice, we evaluate the parameters update using the following proposition.
Proposition 3.2. For  $q_{t}(x) = q(x,\theta)$ , the gradient of reverse KL-divergence can be approximated as
$$
\begin{array}{l} - \nabla_ {\theta} \mathbb {E} _ {q _ {t + \Delta t}} [ \log q (x, \theta) ] = - \int q _ {t + \Delta t} \nabla_ {\theta} \log q (x, \theta) \mathrm {d} x \approx - \int \left[ q _ {t} + \Delta t \frac {\partial}{\partial t} q _ {t} \right] \nabla_ {\theta} \log q (x, \theta) \mathrm {d} x \\ = - \int q _ {t} \nabla_ {\theta} \log q (x, \theta) \mathrm {d} x - \Delta t \int \frac {\partial}{\partial t} q _ {t} \nabla_ {\theta} \log q (x, \theta) \mathrm {d} x \tag {16} \\ \end{array}
$$
Using this approximation, the optimal update from Eq. (15) of parameters becomes
$$
\Delta \theta^ {*} = \underset {\left\| \Delta \theta \right\| _ {2} = 1} {\arg \min } \left\langle \Delta \theta , - \int \frac {\partial}{\partial t} q _ {t} (x) \nabla_ {\theta} \log q (x, \theta) \mathrm {d} x \right\rangle \propto \int \frac {\partial}{\partial t} q _ {t} (x) \nabla_ {\theta} \log q (x, \theta) \mathrm {d} x. \tag {17}
$$
where  $\langle \cdot ,\cdot \rangle$  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.
$$
\Delta \theta^ {*} \propto - \mathbb {E} _ {q _ {t} (x)} \left[ \left(E _ {\mathrm {l o c}} (x, \theta) - \mathbb {E} _ {q _ {t} (x)} \left[ E _ {\mathrm {l o c}} (x, \theta) \right]\right) \nabla_ {\theta} \log q (x, \theta) \right]. \tag {18}
$$
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, \theta)$ , instead of following the update in Eq. (17), we can exploit the information geometry of the statistical manifold of  $q(x, \theta)$ , and define the update using the Fisher information matrix  $\mathcal{F}_{\theta}$
$$
\Delta \theta^{*} = \operatorname *{arg  min}_{\substack{\Delta \theta \\ \| \Delta \theta \|_{\mathcal{F}} = 1}}\Bigg\langle \Delta \theta , - \int \frac{\partial}{\partial t} q_{t}(x)\nabla_{\theta}\log q(x,\theta)  \mathrm{d}x\Bigg\rangle \propto \mathcal{F}_{\theta}^{-1}\int \frac{\partial}{\partial t} q_{t}(x)\nabla_{\theta}\log q(x,\theta)  \mathrm{d}x  ,
$$
$$
\mathcal {F} _ {\theta} = \mathbb {E} _ {q (x, \theta)} \left[ \frac {\partial}{\partial \theta} \log q (x, \theta) \frac {\partial}{\partial \theta} \log q (x, \theta) ^ {\top} \right]. \tag {19}
$$
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  $\theta$  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