A Theoretical Framework for an Efficient Normalizing Flow-Based Solution to the Schrödinger Equation

The Setting   Our overall goal is to compute the ground state wavefunction and energy of a molecule given its molecular parameters and spin multiplicity. Denote $x_{i}=(r_{i},s_{i})$ to be the pair consisting of the position and spin for the $i^{th}$ electron; $x$ will denote the entire ordered list $(x_{1},\dots,x_{n})$, with corresponding definitions for $r$ and $s$. We specify wavefunctions as $\psi(x)$; due to the fact that electrons are Fermions, valid wavefunctions must be antisymmetric, that is if $\pi\in\mathbb{S}_{n}$ is a permutation, then

$$\psi(\pi x)=(-1)^{\pi}\psi(x)$$

(1)

where as usual, $(-1)^{\pi}$ is shorthand for $(-1)^{N(\pi)}$ where $N(\pi)$ is the minimal number of flips to produce $\pi$.

Let $R_{I}$ and $Z_{I}$ denote the position and atomic number of the $I^{th}$ nucleus, and let the Laplacian for the $i^{th}$ electron be $\Delta_{i}=\frac{\partial^{2}}{\partial r_{i1}^{2}}+\frac{\partial^{2}}{\partial r_{i2}^{2}}+\frac{\partial^{2}}{\partial r_{i3}^{2}}$; then the Hamiltonian is given by

$$H=-\frac{1}{2}\underbrace{\sum_{i}\Delta_{i}}_{\Delta}+\underbrace{\sum_{i>j}\frac{1}{\|r_{i}-r_{j}\|}-\sum_{iI}\frac{Z_{I}}{\|r_{i}-R_{I}\|}+\sum_{I>J}\frac{Z_{I}Z_{J}}{\|R_{I}-R_{J}\|}}_{V(x)}=-\tfrac{1}{2}\Delta+V(x)$$

(2)

Our goal is to compute the ground state wavefunction, which we denote as $\psi_{0}(x)$ and corresponding ground state energy $E_{0}$. They may be computed using the variational principle:

$$\psi_{0}=\operatorname*{argmin}_{\psi\in\Psi}\frac{\langle\psi|H|\psi\rangle}{\langle\psi|\psi\rangle}\qquad\text{and}\qquad E_{0}=\frac{\langle\psi_{0}|H|\psi_{0}\rangle}{\langle\psi_{0}|\psi_{0}\rangle}$$

(3)

where $\Psi$ is the set of all possible valid wavefunctions, and $H$ is the Hamiltonian. If we specify the wavefunction ansatz as a neural network with parameters $\theta$, this becomes

$$\theta^{*}=\operatorname*{argmin}_{\theta}\frac{\langle\psi(\cdot;\theta)|H|\psi(\cdot;\theta)\rangle}{\langle\psi(\cdot;\theta)|\psi(\cdot;\theta)\rangle}\qquad\text{and}\qquad E^{*}=\frac{\langle\psi(\cdot;\theta^{*})|H|\psi(\cdot;\theta^{*})\rangle}{\langle\psi(\cdot;\theta^{*})|\psi(\cdot;\theta^{*})\rangle}\geq E_{0}$$

(4)

That is, we compute an upper bound $E^{*}$ to the ground state energy $E_{0}$. The more expressive the ansatz, the tighter the bound will be.

Variational Monte Carlo   The issue with the formulation to this point is the need to compute the inner products in Equations (3) and (4), which correspond to very high-dimensional integrals. A standard solution to this problem is based on a Monte Carlo scheme. To begin with, let us define the local energy as

$$\mathcal{E}(x)\equiv\frac{H\psi(x)}{\psi(x)}=-\frac{\Delta\psi(x)}{2\psi(x)}+V(x)\qquad\text{and}\qquad\mathcal{E}_{r}(x)=\text{Real}\{\mathcal{E}(x)\}$$

(5)

In this case, one can simplify the minimand in Equation (3) (see Appendix A) as

$$\frac{\langle\psi|H|\psi\rangle}{\langle\psi|\psi\rangle}\,=\,\mathbb{E}_{x\sim\rho(\cdot)}\left[\mathcal{E}_{r}(x)\right]\,\approx\,\frac{1}{K}\sum_{k=1}^{K}\mathcal{E}_{r}\left(x^{(k)}\right)\qquad\text{with}\qquad\rho(x)=\frac{|\psi(x)|^{2}}{\langle\psi|\psi\rangle}$$

(6)

where the $x^{(k)}$ are sampled from $\rho(x)=|\psi(x)|^{2}/\langle\psi|\psi\rangle$.

As enumerated in Section 1, a number of recent works have followed the above approach using a variety of neural networks as the ansatz for the wavefunction $\psi(\cdot;\theta)$. In order to do so, one must be able to sample from $\rho(x;\theta)=|\psi(x;\theta)|^{2}/\langle\psi|\psi\rangle$; and as the networks are quite general, the only feasible method for sampling is a Markov Chain Monte Carlo technique such as Langevin Monte Carlo (Umrigar et al., 1993) or any of several variations. These techniques can be time-consuming, as each sample is itself the solution of a stochastic differential equation as time goes to infinity.

A solution to this problem presents itself if we can somehow specify a wavefunction $\psi(x)$ which is easy to sample from. We are interested in wavefunctions which satisfy the following three properties:

1. (W1)

   There is an explicit functional form for the wavefunction $\psi(x)$.

2. (W2)

   $\psi$ is antisymmetric.

3. (W3)

   We can sample non-iteratively (in constant time) from $|\psi(\cdot)|^{2}$.

The first two properties are necessary for any form of Variational Monte Carlo: (W1) allows us to evaluate the local energy $\mathcal{E}_{r}$ in (5) for use in (6); and (W2) is required for valid electronic (Fermionic) wavefunctions. But (W3) is the new ingredient: if we have a family of wavefunctions $\psi$ satisfying (W1)-(W3), then solving the minimization in (4) via the Monte Carlo approach in (6) will be considerably accelerated, as each sample will only require constant time to generate. We add a fourth property, which is not strictly necessary but is both desirable and will prove useful:

1. (W4)

   $\psi$ is normalized, that is $\int|\psi(x)|^{2}dx=1$.

It turns out that generating such wavefunctions is possible using the following procedure:

Proof: See Appendix B.
The general idea expressed in Theorem 1 is that we can build the wavefunction $\psi$ out of an easy-to-sample-from density function satisfying additional properties (D1)-(D2); and a nearly antisymmetric phase function $\kappa$. In what follows, we will show how to construct both of these ingredients. But before doing so, we take a short detour to address the most important practical scenario, that of fixed spin multiplicity.

Notation   As in most approaches to this problem, we assume that the spin multiplicity of the molecule is specified, which is equivalent to fixing the number of spin up and spin down electrons, denoted $n_{u}$ and $n_{d}$ respectively, with $n_{u}+n_{d}=n$. Define the canonical spin vector to be given by $\bar{s}=[\uparrow,\dots,\uparrow,\downarrow,\dots,\downarrow]$, i.e. the first $n_{u}$ are $\uparrow$, the last $n_{d}$ are $\downarrow$. We let the sets of indices of up and down spin electrons for the canonical spin vector be denoted by $\mathcal{N}_{u}=\{1,\dots,n_{u}\}$ and $\mathcal{N}_{d}=\{n_{u}+1,\dots,n\}$. Finally, we will be interested in the subgroup of permutations in which a permutation is applied separately to spin-up and spin-down electrons. We denote this subgroup by

$$\mathbb{G}\equiv\mathbb{S}_{\mathcal{N}_{u}}\times\mathbb{S}_{\mathcal{N}_{d}}\quad\quad(\mathbb{G}\text{ is a subgroup of }\mathbb{S}_{n})$$

(8)

Specification of the Density   In the case of fixed spin multiplicity, the specification of the density $\rho(x)$ is simplified:

Proof: See Appendix C.
To summarize: in the case of fixed spin multiplicity, specifying a wavefunction $\psi$ satisfying our desired conditions (W1)-(W4) is equivalent to specifying a density $\bar{\rho}(r)$ satisfying conditions (R1)-(R2); and then applying the transformations given in Theorems 1 and 2 to map from $\bar{\rho}$ to $\psi$.¹¹1We have for the moment ignored the issue of the phase $\kappa$, which we return to in Sections 3.6 and 4.1. Therefore, henceforth we will focus exclusively on specifying densities $\bar{\rho}(r)$ satisfying conditions (R1)-(R2). To avoid unnecessary notational complexity, we will drop the bars and simply write $\rho(r)$.

Our goal is to use a normalizing flow to construct the density $\rho(r)$. Let $D$ be the ambient dimension (i.e. $D=3$) and $n$ be the number of electrons. The relevant vectors will live in the space $\mathbb{R}^{Dn}$ construed as the Cartesian product $\mathbb{R}^{D}\times\dots\times\mathbb{R}^{D}$ (which is of course isomorphic to $\mathbb{R}^{Dn}$). A normalizing flow will consist of two ingredients: (1) a base random variable $z$, which lives in $\mathbb{R}^{Dn}$, and is described by the density $\rho_{z}(z)$; (2) an invertible transformation $T:\mathbb{R}^{Dn}\to\mathbb{R}^{Dn}$, such that $r=T(z)$. In this case, the density $\rho(r)$ is the push-forward of $\rho_{z}$ along $T$, and is given by the change of variables formula

$$\rho(r)=\rho_{z}(T^{-1}(r))|\det J_{T^{-1}}(r)|$$

(9)

Recall that we would like our density $\rho(r)$ to satisfy conditions (R1)-(R2) laid out in Theorem 2. The following theorem establishes conditions for this to occur:

Proof: See Appendix D.
Armed with this key result, we now set out to design the base density $\rho_{z}$ and transformation $T$ which satisfy the conditions of Theorem 3.

In most cases in machine learning, the base density for a normalizing flow is taken to be a standard distribution, most often a Gaussian. In our case, we require that the base density have certain special properties, namely (R1) and (R2) from Theorem 2. It turns out that Determinantal Point Processes (DPPs) have just the properties we require. In particular, we are interested in the class of DPPs known as Projection DPPs (Gautier et al., 2019; Lavancier et al., 2015), which can be specified as follows. We will let $y$ specify a generic point in $\mathbb{R}^{D}$. Let $h_{k}:\mathbb{R}^{D}\to\mathbb{R}$ for $k=1,\dots,n$ be a set of $n$ functions which are orthogonal, that is $\langle h_{i},h_{j}\rangle=\int_{\mathbb{R}^{D}}h_{i}(y)h_{j}(y)dy=\delta_{ij}$. Let $H(y)$ be the column vector composed by stacking the individual functions $h_{i}(y)$ and define the kernel function as $K(y,y^{\prime})=H(y)^{T}H(y^{\prime})$. Then for a given collection of $n$ points in $\mathbb{R}^{D}$, that is $r=(r_{1},\dots,r_{n})$, we define the $n\times n$ kernel matrix $\mathbf{K}_{n}(r)$, from which the density of the Projection DPP may be specified:

$$\mathbf{K}_{n}(r)=\begin{bmatrix}K(r_{1},r_{1})&\dots&K(r_{1},r_{n})\\ \vdots&\ddots&\vdots\\ K(r_{n},r_{1})&\dots&K(r_{n},r_{n})\end{bmatrix}\qquad\Rightarrow\qquad\rho_{dpp}(r;n)=\frac{1}{n!}\det\mathbf{K}_{n}(r)$$

(10)

Since $\mathbf{K}_{n}(r)$ is positive semi-definite, it follows that its determinant is non-negative so that $\rho_{dpp}(r;n)$ is non-negative, as desired. A proof that $\rho_{dpp}(r;n)$ is properly normalized (i.e. integrates to $1$) can be found, for example, in Proposition 2.10 of (Johansson, 2006).

Given the notion of a Projection DPP, we may define the base density as follows. As above, let the base random variable be $z$, where $z$ can be broken into spin-up and spin-down pieces, denoted $z_{u}$ and $z_{d}$. (Specifically, $z_{u}$ and $z_{d}$ are the parts of $z$ corresponding to electrons in $\mathcal{N}_{u}$ and $\mathcal{N}_{d}$, respectively.) The base density can then be constructed by taking

$$\rho_{z}(z)=\rho_{dpp}(z_{u};n_{u})\rho_{dpp}(z_{d};n_{d})$$

(11)

That is, $z_{u}$ and $z_{d}$ are chosen from two independent Projection DPPs. We then have the following theorem:

Proof: See Appendix E.
We therefore have an explicit form for the base density from Equations (10) and (11). Furthermore, sampling from the base density amounts to sampling from two independent Projection DPPs. A sampling procedure for Projection DPPs is specified in Appendix F.

As noted in Section 3, we require the normalizing flow transformation to be $\mathbb{G}$-equivariant. Of course, chaining together many layers which are each $\mathbb{G}$-equivariant results in an overall transformation which is also $\mathbb{G}$-equivariant. Now, suppose that a particular layer $\ell$ can be written as

$$r^{\ell+1}=T^{\ell}(r^{\ell})$$

(12)

where $r^{\ell}=(r_{1}^{\ell},\dots,r_{n}^{\ell})$ and likewise for $r^{\ell+1}$. We will need to see the action on the spin-up and spin-down electrons separately, so we denote $r_{u}^{\ell}=(r_{i}^{\ell})_{i\in\mathcal{N}_{u}}$ and $r_{d}^{\ell}=(r_{i}^{\ell})_{i\in\mathcal{N}_{d}}$; and we may write

$$r_{u}^{\ell+1}=T_{u}^{\ell}(r_{u}^{\ell},r_{d}^{\ell})\quad\text{and}\quad r_{d}^{\ell+1}=T_{d}^{\ell}(r_{u}^{\ell},r_{d}^{\ell})$$

(13)

For notational convenience, we use $\alpha\in\{u,d\}$ to denote the spin, and the complement of the spin is given by ${\hat{\alpha}}$ (i.e. if $\alpha=u$ then ${\hat{\alpha}}=d$ and vice-versa). Then we have the following theorem:

Proof: See Appendix G.

We now show how to specify continuous and discrete normalizing flows satisfying Theorem 5.

According to Theorem 3, we are required a find a transformation which is $\mathbb{G}$-equivariant. We now show this can be achieved via a continuous normalizing flow. We specify this flow via the ordinary differential equation (ODE)

$$\frac{dv}{dt}=\Gamma_{t}(v),\quad\text{with}\quad v(0)=z\sim\rho_{z}(\cdot)\quad\text{and}\quad r=v(1)$$

(15)

That is, the transformation $r=T(z)$ is derived as follows: the initial condition is sampled from the base density; and $r$ is gotten by integrating the ODE forward to time $t=1$. $\Gamma$’s $t$-dependence is indicated via a subscript for notational convenience. We then have the following theorem:

Proof: See Appendix H.
It therefore suffices to design a $\mathbb{G}$-equivariant function $\Gamma_{t}$. Let us break this down by spin: from Theorem 5, we know that this implies that for all $t$, we have that $\Gamma_{t}(\pi_{\alpha}r_{\alpha},\pi_{\hat{\alpha}}r_{\hat{\alpha}})=\pi_{\alpha}\Gamma_{t}(r_{\alpha},r_{\hat{\alpha}})\text{ for }\alpha\in\{u,d\}$. We show in Appendix K how to implement a layer of $\Gamma$ with a combination of multihead attention, fully connected layers, and linear projections ($\Gamma$ can be composed of many such layers).

Continuous normalizing flows are elegant; however, they can present some numerical difficulties. In particular, the issue of ODE stiffness frequently arises in deep learning pipelines involving continuous normalizing flows. Thus, we now present an alternative method, based on discrete normalizing flows.

Our goal is now to design such functions $T_{u}^{\ell}$ and $T_{d}^{\ell}$ which satisfy Equation (14), and for which the overall transformation $T^{\ell}=(T_{u}^{\ell},T_{d}^{\ell})$ is invertible. The goal of the layer we propose here is to not sacrifice on expressivity, especially when compared to many layers which are designed for discrete normalizing flows. In particular, the main issue will be to show that the expressivity can be retained even with the joint requirements of invertibility and $\mathbb{G}$-equivariance. We note that the kind of transformation we propose below is not generally used for normalizing flows, as the determinant of its Jacobian is not fast to compute; however, this is not an issue in our case, as the dimension of the spaces we are dealing with are relatively small. For a more detailed discussion, see Appendix I.

To solve this problem, we introduce the Split Subspace Layer; we note that this layer may be of broader interest in machine learning, independent of the current setting. As before, we take $D$ to represent the ambient spatial dimension; in our case, $D=3$. A key parameter for the $\ell^{th}$ layer will be the orthogonal matrix $\Lambda_{\alpha}^{\ell}\in O(D)$; in particular, we divide this matrix into 2 pieces

$$\Lambda_{\alpha}^{\ell}=[\beta_{\alpha}^{\ell},\xi_{\alpha}^{\ell}]\quad\text{with }\beta_{\alpha}^{\ell}\in\mathbb{R}^{D\times D_{\beta}}\quad\text{ and }\quad\xi_{\alpha}^{\ell}\in\mathbb{R}^{D\times(D-D_{\beta})}$$

(16)

That is, $\beta_{\alpha}^{\ell}$ represents the first $D_{\beta}$ columns of $\Lambda_{\alpha}^{\ell}$, and $\xi_{\alpha}^{\ell}$ represents the final $D-D_{\beta}$ columns. For each electron $i$, we compute the inner product of its coordinates with $\beta_{\alpha}^{\ell}$, i.e.

$$\gamma_{\alpha,i}^{\ell}=(\beta_{\alpha}^{\ell})^{T}r_{\alpha,i}^{\ell}\quad\text{so that }\gamma_{\alpha,i}^{\ell}\in\mathbb{R}^{D_{\beta}}$$

(17)

We can collect the individual vectors $\gamma_{\alpha,i}^{\ell}$ into a list $\gamma_{\alpha}^{\ell}=(\gamma_{\alpha,i}^{\ell})_{i\in{\mathcal{N}_{\alpha}}}$. Given this, we define the Split Subspace Layer $T_{\alpha}^{\ell}$ on a per-electron basis by

$$r_{\alpha,i}^{\ell+1}=T_{\alpha,i}^{\ell}(r_{\alpha}^{\ell},r_{\hat{\alpha}}^{\ell})=r_{\alpha,i}^{\ell}+\xi_{\alpha}^{\ell}\varphi_{\alpha,i}^{\ell}(\gamma_{\alpha}^{\ell},\gamma_{\hat{\alpha}}^{\ell})\quad\text{with}\quad\varphi_{\alpha,i}^{\ell}(\gamma_{\alpha}^{\ell},\gamma_{\hat{\alpha}}^{\ell})\in\mathbb{R}^{D-D_{\beta}}$$

(18)

where $\varphi_{\alpha}^{\ell}$ is a network, and $\varphi_{\alpha,i}^{\ell}$ is the part of (the output of) $\varphi_{\alpha}^{\ell}$ corresponding to the $i^{th}$ electron. The layer is referred to as the Split Subspace Layer due to the fact that its input is one subspace of $\mathbb{R}^{D}$, given by $\beta_{\alpha}^{\ell}$; whereas its output is in the orthogonal complement of this subspace, given by $\xi_{\alpha}^{\ell}$.

The main ingredient of the layer is the network $\varphi_{\alpha}^{\ell}$. We now show two things: (1) the layer is invertible for any choice of $\varphi_{\alpha}^{\ell}$ (2) we derive conditions on $\varphi_{\alpha}^{\ell}$ to achieve $\mathbb{G}$-equivariance of $T_{\alpha}^{\ell}$.

Proof: See Appendix J.
The Split Subspace Layer therefore depends on implementation of the network $\varphi_{\alpha}^{\ell}$ so that it satisfies Equation (20). We show in Appendix K how $\varphi_{\alpha}^{\ell}$ can be implemented with a combination of multihead attention, fully connected layers, and linear projections. We specify a more general version of the Split Subspace Layer in Appendix L.

Log Domain: Density   In order to avoid numerical issues, it is best to operate in the log domain. Suppose that

$$\psi(r)=e^{q(r)+iw(r)}\quad\Leftrightarrow\quad q(r)=\tfrac{1}{2}\log\rho(r)\quad\text{and}\quad w(r)=\text{atan2}\left(\kappa_{i}(r),\kappa_{r}(r)\right)$$

(21)

where $\kappa_{r}(r)$ and $\kappa_{i}(r)$ are the real and imaginary parts of the phase $\kappa(r)$, respectively; and atan2 is the “full” arctangent.

The log-density $q(r;\theta)$ may be computed for both continuous and discrete normalizing flows, where we now introduce the parameters $\theta$ of the network explicitly. Consider a sample $z$ chosen from the base density $\rho_{z}(z)$, and in analogy to $q(r)$, define $q_{z}(z)=\tfrac{1}{2}\log\rho_{z}(z)$. Now, in the case of a continuous normalizing flow, let $v(t)$ satisfy Equation (15); then $q(r;\theta)$ can be by computed (Chen et al., 2018) by solving the ODE

$$\frac{da}{dt}=-\text{Trace}\left(\frac{\partial\Gamma_{t}}{\partial v}(v(t);\theta)\right)\quad\text{with}\quad\text{ with }a(0)=q_{z}(z)\quad\text{and}\quad q(r;\theta)=a(1)$$

(22)

which is the continuous analogue of the change of variables formula. In the case of a discrete normalizing flow, fix the following notation: $r^{0}=z$, $r=r^{L+1}$, and $T=T^{L}\circ\dots\circ T^{0}$. Then we may use a logarithmic version of the standard change of variables formula (9):

$$q(r;\theta)=q_{z}\left(T^{-1}(r;\theta)\right)+\frac{1}{2}\sum_{\ell=0}^{L}\log\left|\det J_{(T^{\ell})^{-1}}(r^{\ell+1};\theta)\right|$$

(23)