Original: A S ELF -A TTENTION A NSATZ FOR A B - INITIO Q UANTUM C HEMISTRY
Revised: A S ELF -A TTENTION A NSATZ FOR A B - INITIO Q UANTUM C HEMISTRY

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: 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¨odinger equation, the fundamental equation for quantum chemistry and material science.
Revised: 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¨odinger equation, the fundamental equation for quantum chemistry and material science.

Original: This equation can be solved from first principles , requiring no external training data.
Revised: This equation can be solved from first principles , requiring no external training data.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: On larger molecules especially, the ground state energy can be improved by dozens of kcal/mol, a qualitative leap over previous methods.
Revised: On larger molecules especially, the ground state energy can be improved by dozens of kcal/mol, a qualitative leap over previous methods.

Original: 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.
Revised: 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.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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.
Revised: 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.

Original: Although these laws have been known for nearly a century (Schr ¨ odinger, 1926), the fundamental equations are too difficult to solve analytically for all but the simplest systems.
Revised: Although these laws have been known for nearly a century (Schr ¨ odinger, 1926), the fundamental equations are too difficult to solve analytically for all but the simplest systems.

Original: In recent years, tools from deep learning have been used to great effect to improve the quality of computational quantum physics (Carleo & Troyer, 2017).
Revised: In recent years, tools from deep learning have been used to great effect to improve the quality of computational quantum physics (Carleo & Troyer, 2017).

Original: For the study of chemistry in particular, it is the quantum behavior of electrons that matters, which imposes certain constraints on the possible solutions.
Revised: For the study of chemistry in particular, it is the quantum behavior of electrons that matters, which imposes certain constraints on the possible solutions.

Original: 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 a variety of extensions and improvements (Hermann et al., 2022).
Revised: 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 a variety of extensions and improvements (Hermann et al., 2022).

Original: However, follow-up work has mostly focused on applications and iterative improvements to the neural network architectures introduced in the first set of papers.
Revised: However, follow-up work has mostly focused on applications and iterative improvements to the neural network architectures introduced in the first set of papers.

Original: 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.
Revised: 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.

Original: 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., 2021).
Revised: 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., 2021).

Original: The basic selfattention layer is also permutation equivariant, a useful property for applications to chemistry, where physical quantities should be invariant to the ordering of atoms and electrons (Fuchs et al., 2020).
Revised: The basic selfattention layer is also permutation equivariant, a useful property for applications to chemistry, where physical quantities should be invariant to the ordering of atoms and electrons (Fuchs et al., 2020).

Original: 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 mechanics.
Revised: 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 mechanics.

Original: In this work, we introduce a new self-attention neural network, the Wavefunction Transformer (PsiFormer), which can be used as an approximate numerical solution (or Ansatz ) for the fundamental equations of the quantum mechanics of electrons.
Revised: In this work, we introduce a new self-attention neural network, the Wavefunction Transformer (PsiFormer), which can be used as an approximate numerical solution (or Ansatz ) for the fundamental equations of the quantum mechanics of electrons.

Original: We test the PsiFormer on a wide variety of benchmark systems for quantum chemistry and find that it is significantly more accurate than existing neural network Ansatzes of roughly the same size.
Revised: We test the PsiFormer on a wide variety of benchmark systems for quantum chemistry and find that it is significantly more accurate than existing neural network Ansatzes of roughly the same size.

Original: 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 a particularly attractive approach for scaling neural network Ansatzes to larger, more challenging systems.
Revised: 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 a particularly attractive approach for scaling neural network Ansatzes to larger, more challenging systems.

Original: In what follows, we will provide an overview of the variational quantum
Revised: In what follows, we will provide an overview of the variational quantum

Original: Monte Carlo approach to computational quantum mechanics (Sec. 2), introduce the PsiFormer architecture in detail (Sec. 3), present results on a wide variety of atomic and molecular benchmarks (Sec. 4) and wrap up with a discussion of future directions (Sec 5).
Revised: Monte Carlo approach to computational quantum mechanics (Sec. 2), introduce the PsiFormer architecture in detail (Sec. 3), present results on a wide variety of atomic and molecular benchmarks (Sec. 4) and wrap up with a discussion of future directions (Sec. 5).

Original:  B ACKGROUND 2. Q UANTUM M ECHANICS AND C HEMISTRY
Revised:  B ACKGROUND 2. Q UANTUM M ECHANICS AND C HEMISTRY

Original: The fundamental object of study in quantum mechanics is the wavefunction , which represents the state of all possible classical configurations of a system.
Revised: The fundamental object of study in quantum mechanics is the wavefunction , which represents the state of all possible classical configurations of a system.

Original: If the wavefunction is known, then all other properties of a system can be calculated from it.
Revised: If the wavefunction is known, then all other properties of a system can be calculated from it.

Original: While there are multiple ways of representing a wavefunction, we focus on the first quantization approach, where the wavefunction is a map from possible particle states to a complex amplitude.
Revised: While there are multiple ways of representing a wavefunction, we focus on the first quantization approach, where the wavefunction is a map from possible particle states to a complex amplitude.

Original: The state of a single electron x ∈ R 3 × {↑ , ↓} can be represented by its position r ∈ R 3 and spin σ ∈ {↑ , ↓} .
Revised: The state of a single electron x ∈ R 3 × {↑ , ↓} can be represented by its position r ∈ R 3 and spin σ ∈ {↑ , ↓} .

Original: Then the wavefunction for an N-electron system is a function Ψ : (cid:0) R 3 × {↑ , ↓} (cid:1) N → C .
Revised: Then the wavefunction for an N-electron system is a function Ψ : (cid:0) R 3 × {↑ , ↓} (cid:1) N → C .

Original: Let x ≜ x 1 , . .
Revised: Let x ≜ x 1 , . .

Original: , x N denote the set of all electron states.
Revised: , x N denote the set of all electron states.

Original: The wavefunction is constrained to have unit ℓ 2 norm (cid:82) d x | Ψ | 2 ( x ) = 1 , and | Ψ | 2 can be interpreted as the probability of observing a quantum system in a given state when measured.
Revised: The wavefunction is constrained to have unit ℓ 2 norm (cid:82) d x | Ψ | 2 ( x ) = 1 , and | Ψ | 2 can be interpreted as the probability of observing a quantum system in a given state when measured.

Original: Not all functions are valid wavefunctions – particles must be indistinguishable, meaning | Ψ | 2 should be invariant to changes in ordering.
Revised: Not all functions are valid wavefunctions – particles must be indistinguishable, meaning | Ψ | 2 should be invariant to changes in ordering.

Original: Additionally, the Pauli exclusion principle states that the probability of observing any two electrons in the same state must be zero.
Revised: Additionally, the Pauli exclusion principle states that the probability of observing any two electrons in the same state must be zero.

Original: This is enforced by requiring the wavefunction for electronic systems to be antisymmetric .
Revised: This is enforced by requiring the wavefunction for electronic systems to be antisymmetric .

Original: In this paper, we will focus on how to learn an unnormalized approximation to Ψ by representing it with a neural network.
Revised: In this paper, we will focus on how to learn an unnormalized approximation to Ψ by representing it with a neural network.

Original: The physical behavior of non-relativistic quantum systems is described by the Schr¨odinger equation.
Revised: The physical behavior of non-relativistic quantum systems is described by the Schr¨odinger equation.

Original: In its time-independent form, it is an eigenfunction equation
Revised: In its time-independent form, it is an eigenfunction equation

Original: ˆ H Ψ
Revised: ˆ H Ψ

Original: ( x ) = E Ψ
Revised: ( x ) = E Ψ

Original: ( x ) where ˆ H is a
Revised: ( x ) where ˆ H is a

Original: Hermitian linear operator called the Hamiltonian and the scalar eigenvalue E corresponds to the energy of that particular solution.
Revised: Hermitian linear operator called the Hamiltonian and the scalar eigenvalue E corresponds to the energy of that particular solution.

Original: In quantum chemistry, atomic units (a.u.) are typically used, in which the unit of distance is the Bohr radius ( a 0 ), and the unit of energy is Hartree (Ha).
Revised: In quantum chemistry, atomic units (a.u.) are typically used, in which the unit of distance is the Bohr radius ( a 0 ), and the unit of energy is Hartree (Ha).

Original: The physical details of a system are defined through the choice of Hamiltonian.
Revised: The physical details of a system are defined through the choice of Hamiltonian.

Original: For chemical systems, the only details which need to be specified are the locations and charges of the atomic nuclei.
Revised: For chemical systems, the only details which need to be specified are the locations and charges of the atomic nuclei.

Original: In quantum chemistry it is standard to approximate the nuclei as classical particles with fixed positions, known as the Born-Oppenheimer approximation, in which case the Hamiltonian becomes:
Revised: In quantum chemistry it is standard to approximate the nuclei as classical particles with fixed positions, known as the Born-Oppenheimer approximation, in which case the Hamiltonian becomes:

Original: [Equation]
Revised: [Equation]

Original: where ∇ 2 i = (cid:80) 3 j =1 ∂ 2 ∂r 2 ij is the Laplacian w.r.t.
Revised: where ∇ 2 i = (cid:80) 3 j =1 ∂ 2 ∂r 2 ij is the Laplacian w.r.t.

Original: the i th particle and Z I and R I , I ∈ { 1 , ..., N nuc } are the charges and coordinates of the nuclei.
Revised: the i th particle and Z I and R I , I ∈ { 1 , ..., N nuc } are the charges and coordinates of the nuclei.

Original: Two simplifications follow from this.
Revised: Two simplifications follow from this.

Original: First, since ˆ H is a Hermitian operator, solutions Ψ must be real-valued.
Revised: First, since ˆ H is a Hermitian operator, solutions Ψ must be real-valued.

Original: Thus we can restrict our attention to real-valued wavefunctions.
Revised: Thus we can restrict our attention to real-valued wavefunctions.

Original: Second, since the spins σ i do not appear anywhere in Eq.
Revised: Second, since the spins σ i do not appear anywhere in Eq.

Original: 1, we can fix a certain number of electrons to be spin up and the remainder to be spin down before beginning any calculation (Foulkes et al., 2001).
Revised: 1, we can fix a certain number of electrons to be spin up and the remainder to be spin down before beginning any calculation (Foulkes et al., 2001).

Original: The appropriate number for the lowest energy state can usually be guessed by heuristics such as Hund’s rules.
Revised: The appropriate number for the lowest energy state can usually be guessed by heuristics such as Hund’s rules.

Original: While the time-independent Schr¨odinger 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 .
Revised: While the time-independent Schr¨odinger 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 .

Original: Solutions with higher energy, known as excited states , are relevant to photochemistry, but in this paper we will restrict our attention to ground states.
Revised: Solutions with higher energy, known as excited states , are relevant to photochemistry, but in this paper we will restrict our attention to ground states.

Original: For a typical small molecule, the total energy of a system is on the order of hundreds to thousands of
Revised: For a typical small molecule, the total energy of a system is on the order of hundreds to thousands of

Original: Hartrees.
Revised: Hartrees.

Original: However the relevant energy scale for chemical bonds is typically much smaller – on the order of 1 kilocalorie per mole (kcal/mol), or ∼ 1 .
Revised: However the relevant energy scale for chemical bonds is typically much smaller – on the order of 1 kilocalorie per mole (kcal/mol), or ∼ 1 .

Original: 6 mHa – less than one part in one hundred thousand of the total energy.
Revised: 6 mHa – less than one part in one hundred thousand of the total energy.

Original: Calculations within 1 kcal/mol of the ground truth are generally considered “chemically accurate”.
Revised: Calculations within 1 kcal/mol of the ground truth are generally considered “chemically accurate”.

Original: Mean-field methods are typically within about 0.5% of the true total energy.
Revised: Mean-field methods are typically within about 0.5% of the true total energy.

Original: The difference between the mean-field energy and true energy is known as the correlation energy , and chemical accuracy is usually less than 1% of this correlation energy.
Revised: The difference between the mean-field energy and true energy is known as the correlation energy , and chemical accuracy is usually less than 1% of this correlation energy.

Original: [Figure]
Revised: [Figure]

Original: 2. V ARIATIONAL Q UANTUM M ONTE C ARLO
Revised: 2. V ARIATIONAL Q UANTUM M ONTE C ARLO

Original: There are a wide variety of computational techniques to find the ground state solution of the Hamiltonian in Eq 1.
Revised: There are a wide variety of computational techniques to find the ground state solution of the Hamiltonian in Eq 1.

Original: We are particularly interested in solving these equations from first principles ( abinitio ), that is, without any data other than the atomic positions.
Revised: We are particularly interested in solving these equations from first principles ( abinitio ), that is, without any data other than the atomic positions.

Original: The ab-initio method most compatible with the modern deep learning paradigm is variational quantum Monte Carlo (VMC, Foulkes et al.
Revised: The ab-initio method most compatible with the modern deep learning paradigm is variational quantum Monte Carlo (VMC, Foulkes et al.

Original: In VMC, a 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.
Revised: In VMC, a 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.

Original: VMC is variational in the sense that it minimizes an upper bound on the energy of a system.
Revised: VMC is variational in the sense that it minimizes an upper bound on the energy of a system.

Original: Therefore if two VMC solutions give different energies, the one with the lower energy will be closer to the true energy, even if the true energy is not known.
Revised: Therefore if two VMC solutions give different energies, the one with the lower energy will be closer to the true energy, even if the true energy is not known.

Original: In VMC, we start with an unnormalized wavefunction Ansatz Ψ θ : (cid:0) R 3 × {↑ , ↓} (cid:1) N
Revised: In VMC, we start with an unnormalized wavefunction Ansatz Ψ θ : (cid:0) R 3 × {↑ , ↓} (cid:1) N

Original: → R with parameters θ .
Revised: → R with parameters θ .

Original: The expected energy of the system is given by the Rayleigh quotient:
Revised: The expected energy of the system is given by the Rayleigh quotient:

Original: [Equation]
Revised: [Equation]

Original: where we have rewritten the Rayleigh quotient as an expectation over a random variable proportional to Ψ 2 θ on the right hand side.
Revised: where we have rewritten the Rayleigh quotient as an expectation over a random variable proportional to Ψ 2 θ on the right hand side.

Original: The term E L ( x ) =
Revised: The term E L ( x ) =

Original: Ψ − 1 ( x ) ˆ H Ψ
Revised: Ψ − 1 ( x ) ˆ H Ψ

Original: ( x ) is known as the local energy.
Revised: ( x ) is known as the local energy.

Original: Details on how to compute the local energy and unbiased estimates of the gradient of the average energy are given in Sec.
Revised: Details on how to compute the local energy and unbiased estimates of the gradient of the average energy are given in Sec.

Original: A.1 in the appendix.
Revised: A.1 in the appendix.

Original: 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).
Revised: 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).

Original: Notably, the samples x ∼ Ψ 2 θ can be generated from the Ansatz itself , rather than requiring external data.
Revised: Notably, the samples x ∼ Ψ 2 θ can be generated from the Ansatz itself , rather than requiring external data.

Original: The form of Ψ must be restricted to antisymmetric functions to avoid collapsing onto non-physical solutions.
Revised: The form of Ψ must be restricted to antisymmetric functions to avoid collapsing onto non-physical solutions.

Original: This is most commonly done by taking the determinant of a matrix of single-electron functions Ψ( x ) = det [ Φ ( x )] , where Φ ( x ) denotes the matrix with elements ϕ
Revised: This is most commonly done by taking the determinant of a matrix of single-electron functions Ψ( x ) = det [ Φ ( x )] , where Φ ( x ) denotes the matrix with elements ϕ

Original: i ( x j ) , since the determinant is antisymmetric under exchange of rows or columns.
Revised: i ( x j ) , since the determinant is antisymmetric under exchange of rows or columns.

Original: This is known as a Slater determinant , and the minimum-energy wavefunction of this form gives the mean-field solution to the Schr¨odinger equation.
Revised: This is known as a Slater determinant , and the minimum-energy wavefunction of this form gives the mean-field solution to the Schr¨odinger equation.

Original: While ϕ i is a function of one electron in a Slater determinant, any permutation-equivariant function of all electrons can be used as input to Φ and Ψ will still be antisymmetric.
Revised: While ϕ i is a function of one electron in a Slater determinant, any permutation-equivariant function of all electrons can be used as input to Φ and Ψ will still be antisymmetric.

Original: 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 (Kato, 1957).
Revised: 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 (Kato, 1957).

Original: The cusp conditions state that the wavefunction must be non-differentiable at these points, and give exact values for the average derivatives at the cusps.
Revised: The cusp conditions state that the wavefunction must be non-differentiable at these points, and give exact values for the average derivatives at the cusps.

Original: This can be built into an Ansatz by multiplying by a
Revised: This can be built into an Ansatz by multiplying by a

Original: Jastrow factor which satisfies these conditions analytically (Drummond et al., 2004).
Revised: Jastrow factor which satisfies these conditions analytically (Drummond et al., 2004).

Original: 2. R ELATED WORK
Revised: 2. R ELATED WORK

Original: Machine learning has found numerous applications to computational chemistry in recent years, but has mostly focused on problems at the level of classical physics (Sch¨utt et al., 2018;
Revised: Machine learning has found numerous applications to computational chemistry in recent years, but has mostly focused on problems at the level of classical physics (Sch¨utt et al., 2018;

Original: Fuchs et al., 2020; Batzner et al., 2022; Segler et al., 2018; G´omez-Bombarelli et al., 2018), which all rely on learning from large datasets of experiments or ab-initio calculations.
Revised: Fuchs et al., 2020; Batzner et al., 2022; Segler et al., 2018; G´omez-Bombarelli et al., 2018), which all rely on learning from large datasets of experiments or ab-initio calculations.

Original: There is also work on machine learning for density functional theory (DFT) (Nagai et al., 2020; Kirkpatrick et al., 2021), which is 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.
Revised: There is also work on machine learning for density functional theory (DFT) (Nagai et al., 2020; Kirkpatrick et al., 2021), which is 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.

Original: Here instead, we are focused on improving the ab-initio methods themselves.
Revised: Here instead, we are focused on improving the ab-initio methods themselves.

Original: For a thorough introduction to ab-initio quantum chemistry, we recommend Helgaker et al.
Revised: For a thorough introduction to ab-initio quantum chemistry, we recommend Helgaker et al.

Original: and Szabo & Ostlund (2012).
Revised: and Szabo & Ostlund (2012).

Original: Within this field, VMC was considered a 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 diffusion Monte Carlo (DMC) calculations, which are more accurate, but do not produce an explicit functional form for the wavefunction (Foulkes et al., 2001).
Revised: Within this field, VMC was considered a 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 diffusion Monte Carlo (DMC) calculations, which are more accurate, but do not produce an explicit functional form for the wavefunction (Foulkes et al., 2001).

Original: These
Revised: These

Original: VMC calculations typically used a Slater-Jastrow Ansatz (Kwon et al., 1993), which consists of a large linear combination of Slater determinants multiplied by a Jastrow factor.
Revised: VMC calculations typically used a Slater-Jastrow Ansatz (Kwon et al., 1993), which consists of a large linear combination of Slater determinants multiplied by a Jastrow factor.

Original: They also sometimes include backflow , a coordinate transformation that accounts for electron correlations with a fixed functional form (Feynman & Cohen, 1956).
Revised: They also sometimes include backflow , a coordinate transformation that accounts for electron correlations with a fixed functional form (Feynman & Cohen, 1956).

Original: 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 al., 2020;
Revised: 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 al., 2020;

Original: Han et al., 2019; Luo & Clark, 2019; Taddei et al., 2015).
Revised: Han et al., 2019; Luo & Clark, 2019; Taddei et al., 2015).

Original: In first quantization, these Ansatzes used a sum of a small number of determinants to construct antisymmetric functions, but used very general permutation-equivariant functions of all electrons as inputs, rather than the single-electron functions used in Slater determinants.
Revised: In first quantization, these Ansatzes used a sum of a small number of determinants to construct antisymmetric functions, but used very general permutation-equivariant functions of all electrons as inputs, rather than the single-electron functions used in Slater determinants.

Original: Some, like the PauliNet (Hermann et al., 2020), used Jastrow factors, while the FermiNet (Pfau et al., 2020) used non-differentiable input features to learn the cusp conditions.
Revised: Some, like the PauliNet (Hermann et al., 2020), used Jastrow factors, while the FermiNet (Pfau et al., 2020) used non-differentiable input features to learn the cusp conditions.

Original: Most follow-up work, as surveyed in Hermann et al.
Revised: Most follow-up work, as surveyed in Hermann et al.

Original: (2022), integrated these Ansatzes with other methods and applications, but did not significantly alter the architecture of the neural networks.
Revised: (2022), integrated these Ansatzes with other methods and applications, but did not significantly alter the architecture of the neural networks.

Original: The most significant departure in terms of the neural network architecture was Gerard et al.
Revised: The most significant departure in terms of the neural network architecture was Gerard et al.

Original: (2022), which extended the FermiNet – generally recognized to be the most accurate neural network Ansatz up to that point – and integrated it with several details of the PauliNet, especially the continuous-filter convolutions also used by the SchNet (Sch¨utt et al., 2018), claiming to reach even higher accuracy on several challenging systems.
Revised: (2022), which extended the FermiNet – generally recognized to be the most accurate neural network Ansatz up to that point – and integrated it with several details of the PauliNet, especially the continuous-filter convolutions also used by the SchNet (Sch¨utt et al., 2018), claiming to reach even higher accuracy on several challenging systems.

Original: We refer to this architecure as the FermiNet+SchNet.
Revised: We refer to this architecure as the FermiNet+SchNet.

Original:  T HE P SI F ORMER
Revised:  T HE P SI F ORMER

Original: The PsiFormer has the basic form:
Revised: The PsiFormer has the basic form:

Original: [Equation]
Revised: [Equation]

Original: where J θ : ( R 3 × {↑ , ↓} ) N → R and Φ kθ : ( R 3 × {↑ , ↓} ) N → R N × N are functions with learnable parameters θ .
Revised: where J θ : ( R 3 × {↑ , ↓} ) N → R and Φ kθ : ( R 3 × {↑ , ↓} ) N → R N × N are functions with learnable parameters θ .

Original: This is similar to the Slater-Jastrow Ansatz, FermiNet and PauliNet, with the key difference being that in the PsiFormer, Φ kθ consists of a sequence of multiheaded self-attention layers (Vaswani et al., 2017).
Revised: This is similar to the Slater-Jastrow Ansatz, FermiNet and PauliNet, with the key difference being that in the PsiFormer, Φ kθ consists of a sequence of multiheaded self-attention layers (Vaswani et al., 2017).

Original: The high-level structure of the PsiFormer is shown in Fig. 1(b), where it is contrasted with the FermiNet (and SchNet extension) in Fig. 1(a).
Revised: The high-level structure of the PsiFormer is shown in Fig. 1(b), where it is contrasted with the FermiNet (and SchNet extension) in Fig. 1(a).

Original: L i H
Revised: L i H

Original: L i 2 N H 3 C H 4 C O
Revised: L i 2 N H 3 C H 4 C O

Original: N 2 C 2 Hm e t h y la m i n e
Revised: N 2 C 2 Hm e t h y la m i n e

Original: Oe t h a n o l b i c y c l ob u t a n e
Revised: Oe t h a n o l b i c y c l ob u t a n e

Original: [Figure]
Revised: [Figure]

Original: Because a self-attention layer takes a sequence of vectors as input, only features of single electrons are used as input to Φ kθ .
Revised: Because a self-attention layer takes a sequence of vectors as input, only features of single electrons are used as input to Φ kθ .

Original: The input feature vector
Revised: The input feature vector

Original: f 0 i for electron
Revised: f 0 i for electron

Original: i is similar to the one-electron stream of the
Revised: i is similar to the one-electron stream of the

Original: FermiNet, which uses a concatenation of electron-nuclear differences r i − R I and distances | r i − R
Revised: FermiNet, which uses a concatenation of electron-nuclear differences r i − R I and distances | r i − R

Original: I | ,
Revised: I | ,

Original: I = 1 , . .
Revised: I = 1 , . .

Original: , N nuc , with two key differences.
Revised: , N nuc , with two key differences.

Original: First, we found that for systems with widely separated atoms, using FermiNet one-electron features caused self-attention Ansatzes to become unstable, so we rescale the inputs in the PsiFormer by a factor of log(1 + | r
Revised: First, we found that for systems with widely separated atoms, using FermiNet one-electron features caused self-attention Ansatzes to become unstable, so we rescale the inputs in the PsiFormer by a factor of log(1 + | r

Original: i − R
Revised: i − R

Original: I | ) / | r
Revised: I | ) / | r

Original: i − R
Revised: i − R

Original: I | , so that the input vectors grow logarithmically with distance from the nucleus.
Revised: I | , so that the input vectors grow logarithmically with distance from the nucleus.

Original: Second, we concatenate the spin σ i into the input feature vector itself (mapping ↑ to 1 and ↓ to -1).
Revised: Second, we concatenate the spin σ i into the input feature vector itself (mapping ↑ to 1 and ↓ to -1).

Original: While this spin term is kept fixed during training, it breaks the symmetry between spin up and spin down electrons, so that columns of Φ kθ are only equivariant under exchange of same-spin electrons.
Revised: While this spin term is kept fixed during training, it breaks the symmetry between spin up and spin down electrons, so that columns of Φ kθ are only equivariant under exchange of same-spin electrons.

Original: This is a notable departure from the
Revised: This is a notable departure from the

Original: FermiNet and PauliNet, where the difference between spin-up and spin-down electrons is instead built into the architecture.
Revised: FermiNet and PauliNet, where the difference between spin-up and spin-down electrons is instead built into the architecture.

Original: The input features f 

Revised: The input features f

Original: i are next projected into the same dimension as the attention inputs by a linear mapping h 

Revised: i are next projected into the same dimension as the attention inputs by a linear mapping h

Original: i =
Revised: i =

Original: W 0 f 0 i , and then passed into a sequence of multiheaded self-attention layers followed by linear-nonlinear layers, both with residual connections:
Revised: W 0 f 0 i , and then passed into a sequence of multiheaded self-attention layers followed by linear-nonlinear layers, both with residual connections:

Original: [Equation]
Revised: [Equation]

Original: where h indexes the different attention heads, concat h denotes concatenation of the output from different attention heads, and S ELF A TTN denotes standard self-attention:
Revised: where h indexes the different attention heads, concat h denotes concatenation of the output from different attention heads, and S ELF A TTN denotes standard self-attention:

Original: [Equation]
Revised: [Equation]

Original: where d is the output dimension of the key and query weights.
Revised: where d is the output dimension of the key and query weights.

Original: In principle, multiple linear-nonlinear layers could be used in-between self-attention layers, but we found that adding a deeper MLP between self-attention layers was less effective than adding more self-attention layers.
Revised: In principle, multiple linear-nonlinear layers could be used in-between self-attention layers, but we found that adding a deeper MLP between self-attention layers was less effective than adding more self-attention layers.

Original: While a smooth nonlinearity must be used to guarantee that a wavefunction is smooth everywhere except the cusps, we found that using activation functions other than tanh had a marginal impact on performance.
Revised: While a smooth nonlinearity must be used to guarantee that a wavefunction is smooth everywhere except the cusps, we found that using activation functions other than tanh had a marginal impact on performance.

Original: A final linear projection into a NN det dimensional space is applied to the activations, and the output is multiplied by a weighted sum of exponentially-decaying envelopes, Φ kij = Ω kij w kTi h
Revised: A final linear projection into a NN det dimensional space is applied to the activations, and the output is multiplied by a weighted sum of exponentially-decaying envelopes, Φ kij = Ω kij w kTi h

Original: Lj , where Ω k ij = (cid:80) I π k iI exp (cid:0) − σ k iI | r j − R I | (cid:1) .
Revised: Lj , where Ω k ij = (cid:80) I π k iI exp (cid:0) − σ k iI | r j − R I | (cid:1) .

Original: This is so that the boundary condition lim | r |→∞ Ψ θ ( r ) = 0 is
Revised: This is so that the boundary condition lim | r |→∞ Ψ θ ( r ) = 0 is

Original: [Figure]
Revised: [Figure]

Original: enforced, and is the same form as the envelope used by the FermiNet in Spencer et al.
Revised: enforced, and is the same form as the envelope used by the FermiNet in Spencer et al.

Original: (2020a).
Revised: (2020a).

Original: The matrix elements Φ kij are then passed into k determinants and the output summed together.
Revised: The matrix elements Φ kij are then passed into k determinants and the output summed together.

Original: As the distances | r i − R I | are inputs to the PsiFormer, it is capable of learning the electron-nuclear cusp conditions, much like the FermiNet.
Revised: As the distances | r i − R I | are inputs to the PsiFormer, it is capable of learning the electron-nuclear cusp conditions, much like the FermiNet.

Original: 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.
Revised: 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.

Original: Instead, the PsiFormer uses a conventional Jastrow factor only for the electron-electron cusps.
Revised: Instead, the PsiFormer uses a conventional Jastrow factor only for the electron-electron cusps.

Original: We use a particularly simple Jastrow factor:
Revised: We use a particularly simple Jastrow factor:

Original: [Equation]
Revised: [Equation]

Original: which has only two free parameters, α par and α anti , and works well in practice.
Revised: which has only two free parameters, α par and α anti , and works well in practice.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Here we present an evaluation of the PsiFormer on a wide variety of benchmark systems.
Revised: Here we present an evaluation of the PsiFormer on a wide variety of benchmark systems.

Original: Where it is not specified, all results with the FermiNet and FermiNet+SchNet are with our own implementation, forked from Spencer et al.
Revised: Where it is not specified, all results with the FermiNet and FermiNet+SchNet are with our own implementation, forked from Spencer et al.

Original: (2020b).
Revised: (2020b).

Original: We use Kronecker-factored Approximate Curvature (KFAC)
Revised: We use Kronecker-factored Approximate Curvature (KFAC)

Original: (Martens & Grosse, 2015) to optimize the PsiFormer.
Revised: (Martens & Grosse, 2015) to optimize the PsiFormer.

Original: To the best ofour knowledge, this is the first time in the literature that KFAC  has been used with self-attention.
Revised: The use of KFAC for self-attention has been investigated in Zhang et al.

Original: We show the advantage of KFAC, and how it interacts with LayerNorm, in Sec.
Revised: We show the advantage of KFAC, and how it interacts with LayerNorm, in Sec.

Original: A.2.3 of the appendix.
Revised: A.2.3 of the appendix.

Original: We use hyperparameters  similar to the original FermiNet paper, though we have made several modifications which help for larger systems: we pretrain for longer, we generate samples for pretraining from the target 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 to be more robust to outliers.
Revised: We use hyperparameters and a Metropolis-Hastings MCMC algorithm similar to the original FermiNet paper, though we have made several modifications which help for larger systems: we pretrain for longer, we generate samples for pretraining from the target 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 to be more robust to outliers.

Original: Details are given in Sec.
Revised: Details are given in Sec.

Original: A.2.4 of the appendix.
Revised: A.2.4 of the appendix.

Original: 4. R EVISITING S MALL M OLECULES
Revised: 4. R EVISITING S MALL M OLECULES

Original: In Pfau et al.
Revised: In Pfau et al.

Original: (2020), they compared the FermiNet against CCSD(T) extrapolated to the complete basis set (CBS) limit on a number of small molecules (4-30 electrons) from the G3 database (Curtiss et al., 2000).
Revised: (2020), they compared the FermiNet against CCSD(T) extrapolated to the complete basis set (CBS) limit on a number of small molecules (4-30 electrons) from the G3 database (Curtiss et al., 2000).

Original: While the FermiNet captured more than 99% of the correlation energy relative to
Revised: While the FermiNet captured more than 99% of the correlation energy relative to

Original: CCSD(T)/CBS for systems as large as ethene (16 electrons), the quality of the FermiNet calculation began to decline as the system size grew.
Revised: CCSD(T)/CBS for systems as large as ethene (16 electrons), the quality of the FermiNet calculation began to decline as the system size grew.

Original: While some of this discrepancy was reduced simply by changing to a framework with better numerics (Spencer et al., 2020a), the reported difference in energy for bicyclobutane was still greater than 20 mHa.
Revised: While some of this discrepancy was reduced simply by changing to a framework with better numerics (Spencer et al., 2020a), the reported difference in energy for bicyclobutane was still greater than 20 mHa.

Original: Here we revisit these systems, comparing the FermiNet with training improvements against the PsiFormer.
Revised: Here we revisit these systems, comparing the FermiNet with training improvements against the PsiFormer.

Original: [Table]
Revised: [Table]

Original: To investigate whether the PsiFormer performance could be reproduced by simply making the FermiNet larger, we investigated both a “small” and “large” configuration for the PsiFormer and FermiNet.
Revised: To investigate whether the PsiFormer performance could be reproduced by simply making the FermiNet larger, we investigated both a “small” and “large” configuration for the PsiFormer and FermiNet.

Original: The small FermiNet has the same layer dimensions and determinants as the one used in Pfau et al.
Revised: The small FermiNet has the same layer dimensions and determinants as the one used in Pfau et al.

Original: (2020), while the large configuration has twice as many determinants and a one-electron stream twice as wide.
Revised: (2020), while the large configuration has twice as many determinants and a one-electron stream twice as wide, similar to the largest networks in Spencer et al. (2020a).

Original: The PsiFormer configurations have the same number of determinants and the MLP layers are the same width as the FermiNet one-electron stream, though due to the selfattention weights the small PsiFormer has a number of parameters between that of the large and small FermiNet.
Revised: The PsiFormer configurations 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.

Original: Exact details are given in  the appendix.
Revised: Exact details are given in Table 7 in the appendix.

Revised: On these systems, the small PsiFormer ran in a similar amount of time to the large FermiNet, 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).

Original: The results on small molecules can be seen in Fig.
Revised: The results on small molecules can be seen in Fig.

Original: 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.
Revised: 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.

Original: This is true for all systems investigated.
Revised: This is true for all systems investigated.

Original: The improvement from the PsiFormer is particularly dramatic on ozone and bicyclobutane – on ozone, the large PsiFormer is within 1 kcal/mol of CCSD(T)/CBS, while even the largest FermiNet has an error more than 4 times larger than this.
Revised: The improvement from the PsiFormer is particularly dramatic on ozone and bicyclobutane – on ozone, the large PsiFormer is within 1 kcal/mol of CCSD(T)/CBS, while even the largest FermiNet has an error more than 4 times larger than this.

Original: On all molecules, the large
Revised: On all molecules, the large

Original: PsiFormer captures more than 99% of the correlation energy relative to the reference energy.
Revised: PsiFormer captures more than 99% of the correlation energy relative to the reference energy.

Original: 4. C OMPARISON A GAINST F ERMI N ET WITH S CH N ET - LIKE C ONVOLUTIONS
Revised: 4. C OMPARISON A GAINST F ERMI N ET WITH S CH N ET - LIKE C ONVOLUTIONS

Original: 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).
Revised: 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).

Original: The primary innovations of this work were the FermiNet+SchNet architecture and new hyperparameters they claim led to much faster optimization.
Revised: The primary innovations of this work were the FermiNet+SchNet architecture and new hyperparameters they claim led to much faster optimization.

Original: They showed an especially large improvement on heavy atoms like potassium (K) and iron (Fe), and some improvement on larger molecules like benzene.
Revised: They showed an especially large improvement on heavy atoms like potassium (K) and iron (Fe), and some improvement on larger molecules like benzene.

Original: We attempted to reproduce the results of Gerard et al.
Revised: We attempted to reproduce the results of Gerard et al.

Original: (2022) with our own FermiNet+SchNet implementation, with somewhat surprising results in Fig.
Revised: (2022) with our own FermiNet+SchNet implementation, with somewhat surprising results in Fig.

Original: First, the changes in training used for the
Revised: First, the changes in training used for the

Original: FermiNet seem to be enough to close the gap with the published FermiNet+SchNet results on heavy atoms.
Revised: FermiNet seem to be enough to close the gap with the published FermiNet+SchNet results on heavy atoms.

Original: For instance, ablation studies suggested that modifying the hyperparameters plus adding
Revised: For instance, ablation studies suggested that modifying the hyperparameters plus adding

Original: SchNet-like convolutions accounts for a 38.9 mHa improvement in accuracy on the potassium atom at 10 5 iterations, but in our experiments the FermiNet with default hyperparameters is within a few mHa of the published result.
Revised: SchNet-like convolutions accounts for a 38.9 mHa improvement in accuracy on the potassium atom at 10 5 iterations, but in our experiments the FermiNet with default hyperparameters is within a few mHa of the published result.

Original: On benzene, both the SchNet-like convolutions and modified hyperparameters improved the final energy by a few mHa, though still fell slightly short of the published results.
Revised: On benzene, both the SchNet-like convolutions and modified hyperparameters improved the final energy by a few mHa, though still fell slightly short of the published results.

Original: Most importantly, the PsiFormer is clearly either comparable to the best published FermiNet+SchNet results, as on K, or better by a wide margin, as on benzene.
Revised: Most importantly, the PsiFormer is clearly either comparable to the best published FermiNet+SchNet results, as on K, or better by a wide margin, as on benzene.

Original: These results also give us confidence that the FermiNet is as strong a baseline as any other published method for further comparison.
Revised: These results also give us confidence that the FermiNet is as strong a baseline as any other published method for further comparison.

Original: 4. T HIRD -R OW A TOMS
Revised: 4. T HIRD -R OW A TOMS

Original: The FermiNet, FermiNet+SchNet 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.
Revised: The FermiNet, FermiNet+SchNet 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.

Original: Prior work  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.
Revised: Prior work [Table] 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.

Original: Instead, comparison to experimental results is a more practical evaluation metric.
Revised: Instead, comparison to experimental results is a more practical evaluation metric.

Original: The ionization[Table] potential – the amount of energy it takes to remove one electron – is a particularly simple comparison for which good experimental data exists (Koga et al., 1997).
Revised: The ionization potential – the amount of energy it takes to remove one electron – is a particularly simple comparison for which good experimental data exists (Koga et al., 1997).

Original: Here we compare the FermiNet,
Revised: Here we compare the FermiNet,

Original: FermiNet+SchNet and PsiFormer for estimating the ionization potential of potassium, iron and zinc.
Revised: FermiNet+SchNet and PsiFormer for estimating the ionization potential of potassium, iron and zinc.

Original: To compare the relative importance of size and architecture, we also looked at both a small and large configuration of all Ansatzes, with parameters described in  the appendix.
Revised: To compare the relative importance of size and architecture, we also looked at both a small and large configuration of all Ansatzes, with parameters described in Table 6 the appendix.

Original: Results are shown in Table 1.
Revised: Results are shown in Table 1.

Original: 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 results.
Revised: 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 results.

Original: 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.
Revised: 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.

Original: On heavier atoms, the difference between architectures became more pronounced, and the improvement of the
Revised: On heavier atoms, the difference between architectures became more pronounced, and the improvement of the

Original: PsiFormer relative to other models on absolute energy was more robust.
Revised: PsiFormer relative to other models on absolute energy was more robust.

Original: The results on ionization potentials were more mixed, and no Ansatz came within chemical accuracy of the ground truth.
Revised: The results on ionization potentials were more mixed, and no Ansatz came within chemical accuracy of the ground truth.

Original: This shows that even the PsiFormer is not yet converged to the ground truth, though it is the Ansatz closest to reaching it so far.
Revised: This shows that even the PsiFormer is not yet converged to the ground truth, though it is the Ansatz closest to reaching it so far.

Original: 4. L ARGER M OLECULES
Revised: 4. L ARGER M OLECULES

Original: 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.
Revised: 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.

Original: While
Revised: While

Original: CCSD(T) scales with number of electrons as O ( N 7 ) , a single iteration of wavefunction optimization for a fixed neural network size scales as O ( N 4 ) in theory, and in practice is closer to cubic for system sizes of several dozen electrons.
Revised: CCSD(T) scales with number of electrons as O ( N 7 ) , a single iteration of wavefunction optimization for a fixed neural network size scales as O ( N 4 ) in theory, and in practice is closer to cubic for system sizes of several dozen electrons.

Original: 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.
Revised: 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.

Original: Results on systems from benzene (42 electrons) to carbon tetrachloride (CCl 4 , 74 electrons) are given in Table 2.
Revised: Results on systems from benzene (42 electrons) to carbon tetrachloride (CCl 4 , 74 electrons) are given in Table 2.

Original: On these systems, CCSD(T) becomes impractical for us to run without approximations, so we only compare against other QMC results.
Revised: On these systems, CCSD(T) becomes impractical for us to run without approximations, so we only compare against other QMC results.

Original: Due to computational constraints, we were only able to run the small network configurations on these systems.
Revised: Due to computational constraints, we were only able to run the small network configurations on these systems.

Original: Additionally, to help MCMC convergence, we updated half the electron positions at a time in each MCMC move, rather than moving all electrons simultaneously, as all-electron moves become less efficient for larger systems.
Revised: Additionally, to help MCMC convergence, we updated half the electron positions at a time in each MCMC move, rather than moving all electrons simultaneously, as all-electron moves become less efficient for larger systems.

Original: On the three hydrocarbon systems investigated (benzene, toluene and naphthalene), we did not use

Original: LayerNorm, as it seemed to have a limited impact on benzene.

Original: However for CCl 4 we found that

Original: LayerNorm has a signficant impact and the results here make use of it.

Original: In Table 2, it can be seen that the PsiFormer is not only significantly better on benzene than the 
Revised: In Table 2, it can be seen that the PsiFormer is not only significantly better on benzene than the FermiNet, as in Fig.

Original: FermiNet, as in Fig. 3, but it is better than the best previously published DMC energy, a remarkable feat.
Revised: 3, but it is better than the best previously published DMC energy, a remarkable feat.

Original: 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.
Revised: 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.

Original: The PsiFormer continues to outperform the FermiNet by an ever larger margin on larger systems, reaching 120 mHa (75 kcal/mol) on CCl 4 .
Revised: The PsiFormer outperforms the FermiNet by an ever larger margin on larger systems, reaching 120 mHa (75 kcal/mol) on CCl 4 .

Revised: On the three hydrocarbon systems investigated, LayerNorm has only a small impact.

Revised: However for CCl 4 LayerNorm has a signficant impact, accounting for 70 mHa of the total 120 mHa improvement over the FermiNet.

Original: While we 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.
Revised: While we 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.

Original: [Table]
Revised: [Table]

Original: 4. T HE B ENZENE D IMER
Revised: 4. T HE B ENZENE D IMER

Original: Finally
Revised: Finally

Original: [Table]
Revised: [Table]

Original: we look at the benzene dimer, a challenging benchmark system for computational chemistry due to the weak van der
Revised: we look at the benzene dimer, a challenging benchmark system for computational chemistry due to the weak van der

Original: Waals force between the two molecules
Revised: Waals force between the two molecules

Original: [Table]
Revised: [Table]

Original: and the largest molecular system ever investigated using neural network Ansatzes (Ren et al., 2022).
Revised: and the largest molecular system ever investigated using neural network Ansatzes (Ren et al., 2022).

Original: The dimer has several possible equilibrium configurations, many of which have nearly equal energy (Sorella et al., 2007; Azadi &
Revised: The dimer has several possible equilibrium configurations, many of which have nearly equal energy (Sorella et al., 2007; Azadi &

Original: Cohen, 2015), making it additionally challenging to study computationally, but here we restrict our attention to the T-shaped structure (Fig. 4) so that we can directly compare against Ren et al.
Revised: Cohen, 2015), making it additionally challenging to study computationally, but here we restrict our attention to the T-shaped structure (Fig. 4) so that we can directly compare against Ren et al.

Original: Results are given in Table 3.
Revised: Results are given in Table 3.

Original: The results from Ren et al.
Revised: The results from Ren et al.

Original: (2022) are from a small FermiNet ( layers) trained for 2 million iterations.
Revised: (2022) are from a small FermiNet ( layers) trained for 2 million iterations.

Original: 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
Revised: 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

Original: FermiNet trained for 800,000 iterations (Ren et al.
Revised: FermiNet trained for 800,000 iterations (Ren et al.

Original: (2022) Supplementary Figure 7a).
Revised: (2022) Supplementary Figure 7a).

Original: The PsiFormer reaches a significantly lower energy than the FermiNet with the same number of training iterations – within just a few mHa of the DMC result, much like the benzene monomer.
Revised: The PsiFormer reaches a significantly lower energy than the FermiNet with the same number of training iterations, again surpassing the DMC result by a few mHa.

Original: [Figure]
Revised: 4.95 Å

Original: We also tried to estimate the dissociation energy by similar means as Ren et al.
Revised: We also tried to estimate the dissociation energy by similar means as Ren et al.

Original: (2022) – comparing against the energy of the same model trained with a bond length of 10 ˚A, and twice the energy of the same model trained on the monomer.
Revised: (2022) – comparing against the energy of the same model trained with a bond length of 10 ˚A, and twice the energy of the same model trained on the monomer.

Original: Ironically, our FermiNet baseline, which had the worst absolute energy, had the best relative energy between configurations.
Revised: Ironically, our FermiNet baseline, which had the worst absolute energy, had the best relative energy between configurations.

Original: While every model underestimated the energy relative to twice the monomer, the discrepancy was lowest with DMC.
Revised: While every model underestimated the energy relative to twice the monomer, the discrepancy was lowest with DMC.

Original: It should be noted that the zero-point vibrational energy (ZPE) is not included, the dissociation energies from Ren et al.
Revised: It should be noted that the zero-point vibrational energy (ZPE) is not included, the dissociation energies from Ren et al.

Original: (2022) in Table 3 are estimates based on their figures, and there is disagreement over the exact experimental energy (Grover et al., 1987;
Revised: (2022) in Table 3 are estimates based on their figures, and there is disagreement over the exact experimental energy (Grover et al., 1987;

Original: Krause et al., 1991), so this comparison should be considered a rough estimate only.
Revised: Krause et al., 1991), so this comparison should be considered a rough estimate only.

Original: 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 predictingmolecular binding energies.
Revised: 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.

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: We have shown that self-attention networks are capable of learning quantum mechanical properties of electrons far more effectively than comparable methods.
Revised: We have shown that self-attention networks are capable of learning quantum mechanical properties of electrons far more effectively than comparable methods.

Original: The advantage of self-attention networks seems to become most pronounced on large systems, suggesting that these models should be the focus of future efforts to scale to even larger systems.
Revised: The advantage of self-attention networks seems to become most pronounced on large systems, suggesting that these models should be the focus of future efforts to scale to even larger systems.

Original: In addition to the strong empirical results, using standard attention layers means that we can leverage existing work on improving scalability, either by architectural advances (Child et al., 2019;
Revised: In addition to the strong empirical results, using standard attention layers means that we can leverage existing work on improving scalability, either by architectural advances (Child et al., 2019;

Original: 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.
Revised: 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.

Original: This presents a promising path towards studying the most challenging molecules and materials in silico with unprecedented accuracy.
Revised: This presents a promising path towards studying the most challenging molecules and materials in silico with unprecedented accuracy.

Original: E THICS S TATEMENT
Revised: E THICS S TATEMENT

Original: 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.
Revised: 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.

Original: However, in the future, this line of work could lead to 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.
Revised: However, in the future, this line of work could lead to 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.

Original: This could lead to the discovery of new beneficial drugs or industrial chemical processes which are more environmentally friendly.
Revised: This could lead to the discovery of new beneficial drugs or industrial chemical processes which are more environmentally friendly.

Original: 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
Revised: 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

Original: R EPRODUCIBILITY S TATEMENT
Revised: R EPRODUCIBILITY S TATEMENT

Original: Details of training, network parameters, and optimization hyperparameters are given in the appendices for all experiments shown.
Revised: Details of training, network parameters, and optimization hyperparameters are given in the appendices for all experiments shown.

Original: The code will be open sourced under an OSI-approved license by the camera-ready stage.
Revised: The code will be open sourced under an OSI-approved license by the camera-ready stage.

Original: A A PPENDIX A. C ALCULATING E NERGIES AND G RADIENTS
Revised: A A PPENDIX A. C ALCULATING E NERGIES AND G RADIENTS

Original: It is generally more numerically stable to work directly with the log wavefunction, and the local energy can be expressed as
Revised: It is generally more numerically stable to work directly with the log wavefunction, and the local energy can be expressed as

Original: [Equation]
Revised: [Equation]

Original: where V ( x ) is the potential energy (the last three terms of Eq 1).
Revised: where V ( x ) is the potential energy (the last three terms of Eq 1).

Original: The gradient of the energy is given by
Revised: The gradient of the energy is given by

Original: [Equation]
Revised: [Equation]

Original: where the local energy E L ( x )
Revised: where the local energy E L ( x )

Original: =
Revised: =

Original: Ψ − 1 ( x ) ˆ H Ψ
Revised: Ψ − 1 ( x ) ˆ H Ψ

Original: ( x ) .
Revised: ( x ) .

Original: Note that the E L ( x ) − E x ′ ∼ Ψ 

Revised: Note that the E L ( x ) − E x ′ ∼ Ψ

Original: [ E L ( x ′ )] term is the difference between the local energy at x and the average energy over all x ′ .
Revised: [ E L ( x ′ )] term is the difference between the local energy at x and the average energy over all x ′ .

Original: We make a small but critical change to how the gradients are computed relative to Pfau et al.
Revised: We make a small but critical change to how the gradients are computed relative to Pfau et al.

Original: (2020) which stabilizes training for all models considered here.
Revised: (2020) which stabilizes training for all models considered here.

Original: 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).
Revised: 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).

Original: To mitigate this, the local energy is often truncated in practice.
Revised: To mitigate this, the local energy is often truncated in practice.

Original: Let ⟨ E L ⟩ mean denote the mean local energy for one minibatch of walkers and ⟨ E L ⟩ median denote the median.
Revised: Let ⟨ E L ⟩ mean denote the mean local energy for one minibatch of walkers and ⟨ E L ⟩ median denote the median.

Original: Then in Pfau et al.
Revised: Then in Pfau et al.

Original: (2020), the local energies were clipped to be within a constant multiple ρ of the mean absolute deviation around the mean: MAD mean ( E L )
Revised: (2020), the local energies were clipped to be within a constant multiple ρ of the mean absolute deviation around the mean: MAD mean ( E L )

Original: = 1 N (cid:80)
Revised: = 1 N (cid:80)

Original: [Equation]
Revised: [Equation]

Original: mean | .
Revised: mean | .

Original: While this is more robust than the standard deviation, it is still susceptible to large outliers because it is centered at the mean.
Revised: While this is more robust than the standard deviation, it is still susceptible to large outliers because it is centered at the mean.

Original: Instead, we use the mean absolute deviation around the median to determine the window for clipping:
Revised: Instead, we use the mean absolute deviation around the median to determine the window for clipping:

Original: MAD median
Revised: MAD median

Original: [Equation]
Revised: [Equation]

Original: i
Revised: i

Original: [Equation]
Revised: [Equation]

Original: | .
Revised: | .

Original: Secondly, in Pfau et al.
Revised: Secondly, in Pfau et al.

Original: (2020), the average energy term E x ′ ∼ Ψ 2 [ E L ( x ′ )] in the gradient was approximated by the mean local energy ⟨ E L ⟩ mean of a minibatch.
Revised: (2020), the average energy term E x ′ ∼ Ψ 2 [ E L ( x ′ )] in the gradient was approximated by the mean local energy ⟨ E L ⟩ mean of a minibatch.

Original: That meant that outliers were still included in this term.
Revised: That meant that outliers were still included in this term.

Original: We instead use the mean of the clipped local energies.
Revised: We instead use the mean of the clipped local energies.

Original: This guarantees that the mean over a minibatch of the energy difference term is always zero, improving stability during optimization.
Revised: This guarantees that the mean over a minibatch of the energy difference term is always zero, improving stability during optimization.

Original: If we let clip mean / median ( x ; σ ) denote the functions that clip x to be in the range [ ⟨ x ⟩ mean /
Revised: If we let clip mean / median ( x ; σ ) denote the functions that clip x to be in the range [ ⟨ x ⟩ mean /

Original: [Equation]
Revised: [Equation]

Original: − σ, ⟨ x ⟩ mean / median + σ ] , then the gradient for one batch in Pfau et al.
Revised: − σ, ⟨ x ⟩ mean / median + σ ] , then the gradient for one batch in Pfau et al.

Original: (2020) was:
Revised: (2020) was:

Original: [Equation]
Revised: [Equation]

Original: (13) while here we use:
Revised: (13) while here we use:

Original: [Equation]
Revised: [Equation]

Original: While this may seem like a subtle difference, it has a great effect on larger systems, especially heavier atoms.
Revised: While this may seem like a subtle difference, it has a great effect on larger systems, especially heavier atoms.

Original: 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.
Revised: 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.

Original: We found that clipping with proper centering was more effective than removing the walkers entirely.
Revised: We found that clipping with proper centering was more effective than removing the walkers entirely.

Original: A. T RAINING
Revised: A. T RAINING

Original: In this section we give details on training and further differences from the original FermiNet (Pfau et al., 2020).
Revised: In this section we give details on training and further differences from the original FermiNet (Pfau et al., 2020).

Original: A.2. D ENSE DETERMINANTS
Revised: A.2. D ENSE DETERMINANTS

Original: The original FermiNet and PauliNet Ansatzes used block-diagonal determinants for spin-up and spin-down electrons (Pfau et al., 2020; Hermann et al., 2020), such that each determinant could be
Revised: The original FermiNet and PauliNet Ansatzes used block-diagonal determinants for spin-up and spin-down electrons (Pfau et al., 2020; Hermann et al., 2020), such that each determinant could be

Original: [Figure]
Revised: [Figure]

Original: factorized into the product of a determinant of spin-up electrons of size N ↑ × N ↑ and a determinant of spin-down electrons N ↓ × N ↓ , as is the case with conventional VMC ansatzes (Foulkes et al., 2001):
Revised: factorized into the product of a determinant of spin-up electrons of size N ↑ × N ↑ and a determinant of spin-down electrons N ↓ × N ↓ , as is the case with conventional VMC ansatzes (Foulkes et al., 2001):

Original: [Equation]
Revised: [Equation]

Original: (16) where x σ denotes the set of electron states with the specified spin and different functions are used for electrons of different spins.
Revised: (16) where x σ denotes the set of electron states with the specified spin and different functions are used for electrons of different spins.

Original: The FermiNet authors subsequently proposed simply using dense determinants of size N × N , N = N ↑ + N ↓ (Spencer et al., 2020b):
Revised: The FermiNet authors subsequently proposed simply using dense determinants of size N × N , N = N ↑ + N ↓ (Spencer et al., 2020b):

Original: [Equation]
Revised: [Equation]

Original: 17) where now each block is of dimension N × N σ .
Revised: 17) where now each block is of dimension N × N σ .

Original: This has largely become the default choice for
Revised: This has largely become the default choice for

Original: FermiNet architectures and has shown to provide improved accuracy at small additional cost for a myriad of systems (Lin et al., 2021; Cassella et al., 2022; Ren et al., 2022; Gerard et al., 2022; Gao & G ¨ unnemann, 2022).
Revised: FermiNet architectures and has shown to provide improved accuracy at small additional cost for a myriad of systems (Lin et al., 2021; Cassella et al., 2022; Ren et al., 2022; Gerard et al., 2022; Gao & G ¨ unnemann, 2022).

Original: A.2. I MPLEMENTATION OF F ERMI N ET +S CH N ET
Revised: A.2. I MPLEMENTATION OF F ERMI N ET +S CH N ET

Original: To best reproduce the results from Gerard et al.
Revised: To best reproduce the results from Gerard et al.

Original: , we implemented the changes to the FermiNet which they claimed made the largest difference.
Revised: , we implemented the changes to the FermiNet which they claimed made the largest difference.

Original: All experiments in this paper, including the FermiNet, used dense determinants.
Revised: All experiments in this paper, including the FermiNet, used dense determinants.

Original: We added the SchNet-like continuous filter convolutions, as well as the nuclear embedding stream and separate streams for same-spin and different-spin electrons in the two-electron stream.
Revised: We added the SchNet-like continuous filter convolutions, as well as the nuclear embedding stream and separate streams for same-spin and different-spin electrons in the two-electron stream.

Original: We also compared their training hyperparameters against ours, except for batch size, which was kept at 4096 for all experiments.
Revised: We also compared their training hyperparameters against ours, except for batch size, which was kept at 4096 for all experiments.

Original: We did not implement their local input features or envelope initialization, as their results suggested they did not make as significant a difference, and our longer pretraining likely had the same effect as changing the envelope initialization.
Revised: We did not implement their local input features or envelope initialization, as their results suggested they did not make as significant a difference, and our longer pretraining likely had the same effect as changing the envelope initialization.

Original: A.2. C HOICE OF O PTIMIZER AND L AYER N ORM
Revised: A.2. C HOICE OF O PTIMIZER AND L AYER N ORM

Original: The original FermiNet was only able to reach high accuracy when trained with Kronecker-factored approximate curvature (KFAC)
Revised: The original FermiNet was only able to reach high accuracy when trained with Kronecker-factored approximate curvature (KFAC)

Original: (Martens & Grosse, 2015).
Revised: (Martens & Grosse, 2015).

Original: To see if the same holds true for the
Revised: To see if the same holds true for the

Original: PsiFormer, here we compare training with KFAC against ADAM on several systems.
Revised: PsiFormer, here we compare training with KFAC against ADAM on several systems.

Original: Additionally, when trained with ADAM, self-attention layers usually require LayerNorm (Ba et al., 2016).
Revised: Additionally, when trained with ADAM, self-attention layers usually require LayerNorm (Ba et al., 2016).

Original: However, 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.
Revised: However, 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.

Original: Figure 5 shows a comparison for the PsiFormer on the sulphur dimer and the iron atom.
Revised: Figure 5 shows a comparison for the PsiFormer on the sulphur dimer and the iron atom.

Original: Consistent with the FermiNet, the PsiFormer consistently converges faster and to lower energies with KFAC.
Revised: Consistent with the FermiNet, the PsiFormer consistently converges faster and to lower energies with KFAC.

Original: With LayerNorm, the situation is more ambiguous.
Revised: With LayerNorm, the situation is more ambiguous.

Original: On the sulphur dimer, it seems to have marginalimpact, while on the iron atom, it definitely improves training with ADAM, and seems to improve
Revised: On the sulphur dimer, it seems to have marginal

Original: [Table] [Table]
Revised: [Table] [Table]

Original:  molecules (Fig. 2)
Revised: Small molecules (Fig. 2)

Original: [Table]
Revised: [Table]

Original: Third-
Revised: Third-row atoms (Table 1)

Original: [Table]
Revised: [Table]

Original: atoms (Table1)
Revised: Large molecules (Table

Original: [Table]
Revised: [Table]

Original:  stability when using KFAC.As LayerNorm primarily seems to be useful for heavier atoms, in all experiments in the main text we use LayerNorm for systems with atoms beyond the first row (K, Fe, Zn, CCl 4 ) but not for hydrocarbons, unless otherwise specified.
Revised: impact, while on the iron atom, it definitely improves training with ADAM, and seems to improve stability when using KFAC.

Original: A.2. H YPERPARAMETERS
Revised: A.2. H YPERPARAMETERS

Original: Table 4 shows the default hyperparameters used for training all models implemented in this work.
Revised: Table 4 shows the default hyperparameters used for training all models implemented in this work.

Original: Note that Pfau et al.
Revised: Note that Pfau et al.

Original: (2020) took the sum over gradients across the batch on each device and averaged over devices, whereas here the gradients are averaged over the entire batch.
Revised: (2020) took the sum over gradients across the batch on each device and averaged over devices, whereas here the gradients are averaged over the entire batch.

Original: This amounts to a scaling of the learning rate, meaning that the same learning rate as in Pfau et al.
Revised: This amounts to a scaling of the learning rate, meaning that the same learning rate as in Pfau et al.

Original: (2020) cannot be used.
Revised: (2020) cannot be used.

Original: As in previous FermiNet implementations, we pretrain all networks to match Hartree-Fock (HF) orbitals computed using PySCF.
Revised: As in previous FermiNet implementations, we pretrain all networks to match Hartree-Fock (HF) orbitals computed using PySCF.

Revised: Here, we use the LAMB optimizer (You et al., 2020) as the pretraining optimizer.

Original: We find that longer pretraining stabilises training for all models considered.
Revised: We find that longer pretraining stabilises training for all models considered.

Original: In addition, during pretraining, we draw samples from the HF orbitals only, instead of the neural network wavefunction.
Revised: In addition, during pretraining, we draw samples from the HF orbitals only, instead of the neural network wavefunction.

Original: A smaller number of pretraining iterations (20,000) was used for small molecules in Section 4.1, while 100,000 iterations were used for third-row atoms and larger molecules.
Revised: A smaller number of pretraining iterations (20,000) was used for small molecules in Section 4.1, while 100,000 iterations were used for third-row atoms and larger molecules.

Revised: To generate samples from Ψ 2 , we use the Metropolis-Hastings algorithm with symmetric Gaussian proposals, as in Pfau et al.

Original: We increase the number of decorrelation MCMC steps between optimization iterations from 10 in Pfau et al.
Revised: We increase the number of decorrelation MCMC steps between optimization iterations from 10 in Pfau et al.

Original: (2020) to 30.
Revised: (2020) to 30.

Original: Additionally, for larger systems, we do not update all electron positions simultaneously in one Metropolis step.
Revised: Additionally, for larger systems, we do not update all electron positions simultaneously in one Metropolis step.

Original: Instead, we split the electrons into multiple blocks, and iteratively update each block once per step.
Revised: Instead, we split the electrons into multiple blocks, and iteratively update each block once per step.

Original: This is an intermediate between the all-electron and one-electron moves commonly used in VMC.
Revised: This is an intermediate between the all-electron and one-electron moves commonly used in VMC.

Original: Note that while all models were trained for 200,000 optimization iterations, this does not mean that all systems required that many iterations to converge.
Revised: Note that while all models were trained for 200,000 optimization iterations, this does not mean that all systems required that many iterations to converge.

Original: Many smaller systems converged in far fewer iterations than this, but the same number was used for all systems to minimize confusion.
Revised: Many smaller systems converged in far fewer iterations than this, but the same number was used for all systems to minimize confusion.

Original: Not all systems used identical parameters - for larger systems we increased the number of pretraining steps and blocks per MCMC update, and LayerNorm was only usedwith the PsiFormer for atoms heavier than neon.
Revised: Not all systems used identical parameters - for larger systems we increased the number of pretraining steps and blocks per MCMC update, and LayerNorm was not used for small molecules in Table 2.

Original: In Table 5 we specify which systems were trained with which settings.
Revised: In Table 5 we specify which systems were trained with which settings.

Original: In Section 4.1, the performance of “small” and “large” FermiNet and PsiFormer models was investigated on a set of small and medium molecules.
Revised: In Section 4.1, the performance of “small” and “large” FermiNet and PsiFormer models was investigated on a set of small and medium molecules.

Original: Table 6 gives the network parameters used for these
Revised: Table 6 gives the network parameters used for these

Original: [Table] [Table]
Revised: [Table] [Table]

Original: models.
Revised: models.

Original: In all other experiments, unless otherwise specified, the ‘small’ network configurations were used.
Revised: In all other experiments, unless otherwise specified, the ‘small’ network configurations were used.

Original: Table 7 contains the number of parameters used by each model for this set of molecules.
Revised: Table 7 contains the number of parameters used by each model for this set of molecules.

Original: A. C OMPUTATIONAL D ETAILS
Revised: A. C OMPUTATIONAL D ETAILS

Original: 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 parallel using between 16 and 64 A100 GPUs, depending on the system size.
Revised: 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 parallel using between 16 and 64 A100 GPUs, depending on the system size.

Original: All calculations were done at standard single precision, as we found that TensorFloat-32 calculations on A100 had numerical accuracy issues.
Revised: All calculations were done at standard single precision, as we found that TensorFloat-32 calculations on A100 had numerical accuracy issues.

Original: A table of total training time – including pretraining – for several models is given in Table 9.
Revised: A table of total training time – including pretraining – for several models is given in Table 9.

Original: Empirically, the time per iteration scaled roughly cubically, i.e. the benzene dimer required ∼ 8 times the computational resources to run as the benzene molecule, though the number of atoms also played a significant role in timing.
Revised: Empirically, the time per iteration scaled roughly cubically, i.e. the benzene dimer required ∼ 8 times the computational resources to run as the benzene molecule, though the number of atoms also played a significant role in timing.

Original: For systems with very large numbers of atoms, like the benzene dimer, the FermiNet and PsiFormer ran at nearly the same speed.
Revised: For systems with very large numbers of atoms, like the benzene dimer, the FermiNet and PsiFormer ran at nearly the same speed.

Revised: 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.

Original: Geometries are taken from Curtiss et al.
Revised: Geometries are taken from Curtiss et al.

Original: and given in Table 8 or in Pfau et al.
Revised: and given in Table 8 or in Pfau et al.

Original: The benzene dimer geometry is obtained via a rigid translation and rotation of two monomers.
Revised: The benzene dimer geometry is obtained via a rigid translation and rotation of two monomers.

Original: [Table]
Revised: [Table]

Original: H -0. H -0.
Revised: H -0. H -0.

Original: naphthalene
Revised: naphthalene

Original: C 0.
Revised: C 0.

Revised: C 0.

Original: C 0. 
Revised: C 0. C 0. C 0. C 0.

Original: C 0.C 0. C 0. C 0.
Revised: C 0.

Original: C 0.C 0.
Revised: C 0.

Original: H 0. H 0.
Revised: H 0. H 0.

Original: C 0. C 0.
Revised: C 0. C 0.

Original: H 0. H 0.
Revised: H 0. H 0.

Original: H 0. H 0.
Revised: H 0. H 0.

Original: H 0. H 0.
Revised: H 0. H 0.

Original: CCl
Revised: CCl

Original: C 0.
Revised: C 0.

Original: Cl 1.
Revised: Cl 1.

Original: Cl -1. Cl -1.
Revised: Cl -1. Cl -1.

Original: Cl 1.
Revised: Cl 1.

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: A. A BLATION STUDIES
Revised: A. A BLATION STUDIES

Original: Two modifications to the self-attention network proved critical to stability for the PsiFormer - the
Revised: Two modifications to the self-attention network proved critical to stability for the PsiFormer - the

Original: Jastrow factor ensuring correct behaviour at electronic cusps
Revised: Jastrow factor ensuring correct behaviour at electronic cusps

Original: [Figure]
Revised: [Figure]

Original: and rescaling the input distances.
Revised: and rescaling the input distances.

Original: Here we show ablation studies without these modifications.
Revised: Here we show ablation studies without these modifications.

Original: Figure 6 shows the PsiFormer on (a) the nitrogen dimer and (b) ethanol, without a Jastrow factor and without input feature rescaling.
Revised: Figure 6 shows the PsiFormer on (a) the nitrogen dimer and (b) ethanol, without a Jastrow factor and without input feature rescaling.

Original: In the absence of an electron-electron Jastrow factor to enforce cusp conditions, the energy is very noisy.
Revised: In the absence of an electron-electron Jastrow factor to enforce cusp conditions, the energy is very noisy.

Original: If the input features are not rescaled, the energy is often unstable.
Revised: If the input features are not rescaled, the energy is often unstable.

Original: For smaller systems, as in the nitrogen dimer shown, training may recover from the instability, but especially for medium or larger systems, the energy often diverges.
Revised: For smaller systems, as in the nitrogen dimer shown, training may recover from the instability, but especially for medium or larger systems, the energy often diverges.

Original: A. A TTENTION M APS
Revised: A. A TTENTION M APS

Original: Figures 7 and 8 show attention maps of the PsiFormer for a sample electron configuration, for the benzene dimer at equilibrium and dissociated geometry respectively.
Revised: Figures 7 and 8 show attention maps of the PsiFormer for a sample electron configuration, for the benzene dimer at equilibrium and dissociated geometry respectively.

Original: Each attention map shows the output of the dot product between key and query vectors before the softmax operation.
Revised: Each attention map shows the output of the dot product between key and query vectors before the softmax operation.

Original: Four attention maps are shown for each network layer, one per attention head.
Revised: Four attention maps are shown for each network layer, one per attention head.

Original: 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.
Revised: 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.

Original: The block structure of the attention weights is clearly visible, where electrons attend more to other electrons around the same molecule.
Revised: The block structure of the attention weights is clearly visible, where electrons attend more to other electrons around the same molecule.

Original: A tt e n t i o n
Revised: A tt e n t i o n

Original: H e a d
Revised: H e a d

Original: [Figure]
Revised: [Figure]

Original: A tt e n t i o n
Revised: A tt e n t i o n

Original: H e a d
Revised: H e a d

Original: [Figure]
Revised: [Figure]