On Representing Electronic Wave Functions with Sign Equivariant Neural Networks

Authors: nicholas  gao
stephan g ünnemann
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-05 09:30:38

Abstract: Recent neural networks demonstrated impressively accurate approximations of electronic ground-state wave functions. Such neural networks typically consist of a permutation-equivariant neural network followed by a permutation-antisymmetric operation to enforce the electronic exchange symmetry. While accurate, such neural networks are computationally expensive. In this work, we explore the flipped approach, where we first compute antisymmetric quantities based on the electronic coordinates and then apply sign equivariant neural networks to preserve the antisymmetry. While this approach promises acceleration thanks to the lower-dimensional representation, we demonstrate that it reduces to a Jastrow factor, a commonly used permutation-invariant multiplicative factor in the wave function. Our empirical results support this further, finding little to no improvements over baselines. We conclude with neither theoretical nor empirical advantages of sign equivariant functions for representing electronic wave functions within the evaluation of this work.

Main paper: 
 introduction: 
 The stationary Schrödinger equation (Schrödinger, 1926) at the heart of quantum chemistry is a partial differential equation in 3N dimensions, where N is the number of electrons in the system:Hψ = Eψ (1)where ψ : R N ×3 → R is the wave function, E the energy and H the Hamiltonian of the system, see Appendix F. Electronic wave functions must obey the exchange antisymmetry ψ(..., r i , ..., r j , ...) = -ψ(..., r j , ..., r i , ...). Typically, one enforces this by using a so-called Slater determinant ψ(r) = det [φ j (r i )] i,j∈{0,...,N } = det Φ(r).(2)In practice, one often uses linear combinations of determinants ψ(r) = K k=1 w k det Φ k (r). Further, note that the antisymmetric property is preserved by multiplying the antisymmetric function with a symmetric function J : R N ×3 → R, leading to the so-called Slater-Jastrow wave functionψ(r) = J(r) K k=1 w k det Φ k (r).(3)Recent works improved these Ansätze by replacing the so-called orbital functions φ i in Equation 2 by neural networks (Hermann et al., 2023) . While this has shown remarkably accurate approximations of the electronic wave function of various molecules, it comes at a significant computational cost due to explicit electron-electron interactions. In classical quantum chemistry, one avoids this high cost by using easily integrable Gaussian basis functions to construct the orbital functions φ i .Unfortunately, due to the simple structure of the orbital functions, one may typically require thousands to millions of determinants to capture electronic correlations correctly (Gao et al., 2024) .In this work, we aim to reduce the classical large number of determinants via neural networks. But, instead of defining the orbital functions via neural networks, we explore non-linear combinations of determinants, i.e., we are reformulating the classic Slater-Jastrow wave function toEQUATIONwhere f (x, y) = y K k=1 w k x k recovers the classical case. To accomplish this, we define a set of sign equivariant operations to preserve the antisymmetry of the determinants in the final output. Our theoretical analysis finds that such an algebraic construction is identical to a Jastrow factor and, thus, cannot shrink the zero set of the wave function. We support these findings via empirical results. 
related work: 
 Solving the stationary Schrödinger equation accurately for many systems opens the path to fast and accurate machine-learned force fields (Schütt et al., 2018; Kosmala et al., 2023; Wollschläger et al., 2023) . The machine learning approach to this has been, since the first work by Lou et al. (2023) , to parameterize the orbitals from from Equation 2 with neural networks (Hermann et al., 2023; Zhang et al., 2023) . While subsequent works tweak architectures (Gerard et al., 2022; von Glehn et al., 2023) , explore new applications (Cassella et al., 2023; Kim et al., 2023; Lou et al., 2023; Wilson et al., 2022; Pescia et al., 2023) , compute excited states (Pfau et al., 2023; Entwistle et al., 2022) , generalize across molecules (Scherbela et al., 2022; 2024; 2023; Gao & Günnemann, 2022; 2023b; a) , or explore the use of Diffusion Monte Carlo (Wilson et al., 2021; Ren et al., 2023) , the underlying structure of the wave function remained the same. This work explores a different approach by parametrizing the wave function via non-linear combinations of determinants. 
sign equivariant functions: 
 Throughout this work, we will mainly refer to two different symmetries. Firstly, the fermionic antisymmetry to permutations, i.e., ψ(π(r)) = sign(π)ψ(r), and, secondly, to odd function, more formally functions that are equivariant to the cyclic group C 2 , i.e., f (-x) = -f (x). The framework of equivariance allows us to define these symmetries more generally. Definition 1. A function f : X → Y on real vector spaces X , Y is said to be equivariant under groupG iff ∀g ∈ G f (G X g x) = G Y g f (x) where G X g , G Yg are the group representations of g acting on the vector spaces X and Y, respectively.As a special case of equivariance, one can define invariance where the result of a function does not change under group actions, i.e., G Y g is the identity for all g ∈ G. Given these definitions, we can now concretely describe antisymmetric and odd functions. Further, to aid later discussion, we also introduce symmetric and even functions as counterparts. Definition 2. A function f : X → Y is antisymmetric iff f is equivariant under the symmetric group S n and the group acts on Y as y → sign(π)y, ∀π ∈ S n , y ∈ Y.Definition 3. A function f : X → Y is symmetric iff f is invariant under the symmetric group S n . Definition 4. A function f : X → Y is odd iff f is equivariant under the cyclic group C 2 = {-1, 1} and the group acts on Y as y → gy, ∀g ∈ C 2 , y ∈ Y. Definition 5. A function f : X → Y is even iff f is invariant under the cyclic group C 2 .Implicit Odd Functions A function f acting on the vector of determinants x = [det Φ k (r)] K k=1 must be odd, i.e., f (-x) = -f (x), to preserve the fermionic antisymmetry. Let f 1 and f 2 be odd functions and g an even function. The following functions are odd: (1) multiplication with a constant αf (x), (2) addition f 1 (x) + f 2 (x), (3) element-wise odd functions, e.g., f (x) = [tanh(x i )] K k=1 , (4) multiplication with an even function g, i.e., f (x)g(x), and (5) chaining f 2 (f 1 (x)). With (5), we can compose complex functions by chaining simpler ones.Combining (1) and (2) yields linear combinations, i.e., we can construct linear layers without bias terms. We construct neural networks f (T ) by combining this with non-linear activation functions (3):f (t+1) (x (t) ) = tanh x (t) W (t) * J (t)(5)where J (t) is obtained from a Jastrow factor, i.e., a symmetric function of the electronic coordinates.Explicit Odd Functions. Alternatively, given an arbitrary function g : X → Y, one can construct an odd function f : X → Y via f (x) = g(x)g(-x). Inversely, as proven in Appendix A, every odd function f can be expressed by a non-odd function g: Theorem 1. Any odd function f : X → Y on real vector spaces X , Y can be represented asf (x) = g(x) -g(-x) where g : X → Y, x ∈ X .In our experiments, we implement g via an MLP with Jastrow factors J (0) and J (T ) :EQUATIONLinear-logarithmic Domain. One may notice that the distribution of ψ for samples drawn from ψ 2 varies by several orders of magnitude. We illustrate this in Appendix C for LiH. While well separated in the logarithmic domain, the data is crammed into a small region when viewed in the linear domain.Unfortunately, we cannot directly work in the log-domain as it prohibits any intermediate values from being zero. We address this by working in a domain that resembles a linear relationship around zero and a logarithmic relationship far from zero. We define the linlog transformation as linlog α (x) = sign(x) log(|x|e α + 1) (7) where α controls a shift of the data towards the linear or logarithmic region. We depict this function and discuss its stable implementation in Appendix D. 
theoretical results: 
 Equivalence to Jastrow Factor. Before empirically analyzing the sign equivariant functions, we point out an equivalence between odd functions and a Jastrow factor in the context of electronic wave functions. Interestingly, one can show that all functions of the form in Equation 4 can be represented by functions in the classical Slater-Jastrow form in Equation 3. Theorem 2. Let R, X , Y, Z be real vector spaces, φ : R → X an antisymmetric function, J : R → Y a symmetric function, and f : X × Y → Z an in the first argument odd function, i.e., f (-x, y) = -f (x, y). Any antisymmetric function ψ : R → Z; ψ(r) = f (φ(r), J(r)) can be expressed a.e. as ψ(r) = (φ(r) T x) Ĵ(r) with x ∈ X \ {0} and a symmetric function Ĵ : R → Z.We prove this theorem in Appendix B. In the context of quantum chemistry, this result implies that a linear combination and a non-positive-constrained Jastrow factor can equally represent any odd non-linear combination of the antisymmetric function. While this result guarantees the existence of such a Jastrow factor, it provides no statement about the ease of finding such a solution.Cusp Conditions. Thanks to Kato (1957) 's theorem, we know that the solutions to Equation 1 must fulfill the cusp condition lim r→Rm -1 ψ(r) ∂ψ(r) ∂r= Z m where R m , Z m are the position and charge of the m-th nucleus, respectively. However, one cannot fulfill the cusp conditions if one chooses Gaussian-type orbitals as basis functions φ i in Equation 2. In fact, all extrema of the original wave function are preserved. One can verify this via the chain rule on ψ(r) = f ([det Φ k (r)] K k=1 ): ∂ψ(r) ∂r = ∂ψ(r) ∂ [det Φ k (r)] K k=1 ∂ [det Φ k (r)] K k=1∂r .(8) Thus, anytime the orbital derivatives are 0, the derivatives of the final wave function will be 0. As the derivatives of Gaussian-type orbitals are zero at the nuclei, the cusp conditions cannot be fulfilled. 
empirical results: 
 While Theorem 2 guarantees the existence of equivalent Jastrows, it does not provide insights into finding one. Thus, we experimentally evaluate the impact of odd functions on the energy of wave functions. We test odd functions on four different systems, LiH, Li2, and two states of N2, as detailed in Appendix E. We use two different approaches to get the initial antisymmetric input for the odd functions: (1) CASSCF wave functions, a series of Slater determinants from standard quantum chemistry (Szabo & Ostlund, 2012) , and (2) FermiNet, a neural network wave functions (Pfau et al., 2020) . For optimization, we use the variational Monte Carlo (VMC) framework, which we highlight in Appendix F, with K-FAC (Martens & Grosse, 2015) . Ablations with Prodigy (Mishchenko & Defazio, 2023 ) can be found in Appendix H and the exact hyperparameters in Appendix G.CASSCF. We first use the cheaper CASSCF wave functions to test the hyperparameter α ∈ {-2, 0, 2}'s impact from Equation 7 on the final energy. Additionally, we repeat the experiment in the linear and linlog domain. Further, we compare a classical linear readout, implicit odd functions from Equation 5, and the explicit odd functions from Equation 6. The final energies are plotted in Figure 1 . On small structures like LiH and Li2, the explicit odd functions improve energies by up to 8 mE h and 4 mE h , respectively. Unfortunately, this advantage does not carry over to larger, more challenging molecules like the nitrogen dimer. We observe significant numerical instabilities with implicit or explicit odd functions in the optimization process; we explore this further in Appendix H.FermiNet. For FermiNet we picked the best performing α from the CASSCF experiments, i.e., α = -2. We evaluate the energy change by including a Jastrow factor. We consider three types of Jastrows, (1) none, (2) a traditional Jastrow, Equation 3, (Jastrow), and (3) a Jastrow included the odd function, Equations 5 and 6, (SymmetricOdd). Appendix G details the exact Jastrow function.In contrast to CASSCF wave functions, FermiNet can closely recover the ground state thanks to the explicit inclusion of electron correlation in the orbitals. We plot the final energies in Figure 2 . However, neither odd function positively impacts the optimization but worsens energies. Adding a Jastrow factor leads to improvements in the more challenging nitrogen dimer. 
conclusion: 
 Moving neural network complexity from the orbital functions to a non-linear combination of simpler basis functions promises faster quantum chemistry methods than deep neural networks applied to electronic coordinates. To design such functions, we present two approaches: implicit constructions by combining simple odd functions and explicitly odd functions where one enforces the antisymmetry at the end. While our experimental results show that such odd functions combined with classical determinants can yield better results on small structures, optimizing such odd functions proves difficult. At larger system sizes, we frequently observe numerical instabilities and degraded performance. In combination with neural network wave functions, we found such odd functions to worsen energies. From a theoretical point of view, we have shown the limitations of odd functions. Specifically, they correspond to a subset of the set of functions obtainable via Jastrow factors. Given the theoretical and empirical results, we find little evidence for computational or accuracy advantages of non-linear combinations of Slater determinants for machine-learning quantum chemistry. 
 
 Appendix: Proof of Theorem 1. DefineEQUATIONfor some arbitrary non-zero vector v ∈ X /0 . This construction yieldsEQUATIONEQUATIONB PROOF OF THEOREM 2Lemma 1. Every odd function f : X → Y on real vector spaces X , Y can be represented almost everywhere as a product of a linear combination x T v, v ∈ X \ {0} and an even function h :EQUATION)Proof of Lemma 1. We start with proving that h(x) = f (x)x T v for some v ∈ X /{0} is an even functionEQUATIONPlugging h into Equation 12 yieldsEQUATION)Proof of Theorem 2. Directly applying Lemma 1, we getEQUATION)If we set Ĵ(r) = f (φ(r),J(r))φ(r) T v, it remains to show that Ĵ is symmetric:EQUATIONEQUATIONC DISTRIBUTION OF WAVE FUNCTION AMPLITUDESIn Figure 3 , we plot the distribution of wave function amplitudes for LiH in the logarithmic, linear, and linlog domains. We drew samples from ψ 2 and plotted the distribution of the amplitudes ψ.While the data looks well distributed in the logarithmic domain, it is crammed into a small region in the linear domain. Our linlog transformation from Equation 7 addresses this issue by viewing part of the data in the linear domain and part in the logarithmic domain. The scaling factor α controls the shift between the two domains. It plays a crucial role as a hyperparameter in our experiments. A small α leads to a spiky distribution like in the linear domain, and a large α leads to a flat distribution. Our proposed linlog transformation from Equation 7 is C ∞ smooth and can be numerically stably implemented to transform from and to the logarithmic domain directly via linlog α (sign(x), log(|x|)) = sign(x)softplus(log(|x|) + α).(20) To avoid the spiky distribution one typically observed in the linear domain, we initialize α = α init + median i {max k {det Φ k (r i )}} via a hyperparameter α init and the median of a batch of electronic configurations sampled from ψ 2 with a linear readout instead of the learnable odd function. We plot the transformation relative to the linear domain in Figure 4 .We use four different systems for our empirical analysis: LiH, Li2, N2, and a distorted N2. For LiH and Li2, we pick the equilibrium structures from Pfau et al. (2020) . For N2, we use the equilibrium structure from Pfau et al. (2020) as well a highly distorted structure where the error of FermiNet from Pfau et al. (2020) Variational Monte Carlo (VMC) is a method to approximate solutions to the Schrödinger equation from Equation 1. In this work, we are interested in molecular systems where the Hamiltonian takes the following form in the Born-Oppenheimer approximation:EQUATION) with r i being the position of the i-th electron, R m the position of the m-th nucleus, and Z m the charge of the m-th nucleus. In linear algebra, Equation 1 is an eigenvalue problem, i.e., we are looking for the lowest eigenvalue E 0 and corresponding eigenvector ψ 0 . To accomplish this, we use the variational principle, i.e., we approximate the ground state by a trial wave function ψ T and minimize the energy expectation value. The variational principle (Szabo & Ostlund, 2012) states that the energy expectation value of any trial wave function ψ θ is an upper bound to the ground state energy E 0 :EQUATION)If we now define the probability density function (PDF) p(r) = ψ 2 (r) ψ 2 (r) d r , we can rewrite Equation 22 asEQUATIONwhere the right-hand side is the so-called VMC energy. By taking gradients of the VMC energy to the parameters θ of the trial wave function ψ θ , we can optimize the parameters to minimize the energy. These gradients can be computed as∇ θ = E p(r) (E L (r) -E p(r) [E L (r)])∇ θ log ψ θ (r) .(24) By approximating the expectation values via Monte Carlo sampling, we can compute the gradients via Equation 24 and optimize the parameters via gradient descent (Ceperley et al., 1977) .This section details the setup of our experiments.As Jastrow factors, we use an MLP on the averaged electronic coordinates:EQUATIONWe picked this specific formulation as it provides a similar computational cost to the odd function defined in Section 3 as it does not explicitly consider the individual electronic coordinates. We use this Jastrow as a standalone Jastrow or the symmetric function in Equation 5 and 6.As standard in the field, we pretrain the FermiNet on a Hartree-Fock wave function and then optimize the odd function with the pretrained FermiNet within the VMC framework (Pfau et al., 2020) .We implement everything in JAX (Bradbury et al., 2018) . To compute the laplacian from Equation 21, we use the forward laplacian algorithm from Li et al. (2023) implemented by Gao et al. (2023) .We use the hyperparameters listed in Table 2 for our experiments.In addition to the experiments in the main body, we present additional empirical data here. As an alternative to the KFAC optimizer, we present results with prodigy (Mishchenko & Defazio, 2023) , a learning-rate-free version of Adam (Kingma & Ba, 2014) , in Figure 5 . However, the conclusion remains the same: The odd functions improve energies on small structures and worsen them on large ones. Further, we present CASSCF experiments with Jastrow factors in Figure 6 .Here, we observe that including a Jastrow factor closes the gap between the linear and non-linear combinations. Finally, we present CASSCF experiments in float64 precision in Figure 7 to show that the numerical instabilities are not due to numerical precision. 


Decision: Thanks to reviewers for their careful review and nice suggestions and questions. Both reviewers seem to agree with the acceptance oft his paper. I also went through their points and the paper, and recommend authors to address them in their final revision.

Reviews: review: title: title: Review
main_review: main_review: **Summary**
The paper investigates the utility of sign equivariant neural networks that combine determinants *non-linearly* for representing electronic wave functions. This approach is different from existing neural-network-based approximations to the wave function that rely on *linear* combinations of determinants (so-called Slater-Jastrow wave function).

Specifically, the paper contains two main contributions:
1. A theoretical result, showing that any sign equivariant non-linear combination of determinants can be equivalently represented using the classical "linear" Slater-Jastrow form.
2. An empirical evaluation of two classes of wave functions based on sign equivariant neural networks. The evaluation shows that the sign equivariant architectures do not provide a consistent advantage over the traditional methods based on the Slater-Jastrow form.

**Strengths/weaknesses**

The paper studies an important ML problem, provides a clear presentation of the results, and examines an interesting question overlooked by prior work. The findings of this paper may be important in shaping the direction of future work in this ML area.

My only criticism of the paper is that the conclusion in the abstract ("sign equivariant functions are unsuitable for representing electronic wave functions") might be too strong given the evidence presented in the paper. 

The theoretical result shows that the non-linear component can be in theory "absorbed" into the permutation invariant Jastrow factor $J(r)$, while the $\det \Phi(r)$ terms are used to capture the fermionic antisymmetry to permutations. The experimental evaluation studies two specific sign equivariant NN architectures. Neither of these statements excludes the possibility that different architectures for $J(r)$ or the sign equivariant aggregation function $f(x, y)$ would lead to a different conclusion.  As a somewhat far-fetched example, the strong conclusion presented in the paper is akin to concluding that neural networks are unsuitable for learning functions of images by 1/ invoking the universal approximation property of MLPs and 2/ showing their poor results on some image-related ML tasks. I would recommend to add more nuance to the statement, taking the limitations of the experimental setup into account.

**Minor suggestions:**
- Typos in Definition 2, should instead be "the group acts on $\mathcal{Y}$ as $\pi y \mapsto \operatorname{sign}(\pi) y$"
- The description of "Linear-logarithmic domain" is Section 3 feels out of place. The statement about samples drawn from $\phi^2$ refers to the experiments that come much later. Also, previous paragraphs use $x$ to denote the vector of determinants, but in the definition of $\operatorname{linglog}$, $x$ is a scalar
- Potential typo: $f^{(t)}$ instead of $f^{(T)}$ before equation 5.

Disclaimer: I have limited experience in ML for quantum chemistry, so please feel free to downweight my feedback in favor of reviewers who are experts in this domain.
title: title: Strong paper which is focused more on the application than on machine learning
main_review: main_review: ## Summary
The authors present a method to approximate the electronic ground-state wave functions in molecules. To this end, related work has mostly proposed to approximate individual wave functions within the Slater-determinant using neural networks. These determinants are then combined linearly and multiplied with a Jastrow factor. In contrast, this paper suggests to replace this linear combination with a neural network which is constructed to guarantee the anti-symmetric property of the wave function. The authors provide proof that the network structure guarantees the anti-symmetric property. They go on to show, that for every function that can be represented by their neural network there is a Jastrow-Slater-Wavefunction that is identical to it wherever (phi(r)^T x) is non-zero. In empirical experiments they do not find a clear advantage of their method.

## Suitability for this workshop
I think the results are interesting to the workshop as the authors build a specialized neural network to ensure certain properties of the solution.

## Comments
While I would have wished for more details on the neural networks for this venue, I think that given the page limit, the selection of information is comprehensible. I think that the paper is clearly written. It was sometimes difficult to follow for me as the paper requires some background in quantum mechanics that I lack. However, given the page limit I do not think a longer introduction to this topic is feasible. I think the discussion of the results is good, only the final sentence of the abstract might be a slight over-claim because it is difficult to claim it being impossible from it not showing in one experiment.

##  Justification of Rating
I think this is a very good work as it presents empirical experiments that are well motivated and further gives good theoretical insights why these empirical experiments turned out the way they did. The only downside in my opinion is that the focus is much more on quantum mechanics and less on learning representations. Nevertheless, I think this paper should clearly be accepted.