A SCORE-BASED MODEL FOR LEARNING NEURAL WAVEFUNCTIONS

Abstract

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

1. INTRODUCTION

Understanding the properties of quantum systems lies at the core of many scientific disciplines, such as condensed matter physics, material science, and quantum chemistry. A quantum system is characterized by its ground state wavefunction, formally obtained by solving the Schrödinger equation. However, directly solving the Schrödinger equation for quantum systems with many particles is impractical due to the exponentially large Hilbert space. Owning to its strong dimension reduction capabilities, deep learning methods have been used as a strong candidate to approximately solve the Schrödinger equation and extract properties of quantum systems with the desired accuracy. For example, under the supervised learning setting, deep learning methods have been successfully applied to predict the quantum properties of molecular systems based on training data generated from density functional theory (DFT) calculation (Schütt et al., 2017; Gasteiger et al., 2020; Liu et al., 2022; Wang et al., 2022) . However, supervised methods rely on expensive computational simulations to generate a large amount of training data, and the accuracy of these methods is fundamentally limited by the data quality. Furthermore, DFT calculations involve various approximations and are not guaranteed to reach true ground states. A common scheme for approximately solving the Schrödinger equation is the variational principle, which optimizes a trial wavefunction to reach the ground state by minimizing its energy as much as possible via quantum Monte Carlo (QMC). Such a method is called variational Monte Carlo, whose accuracy relies on the expressive power of the trial wavefunction. Recently, deep learning methods coupled with variational Monte Carlo have unleashed the potential of both methods (Carleo & Troyer, 2017; Hermann et al., 2022) . Powered by the efficient sampling and optimization framework of quantum Monte Carlo and the universal approximation capability of deep neural networks, neural wavefunctions can accurately model quantum states, and dramatic improvements have been achieved (Pfau et al., 2020; Hermann et al., 2020) . Modeling a wavefunction is conceptually similar to modeling a probability density. Existing methods model the wavefunction explicitly by training a neural network to directly output the wavefunction values. However, numerous examples in machine learning have shown that implicitly modeling data distributions provides better representations (Kingma & Welling, 2014; Goodfellow et al., 2014; Ho et al., 2020) . As our direct reference, score-based methods have demonstrated their strong suc-cess in generative modeling (Song & Ermon, 2019; Song et al., 2020) . A score is defined as the gradient of the log probability. For example, realistic images can be generated from random noise by following dynamics defined by scores. In this paper, we show that the quantum wavefunction can be represented by score models and be optimized within the QMC framework. Our motivation to relate score-based models with QMC is based on an interesting connection between energy computations and score-based formulations. In QMC, the energy of a system is averaged over local energy of plausible quantum states. Our observation is that local energy only involves gradients of the logarithmic wavefunction, which we define as the score of the wavefunction. As a result, to minimize energy, the score must be explicitly computed. On the other hand, the actual wavefunction values are only used for sampling and optimization. To this end, we propose a new optimization framework for QMC where sampling and optimization are also achieved by using score functions alone, eliminating the need to explicitly compute the wavefunction value. In our proposed score-based framework, the sampling is done via Langevin dynamics and optimization is done through a new loss function inspired by diffusion Monte Carlo. Our score-based method enables the possibility of performing QMC computation with only score functions, which is infeasible in existing optimization frameworks. A direct benefit is that, by predicting gradients, we avoid the need to recompute it from the wavefunction. Moreover, score functions can be interpreted as the force of quantum systems, by implicitly modeling distributions with score functions, the dynamics of quantum systems could be better captured. Our experimental results show that with our score-based optimization framework, ground states of quantum systems can be accurately learned.

2.1. QUANTUM MANY-BODY WAVEFUNCTION IN CONTINUOUS SPACE

We use x ∈ R N ×d to denote the coordinates of N particles in d-dimension. The quantum state of a system is defined by its wavefunction ψ : R N ×d → R. By definition ψ is normalized ( x |ψ(x)| 2 = 1) and |ψ(x)| 2 gives the probability density of observing x. Any wavefunction ψ can be expressed as linear combination of eigenfunctions ψ n , which are solutions to the time-independent Schrödinger equation Ĥψ n (x) = E n ψ n (x), where Ĥ is an linear operator known as the Hamiltonian, and E n is a scalar giving the energy of the n-th eigen state. The Hamiltonian is defined as Ĥψ(x) = - 1 2 i ∇ 2 i ψ(x) + V (x)ψ(x), where the index i runs over all of the N × d dimensions in the summation. The first term in the Hamiltonian takes the sum of the second-order partial derivatives of the wavefunction and is related to the kinetic energy of the system. The second term in the Hamiltonian multiplies the wavefunction by a scalar value and is related to the potential energy of the system. The kinetic term is intrinsic to the Schrödinger equation and always takes the same form, whereas the potential function V : R N ×d → R varies for different physics problems. Note that although a wavefunction can be complex-valued in general, we can let ψ be real-valued because Ĥ is real. Our objective is to find the ground state ψ 0 , which is the eigen state associated with the lowest energy E 0 . In our notation, the coordinates x can be either viewed as N d-dimensional vectors or as a flattened N • d dimensional vector. In the rest of this paper, depending on the context, we may use bold x i to denote the i-th particle or use the regular font x i to denote the i-th scalar component of the flattened vector.

2.2. VARIATIONAL MONTE CARLO

The variational Monte Carlo (VMC) method uses a parameterized function ψ θ : R N ×d → R (called the Ansatz) to model a wavefunction, where θ denotes the parameters to be optimized. The normalization of ψ θ is not required. The energy expectation of ψ θ is computed as: L(θ) = ψ θ (x) Ĥψ θ (x)dx ψ θ (x)ψ θ (x)dx = ψ θ (x) 2 Ĥψ θ (x) ψ θ (x) dx ψ θ (x) 2 dx = E x∼ψ 2 θ ψ 2 θ E L (x; θ),

