A Self-Attention Ansatz for Ab-initio Quantum Chemistry
Authors: Ingrid von Glehn
James S Spencer
David Pfau
Venue: ICLR 2023 poster (2023)
Keywords: Machine learning for science
attention
Transformers
Monte Carlo
MCMC
self-generative learning
quantum physics
chemistry
machine learning for physics
machine learning for molecules
machine learning for chemistry

Abstract: We present a novel neural network architecture using self-attention, the Wavefunction Transformer (PsiFormer), which can be used as an approximation (or "Ansatz") for solving the many-electron Schrödinger equation, the fundamental equation for quantum chemistry and material science. This equation can be solved *from first principles*, requiring no external training data. In recent years, deep neural networks like the FermiNet and PauliNet have been used to significantly improve the accuracy of these first-principle calculations, but they lack an attention-like mechanism for gating interactions between electrons. Here we show that the PsiFormer can be used as a drop-in replacement for these other neural networks, often dramatically improving the accuracy of the calculations. On larger molecules especially, the ground state energy can be improved by dozens of kcal/mol, a qualitative leap over previous methods. This demonstrates that self-attention networks can learn complex quantum mechanical correlations between electrons, and are a promising route to reaching unprecedented accuracy in chemical calculations on larger systems.

Published as a conference paper at ICLR 2023
A SELF-ATTENTION ANSATZ FOR AB-INITIO QUANTUM CHEMISTRY
Ingrid von Glehn, James S. Spencer & David Pfau (ingridvy, jame ssspencer, pfau @deeeedde.eee.coomm com Sarrggaames
ABSTRACT
We present novel neural network architecture using self-attention, the Wavefune- tion Transformer (Psiformer), which can be used as approximation (or Ansatz) for solving the many-electron Schrödinger equation, the fundamental equation for quantum chemistry and material science. This equation be solved from first principles, requiring no external training data. recent years, deep neural net- works like the FermiNet and PauliNet have been used significantly improve the accuracy of these first-principle calculations, but they lack an attention-like mechanism for gating interactions between electrons. Here show that the Psi- former can be used as a drop-in replacement for these other neural networks, often dramatically improving the accuracy of the calculations. On larger molecules es- pecially, the ground state energy can be improved by dozens of kcal/mol, quali- tative leap over previous methods. This demonstrates that self-attention networks can learn complex quantum mechanical correlations between electrons, and are a promising route to reaching unprecedented accuracy chemical calculations on larger systems.
INTRODUCTION
The laws of quantum mechanics describe the nature of matter at the microscopic level, and underpin the study of chemistry, condensed matter physics and material science. Although these laws have been known for nearly a century (Schrödinger) (9266), the fundamental equations are too difficult to solve analytically for all but the simplest systems. In recent years, tools from deep learning have been used to great effect to improve the quality of computational quantum physics (Carleo & Troyer, 2017). For the study of chemistry in particular, it is the quantum behavior of electrons that matters, which imposes certain constraints on the possible solutions. The use of deep neural networks for successfully computing the quantum behavior of molecules was introduced almost simultaneously by several groups (Pfau et al., 2020, Hermann et al.) 2020; Choo et al.) 2020), and has since led to variety of extensions and improvements (Hermann al.] (022). However, follow-up work has mostly focused on applications and iterative improvements to the neural network architectures introduced in the first set of papers.
At the same time, neural networks using self-attention layers, like the Transformer (Vaswani et al., 2017), have had a profound impact on much of machine learning. They have led to breakthroughs in natural language processing (Devlin et al., 2018). language modeling (Brown et al. 2020), image recognition (Dosovitskiy et al. 2020), and protein folding (Jumper et al. [021)) The basic self- attention layer is also permutation equivariant, a useful property for applications to chemistry, where physical quantities should be invariant the ordering of atoms and electrons (Fuchs et al., 2020). Despite the manifest successes in other fields, no one has yet investigated whether self-attention neural networks are appropriate for approximating solutions in computational quantum mechanies.
In this work, we introduce a new self-attention neural network, the Wavefunction Transformer (Psi- former), which can be used as an approximate numerical solution (or Ansatz) for the fundamental equations of the quantum mechanics of electrons. We test the Psiformer on a wide variety of bench- mark systems for quantum chemistry and find that is significantly more accurate than existing neural network Ansatzes of roughly the same size. The increase in accuracy is more pronounced the larger the system is as much as 75 times the normal standard for "chemical accuracy" - suggesting that the Psiformer is particularly attractive approach for scaling neural network Ansatzes to larger,

---

Published as a conference paper at ICLR 2023
more challenging systems. In what follows, we will provide overview of the variational quantum Monte Carlo approach to computational quantum mechanics (Sec. 2 introduce the Psiformer ar- chitecture in detail (See. 3), present results wide variety of atomic and molecular benchmarks (Sec. 3 and wrap up with a discussion of future directions (Sec.
BACKGROUND
2.1
QUANTUM MECHANICS AND CHEMISTRY
The fundamental object of study in quantum mechanics is the wavefumction, which represents the state of all possible classical configurations of system. If the wavefunction known, then all other properties of a system can be calculated from it. While there are multiple ways of representing wavefunction, we focus on the first quantization approach, where the wavefunction is a map from possible particle states to complex amplitude. The state of a single electron x E R x (1, can be represented by its position r IR's and spin a [t, Then the wavefunction for an N-electron system is a function V/ (I33 x (n.U)) - C. Let æ _ Xi..... XN denote the set of all electron states. The wavefunction is constrained to have unit & norm J dæ ( 1, and [ can be interpreted as the probability of observing a quantum system in given state when measured.
Not all functions valid wavefunctions - particles must be indistinguishable, meaning [1/2 should be invariant to changes in ordering. Additionally, the Pauli exclusion principle states that the prob- ability of observing any two electrons in the same state must be zero. This enforced by requiring the wavefunction for electronic systems to be antisymmetric. In this paper, we will focus on how to learn an unnormalized approximation to by representing it with a neural network.
The physical behavior of non- relateiatiiistic quantum systems is described by the Schrödinger equation. In its time-in ineeeindepeneent form, it is an eigenfunction equation (x)) = (xx where F is Hermitian linear operator called the Hamiltonian and the scalar eigenvalue E corresponds to the energy of that particular solution. In quantum chemistry, atomic units (a.u.) are typically used, in which the unit of distance the Bohr radius (oo), and the unit energy is Hartree (Ha).
The physical details of system are defined through the choice of Hamiltonian. For chemical sys- tems, the only details which need to be specified are the locations and charges of the atomic nuclei. In quantum chemistry it is standard to approximate the nuclei classical particles with fixed posi- tions, known as the Born- Noon-Oppenheimer approximation, in which case the Hamiltonian becomes:
Zr
212,
A = 2 V
-
+
Σ TJJ
(1)
R,
R, - RR
27
where V Σ
0² is the Laplacian w.r.t. the ith particle and Z₁ and R,, E (1, N1.....Nudd Nuue } are
3r the charges and coordinates of the nuclei.
Two simplifications follow from this. First, since Á is a Hermitian operator, solutions Vi must be restrict our attention real-valued wavefunctions. Second, since the spins
real-valued. Thus o, do not appear anywhere in Eq. we can fix a certain number of electrons to be spin up and the remainder be spin down before beginning any calculation (Foulkes et al. 2001). The appropriate number for the lowest energy state can usually be guessed by heuristics such as Hund's rules.
While the time-independent Schrödinger equation defines the possible solutions of constant energy, at the energy scales relevant for most chemistry the electrons are almost always found near the lowest energy state, known as the ground state. Solutions with higher energy. known as excited states, are relevant to photochemistry, but in this paper we will restrict our attention to ground states.
For typical small molecule, the total energy of a system is on the order of hundreds to thousands of Hartrees. However the relevant energy scale for chemical bonds is typically much smaller on the order of I kilocalorie per mole (kcal/mol), or 1.6 mHa - less than one part in one hundred thousand of the total energy. Calculations within I kcal/mol of the ground truth are generally considered "chemically accurate". Mean-fied field methods are typically within about 0.5% of the true total energy. The difference between the mean mean-field energy and true energy is known the correlation energy, and chemical accuracy is usually less than 1 1% of this correlation energy. For example, the binding energy of the benzene dimer (investigated in Section 4.5) is only ~4 mHa.
2

---

Published as a conference paper at ICLR 2023
V
Legend
V
learnable function
function
det
det
Jastrow
input array
envelope
projection
envelope
projection
concel diconv
MLP
self attention
XL
xL
afcony
MLP
MLPP
self- aoff-tttenioon
elecronnnuclest
elecron-nu clear
electron fealures
fealuras
(a) FermiNet eehhiNe
(b) Psiformer
Figure 1: Comparison of (a) FermiNet and SermiNet++ShhNet and (b) the Psiformer. The Fer- miNet variants have two streams, acting on electron-nucleat and electron-electron features, which merged via concatenation continuous-fiter convolution operations. In contrast, the Psiformer dlectron-nuclear uses a single stream of self-attention layers, acting on electron nuclear features only. Electron- electron features appear only via the Jastrow factor. The FermiNet+SchNe also includes nuclear embedding stream and separate spin-dependant electron-electron streams, not pictured here.
2.2 VARIATIONAL QUANTUM MONTE CARLO
There wide variety of computational techniques to find the ground state solution of the Hamil- tonian in Eq I We are particularly interested solving these equations from first principles (ab- initio), that is, without any data other than the atomic positions. The ab-initio method most compat- ible with the modern deep learning paradigm is variational quanium Monte Carlo (VMC, Foulkes) al. (200T)) In VMC, parametric wavefunction approximation (or Ansatz) is optimized using samples from the Ansatz itself, in much the same way that deep neural networks are optimized by gradient descent on stochastic minibatches. VMC variational in the sense that it minimizes an upper bound on the energy of system. Therefore two VMC solutions give different energies, the one with the lower energy will be closer to the true energy, even the true energy not known.
In VMC, we start with an unnormalized wavefunction Ansatz Vo
(R3 x (t.4)) N > R with parameters The expected energy of the system given by the Rayleigh quotient:
Co (VgHWeo J dax (x) HV)())) (1)))
Ex-ez
En [EI())
(2)
J (x
where we have rewritten the Rayleigh quotient an expectation over random variable proportional to V/ the right hand side. The term (x) -x= () is known as the local energy. Details on how to compute the local energy and unbiased estimates of the gradient of the average energy are given in Sec. Sec.AA.Ilin the appendix.
In VMC, samples from the distribution proportional to 2 are generated by Monte Carlo methods, and unbiased estimates of the gradient are used to optimize the Ansatz, either by standard stochastic gradient methods, or more advanced methods (Umrigar et al.) 2007): Sorella, 1998). Notably, the samples V can be generated from the Ansatz itself, rather than requiring external data.
The form of L must be restricted to antisymmetric functions to avoid collapsing onto non non-physical solutions. This is most commonly done by taking the determinant of matrix of single-electron
3

---

Published as a conference paper at ICLR 2023
functions (x) det [o(x), where ((x) denotes the matrix with elements (: since the deter- minant is antisymmetric under exchange of rows or columns. This is known as a Slater determinant, and the minimuum-energy wavefunction of this form gives the mean-field solution to the Schrödinger equation. While Φ is a function of one electron in a Slater determinant, any permutation-equivariann function of all electrons can be used input to + and V will still be antisymmetric.
The potential energy becomes infinite when particles overlap, which places strict constraints on the form of the wavefunction at these points, known as the Kato cusp conditions (Katol (957). The cusp conditions state that the wavefunction must be non-differentiable at these points, and give exact values for the average derivatives the cusps. This can be built into an Ansatz by multiplying by a Jastrow factor which satisfies these conditions analytically (Drummond et al. 2004).
2.3 RELATED WORK
Machine learning has found numerous applications to computational chemistry in recent years, but has mostly focused on problems at the level of classical physies (Schütt et al.) [2018] Fuchs et al. et al. 2018), which all rely on learning from large datasets of experiments or ab-initio calculations. There is also work on machine learning for density functional theory (DFT) (Nagai et al. 2020; Kirkpatrick et al.) (02T), which is
2020
Batzner et al. 2022] Segler et al. 2018 Jomezz-Bomarne
an intermediate between classical and all-electron quantum chemistry, but even this is still primarily a supervised learning problem and relies on ab-initio calculations for data. Here instead, we are focused on improving the ab-initio methods themselves.
For thorough introduction to ab-initio quantum chemistry, we recommend Helgaker et al. (2014) and Szabo & Ostlund (2012). Within this field, VMC was considered simple but low-accuracy method, failing to match the performance of sophisticated methods (Motta et al.) 2020), but often used as a starting point for diffusioo Monte Carlo (DMC) calculations, which are more accurate, but do not produce an explicit functional form for the wavefunction (Foulkes et al. 2001). These VMC calculations typically used a Slater-Justrow Ansatz (Kwon et al., 1993). which consists of a large linear combination of Slater determinants multiplied by Jastrow factor. They also sometimes include backflow, a coordinate transformation that accounts for electron correlations with a fixed functional form (Feynman & Cohen 1956).
Recently, the use of neural network Ansatzes in VMC was shown to greatly improve the accuracy of many-electron calculations, often making them competitive with, or in some circumstances superior to, methods like DMC or coupled cluster (Pfau et al.
2020 Hermann et al. 2020 Choo et af.) 2020 Han et al. 2019; Luo & Clark] 2019 Taddei et al. 2015). In first quantization, these Ansatzes used a sum of a small number of determinants to construct antisymmetric functions, but used very general
functions of all electrons as inputs, rather than the single-electron functions used in Slater determinants. Some, like the PauliNet (Hermann et al.) 2020), used Jastrow factors, while the FermiNet (Přau et al. 2020) used con-differentiable input features learn the cusp conditions.
Most follow-up work, as surveyed in Hermann et al. 2022), integrated these Ansatzes with other methods and applications, but did not significantly alter the architecture of the neural networks. The most significant departure terms the neural network architecture was Gerard et al. (022)) which extended the FermiNet generally recognized to be the most accurate neural network Ansatz up to that point- and integrated with several details of the PauliNet, especially the continuous-fi convolutions also used by the SchNet (Schütt et al. 20188), claiming to reach even higher accuracy on several challenging systems. We refer to this architecure as the FerniNet+SchNNet.
3 THE PSIFORMER
The Psiformer has the basic form:
New
Vo(x)) exp (Jo(z:) Σ det ())
(3)
where To : : x (t. + R and Φ (R' x ()) - NXN are functions with learnable
parameters 0. This is similar the Slater-Jastrow Ansatz, FermiNet and PauliNet, with the key
4

---

Published as a conference paper at ICLR 2023
difference being that in the Psiformer, Ø consists of a sequence of multiheaded self-attention layers Vaswani et al. 2017). The key motivation for this is that the electron-electron dependence in the Hamiltonian introduces subtle and complex dependence in the wavefunction. Self-attention is one way of introducing this without fixed functional form. The high-level structure of the Psiformer is shown in Fig. [b), where is contrasted with the FermiNet (and SchNet extension) in Fig. Ta).
Because a self-attention layer takes sequence of vectors input, only features of single electrons are used as input to $ The imput fea- ture vector f, for electron i is simi-
0.03
FermiNet Pfau et al. (2020)
FermiNet (small)
FermiNet (large)
Psifarmer (small)
lar to the one-electron stream of the FermiNet, which uses concatena- tion of electron-nuclear differences R and distances r, Ryl.
0.02
E E-Ere Erer (a
Psiformer (large)
0.01
== I,..,, with two key dif- ferences. First, found that for systems with widely separated atoms, using FermiNet one-electron features caused self-attention Ansatzes to be- come unstable, we rescale the in- puts in the Psiformer by a factor of R1). so
0.00
LiH
Li2
NH3
CH4
CO
N2
C2H4
methyl-
amine 03
ethanol
bicyclo-
butane
log(1 + Ir; Rrl)/Ir; that the input vectors grow logarith- mically with distance from the nu- cleus. Second, we concatenate the spin o, into the input feature vector itself (mapping to I and J to -1). While this spin term is kept fixed dur- ing training, breaks the symmetry between spin up and spin down electrons, that columns of T are only equivariant under exchange of same-spin electrons. This a notable departure from the FermiNet and PauliNet, where the difference between spin-up and spin-down electrons is instead built into the architecture.
Figure 2: FermiNet and Psiformer accuracy on small molecules. Geometries and CCSD(T)/CBS reference ener- gies are taken from Pfau et al. (002)). The grey region indi- cates chemical accuracy (1 kcal/mol or 1.6 mHa) relative to the reference energies.
The input features f are next projected into the same dimension the attention inputs by linear mapping h = o 4, and then passed into a sequence of multiheaded self-attention layers followed by mnnaien layers, both with residual connections:
tt+1 = h W concatt, (S hyy W w w
h.t1 - = p + tanh (W + b
(5)
where h indexes the different attention heads, concat denotes concatenation of the output from different attention heads, and SELFATTN denotes standard self-attention:
(E hN; W WA, W.) Vd
Σ oj (qf ...... Kkik vj
(6)
k. Weh, q W W,h,, W.h;
exp(x;:
Oz(21...... .N) E, exp(23)
(8)
where dl is the output dimension of the key and query weights. In principle, multiple linear-nonlinear layers could be used in-between self-attention layers, but we found that adding deeper MLP be- tween self-attention layers less effective than adding more self-attention layers. While smooth nonlinearity must used to guarantee that a wavefunction smooth everywhere except the cusps, found that using activation functions other than tanh had marginal impact on performance.
final linear projection into a N Naet dimensional space is applied to the activations, and the output is multiplied by a weighted of sxponentially-decayae envelopes, T 2 wt hy, where
5

---

Published as a conference paper at ICLR 2023
Benzene
-599.000
-1263.630
-232.210
-1263.035
-232.215
-99.9005
Energy (a. v....)
-1263.640
-232.220
LN, Pfau
VF. LN, Gerard
-599.15
-1263.650
VF, no LN, Gerard
FN, Pfau
12563.555
-322.335
-599.920
-1263.660
-232.240
FN-SN, Pfau
FN+SN, Gerard
599.925 0 50 100 150 200
1263 665
100 150 200
-322.245
200 Iterations (thousand)
50
100
150
FN-SN (original)
lterations (thousand)
Iterations (thousand)
Figure 3: Comparison of the Psiformer (V/F), with and without LayerNorm (LN) against the Fer- miNet (FN) and FermiNet++ (FN+SN) using training hyperparameters from Pfau et al. (2020) and Gerard al. 20022). Learning curves from the original FN+SN implementation in Gerard al. 2022) are in green. Energies are smoothed over 4000 iterations.
& exp (-ofylr - Ryl). This is so that the boundary condition lim -x (T( = is enforced, and is the same form the envelope used by the FermiNet in |Spencer al. (2020a). The matrix elements $ are then passed into determinants and the output summed together.
"i =, J
As the distances r R,y are inputs to the Psiformer, is capable of learning the electron- cusp conditions, much like the FermiNet. However, the self-attention part of the Psiformer does not take pairwise electron distances as inputs, so it cannot learn the electron-electron cusp conditions. Instead, the Psiformer uses conventional Jastrow factor only for the electron-electree cusps. We use particularly simple Jastrow factor:
Io(x)=
Σ
1
paz
Σ
1
(9)
4 O'par
r; r
2 Olanti +
rr rg
i<ji;=a;
i.joo
which has only two free parameters, Opar and Clantis and works well in practice.
4 EXPERIMENTS
Here present evaluation of the Psiformer wide variety of benchmark systems. Where it is not specified, all results with the FermiNet and SermiNet+SchNett are with our own implementation, forked from Spencer et al. (2020b). We
fantecer-facctoed Approximate Curvature (KFAC) (Martens & Grosse| 2015) to optimize the Psiformer. The use of KFAC for self-attention has been investigated in Zhang et al. (011)). We show the advantage of KFAC, and how interacts with Lay- erNorm, in Sec. A.2.3 of the appendix. We use hyperparameters and a Metropolis-Hastingg algorithm similar the original FermiNet paper, though we have made several modifications which help for larger systems: we pretrain for longer, generate samples for pretraining from the tar- get wavefunction, we take more MCMC steps between parameter updates, we propose updates for subsets of electrons rather than all electrons simultaneously, and we slightly change the gradient computation be more robust to outliers. Details are given in Sec. A.c.44.2.4ff the appendix.
MCMC
4.1 REVISITING SMALL MOLECULES
In Pfau et al. (220). they compared the FermiNet against CCSD(T) extrapolated the complete basis set (CBS) limit on number of small molecules (4-30 electrons) from the G3 database (Curtiss) ct al. 2000). While the FermiNet captured more than 99% the correlation energy relative to CCSD(T)CBS for systems large as ethene (16 electrons), the quality of the FermiNet calculation began to decline as the system size grew. While some of this discrepancy was reduced simply by changing framework with better numerics (Spencer al., 2020a), the reported difference in energy for bicyclobutane was still greater than 20 mHa. Here we revisit these systems, comparing the FermiNet with training improvements against the Psiformer.

---

Published as a conference paper at ICLR 2023
FermiNet
CermiNet+Schhee
Paiformer
System
Expl.
Small
Large
Small
Large
Small
Large
K (Ha)
-599.11332)
-599.114922)
-599.9153(2)
-599.9175(3)
-509.9202(2)
$99,9205(1)
KT (Ha)
-599.7561())
-59.7677122)
-5997777(22)
-599.700((2)
-599.7589(1)
-599,7617(1)
Ix (eV)
4,278(9)
4.000g9)
4,286(8)
4,28(1)
4.388(6)
4.321(5)
4,32631
Fe (Ha)
-1263.6363(5)
-1263.6381(4)
-1263.6464(3
-1263.6385(()
(263.6595(()
-1263.0613(3)
Fet (Ha)
-1263.35705))
-1263.3616(4)
-1663362004
-1263,3676(4)
1263,3776((3
-126333768(3)
1Pr (eV)
7,60(2)
7,52(2)
7,74(2)
7,37(2)
7.75(1)
7.74(1)
7,84194
7n (Ha)
-1779
-17794131(66
-1779419556)
-17794267((5)
-37794304(4)
-1779,4365(5)
Za (Ha)
-1779 077816)
-1779.0846(6)
-1779.0843(7)
-1779.0909(5)
-1779,102(1)
-1779.1054(3)
1P (eV)
9.04(2)
8,94()
9.12(2)
9,14(2)
8.94(3)
9.01()
9.24695
Table 1: Energies of third-row neutral atoms and cations. Ionization potentials are compared against experimental results from from|Koga et al. a (1997). Total energies are in Hartree while ionization potentials are in eV. Chemical accuracy is 0.043 eV.
To investigate whether the Psiformer performance could he reproduced by simply making the Fer- miNet larger, investigated both "small" and "large" configuration for the Psiformer and Fer- miNet. The small FermiNet has the same layer dimensions and determinants as the one used in Pfau al. 22200), while the large configuration has twice many determinants and one-electron stream Twice as wide, similar to the largest networks in Spencer et al. (2020a). The Psiformer configura- tions have the same number of determinants and the MLP layers are the same width as the FermiNet one-electron stream, though due to the self-attention weights the small Psiformer has a number of parameters between that of the large and small FermiNet. Exact details are given in Table 7 in the appendix. On these systems, the small Psiformer ran in similar amount of time to the large Fer- miNet, though for even larger systems, the difference in wall time between the small FermiNet and small Psiformer became much smaller (see Table 9 in the appendix).
The results on small molecules can seen in Fig. 2 While increasing the size of the FermiNet increases the accuracy somewhat, the small Psiformer is more accurate than the large FermiNet, despite having fewer parameters, while the large Psiformer is the most accurate of all. This is true for all systems investigated. The improvement from the Psiformer is particularly dramatic ozone and bicyclobutane on ozone, the large Psiformer is within kcal/mol of CCSD(TYCBS, while even the largest FermiNet has error more than 4 times larger than this. On all molecules, the large Psiformer captures more than 99% of the correlation energy relative to the reference energy.
4.2 COMPARISON AGAINST FERMINET WITH SCHNET-LIKE CONVOLUTIONS
While the performance of the Psiformer relative to the FermiNet is impressive, recent work has proposed several innovations to reach even lower energies (Gerard et al., 2022). The primary inno- vations of this work were the SchniNettSchNett architecture and new hyperparameters they claim led to much faster optimization. They showed an especially large improvement on heavy atoms like potassium (K) and iron (Fe), and some improvement on larger molecules like benzene.
We attempted reproduce the results of Gerard et al. (022) with our own SerniNet+SchNeet im- plementation, with somewhat surprising results in Fig. First, the changes in training used for the FermiNet seem to be enough to close the gap with the published SermiNet+SchNet results heavy atoms. For instance, ablation studies suggested that modifying the hyperparameters plus adding SchNet-like convolutions accounts for 38.9 mHa improvement in accuracy on the potassium atom at 10th iterations, but in our experiments the FermiNet with default hyperparameters is within a few mHa of the published result. On benzene, both the SchNet-like convolutions and modified hyper- parameters improved the final energy by a few mHa, though still fell slightly short of the published results. Most importantly, the Psiformer is clearly either comparable to the best published Fer- SeNet+SchNet results, as on K, or better by wide margin, as on benzene. These results also give us confidence that the FermiNet as strong a baseline as any other published method for further comparison.
4.3 THIRD-ROW ATOMS
The FermiNet, FerniNet+SchNett and Psiformer all seem to perform well on the potassium and iron atoms, but it is difficult to judge how close to the ground truth these results are. Prior work
7

---

Published as a conference paper at ICLR 2023
Ours
Best Published
System
Psiformer
FermiNet
VMC
DMC
No LayerNorm
LayerNorm
Benzene
-232.2205(2)
-232.2400(1)
2232.2393(1)
-232.2267
-232.2370(3)
Toluene
-271.5274(2)
-271.5494()
-271.5538(1
Naphthalene
-385.8147(4)
-385.8679(2)
-385.8685(2)
CCI,
-1878.684(1)
-1878.734(1)
-1878.804(1)
Table Energies for molecules with between 42 and 74 electrons. For benzene, the best published results using neural network Ansatzes are either from GGearr et al. (2022) or "Ren et al. (2022).
had shown that the FermiNet can achieve energies within chemical accuracy of exact results for atoms up to argon (Spencer et al. 2020a), but exact calculations are impractical for third-row atoms. Instead, comparison to experimental results is more practical evaluation metric. The ionization potential the amount of energy it takes to remove one electron is a particularly simple compar- ison for which good experimental data exists (Koga et al. 1997). Here compare the FermiNet, FermiNet+SchNNt and Psiformer for estimating the ionization potential of potassium, iron and zinc To compare the relative importance of size and architecture, also looked both small and large configuration of all Ansatzes, with parameters described in Table the appendix.
Results are shown in Table For atoms this heavy, all methods showed some amount of run-to-run variability, usually small, but on occasion as large as 10 mHa, which may explain some outlier re- sults. On potassium, the difference between Ansatzes was within the range of run-to-run variability, but the Psiformer still did quite well, and reached the most accurate ionization potential. On heavier atoms, the difference between architectures became more pronounced, and the improvement of the Psiformer relative to other models on absolute energy was more robust. The results on ionization potentials were more mixed, and no Ansatz came within chemical accuracy of the ground truth. This shows that even the Psiformer is not yet converged to the ground truth, though it is the Ansatz closest to reaching so far.
4.4 LARGER MOLECULES
Much of the promise for deep neural networks for QMC comes from the fact that they can, in theory, scale much better than other all-electron methods, though this promise has yet to be realized. While CCSD(T) scales with number of electrons as O(N)), ). single iteration of wavefunction optimization for a fixed neural network size scales as O(N4) in theory, and in practice is closer to cubic for system sizes of several dozen electrons. Self-attention has been especially powerful when scaling to extremely large problems in machine learning - here we investigate whether the same holds true for QMC, by applying both the FermiNet and Psiformer to single molecules much larger than those which have been investigated in most prior work.
Results systems from benzene (42 electrons) carbon tetrachloride (CCl,, 74 electrons) given in Table[2] On these systems, CCSD(T) becomes impractical for us to run without approxima- tions, so we only compare against other QMC results. Due to computational constraints, we were only able to run the small network configurations on these systems. Additionally, to help MCMC convergence,
updated half the electron positions moving all electrons simultaneously, as all-electron moves become less efficient for larger systems.
time in each MCMC move, rather than
In Table 2 it can be seen that the Psiformer is not only significantly better on benzene than the Fer- miNet, as in Fig. 3 but it better than the best previously published DMC energy. a remarkable feat. For even larger molecules, there are no results in the literature with neural network wavefunctions to compare against, so we only compare the Psiformer and FermiNet directly. The Psiformer outper- forms the FermiNet by an ever larger margin on larger systems, reaching 120 mHa (75 kcal/mol) on CCI.. On the three hydrocarbon systems investigated, LayerNorm has only small impact. How- ever for CCI, LayerNorm has a signficant impact, accounting for 70 mHa of the total 120 mHa improvement over the FermiNet. While do not claim these are the best variational results in the literature, it is clear that on larger molecules the Psiformer is a significant improvement over the FermiNet, which is itself the most accurate Ansatz for many smaller systems.
8

---

Published as a conference paper at ICLR 2023
Ours
Ren et al. 2022
Paiformer
Expt.
NermiNet
VMC
DMC
No LayerNorm
LayerNorm
Equilibrium
-464.3770(5)
464.4624(2
464.4667(2)
-664.406((3
-661,4610(2)
Dissociated
464.3724(6)
464.4674(2)
464.4660(2)
A Amoono -0.064066)
-001766(2)
-0.0119(2)
0.0219(3)
-0.0020:5)
0.0038866
0.00446(8)
0.0050(3)
0.0007(3)
0.0183(5)
0.0092(()
0.0038/6)
Table 3: Energies for the benzene dimer, with center separation of 4.95 Á (cquilibrium) and 10 (dissociated), and the estimated dissociation energy from taking the difference of the equilibrium energy from twice the mononer energy (A Emono)) and the dissociated energy (&E, 1oAA All A Emuono results are based on comparing like-for-like dimer and monomer calculations. AE in results from Ren et al. 2022) are estimated from figures. Experimental energies are from Grover et al. (1987)
4.5 THE BENZENE DIMER
Finally, we look at the benzene dimer, a challenging benchmark system for computational chemistry due to the weak der Waals force between the two molecules, and the largest molecular system ever investigated using neural network Ansatzes (Ren et al. 2022). The dimer has several possible equilibrium configurations, many of which have nearly equal energy (Sorella al. 2007 Azadi &
Cohen 2013), making it additionally challenging to study computationally. but here We restrict our attention the T-shaped structure (Fig. 4 so that we can directly compare against Ren
(2022).
Results are given in Table The results from Ren et al. (0022) are from a small FermiNet (3 layers) trained for 2 million Ierations. Even our "small" 4-layer FermiNet baseline (which is larger than their FermiNet) trained for 200,000 iterations is able to reach the same accuracy as their small FermiNet trained for 800,000 iterations (Ren al. (2022) Supplementary Figure 7a). The Psiformer reaches significantly lower energy than the FermiNet with the same number of training iterations, again surpassing the DMC result by a few mHa.
We also tried to estimate the dissociation energy by similar means as Ren et al. (2022) comparing against the energy of the same model trained with bond length of 10 Á, and twice the energy of the same model trained on the monomer. Ironically, our FermiNet baseline, which had the worst absolute energy, had the best relative energy between configurations. While every model underestimated the energy relative twice the monomer, the discrepancy est with DMC. It should be noted that the zero-point vibrational energy (ZPE) is not included, the dissociation energies from Ren jet al. (2022) in Table 3 are estimates based on their figures, and there disagreement over the exact experimental energy (Grover et al.] 1987 Krause et al., 1991). so this comparison should be con- sidered a rough estimate only. Our main result is that the absolute energy of the Psiformer is a vast improvement over the FermiNet, and we leave it to future work to properly apply the Psiformer to predicting binding energies.
low-
4.95 A
Figure 4: The T-shaped ben- zene dimer equilibrium.
5 DISCUSSION
We have shown that self-attention networks are capable of learning quantum mechanical properties of electrons far more effectively than comparable methods. The advantage of self-attention networks seems to become most pronounced on large systems, suggesting that these models should be the fo- cus of future efforts to scale to even larger systems. In addition to the strong empirical results, using standard attention layers means that we can leverage existing work improving scalability, either by architectural advances (Child et al. 2019 Wang et al. 2020: Xiong et al. 2021 Jaegle et al. 2021a/b) or software implementations optimized for specialized hardware (Dao et al. 2022). which
could make it possible to scale these models even further. This presents a promising path towards studying the most challenging molecules and materials in silico with unprecedented accuracy.

---

Published as a conference paper at ICLR 2023
ETHICS STATEMENT
The work presented here focuses on fundamental questions in computational chemistry, and is not yet at the stage where it is likely to be widely adopted by experimental chemists. However, in the fu- ture, this line of work could lead computational chemistry becoming much more accurate, making it easier for chemists and material scientists to make new discoveries without requiring cumbersome trial-and-error physical experiments. This could lead to the discovery of new beneficial drugs or industrial chemical processes which are more environmentally friendly. The field of experimental chemistry already has robust ethical standards and processes for preventing harmful applications, and we are confident that more accurate computational methods will not in any way make it easier to circumvent these safeguards.
REPRODUCIBILITY STATEMENT
Details of training, network parameters, and optimization hyperparameters are given the appen- dices for all experiments shown. The code is available under the Apache License 2.0 part of the FermiNet repo at https: /tte:://iitttt com /dee//eeeppmm//e hdd fereerrdddieee/e/////m fermfrrrrnnnee net.
ACKNOWLEDGMENTS
We would like to thank Alex G. de G. Matthews for suggesting the model name, Alex Botev for assistance with KFAC, Michael Scherbela and Leon Gerard for providing data for figures, and James Kirkpatrick for support and encouragement.
REFERENCES
Sam Azadi and R. E. Cohen. Chemical Accuracy from Quantum Monte Carlo for the Benzene Dimer. The Journal of Chemical Physics, 143(10):10430,, 2015.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer Normalization. arXiv preprint arXiv: N000://0.00000 060400645 2016.
Simon Batzner, Albert Musaclian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Ko- mbluthh Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E(3)-Equivariant Graph Neural Networks for Data- Efficiiff and Accurate Interatomic Potentials. Nature Communications, 13 (1):1-11, 2022.
Aleksandar Botev and James Martens. KFAC-JAX, 2022. URL http::/github..cm com/
deepmi nd/ facaa-ca jax]
James Bradbury, Roy Frostig, Peter Hawkins. Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Mo Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
//github. com/ poom//gogle/jax
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D. Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language Models are Few-Shot Learners. Advances Neural Information Processing Systems, 33:1877-1901, 2020.
Giuseppe Carlco and Matthias Troyer. Solving the Quantum Many-Body Problem with Artificial 2017.
Neural Networks, Science, (55(6325):602--
Gino Cassella, Halvard Sutterud, Sam Azadi, N. D. Drummond, David Pfau, James S. Spencer, and W. M. C. Foulkes. Discovering Quantum Phase Transitions with Fermionic Neural Networks. arXiv preprint arXiv:2202.05188. 2022.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating Long Sequences with Sparse Transformers. arXiv preprint arXiv:1904. 10509, 2019.
Kenny Choo, Antonio Mezzacapo, and Giuseppe Carleo. Fermionic Neural-Network States for Ab-initio Electronic Structure. Nature Communications, 1(1):1-7, 2020.
10

---

Published as a conference paper at ICLR 2023
Larry A Curtiss, Krishnan Raghavachari, Paul C Redfern, and John A Pople. Assessment of gaussian-3 and density functional theories for a larger experimental test set. Chem. Phys., 112(17):7374-788, May 2000.
Tri Dao, Daniel Y Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and Memory-Efficient Exact Attention with 10-Awareness. arXiv preprini (iti::2205.1433333 2022.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding, arXiv preprint arXiv: arXiv: 1810,04805, 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. arXiv preprini arXiv:2010. 1100000.199999 1929, 2020.
N. D. Drummond, M. D. Towler, and R. Needs. Jastrow Correlation Factor for Atoms, Molecules, and Solids. Physical Review B, 70(23):235119, 2004.
RP Feymman and Michael Cohen. Energy spectrum of the excitations in liquid helium. Physical Review, 102(5):1189, 1956.
W. M. C. Foulkes, Lubos Mitas, J. Needs, and Guna Rajagopal. Quantum Monte Carlo Simula- tions of Solids. Reviews of Modern Physics, 73(1):33, 2001.
Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Transforanfso:mrr:: 3D Roto-
Translation Equivariant Attention Networks. Advances in Neural Information Processing Sys- tems, 1991999999811 2020.
Nicholas Gao and Stephan Günnemann. Sampling-Free Inference for Ab-initio Potential Energy Surface Networks. arXiv preprint arXiv:2205. 14962, 2022.
Leon Gerard, Michael Scherbela, Philipp Marquetand, and Philipp Grohs. Gold-Standard Solutions to the Schrödinger Equation using Deep Learning: How Much Physics Do We Need? Advances in Neural Information Processing Systems, 2022.
Rafael Jommez-Bombarell.. Jennifer N. Wei. David Duvenaud, José Miguel Hernández-Lobato, Benjamín Sánchez- Sánchez-Lengeling. Dennis Sheherla, Jorge Aguilera-parraguire, Timothy D. Hirzel, Ryan P. Adams, and Alán Aspuru-Guzik. Automatic Chemical Design Using Data-Driven Con- tinuous Representation of Molecules. ACS Central Science, 4(2):268-276, 2018.
J. R. Grover, E. A. Walters, and E. T. Hui. Dissociation Energies of the Benzene Dimer and Dimer Cation. Journal of Physical Chemistry, 91(12):3223-3237, 1987.
Jiequn Han, Linfeng Zhang, and E Weinan. Solving Many-Electron Schrödinger Equation Using Deep Neural Networks. Journal of Computational Physics, 399:108929, 2019.
Trygve Helgaker, Poul Jorgensen, and Jeppe Olsen. Molecular Flectronic-Structuue Theory. John Wiley & Sons, 2014.
Jan Hermann, Zeno Schätzle, and Frank Noé. Deep-Neural-Network Solution of the Electronic Schrödinger Equation. Nature Chemistry, (3(10):911887,,000
Jan Hermann, James Spencer, Kenny Choo, Antonio Mezzacapo, W. M. C. Foulkes, David Pfau, Giuseppe Carleo, and Frank Noé. Ab-initio Quantum Chemistry with Neural- Wave- functions. arXiv preprint ar0iv:2208. 102000..25990 /2590, 2022.
Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Tonescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, et al. PerceiverIO: A General Architecture for Structured Inputs & Outputs. arXiv preprint arXiv:2107. 2021a.
Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and Joao Carreira. Perceiver: General Perception with Iterative Attention. In International conference on machine learning, pp. 4651 16614466. PMLR, 2021b.
11

---

Published as a conference paper at ICLR 2023
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly Accurate Protein Structure Prediction with AlphaFold. Nature, 596(7873):883-58) 2021.
Tosio Kato. On the Eigenfunctions Many- Systems in Quantum Mechanics. Communica- tions on Pure and Applied Mathematics, 10(2):151-177 1957.
James Kirkpatrick, Brendan McMorrow, David HP Turban, Alexander L. Gaunt, James S. Spencer, Alexander G. D. G. Matthews, Annette Obika, Louis Thiry, Meire Fortunato, David Pfau, et al. Pushing the Frontiers of Density Functionals by Solving the Fractional Electron Problem. Science, 374(6573):1385 1389, 2021.
Toshikatsu Koga, Hidenori Aoki, JM de la Vega, and Hiroshi Tatewaki. Atomic Ionization Potentials and Electron Affinities with Relativistic and Mass Corrections. Theoretical Chemistry Accounts, 96(4):248-25) 1997.
H. Krause, B. Ernstberger, and J. Neusser. Binding Energies of Small Benzene Clusters. Chemical Physics Letters, 184(5-6):411-411 1991.
Yongkyung Kwon, DM Ceperley, and Richard M Martin. Effects of Three-Body and Backflow Correlations in the Two-Dimensional Electron Gas. Physical Review B, 48(16):12037, 1993.
Xiang Li, Cunwei Fan, Weiluo Ren, and Ji Chen. Fermionic Neural Network with Effective Core Potential. Physical Review Research, 4(1):01302,, 2022.
Jeffimin Lin, Gil Goldshlager, and Lin Lin. Explicitly Antisymmetrried Neural Network Layers for Variational Monte Carlo Simulation. arXiv preprint arXiv:2112.0349) 2021.
Di Luo and Bryan K Clark. Backflow Transformations via Neural Networks for Quantum Many- Body Wave Functions. Physical Review Letters, 122(22):226401, 2019.
James Martens and Roger Grosse. Optimizing Neural Networks with Kronecker-factored Approx- imate Curvature. In Proceedings of the 32nd International Conference on Machine Learning Volume 37, pp. 2408 2017. PMLR, 2015.
James Martens, Andy Ballard, Guillaume Desjardins, Grzegorz Swirszez, Valentin Dalibard, Jascha Sohl-Dickstein, and Samuel S Schoenholz. Rapid Training of Deep Neural Networks With- Skip Connections Normalization Layers Using Deep Kernel Shaping. arXiv preprint arXiv:21 10.01 765, 2021.
Mario Motta, Claudio Genovese, Fengjie Ma, Zhi-Hao Cui, Randy Sawaya, Garnet Kin-Lic Chan, Natalia Chepiga, Phillip Helms, Carlos Jiménez-Hoyoon Andrew J Millis, al. Ground-Strre Properties of the Hydrogen Chain: Dimerization, Asulator-to-Mette Transition, and Magnetic
Phases. Physical Review X, 10(3):03106 1058, 2020.
Ryo Nagai, Ryosuke Akashi, and Osamu Sugino. Completing Density Functional Theory by Machine- Learning Hidden Messages from Molecules. npj Computational Materials, 6(1):43, 2020. ISSN 2057-3960. doi: 10. 10.1038//41524-02-0011111
Shivesh Pathak and Lucas K. Wagner. A Light Weight Regularization for Wave Function Parameter Gradients in Quantum Monte Carlo. AIP Advances, 10(8):085213, 2020.
David Pfau, James S. Spencer, Alexander G. G. Matthews, and W. M. C. Foulkes. Ab-initio Solu- tion of the Many-Flectron Schrödinger Equation with Deep Neural Networks. Physical Review Research, 2(3):033429, 2020.
Weiluo Ren, Weizhong Fu, and Ji Chen. Towards the Ground State of Molecules via Diffusion Monte Carlo on Neural Networks. arXiv preprint arXiv:2204.13903 2022.
Erwin Schrödinger: An Undulatory Theory of the Mechanics of Atoms and Molecules. Physical Review, 1846). 1926.
Kristof T Schütt, Huziel E Sauceda, P-J Kindermans, Alexandre Tkatchenko. and K-R Müller. Schnet-a Deep Learning Architecture for Molecules and Materials. The Journal of Chemical Physics, 148(24)2441222, 2018.
12

---

Published as a conference paper at ICLR 2023
Marwin H. S. Segler, Mike Preuss, and Mark P. Waller. Planning Chemical Syntheses with Deep Neural Networks and Symbolic AI. Nature, 555(7698):604 610, 2018.
Sandro Sorella. Green Function Monte Carlo with Stochastic Reconfiguration. Physical Review Letters, 80(20):455,, 1998.
Sandro Sorella, Michele Casula, and Dario Rocca. Weak Binding Between Two Aromatic Rings: Feeling the Van der Waals Attraction by Quantum Monte Carlo Methods. The Journal of Chemical Physics, (27(1):0144)) 2007.
James Spencer, David Pfau, Alcksandar Botev, and W. M. C. Foulkes. Better, Faster Fermionic Neural Networks, arXiv preprint 10.0x:20/1.07/25,5 2020a.
James S. Spencer, David Pfau, and FermiNet Contributors. FermiNet, 2020b. URL http://
github. h/tp://p//////////c//cccccccc cttpreppminn///emmieeeeeee
Attila Szabo and Neil S. Ostlund. Modern Quantum Chemistry: Introduction to Advanced Electronic Structure Theory. Courier Corporation, 2012.
Michele Taddei, Michele Ruggeri, Saverio Moroni, and Markus Holzmann. Iterative Backflow Renormalization Procedure for Many-Body Ground-Star Wave Functions of Strongly Interacting Normal Fermi Liquids. Physical Review B, 91(11):115106, 2015.
C. Umrigar, Julien Toulouse, Claudia Filippi, Sandro Sorella, and Richard G. Hennig, Alleviation of the Fermion-Sign Problem by Optimization of Many-Body Wave Functions. Physical Review Letters, 98(1 10:001, 2007.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention All You Need. Advances in Neural Information Processing Systems, 30, 2017.
Sinong Wang, Belinda 7. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-Attention with Linear Complexity. arXiv preprint arXiv:2006. 04768, 2020.
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A Nyström- -aaedd Algorithm for Approximating Self-AAtention. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 14138-14148, 2021.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song. James Demmel. Kurt Keutzer, and Cho-Jui Hsich. Large Batch Optimization for Deep Learning: Training BERT in 76 Minutes. Eighth International Conference on Learning Repre- sentations (ICLR), 2020.
Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George Dahl, Chris Shallue, and Roger B Grosse. Which algorithmic choices matter at which batch sizes? Insights from noisy quadratic model. Advances Neural Information Processing Systems (NeurIPS), 32, 2019.
13

---

Published as a conference paper at ICLR 2023
A APPENDIX
A.1 CALCULATING ENERGIES AND GRADIENTS
It is generally more numerically stable to work directly with the log wavefunction, and the local energy can be expressed as
N 3
1g (()))
+ (
27 + V(())
E (z
ΣΣ 2 5=1
(10)
dr
Öriy
where V(z) is the potential energy (the last three terms of Eq The gradient of the energy is given by
Vi w [x)( 2B [()(()) Ev (Ez(z))) (= Vlog (x))
(1I)
where the local energy (x) = -p - (x=) the difference between the local energy at x and the average energy over all x'.
Note that the (x) - Em (z)2')) term is
We make a small but critical change to how the gradients are computed relative to Pfau et al. (2020) which stabilizes training for all models considered here. The local energy often has very large tails due to numerical issues, especially near cusps (where two particles overlap) and nodes (where the wavefunction goes to zero) (Pathak & Wagner, 2020). To mitigate this, the local en- ergy is often truncated in practice. Let (EL) mean denote the mean local energy for one mini- batch of walkers and (EL) median denote the median. Then Pfau et al. (2020), the local ener- gies were clipped to be within a constant multiple p of the mean absolute deviation around the >>>> E, (x;)):; - (F1..) me..n I. While this is more robust than the stan- dard deviation, it is still susceptible to large outliers because it is centered at the mean. Instead, we use the mean absolute deviation around the median to determine the window for clipping: M.D (E) N Σ; (x;) (EL) median
mean: MADmean FE)) =
Secondly, in Pfau et al. 2020). the average energy term Eg'ny> [E,, [E,, (2')] in the gradient was ap- proximated by the mean local energy (Ez) of a minibatch. That meant that outliers were still included in this term. We instead use the mean of the clipped local energies. This guarantees that the mean over minibatch of the energy difference term always zero, improving stability during optimization. If we let clipmean/ clippean/moedian (ar; denote the functions that clip to be in the range
((x) aucan/median o, (z) mean, 'median + o). then the gradient for one batch in Pfau et al. (2020) was:
(x)) clipmean (FI,(ar); pMAD MAMDmaan mean (FL))
(12)
((Br(1)) - (Er.) mean) Vlog V (z))) mean
(13)
while here we use:
Br(x)= = clipmediaa (Zr( PMADmedian (Er.)) (()((z)) (Bx)
median
(14)
(Ei) mean
Vlog v (x))
(15)
mean
While this may seem like a subtle difference, it has a great effect on larger systems, especially heavier atoms. A similar technique was used to stabilize training the FermiNet with pseudopotentials
(Li et al., 2022), but instead of clipping outlier local energies, the outlier walkers were removed
entirely for that minibatch. We found that clipping with proper centering was more effective than removing the walkers entirely.
A.2 TRAINING
In this section we give details on training and further differences from the original FermiNet (Pauu et al. 2020).
A.2.1
DENSE DETERMINANTS
The original FermiNet and PauliNet Ansatzes used block-diagonal determinants for spin-up and spin-down electrons (Ptau al. 2020 Hermann al. 2020), such that each determinant could be
14

---

Published as a conference paper at ICLR 2023
796.34
-1253.62
-1263.63
776.35
Energy (a.u.)
Energy (a.u.)
-1263.64
-96.36
-1263.85
-796.37
ADAM LN
UN
-1263.66
KFAC. LN
KFAC, LN
KFAC, no LN
796.38
-1263.67
100 Iterations (thousand)
200
iterations (thousand)
(a) S2, R=1.9202
(b) Fe atom
Figure 5: Comparison of different optimization algorithms (KFAC and ADAM) for the Psiformer on (a) the sulphur dimer and (b) the iron atom, with and without LayerNorm. While there is clear advantage using KFAC on both systems, LayerNorm has a marginal impact on the sulphur dimer. On the iron atom, however, LayerNorm improves the accuracy of ADAM and the stability of KFAC. learning rate of 3e-4 was used for ADAM. rolling mean of the last 1000 iterations is used to smooth the energy.
factorized into the product of a determinant of spin-up electrons of size Nj x N, and determinant of spin-down electrons N₂ x N1, as the case with conventional VMC ansatzes (Foulkes et al. 2001:
det (xx) = o (x++) 0
2 (x1) = det [#+( det (x.)).
(16)
where H, denotes the set of electron states with the specified spin and different functions are used for electrons of different spins. The FermiNet authors subsequently proposed simply using dense determinants of size N x N. N N4 + N, (Spencer et al. 2020b):
det + [xx) (2+(x+)
(17)
where now each block is of dimension N x Ng. This has largely become the default choice for FermiNet architectures and has shown to provide improved accuracy at small additional cost for a
myriad of systems (Lin al. 2021 Cassella et al.) 2022 Ren et al. 2022 Gerard et al.
2022
Gao
& Günnemann 2022).
A.2.2 IMPLEMENTATION OF FERMINET+SCHNET
To best reproduce the results from Gerard et al. (2022). we implemented the changes to the FermiNet which they claimed made the largest difference. All experiments in this paper, including the Fer- miNet, used dense determinants. We added the SchNet-like continuous filter convolutions, well as the nuclear embedding stream and separate streams for same-spin and different-spin electrons in the two-electron stream. We also compared their training hyperparameters against ours, except for batch size, which was kept at 4096 for all experiments. We did not implement their local input features or envelope initialization, their results suggested they did not make as significant difference, and our longer pretraining likely had the same effect as changing the envelope initialization.
A.2.3 CHOICE OF OPTIMIZER AND LAYERNORM
The original FermiNet was only able to reach high accuracy when trained with Kronecker-factored approximate curvature (KFAC) (Martens Grosse 2015). see the same holds true for the Psiformer, here we compare training with KFAC against ADAM on several systems. Additionally, when trained with ADAM, self-attention layers usually require LayerNorm (Ba et al., 2016). How- other work has suggested that in some contexts normalization may not be necessary when using KFAC (Martens et al. 2021), ), so we also compare the Psiformer with and without LayerNorm. Figure shows comparison for the Psiformer on the sulphur dimer and the iron atom. Consistent with the FermiNet, the Psiformer consistently converges faster and to lower energies with KFAC. With LayerNorm, the situation is more ambiguous. On the sulphur dimer, it seems to have marginal
15

---

Published as a conference paper at ICLR 2023
Parameter
Value
Training iterations
2e5
Training
Learning rate time
1ro(1 + to
Initial learning rate Iro
0.05
Learning rate decay to
le5
Local energy clipping p
5.0
Pretraining optimizer
LAMB
Pretraining
Pretraining iterations
2e4 or le5
Pretraining basis set
STO-6G
Markov Chain
Batch size
4096
Monte Carlo
Decorrelation steps
30
Norm constraint Damping
1e-3
KFAC
le-3
Table 4: Table of default hyperparameters used.
Systems
Pretraining Steps
MCMC Blocks
LayerNorm
Small molecules (Fig. Third-row atoms (Table
20,000
1
No
100,000
2
Yes
Large molecules (Table
100,000
2
Both
Benzene dimer
100,000
4
Both
Table 5: Variations in hyperparameters between experiments.
impact, while on the iron atom, it definitely improves training with ADAM, and seems to improve stability when using KFAC.
A.2.4 HYPERPARAMETERS
Table 4 shows the default hyperparameters used for training all models implemented in this work. Note that [Pau et al. (2020) took the sum over gradients across the batch on each device and averaged over devices, whereas here the gradients are averaged of the learning rate, meaning that the same learning rate as in Pfau et al. (2020) cannot be used.
the entire batch. This amounts scaling
As in previous FermiNet implementations, we pretrain all networks to match Hartree-Fock (HF) bitals computed using PySCF. Here, we use the LAMB optimizer (You et al. (0200) as the pretraining optimizer. We find that longer pretraining stabilises training for all models considered. In addition, during pretraining, we draw samples from the HF orbitals only, instead of the neural network wave- function. smaller number of pretraining iterations (20,000) was used for small molecules in Section while 100,000 iterations were used for third-row atoms and larger molecules.
To generate samples from 2 we use the Metropolis-Hasting algorithm with symmetric Gaussian proposals, as in Pfau et al.] (022)). We increase the number of decorrelation MCMC steps between optimization iterations from 10 in Pfau et al. (2020) 30. Additionally, for larger systems, we do not update all electron positions simultaneously one Metropolis step. Instead, split the electrons into multiple blocks, and iteratively update each block once per step. This an intermediate between the all-electron and one-electron moves commonly used in VMC.
Note that while all models were trained for 200,000 optimization iterations, this does not mean that all systems required that many iterations converge. Many smaller systems converged in far fewer iterations than this, but the same number was used for all systems to minimize confusion.
Not all systems used identical parameters for larger systems increased the number of pretraining steps and blocks per MCMC update, and LayerNorm was not used for small molecules in Table In Table we specify which systems were trained with which settings.
In Section 4.1 the performance of "small" and "large" FermiNet and Psiformer models was investi- gated set of small and medium molecules. Table Tlives the network parameters used for these
16

---

Published as a conference paper at ICLR 2023
FermiNet
Psiformer
Parameter
Small
Large
Small
Large
Determinants
16
32
16
32
Network layers
4
4
Attention heads
N/A
N/A
4
8
Attention dimension
N/A
N/A
64
64
MLP hidden dims: one-electron
256
512
256
512
MLP hidden dims: two-electron
32
32
N/A
N/A
Table 6: Table of network parameters used for the Psiformer and FermiNet.
System
FermiNet
Psiformer
Pfau et 22220)
Small
Large
Small
Large
LiH
Energy
-8.0705(T)
8.07755XX)
-8.0705 (5705
-8.07052815)
8,07053664)
# parameters
668128
683744
2610400
1595858
6366210
L12
Energy
14.99475(1)
14.99480K2)
-14.99484:2)
-14,99466(1)
-14 14.99485(2)
paramelers
2676448
1602178
6399234
NH3
Energy
(6.56295(8)
36.5634744)
56.56388(3)
56.5636712)
36.3612)
# paramelers
703968
741088
2823399
1621506
6470658
CH4
Energy
40.51400(7)
-0.51443033
30.51455033
40.5145412)
40.51461(()
# parameters
708640
744800
2830816 16
1622850
6473346
CO
Energy
11333218(1)
-13.32354(7)
113.32444(5)
113.32416(4)
113.32466(3)
# parameters
712288
766944
2940640
1635458
6531330
N2
Energy
(89.538((1)
109 (09.54046/6)
109.54
(09.5414866)
109.
(09.54137(4)
109 (017941))(4)
parameters
712288
766944
2940640
6531330
C2H4
Energy
18.5844(1)
-78 58604(5)
78.5877213)
-78 5879413)
# parameters
743648
799968
3039456
1649922
6376642
methylamine
Energy
95,855((2)
95
-5.88 860000555
8605004)
86096(3)
# parameters
759712
821344
3114976
1660098
6613378
03
Energy
225,4145(3)
222.4226(2)
225.4268(1)
225 (323118))
parameters
763360
854752
3280096
1678850
6700034
othanol
Energy
135.0308(3)
155.0419(1)
11550045(()
155.04650()
155.04759166
# paramelers
815904
899936
3403232
1698370
6755458
bicyclobutane
Energy
155.9263(6)
(35.9388(()
-155
155,9461918)
-155 948367)
# parametets
845920
940000
3548896
1717890
6827266
Table 7: Energies and number of parameters in network for set of small molecules studied in Pfau al. (2020.. Note Přau et al. (2020) used block diagonal determinants and did not report precise parameter counts; parameters counts for their results are estimated using their published settings and the FermiNet JAX implementation (Spencer et al. 2020b). Our FermiNet experiments used the en- velope function proposed in /Spencer et al. (2020a) and dense determinants, resulting in slightly dif- ferent numbers of parameters between the networks of Pfau et al. (2020) and our FermiNet (Small) networks.
models. In all other experiments, unless otherwise specified, the 'small' network configurations were used. Table Table contains the number of parameters used by each model for this set molecules.
A.3 COMPUTATIONAL DETAILS
All models were implemented in JAX (Bradbury et al. 2018) based upon the public FermiNet (Spencer et al. 2020b) and KFAC implementations (Botev & Martens 2022), and trained in par- allel using between 16 and 64 A 1000 GPUs, depending on the system size. All calculations were done standard single precision, as we found that TensorFloat-33 calculations on A100 had numer- ical accuracy issues. A table of total training time including pretraining for several models is given Table 9 Empirically, the time per iteration scaled roughly cubically, i.e. the benzene dimer required ~88 times the computational resources to run as the benzene molecule, though the number of atoms also played a significant role in timing. For systems with very large numbers of atoms, like the benzene dimer, the FermiNet and Psiformer ran nearly the same speed. This is likely because the determinants and envelope, which are similar between the FermiNet and Psiformer, make up a larger share of the total run time as the systems become larger. Geometries taken from Curtiss jet al. (2000) and given in Table 3 or in Pfau et al. (2020). The benzene dimer geometry is obtained via rigid translation and rotation of two monomers.
17

---

Published as a conference paper at ICLR 2023
System benzene
atom
position (ao)
C
0.00000
2.63664
0.00000
C
2.28339
1.31832
0.00000
C
2.28339
-1.31832
0.00000
C
0.00000
-2.63664
0.00000
C
22.28339
-1.31832
0.00000
C
-2.28339
1.31832
0.00000
H
0.00000
4.69096
0.00000
H
4.06250
2.34549
0.00000
H
4.06250
-2.45549
0.00000
H
0.00000
4.69096
0.00000
H
-4.62550
-2.34549
0.00000
H
-4,06250
2.34549
0.00000
toluene
C
-0.01831
1.72486
0.00000
C
-0.01297
0.37035
2.27060
C
-0.01297
0.37035
-2227060
C
-0.01297
-2.26452
2.27703
C
-0.01297
-2.26452
-.277003
C
-0,00900
3359166
0.00000
C
0.05583
4.56877
0.00000
H
2.00464
5.26559
0.00000
H
0088281
5.33834
-1.67217
H -0.88281
5,33834
1.67217
H -0.01841
-3.28187
4.06225
H
-0.01841
-3.28187
-4.06225
H
-0.01415
-5.64576
0.00000
H
-0.02402
1,39270
4.05592
H
-0.02402
1.39270
4.05592
naphthalene
C
0.00000
0.00000
1.35203
C
0.00000
0.00000
-1.35203
C
0.00000
2.34349
2.64930
C
0.00000
-2.34349
2.64930
C
0.00000
2.34349
-2.64930
C
0.00000
-2.34349
-2.64930
C
0.00000
-4.59147
1.33509
C
0.00000
4.59147
1.33509
H
0.00000
2.34107
4.70689
H
0.00000
-2.34107
4.70689
C
0,00000
-4.59147
-1.33509
C
0.00000
4.59147
-1.33009
H
0.00000
2.34107
4.70689
H
0.00000
-2.4107
-4.70689
H
0.00000
-6.37654
2.35261
H 0.00000
6.37654
2,35261
H
0.00000
-6.37654
-.35261
0.00000 H
6.37654
-2.35261
CCI,
C
0.00000
0.00000
0.00000
Cl
1.93005
1.93005
1.93005
Cl
-1,93005
-1.93005
1.93005
Cl
-1.93005
1.93005
-1.93005
Cl
1.93005
-1.93005
-1.93005
Table 8: Molecular geometries in Bohr.
18

---

Published as a conference paper at ICLR 2023
FermiNet
SerniNettSchNNet
Psiformer
System
Small
Large
Small
Large
Small
Large
Bicyclobutane
484
789
712
1412
7n
592
1200
808
1365
1317
1600
Benzene
1224
1608
1768
Benzene Dimer
10576
10695
Table 9: Number of A100 GPU hours required to train several systems.
VT, no jastrow UF. no input
ut. Jastrow
-109.535
-155.02
UF, input rescalinng
Energy (a.u.)
Energy (a.u.)
-109.54
-109.545
-155.05
0
50 Iterations (thousand)
100
Iterations (thousand)
(a) N2, R=1.094 Å
(b) ethanol
Figure 6: Ablation studies for the Psiformer without a Jastrow factor or input rescaling on (a) the nitrogen dimer and (b) ethanol. A rolling mean of the last 1000 iterations is used to smooth the energy. The divergence of the energy without using input rescaling, as shown here for ethanol, is typical for medium or large systems.
A.4 ABLATION STUDIES
Two modifications to the self-attention network proved critical to stability for the Psiformer the Jastrow factor ensuring correct behaviour at electronic cusps, and rescaling the input distances. Here we show ablation studies without these modifications.
Figure the Psiformer on (a) the nitrogen dimer and (b) ethanol, without Jastrow factor and without input feature rescaling. the absence of an electron-electron Jastrow factor to enforce cusp conditions, the energy is very noisy. If the input features are not rescaled, the energy is often unsta- ble. For smaller systems, as in the nitrogen dimer shown, training may recover from the instability, but especially for medium larger systems, the energy often diverges.
A.5 ATTENTION MAPS
Figures 7 and 8 show attention maps of the Psiformer for sample electron configuration, for the benzene dimer equilibrium and dissociated geometry respectively. Each attention map shows the output of the dot product between key and query vectors before the softmax operation. Four attention maps are shown for each network layer, one per attention head. The ordering of the electrons is chosen to highlight the structure of the attention maps in relation to electron distances: the electrons are grouped according to their closest atom, where the atoms of each benzene molecule are ordered around the ring. The block structure of the attention weights is clearly visible, where electrons attend more to other electrons around the same molecule.
19

---

Published as a conference paper at ICLR 2023
Network Layer
2
3
4
Attention Head
1
4
0
Figure/ 7: Attention maps for the Psiformer for a sample electron configuration from the benzene dimer at equilibrium geometry. For each network layer, the attention maps for attention heads shown. The electrons are grouped by nearest atom.
20

---

Published as a conference paper at ICLR 2023
Network Layer
2
3
4
Attention Head
4
0
Figure 8: Attention maps for the Psiformer for a sample electron configuration from the benzene dimer at dissociated geometry. For each network layer, the attention maps for 4 attention heads shown. The electrons are grouped by nearest atom.
21


Reviews:
[{"review_id": "xX9zOEXnx1E", "reviewer": "Reviewer_GTMe", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposed to use self-attention as ansatz in FermiNet. There are some minor changes, including using log distance and spin in the input feature. The proposed model is proved better than FermiNet in several experiments including third-row atoms, large molecules and benzene dimer.", "summary_of_the_review": "This paper introduce self-attention in the ansatz in FermiNet. I guess the idea may have been studied by many researchers, and this paper is the first one to make it work, which is a good accomplishment.", "strength_and_weaknesses": "Strength: the main strength of this paper is that --- it really works. I believe there are many researchers actively working on incorporating complicated model structure in FermiNet, and many of them may have the idea to use self-attention. This paper is the first one to make this work, which is not a naive accomplishment. The runtime scaled cubically, which means for large systems, FermiNet and PsiFormer ran at nearly the same speed. Also the experiments are very detailed.\n\nWeaknesses: I don't see any major weaknesses, but I have several questions or comments below.\n\n1. Why do you only consider features of electrons? Can we include the features of atoms in the self-attention mechanism? (DeepWF https://arxiv.org/pdf/1807.07014.pdf did something like this.)\n2. More ablation study would be helpful for understanding how the number of determinants, number of attention heads, attention dimensions and etc (basically the hyperparameters listed in table 6) make contributions to the accuracy of calculation.\n3. The divergence of training energy as in Figure 6 without input rescaling seems quite outrageous. Does the authors see similar issue in FermiNet and FermiNet SchNet?\n4. The authors claim that this work is the first time that KFAC is applied to self-attention layer (in Section 4). However, https://arxiv.org/abs/1907.04164 (from KFAC authors) already claimed that this is the first work where KFAC is applied to transformer model which include self-attention.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "LjblF9qZJ0", "reviewer": "Reviewer_58NZ", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This is an interdisciplinary research paper on AI for science, it proposes a neural network framework based on self-attention mechanism, which can get an Anstazes solving the many-electron Schrodinger equation. The authors scaled the input features and concatenate the numerical spin into the input feature vector. Then they project input features to the same dimension of attention inputs by linear mapping to add self-attention mechanism into neural network framework. ", "summary_of_the_review": "AI for science is novel for our community. However, the contribution of this work is to introduce self-attention. It is an incremental contribution.", "strength_and_weaknesses": "The idea is novel for AI+Science. However, there are some incomplete or vague expressions in this paper. Please improve the language presentation carefully. The motivation for joining self-attention is not clear enough. How does the attention mechanism work on gating interactions between electrons and make the framework better? ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "4X22z0c3uQ", "reviewer": "Reviewer_x9Gw", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a self-attention based architecture called PsiFormer for approximating many-electron Schrödinger equation. The architecture uses self-attention transformer to learn the embedding of electrons and the pairwise features are incorporated through Jastrow factor. The authors did a series of experiments on multiple molecular systems and compared PsiFormer with FermiNet, the previous state-of-the-art system for this task.", "summary_of_the_review": "Overall, it is unclear to me how significant the results are and hard for me to draw a conclusion. I am not an expert in the field of quantum chemistry calculations so I am not confident in my judgment. I vote for weak rejection for now but open to suggestions during author feedback and discussion.", "strength_and_weaknesses": "The strength of this paper is that PsiFormer shows improvement over the FermiNet, especially on larger atoms with more electrons. The weakness, however, is mostly about the experiment section. There are many parts of the experiment that needs further clarification. \n1. It's quite unclear whether the results are significant. The difference between FermiNet and PsiFormer is very close (within 0.01 for smaller atoms and 0.1 for larger atoms). Also, what's the ground truth for each electron system? Is it the result coming from VMC or DMC? If so, shouldn't you report the mean absolute error (MAE) between FermiNet/PsiFormer's prediction and VMC/DMC? Should a system get a lower MAE instead?\n2. What's the level of accuracy (or precision) needed for an algorithm to be useful? Is FermiNet good enough already? (again, improving results by 0.01 seems very small but I can be wrong).\n3. Why do you only report $\\Delta E_{mono}$ and $\\Delta E_{10A}$ for Benzene Dimer but not third-row atoms?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "1cxzhBGqSg5", "reviewer": "Reviewer_RSBm", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a new attention-based neural architecture for solving many-electron Schrodinger equation by optimizing the variational quantum Monte Carlo objective. While previous works proposed architectures like FermiNet, the proposed PsiFormer differs by (1) using self-attention between $O(N)$ electron-nuclear features and (2) incorporating $O(N^{2})$ electron-electron features through (parameterized) Jastrow factor. The proposed algorithm empirically outperforms the prior works in various experiments.", "summary_of_the_review": "This paper proposes a simple idea for solving an important problem. I believe significance and complexity of the problem being solved is enough to compensate for the (arguable) lack of novelty. I am concerned about reproducibility of this paper. I do not think the follow-up works will be able to reproduce this paper without allocating a significant amount of time.", "strength_and_weaknesses": "Strength 1: This paper studies a very important problem in computational chemistry: solving the many-electron Schrodinger equation. Even a small progress in this area is likely to be impactful for future research.\n\nStrength 2: The proposed self-attention mechanism is reasonable since the self-attention mechanism has shown success for many machine learning applications. \n\nStrength 3: The experiments are thorough and it looks like the authors did a good job at solving a complex problem that requires a lot of engineering. \n\nWeakness 1: The proposed idea itself is not very novel. To my knowledge, the novelty of this paper is limited to the self-attention mechanism and the Jastrow factor-based parameterization of electron-electron features. Both of them has already been considered in prior works, like the FermiNet.\n\nWeakness 2: I hope the authors could further highlight the running time of their algorithm in the main text. I think this is importance since (1) running time is important for this type of problem and (2) the proposed self-attention mechanism is known to be quite slow. Currently, I find that the authors report their training time in Table 9 of the appendix. I guess this is the time required to solve the Schrodinger equation for a molecule. How does this time compare to the existing DNN-based or non-DNN-based solvers?\n\nWeakness 3: I would like to have seen some justification behind using the self-attention mechanism. Besides the additional expressive power, is there any reason that self-attention is more useful for parameterization of the energy function? Is there any gating mechanism in “ground truth” for solving the Schrodinger equation?\n\nWeakness 4: There is no explanation on which Markov chain Monte Carlo algorithm is used for the variational quantum Monte Carlo. \n\nWeakness 5: Since the authors do not provide a code repository for their paper, given complexity of the problem being solved, the possibility of reproducing this paper’s experiment is very low.\n\nQuestion: There are some works on graph-based Transformer to incorporate pairwise features to self-attention like Graphormer. I wonder if the authors have already considered this type of incorporation for their work, but discarded due to poor performance.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster