# Message-Passing Neural Quantum States for the Homogeneous Electron Gas

Gabriel Pescia  Institute of Physics, École Polytechnique Fédérale de Lausanne
(EPFL), CH-1015 Lausanne, Switzerland  Center for Quantum Science and
Engineering, École Polytechnique Fédérale de Lausanne (EPFL), CH-1015
Lausanne, Switzerland  Jannes Nys  Institute of Physics, École Polytechnique
Fédérale de Lausanne (EPFL), CH-1015 Lausanne, Switzerland  Center for Quantum
Science and Engineering, École Polytechnique Fédérale de Lausanne (EPFL),
CH-1015 Lausanne, Switzerland  Jane Kim  Department of Physics and Astronomy
and Facility for Rare Isotope Beams, Michigan State University, East Lansing,
Michigan 48824, USA  Alessandro Lovato  Physics Division, Argonne National
Laboratory, Argonne, Illinois 60439  Computational Science Division, Argonne
National Laboratory, Argonne, Illinois 60439  INFN-TIFPA Trento Institute for
Fundamental Physics and Applications, 38123 Trento, Italy  Giuseppe Carleo
Institute of Physics, École Polytechnique Fédérale de Lausanne (EPFL), CH-1015
Lausanne, Switzerland  Center for Quantum Science and Engineering, École
Polytechnique Fédérale de Lausanne (EPFL), CH-1015 Lausanne, Switzerland

(August 30, 2024)

###### Abstract

We introduce a message-passing-neural-network-based wave function Ansatz to
simulate extended, strongly interacting fermions in continuous space. Symmetry
constraints, such as continuous translation symmetries, can be readily
embedded in the model. We demonstrate its accuracy by simulating the ground
state of the homogeneous electron gas in three spatial dimensions at different
densities and system sizes. With orders of magnitude fewer parameters than
state-of-the-art neural-network wave functions, we demonstrate better or
comparable ground-state energies. Reducing the parameter complexity allows
scaling to $N=128$ electrons, previously inaccessible to neural-network wave
functions in continuous space, enabling future work on finite-size
extrapolations to the thermodynamic limit. We also show the Ansatz’s
capability of quantitatively representing different phases of matter.

##  I Introduction

Predicting emergent physical phenomena and system properties from the _ab-
initio_ description of the system’s constituents is notoriously difficult [1,
2]. Fermionic systems, in particular, can exhibit strong correlations among
the particles, leading to collective phenomena in the form of exotic phases of
matter, e.g. superconductivity and superfluidity [3, 4].

In recent years, progress in numerical simulations of strongly correlated
systems was triggered by the development of increasingly precise machine-
learning approximation techniques. Most notably, artificial neural network
(NN) architectures, in combination with Variational Monte Carlo (VMC), have
shown great promise in representing ground states of quantum spin systems,
especially in more than one spatial dimension [5, 6, 7, 8, 9, 10, 11, 12]. Due
to the universal approximation property of NNs, neural-network quantum states
(NQS) can, in theory, accurately represent any quantum many-body state [13,
14]. NQS have been extended to fermionic degrees of freedom in a discrete
basis [15, 16, 17], by incorporating the indistinguishability of quantum
particles. More recently, advancements to ground and excited state searches
for fermionic and bosonic continuous degrees of freedom with open [18, 19, 20]
and periodic boundary conditions (PBCs) [21, 22, 23], have been introduced.

The flexibility of NQS, compared to more traditional models, allows for the
representation of multiple phases of matter, and even different physical
systems, with a single Ansatz. To exemplify this point, we refer to NQS
studies on the ground state of molecular systems [18, 19], solutions to
effective field theory Hamiltonians describing atomic nuclei [24, 25, 26],
bulk studies of fermionic and bosonic extended systems [21, 27, 23, 28], as
well as NQS simulations of low-density neutron matter found in neutron stars
[29]. The downside of this flexibility, especially in continuous space, is
that NQS typically need a significant amount of variational parameters to
reach a given accuracy. This makes optimization challenging and costly,
preventing the usage of refined optimization schemes, e.g. second order
optimization procedures [30, 31]. As a result, the accessible system sizes are
limited to a few tens of particles. However, studying larger system sizes is
of utmost importance to estimate physical properties in the thermodynamic
limit [32, 33, 34, 35]. To remedy the situation, novel NQS architectures must
be developed that significantly reduce the parameter complexity while
retaining high accuracy.

This work introduces a neural-network wave function suitable for simulating
strongly interacting fermionic quantum systems in continuous space with _one
to two_ orders of magnitude fewer parameters than current state-of-the-art
NQS. The general form of the Ansatz is motivated by an analytical argument,
relating the exact ground-state wave function to a many-body coordinate
transformation of the electronic coordinates. It uses a permutation-
equivariant message-passing architecture on a graph, inherently implementing
the indistinguishability of same-species quantum particles [36]. As an
application, we study the Homogeneous Electron Gas (HEG) in three spatial
dimensions without explicitly breaking any of the fundamental symmetries of
the system, such as translations and spin-inversion symmetry. This allows
characterizing, from first principles, the onset of Wigner crystallization at
low densities.

##  II Methods

Throughout this work, we consider a non-relativistic Hamiltonian of identical
particles with mass $m$ in $d$ spatial dimensions:

$\displaystyle H$
$\displaystyle=-\frac{\hbar^{2}}{2m}\sum_{i}^{N}\nabla_{\mathbf{r}_{i}}^{2}+V(\mathbf{X}),$
(1)

where the potential and interaction energy, $V$, is assumed to be diagonal in
position representation, defined by the particle coordinates
$\mathbf{X}=(\mathbf{r}_{1},...,\mathbf{r}_{N})$,
$\mathbf{r}_{i}\in\mathbb{R}^{d}$. In the following, we derive an analytic
functional form of the ground-state wave function and relate it to our
variational Ansatz.

###  II.1 Exact Backflow Transformations

Given a suitable reference state $|\Phi_{0}\rangle$, as initial condition for
the imaginary-time ($\tau$) evolution induced by the Hamiltonian,
$\Phi_{\tau}(\mathbf{X})=\langle\mathbf{X}|e^{-\tau H}|\Phi_{0}\rangle$, the
exact ground-state is obtained in the large imaginary-time limit:
$\lim_{\tau\rightarrow\infty}\Phi_{\tau}(\mathbf{X})\propto\Psi_{0}(\mathbf{X})$,
provided $|\Phi_{0}\rangle$ is non-orthogonal to the exact ground state,
$|\langle\Psi_{0}|\Phi_{0}\rangle|>0$. For fermions, non-orthogonality implies
that the wave function must be at least antisymmetric w.r.t. the exchange of
two particles i.e.
$\Phi_{0}(\mathcal{P}_{ij}(\mathbf{X}))=-\Phi_{0}(\mathbf{X})$
($\mathcal{P}_{ij}$ permutes particles $i$ and $j$). Assuming a valid quantum
reference state (twice differentiable and integrable) and finite matrix
elements of the propagator [37], we can apply the mean-value theorem to the
imaginary-time evolved state, yielding:

$\displaystyle\Phi_{\tau}(\mathbf{X})$ $\displaystyle=$
$\displaystyle\int_{\Omega}d\mathbf{X}^{\prime}G_{\tau}(\mathbf{X},\mathbf{X}^{\prime})\Phi_{0}(\mathbf{X}^{\prime})$
(2) $\displaystyle=$ $\displaystyle\text{Vol}(\Omega)\times
G_{\tau}\left(\mathbf{X},\mathbf{Y}(\mathbf{X})\right)\Phi_{0}(\mathbf{Y}(\mathbf{X})),$
(3)

where $\Omega$ is the (convex) integration domain for the positional degrees
of freedom, and
$G_{\tau}(\mathbf{X},\mathbf{X}^{\prime})=\langle\mathbf{X}|e^{-\tau
H}|\mathbf{X}^{\prime}\rangle$ is the matrix element of the imaginary-time
propagator, which is bounded for finite $\tau$. In Eq. (3), we introduced the
mean-value point
$\mathbf{Y}(\mathbf{X})=(\mathbf{y}_{1}(\mathbf{X}),...,\mathbf{y}_{N}(\mathbf{X}))\in\Omega$,
depending parametrically on the coordinates $\mathbf{X}$.

For general Hamiltonians in the form of Eq. (1), we have
$G_{\tau}(\mathbf{X},\mathbf{X}^{\prime})\geq 0$, for all
$\mathbf{X},\mathbf{X}^{\prime}$. Moreover,
$G_{\tau}(\mathbf{X},\mathbf{X}^{\prime})$ is invariant under the exchange of
particle coordinates:
$G_{\tau}(\mathcal{P}_{ij}(\mathbf{X}),\mathcal{P}_{ij}(\mathbf{X}^{\prime}))=G_{\tau}(\mathbf{X},\mathbf{X}^{\prime})$.
In the fermionic case, the latter implies that $\mathbf{Y}(\mathbf{X})$ must
be equivariant under particle exchange,
$\mathbf{Y}(\mathcal{P}_{ij}(\mathbf{X}))=\mathcal{P}_{ij}(\mathbf{Y}(\mathbf{\mathbf{X}}))$,
to ensure antisymmetry of the total wave-function. Eq. (3) therefore yields
the product between a permutation symmetric, positive semi-definite function
$J(\mathbf{X})=G_{\tau}(\mathbf{X},\mathbf{Y}(\mathbf{X}))\times\mathrm{Vol}(\Omega)$
and a reference state computed at modified coordinates
$\mathbf{Y}(\mathbf{X})$:

$\displaystyle\Phi_{\tau}(\mathbf{X})=J(\mathbf{X})\times\Phi_{0}(\mathbf{Y}(\mathbf{X})).$
(4)

Identification of the mean-value point $Y(X)$ with a many-body coordinate
transformation gives an alternative justification for the backflow
transformations [38] of single-particle coordinates. With a Slater determinant
of given spin orbitals $\phi_{\mu}(\mathbf{r}_{i})$ as initial state,
$\Phi_{0}(\mathbf{X})=\det{\phi_{\mu}(\mathbf{r}_{i})}/{\sqrt{N!}}$, Eq. (4)
is structurally related to the heuristic Jastrow-Backflow variational form
[39, 40]. We remark that the symmetric contribution, $J(X)$, can be
incorporated to the determinant:

$\displaystyle\Phi_{\tau}(\mathbf{X})=\mathcal{K}\times\det{\varphi_{\mu}(\mathbf{y}_{i}(\mathbf{X}))},$
(5)

with
$\varphi_{\mu}(\mathbf{y}_{i}(\mathbf{X}))=\phi_{\mu}(\mathbf{y}_{i}(\mathbf{X}))\times\sqrt[N]{J(\mathbf{X})}$,
and $\mathcal{K}$ a normalization constant.

The functional form (4),(5) is exact, provided that the symmetric factor
$J(\mathbf{X})$ and the mean-value coordinates $\mathbf{Y}(\mathbf{X})$
satisfy Eq. (3), and the reference state is not orthogonal to the exact ground
state. An approximate but explicit form for the coordinate transformation
$\mathbf{Y}(\mathbf{X})$ can be obtained by repeatedly applying the imaginary-
time propagator to the reference state in the limit of small $\tau$. This
process gives rise to the iterative backflow transformation, as introduced in
Ref. [41, 42].

###  II.2 Message-Passing Neural Quantum States

Motivated by Eq. (5), we use single-particle orbitals,
$\\{\phi_{\mu}\\}_{\mu=1}^{N}$, evaluated at many-body backflow coordinates,
$\mathbf{Y}(\mathbf{X})$, to construct the variational Ansatz. The backflow
transformation is parameterized with permutation-equivariant message-passing
NNs (MPNN) [36] (see Fig. 1), hence we name it Message-Passing Neural Quantum
State (MP-NQS).

Figure 1:  Illustration of backflow transformation via MPNN. Single-particle
coordinates $\mathbf{r}_{i}\in\mathbb{R}^{d}$ (black dots, top left) are
mapped to quasi-particle coordinates
$\delta\mathbf{r}_{i}\in\mathbb{C}^{d}$(black/white dots for real/imaginary
part, top right). From an initial graph consisting of an initial feature
vector (dark grey) and a hidden state (green), a new graph, consisting of the
initial features and an updated hidden state (indicated by different
coloring), is constructed via messages. The final node states are linearly
transformed back to $\delta\mathbf{r}_{i}=W\cdot\mathbf{g}_{i}^{(T)}$,
containing information about all particles ($D$ is the dimension of the last
graph’s nodes).

In the MPNN, an all-to-all connected graph, encoding effective particle
positions (nodes) and their interactions (edges):

$\displaystyle\mathbf{g}_{i}^{(t)}$
$\displaystyle=[\mathbf{x}_{i}^{(0)},\mathbf{h}_{i}^{(t)}]\in\mathbb{R}^{D_{1}}$
(6) $\displaystyle\mathbf{g}_{ij}^{(t)}$
$\displaystyle=[\mathbf{x}_{ij}^{(0)},\mathbf{h}_{ij}^{(t)}]\in\mathbb{R}^{D_{2}},$
(7)

is updated iteratively. Here, we have introduced a discrete step index $t\geq
0$, initial feature vectors $\mathbf{x}_{i}^{(0)}\in\mathbb{R}^{D_{1}^{(0)}}$,
$\mathbf{x}_{ij}^{(0)}\in\mathbb{R}^{D_{2}^{(0)}}$, and we denote
concatenation with $[\cdot,\cdot]$. The iterative equations above also contain
auxiliary variables known as hidden states
($\mathbf{h}_{i}^{(t)}\in\mathbb{R}^{D^{h}_{1}}$,
$\mathbf{h}_{ij}^{(t)}\in\mathbb{R}^{D^{h}_{2}}$) with suitably chosen
_feature_ dimensions such that $D_{1}^{h}+D_{1}^{(0)}=D_{1}$, and
$D_{2}^{h}+D_{2}^{(0)}=D_{2}$. Construction of the initial feature vectors
($\mathbf{x}_{i}^{(0)}$, $\mathbf{x}_{ij}^{(0)}$) is system dependent and will
be discussed in detail later.

The hidden states are initialized, independently of their index, with vectors
of learnable parameters. They are updated using permutation-equivariant
_messages_ , obtained from a variant of the attention mechanism [43], we dub
here particle attention. Specifically, the messages are given by weighted
transformations of the edges $\mathbf{g}_{ij}^{(t)}$:
$\mathbf{m}_{ij}^{(t+1)}=\boldsymbol{\omega}^{(t)}_{ij}(\mathbf{g}_{ij}^{(t)})\odot\boldsymbol{\phi}(\mathbf{g}_{ij}^{(t)})$,
where $\odot$ represents element-wise multiplication along the feature
dimension, and $\boldsymbol{\omega}^{(t)}_{ij}\in\mathbb{R}^{D_{2}}$ are
weight vectors. The weights are obtained using _query_ /_key_ matrices given
by $\mathbf{Q}_{ij}^{(t)}=W_{Q}\cdot\mathbf{g}_{ij}^{(t)}$ and
$\mathbf{K}_{ij}^{(t)}=W_{K}\cdot\mathbf{g}_{ij}^{(t)}$, with weight matrices
$W_{Q},W_{K}\in\mathbb{R}^{D_{2}\times D_{2}}$. Applying an element-wise GELU
non-linearity [44] to the overlap between queries and keys along the _particle
dimension_ (as opposed to the feature dimension [43, 45]), results in
permutation-equivariant weights

$\displaystyle\boldsymbol{\omega}^{(t)}_{ij}$
$\displaystyle=\mathrm{GELU}\left(\sum_{l}\mathbf{Q}_{il}^{(t)}\mathbf{K}_{lj}^{(t)}\right).$
(8)

The particle attention compares environments of particles $i$ and $j$, and
effectively increases the order of correlations that can be embedded in a
single iteration of the network. This is crucial to reduce the total number of
network iterations (parameters) and capture many-body effects. The hidden
states are updated using the current graph and messages:

$\displaystyle\mathbf{h}_{i}^{(t+1)}$
$\displaystyle=\boldsymbol{f}\left(\mathbf{g}_{i}^{(t)},\sum_{j\neq
i}\mathbf{m}_{ij}^{(t+1)}\right)$ (9) $\displaystyle\mathbf{h}_{ij}^{(t+1)}$
$\displaystyle=\boldsymbol{\tilde{f}}\left(\mathbf{g}_{ij}^{(t)},\mathbf{m}_{ij}^{(t+1)}\right)$
(10)

The functions $\boldsymbol{\phi}$, $\boldsymbol{f}$, and
$\boldsymbol{\tilde{f}}$ are parameterized by Multilayer Perceptrons (MLPs).
The updated graph then has the same structure as the former:
$\mathbf{g}_{i}^{(t+1)}=[\mathbf{x}_{i}^{(0)},\mathbf{h}_{i}^{(t+1)}]\in\mathbb{R}^{D_{1}}$,
$\mathbf{g}_{ij}^{(t+1)}=[\mathbf{x}_{ij}^{(0)},\mathbf{h}_{ij}^{(t+1)}]\in\mathbb{R}^{D_{2}}$.
Inclusion of the initial inputs, referred to as a “skip connection” in ML
literature [46], mitigates the vanishing gradient problem and allows a more
efficient capture of correlations.

The final backflow coordinates are constructed as
$\mathbf{y}_{i}(\mathbf{X})=\mathbf{r}_{i}+\delta\mathbf{r}_{i}(\mathbf{X})$,
where the displacements, $\delta\mathbf{r}_{i}\left(\mathbf{X}\right)$, are
obtained via a linear transformation of the final node states to $d$
dimensions, i.e.
$\delta\mathbf{r}_{i}\left(\mathbf{X}\right)=W\cdot\mathbf{g}_{i}^{(T)}$ with
$W\in\mathbb{C}^{d\crossproduct D_{1}}$. The complex-valued backflow
transformation allows changing the degree of localization, determined by the
chosen single-particle orbitals, and representing complex-valued wave
functions in general.

Following (5), we further augment the orbitals with a permutation-invariant
factor $J$, of the form

$\displaystyle J(\mathbf{Y},\mu)=\sum_{i}j(\mathbf{y}_{i},\mu)$ (11)

where $j$ is parameterized with a MLP, and $\mu$ denotes the quantum numbers
of the orbitals, yielding:

$\displaystyle\Psi(\mathbf{X})=\det{\varphi_{\mu}(\mathbf{y}_{i}(\mathbf{X}))},$
(12)

with
$\varphi_{\mu}(\mathbf{y}_{i})=\exp\left[J(\mathbf{Y},\mu)\right]\times\phi_{\mu}(\mathbf{y}_{i})$.

###  II.3 Electron Gas

We now study the case of the homogeneous electron gas (HEG) in $d=3$ spatial
dimensions, a prototypical model for the electronic structure in solids. It
includes Coulomb interactions among the solids’ electrons while treating its
positively charged ions as uniform, static, positive background [47]. Despite
this simplification, the HEG exhibits different phases of matter and captures
properties of real solids, particularly of Alkali metals. The Hamiltonian (in
units of Hartree), for a system of $N$ electrons with uniform density
$n=\frac{N}{V}$, is given by:

$\displaystyle H$
$\displaystyle=-\frac{1}{2r_{s}^{2}}\sum_{i}^{N}\nabla_{\mathbf{r}_{i}}^{2}+\frac{1}{r_{s}}\sum_{i<j}^{N}\frac{1}{\norm{\mathbf{r}_{i}-\mathbf{r}_{j}}}+\mathrm{const.}$
(13)

where we introduced the Wigner-Seitz radius $r_{s}=\sqrt[3]{3/(4\pi n)}$, and
a constant arising from the electron-background interaction [39]. The
conditionally convergent series of pairwise Coulomb interactions is evaluated
using the Ewald summation technique, as is standard for extended systems in
QMC [48, 49, 50]. We will assume a fixed spin-polarization with
$N=N_{\uparrow}+N_{\downarrow}$, where $N_{\uparrow/\downarrow}$ denotes the
number of up/down spins. Additionally, $s_{i}\in\\{\uparrow,\downarrow\\}$
denotes the spin of the $i$-th electron. We equip the cubic simulation cell of
side length $L$ with periodic boundary conditions (PBCs) in all spatial
directions to access the bulk of the system.

As in (5), we use a single Slater determinant as a reference state,
$\Phi_{0}(\mathbf{X})$. We expect a liquid-crystal phase transition for the
HEG, as a function of the density $n$. The dominating kinetic energy in Eq.
(13) ($\sim 1/r_{s}^{2}$) for large $n$ leads to the well-known Fermi liquid
behavior. For small $n$, the potential energy ($\sim 1/r_{s}$) dominates and
enforces a crystalline BCC structure among the electrons, known as Wigner
crystal [51]. Plane-wave orbitals are a natural and physically-motivated
choice to model the liquid, translation invariant phase:
$\phi_{\mathbf{k}}(\mathbf{r})=\exp\left[i\mathbf{k}\cdot\mathbf{r}\right]$
with $\mathbf{k}=\frac{2\pi}{L}\mathbf{n}$ where
$\mathbf{n}\in\mathbb{Z}^{d}$. These orbitals allow modeling the translation
invariant system at fixed total momentum
$\mathbf{k}_{\textrm{tot}}=\sum_{i=1}^{N}\mathbf{k}_{i}$. To account for the
spin $s$ of a particle located at $\mathbf{r}$, we use spin-orbitals
$\phi_{\mu}(\mathbf{r},s)=\phi_{\mathbf{k}_{\mu}}(\mathbf{r})\delta_{s_{\mu},s}$,
where each spin-orbital is characterized by the quantum numbers
$\mu=(\mathbf{k}_{\mu},s_{\mu})$. This choice of spin-orbitals lets the
determinant factorize into a product of determinants of up and down spin
orbitals. To study the Wigner crystal phase, we employ localized Gaussian
orbitals centered around the BCC lattice sites $\mathbf{R}_{\mu}$:
$\phi_{\mu}(\mathbf{r},s)=\sum_{\mathbf{R}_{n}}\exp\left[-\alpha(\mathbf{r}-\mathbf{R}_{\mu}+\mathbf{R}_{n})\right])\delta_{s_{\mu},s}$,
where $\alpha$ is a variational parameter, $\mu=(\mathbf{R}_{\mu},s_{\mu})$,
and we sum over simulation cell lattice vectors $\mathbf{R}_{n}$ to ensure
periodicity. For all densities we use a simple cubic simulation cell to not
bias towards any one of the phases.

We further specialize the MP-NQS to the HEG by defining initial feature
vectors. Respecting the spin inversion and translation symmetries of the HEG
requires us to ignore single-particle positions and spins. We, therefore,
initialize the nodes to a learnable embedding vector
$\mathbf{e}\in\mathbb{R}^{D_{1}-D_{1}^{h}}$, that does not depend on the
particle index $i$. For the edge features, we use the translation invariant
particle-distances $\mathbf{r}_{ij}=\mathbf{r}_{i}-\mathbf{r}_{j}$ and their
norm. Same- and opposite-spin pairs are distinguished using products of the
form $s_{i}\cdot s_{j}=\pm 1$ in the edge features. Overall, we obtain the
following initial features:

$\displaystyle\mathbf{x}_{i}^{(0)}=\mathbf{e}\mathrm{,}\qquad\mathbf{x}_{ij}^{(0)}=[\mathbf{r}_{ij},\norm{\mathbf{r}_{ij}},s_{i}\cdot
s_{j}].$ (14)

Notice that this choice preserves the spin quantum number of each particle.

The PBCs of the simulation box are incorporated by mapping the components of
vectors $\mathbf{r}\in\mathbb{R}^{d}$ (where $\mathbf{r}=\mathbf{r}_{i}$ or
$\mathbf{r}=\mathbf{r}_{ij}$) to a Fourier basis
$\mathbf{r}\mapsto\left[\sin(\frac{2\pi}{L}\mathbf{r}),\cos(\frac{2\pi}{L}\mathbf{r})\right]\in\mathbb{R}^{2d}$,
and the norm of the distance between two particles, $\norm{\mathbf{r}_{ij}}$,
to a periodic surrogate
$\norm{\mathbf{r}_{ij}}\mapsto\norm{\sin(\frac{\pi}{L}\mathbf{r}_{ij})}$, as
in Ref. [21]. In summary, our Ansatz exhibits translation and spin-inversion
invariance and maintains a fixed total momentum, $\mathbf{k}_{\textrm{tot}}$.
Its number of parameters is independent of system size (here $\sim 19000$)
and, using Stochastic Reconfiguration (SR) [31], only $\mathcal{O}(10^{3})$
optimization steps are needed to reach convergence. A comparison to other NQS
approaches is given in the Supplemental Material.

##  III Results

We study the fully spin-polarized and unpolarized HEG in different density
regimes $r_{s}\in[1,200]$ and up to system sizes of $N=128$ electrons. We
compare our ground-state energies against state-of-the-art NQS architectures –
FermiNet [23] and WAPNet [22] – for small system sizes $N\in\\{14,19\\}$ and
against _Diffusion Monte Carlo_ (DMC) with backflow (BF-DMC) [52, 53] for
$N=54$ electrons. Additional benchmarks w.r.t. common quantum chemistry
methods, including transcorrelation augmented full configuration interaction
method (FCI) and distinguishable clusters with doubles (DCD) method [34], are
included in the Supplementary Material. The effect of the backflow
transformation on the nodal surface is studied by comparing to fixed-node DMC
(FN-DMC) results. We use an energy of $1.5$ mHa per particle (chemical
accuracy) to assess the significance of energy differences between the
different methods. An overview of all results and benchmarks for the various
system sizes and densities is provided in the Supplemental Material.

###  III.1 Small Systems.–

The best available results for $N=14$ are obtained with the FCI method,
compared to which we find an energy difference of less than $1.5$ mHa per
particle. State-of-the-art NQS architectures perform comparably to the MP-NQS:
The unrestricted FermiNet performs slightly better
($\mathcal{O}\left(10^{-5}\right)$ Ha/N) than both MP-NQS and WAPNet for
$r_{s}\leq 2$, while the MP-NQS and WAPNet improve over this version of
FermiNet for $r_{s}=5$. The restricted FermiNet yields worse ground-state
energies than the MP-NQS over all probed densities [23] (see Fig. 2). It
should be noted that the energy differences for this system are of the order
of machine precision. For $r_{s}\geq 5$, we find slightly better performance
than WAPNet for all of the reported densities. All differences lie within a
range of $1.5$ mHa per particle. Similarly, for $N=19$, the MP-NQS obtains
slightly higher energies than WAPNet for large densities ($r_{s}\leq 5)$ and
marginally lower ones at smaller densities ($r_{s}\geq 5$) with differences
lower than $1.5$ mHa per particle.

Figure 2: Comparison of ground-state energies obtained with different methods
for various densities, polarizations, and system sizes, in mHa per particle.
Each line shows the difference in energy of the respective method to the MP-
NQS energy with plane-wave orbitals $E_{\mathrm{MPNN}}$ (values above zero are
worse than the MP-NQS baseline). (Top) $N=14,19$, (Bottom): $N=54$ particles.
Error bars are too small to be visible for most densities. The corresponding
numerical data can be found in the Supplemental Material.

###  III.2 Large Systems.–

For $N=54$ particles, accurate results are obtained with the FCI method [34].
At this system size, the energy differences per particle are smaller than for
$N=14$, suggesting size-consistency of our Ansatz, provided that the FCI
method is size consistent as well. When compared to purely variational
methods, we obtain significantly better ground-state energies than FN-DMC and
BF-DMC, especially at high densities. This is in stark contrast to the
(FermiNet-based) architecture of [27] (dubbed LiNet in the following), which
does not improve upon BF-DMC energies over the whole density regime (see Fig.
2, bottom). At $r_{s}\geq 50$, we observe improved results using the MP-NQS
with Gaussian orbitals compared to plane-waves (see Fig. 2, bottom panel,
yellow line). This strongly suggests a transition from a de-localized Fermi
liquid to a localized crystalline phase as expected from previous studies [54,
35]. To analyze finite size effects, we also simulate a larger system of
$N=128$ electrons at $r_{s}=50,110,200$, and confirm that Gaussian orbitals
lead to lower ground-state energies, compared to the plane-waves for
$r_{s}>50$ (see Supplementary Material). Furthermore, for $r_{s}=110,200$ the
crystalline character of the variational state can be clearly seen in the
radial distribution functions and corresponding structure factors displayed in
Fig. 3. The prominent peak in the structure factor and the pronounced density
fluctuations in the radial distribution function up to the maximum distance of
$L/2$, indicate the crystalline nature of the represented state. Note that
these are absent for $r_{s}=50$, suggesting a fluid state.

Figure 3:  Spin-averaged radial distribution function (top) and corresponding
structure factor (bottom) for the HEG system with $N=128$ electrons at
$r_{s}=50,110,200$. For $r_{s}=50$ we used plane-waves as reference state
while for $r_{s}=110,200$ Gaussian orbitals centered at BCC lattice sites were
used. Error bars are smaller than the symbols. The crystal and liquid
benchmarks are obtained from [54] for $r_{s}=110$.

##  IV Conclusions.–

We have introduced MP-NQS, a novel NQS architecture that leverages MPNNs to
build highly expressive backflow coordinates. We demonstrate its power on the
HEG system, reducing the number of parameters by orders of magnitudes compared
to state-of-the-art NQS in continuous space while reaching at par or better
accuracy. We also show improvement upon state-of-the-art BF-DMC results on
large systems. The favorable scaling allows us to accurately simulate large
periodic electronic systems, previously inaccessible to state-of-the-art NQS
models. We increase the available system sizes from $N=27$ and $N=54$
electrons in periodic systems [27, 23, 22] to $N=128$ electrons in this work.
Hence we open the door to extrapolation methods to the thermodynamic limit for
extended systems. Furthermore, we reproduce the liquid-crystal phase
transition of the HEG around $r_{s}=100$, matching previous predictions on the
transition density [55, 54, 35], showing the MP-NQS capability to describe
different phases of matter quantitatively better than previous studies of the
HEG based on neural quantum states [23]. In addition to the numerical results,
we also introduced an analytical argument, justifying commonly adopted
backflow transformations. Our argument shows that a backflow transformation
over a reference state is sufficient to obtain the exact ground-state wave
function. It will be of particular interest to characterize the geometrical
properties of these transformations and understand in what cases neural-
network parameterizations can efficiently describe them.

###### Acknowledgements.

We acknowledge stimulating discussions with Markus Holzmann, and Francesco
Pederiva. The numerical tools used in this work are based on the open-source
software NetKet[56, 57, 58] version 3. The present research is supported by
the Swiss National Science Foundation under Grant No. 200021_200336, by the
NCCR MARVEL, a National Centre of Competence in Research, funded by the Swiss
National Science Foundation (grant number 205602), Microsoft Research, by the
U.S. Department of Energy, Office of Science, Office of Nuclear Physics, under
contracts DE-AC02-06CH11357 (A.L), by the 2020 DOE Early Career Award program
(A.L.), the NUCLEI SciDAC program (A.L.), and Argonne LDRD awards (A. L.). J.
K. is supported by the U.S. Department of Energy, Office of Science, office of
Nuclear Physics under grant No. DE-SC0021152 and U.S. National Science
Foundation Grants No. PHY-1404159 and PHY-2013047.

## References

  * Austin  _et al._ [2012] B. M. Austin, D. Y. Zubarev, and W. A. Lester Jr, Quantum monte carlo and related approaches, Chemical reviews 112, 263 (2012).
  * Carlson  _et al._ [2015] J. Carlson, S. Gandolfi, F. Pederiva, S. C. Pieper, R. Schiavilla, K. Schmidt, and R. B. Wiringa, Quantum monte carlo methods for nuclear physics, Reviews of Modern Physics 87, 1067 (2015).
  * Leggett [1975] A. J. Leggett, A theoretical description of the new phases of liquid ${}^{3}\mathrm{He}$, Rev. Mod. Phys. 47, 331 (1975).
  * Tinkham [2004] M. Tinkham,  _Introduction to superconductivity_ (Courier Corporation, 2004).
  * Carleo and Troyer [2017] G. Carleo and M. Troyer, Solving the quantum many-body problem with artificial neural networks, Science 355, 602 (2017).
  * Carrasquilla [2020] J. Carrasquilla, Machine learning for quantum matter, Advances in Physics: X 5, 1797528 (2020).
  * Choo  _et al._ [2018] K. Choo, G. Carleo, N. Regnault, and T. Neupert, Symmetries and many-body excitations with neural-network quantum states, Phys. Rev. Lett. 121, 167204 (2018).
  * Ferrari  _et al._ [2019] F. Ferrari, F. Becca, and J. Carrasquilla, Neural gutzwiller-projected variational wave functions, Phys. Rev. B 100, 125131 (2019).
  * Hibat-Allah  _et al._ [2020] M. Hibat-Allah, M. Ganahl, L. E. Hayward, R. G. Melko, and J. Carrasquilla, Recurrent neural network wave functions, Phys. Rev. Res. 2, 023358 (2020).
  * Bukov  _et al._ [2021] M. Bukov, M. Schmitt, and M. Dupont, Learning the ground state of a non-stoquastic quantum Hamiltonian in a rugged neural network landscape, SciPost Phys. 10, 147 (2021).
  * Nomura and Imada [2021] Y. Nomura and M. Imada, Dirac-type nodal spin liquid revealed by refined quantum many-body solver using neural-network wave function, correlation ratio, and level spectroscopy, Phys. Rev. X 11, 031034 (2021).
  * Astrakhantsev  _et al._ [2021] N. Astrakhantsev, T. Westerhout, A. Tiwari, K. Choo, A. Chen, M. H. Fischer, G. Carleo, and T. Neupert, Broken-symmetry ground states of the heisenberg model on the pyrochlore lattice, Phys. Rev. X 11, 041021 (2021).
  * Deng  _et al._ [2017] D.-L. Deng, X. Li, and S. D. Sarma, Quantum entanglement in neural network states, Physical Review X 7, 021021 (2017).
  * Levine  _et al._ [2019] Y. Levine, O. Sharir, N. Cohen, and A. Shashua, Quantum entanglement in deep learning architectures, Phys. Rev. Lett. 122, 065301 (2019).
  * Luo and Clark [2019] D. Luo and B. K. Clark, Backflow transformations via neural networks for quantum many-body wave functions, Phys. Rev. Lett. 122, 226401 (2019).
  * Choo  _et al._ [2020] K. Choo, A. Mezzacapo, and G. Carleo, Fermionic neural-network states for ab-initio electronic structure, Nature communications 11, 1 (2020).
  * Robledo Moreno  _et al._ [2022] J. Robledo Moreno, G. Carleo, A. Georges, and J. Stokes, Fermionic wave functions from neural-network constrained hidden states, Proceedings of the National Academy of Sciences 119, e2122059119 (2022).
  * Pfau  _et al._ [2020] D. Pfau, J. S. Spencer, A. G. D. G. Matthews, and W. M. C. Foulkes, Ab initio solution of the many-electron schrödinger equation with deep neural networks, Phys. Rev. Res. 2, 033429 (2020).
  * Hermann  _et al._ [2020] J. Hermann, Z. Schätzle, and F. Noé, Deep-neural-network solution of the electronic schrödinger equation, Nature Chemistry 12, 891 (2020).
  * Entwistle  _et al._ [2023] M. Entwistle, Z. Schätzle, P. A. Erdman, J. Hermann, and F. Noé, Electronic excited states in deep variational monte carlo, Nature Communications 14, 274 (2023).
  * Pescia  _et al._ [2022] G. Pescia, J. Han, A. Lovato, J. Lu, and G. Carleo, Neural-network quantum states for periodic systems in continuous space, Physical Review Research 4, 023138 (2022).
  * Wilson  _et al._ [2022] M. Wilson, S. Moroni, M. Holzmann, N. Gao, F. Wudarski, T. Vegge, and A. Bhowmik, Wave function ansatz (but periodic) networks and the homogeneous electron gas, arXiv preprint arXiv:2202.04622 (2022).
  * Cassella  _et al._ [2023] G. Cassella, H. Sutterud, S. Azadi, N. D. Drummond, D. Pfau, J. S. Spencer, and W. M. C. Foulkes, Discovering quantum phase transitions with fermionic neural networks, Phys. Rev. Lett. 130, 036401 (2023).
  * Adams  _et al._ [2021] C. Adams, G. Carleo, A. Lovato, and N. Rocco, Variational monte carlo calculations of $a\leq 4$ nuclei with an artificial neural-network correlator ansatz, Phys. Rev. Lett. 127, 022502 (2021).
  * Gnech  _et al._ [2022] A. Gnech, C. Adams, N. Brawand, G. Carleo, A. Lovato, and N. Rocco, Nuclei with up to $\backslash$ $\\{$A= 6$\\}a=6$ nucleons with artificial neural network wave functions, Few-Body Systems 63, 1 (2022).
  * Lovato  _et al._ [2022] A. Lovato, C. Adams, G. Carleo, and N. Rocco, Hidden-nucleons neural-network quantum states for the nuclear many-body problem, Physical Review Research 4, 043178 (2022).
  * Li  _et al._ [2022] X. Li, Z. Li, and J. Chen, Ab initio calculation of real solids via neural network ansatz, Nature Communications 13, 7895 (2022).
  * Xie  _et al._ [2022] H. Xie, L. Zhang, and L. Wang, $m^\ast$ of two-dimensional electron gas: a neural canonical transformation study (2022), arXiv:2201.03156 [cond-mat, physics:physics].
  * Fore  _et al._ [2022] B. Fore, J. M. Kim, G. Carleo, M. Hjorth-Jensen, and A. Lovato, Dilute neutron star matter from neural-network quantum states, arXiv preprint arXiv:2212.04436 (2022).
  * Toulouse and Umrigar [2007] J. Toulouse and C. J. Umrigar, Optimization of quantum Monte Carlo wave functions by energy minimization, J. Chem. Phys. 126, 084102 (2007), arXiv:physics/0701039 [physics.chem-ph] .
  * Sorella [1998] S. Sorella, Green function monte carlo with stochastic reconfiguration, Phys. Rev. Lett. 80, 4558 (1998).
  * Lin  _et al._ [2001] C. Lin, F. Zong, and D. M. Ceperley, Twist-averaged boundary conditions in continuum quantum monte carlo algorithms, Physical Review E 64, 016702 (2001).
  * Drummond and Needs [2013] N. Drummond and R. Needs, Diffusion quantum monte carlo calculation of the quasiparticle effective mass of the two-dimensional homogeneous electron gas, Physical Review B 87, 045131 (2013).
  * Liao  _et al._ [2021] K. Liao, T. Schraivogel, H. Luo, D. Kats, and A. Alavi, Towards efficient and accurate ab initio solutions to periodic systems via transcorrelation and coupled cluster theory, Phys. Rev. Research 3, 033072 (2021).
  * Azadi  _et al._ [2022] S. Azadi, N. Drummond, and S. Vinko, Correlation energy of the paramagnetic electron gas at the thermodynamic limit, arXiv preprint arXiv:2209.10227 (2022).
  * Gilmer  _et al._ [2017] J. Gilmer, S. S. Schoenholz, P. F. Riley, O. Vinyals, and G. E. Dahl, Neural message passing for quantum chemistry, in  _International conference on machine learning_ (PMLR, 2017) pp. 1263–1272.
  * Apostol [1974] T. M. Apostol,  _Mathematical analysis_ (1974).
  * Feynman and Cohen [1956] R. P. Feynman and M. Cohen, Energy spectrum of the excitations in liquid helium, Phys. Rev. 102, 1189 (1956).
  * Kwon  _et al._ [1998] Y. Kwon, D. M. Ceperley, and R. M. Martin, Effects of backflow correlation in the three-dimensional electron gas: Quantum monte carlo study, Phys. Rev. B 58, 6800 (1998).
  * Kwon  _et al._ [1993] Y. Kwon, D. Ceperley, and R. M. Martin, Effects of three-body and backflow correlations in the two-dimensional electron gas, Physical Review B 48, 12037 (1993).
  * Taddei  _et al._ [2015] M. Taddei, M. Ruggeri, S. Moroni, and M. Holzmann, Iterative backflow renormalization procedure for many-body ground-state wave functions of strongly interacting normal fermi liquids, Physical Review B 91, 115106 (2015).
  * Ruggeri  _et al._ [2018] M. Ruggeri, S. Moroni, and M. Holzmann, Nonlinear network description for many-body quantum systems in continuous space, Phys. Rev. Lett. 120, 205302 (2018).
  * Vaswani  _et al._ [2017] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin, Attention is all you need, Advances in neural information processing systems 30 (2017).
  * Hendrycks and Gimpel [2016] D. Hendrycks and K. Gimpel, Gaussian error linear units (gelus), arXiv preprint arXiv:1606.08415 (2016).
  * von Glehn  _et al._ [2022] I. von Glehn, J. S. Spencer, and D. Pfau, A self-attention ansatz for ab-initio quantum chemistry, arXiv preprint arXiv:2211.13672 (2022).
  * He  _et al._ [2016] K. He, X. Zhang, S. Ren, and J. Sun, Deep residual learning for image recognition, in  _Proceedings of the IEEE conference on computer vision and pattern recognition_ (2016) pp. 770–778.
  * Pines [1963] D. Pines, Elementary excitations in solids, edited by jd jackson and d. pines (1963).
  * Ewald [1921] P. P. Ewald, Die berechnung optischer und elektrostatischer gitterpotentiale, Annalen der physik 369, 253 (1921).
  * Fraser  _et al._ [1996] L. M. Fraser, W. M. C. Foulkes, G. Rajagopal, R. J. Needs, S. D. Kenny, and A. J. Williamson, Finite-size effects and coulomb interactions in quantum monte carlo calculations for homogeneous systems with periodic boundary conditions, Phys. Rev. B 53, 1814 (1996).
  * Toukmaji and Board Jr [1996] A. Y. Toukmaji and J. A. Board Jr, Ewald summation techniques in perspective: a survey, Computer physics communications 95, 73 (1996).
  * Wigner [1934] E. Wigner, On the interaction of electrons in metals, Phys. Rev. 46, 1002 (1934).
  * López Ríos  _et al._ [2006] P. López Ríos, A. Ma, N. D. Drummond, M. D. Towler, and R. J. Needs, Inhomogeneous backflow transformations in quantum monte carlo calculations, Phys. Rev. E 74, 066701 (2006).
  * Azadi and Drummond [2022] S. Azadi and N. D. Drummond, Low-density phase diagram of the three-dimensional electron gas, Phys. Rev. B 105, 245135 (2022).
  * Drummond  _et al._ [2004] N. Drummond, Z. Radnai, J. Trail, M. Towler, and R. Needs, Diffusion quantum monte carlo study of three-dimensional wigner crystals, Physical Review B 69, 085116 (2004).
  * Ceperley and Alder [1980] D. M. Ceperley and B. J. Alder, Ground state of the electron gas by a stochastic method, Phys. Rev. Lett. 45, 566 (1980).
  * Carleo  _et al._ [2019] G. Carleo, K. Choo, D. Hofmann, J. E. Smith, T. Westerhout, F. Alet, E. J. Davis, S. Efthymiou, I. Glasser, S.-H. Lin, _et al._ , Netket: a machine learning toolkit for many-body quantum systems, SoftwareX 10, 100311 (2019).
  * Häfner and Vicentini [2021] D. Häfner and F. Vicentini, mpi4jax: Zero-copy mpi communication of jax arrays, Journal of Open Source Software 6, 3419 (2021).
  * Vicentini  _et al._ [2022] F. Vicentini, D. Hofmann, A. Szabó, D. Wu, C. Roth, C. Giuliani, G. Pescia, J. Nys, V. Vargas-Calderón, N. Astrakhantsev, _et al._ , Netket 3: Machine learning toolbox for many-body quantum systems, SciPost Physics Codebases , 007 (2022).
  * Neuscamman  _et al._ [2012] E. Neuscamman, C. J. Umrigar, and G. K.-L. Chan, Optimizing large parameter sets in variational quantum monte carlo, Phys. Rev. B 85, 045103 (2012).
  * Batatia  _et al._ [2022] I. Batatia, D. P. Kovacs, G. Simm, C. Ortner, and G. Csányi, Mace: Higher order equivariant message passing neural networks for fast and accurate force fields, Advances in Neural Information Processing Systems 35, 11423 (2022).
  * Whitehead  _et al._ [2016] T. M. Whitehead, M. H. Michael, and G. J. Conduit, Jastrow correlation factor for periodic systems, Phys. Rev. B 94, 035157 (2016), arXiv:1607.05921 [cond-mat.str-el] .
  * Pudliner  _et al._ [1997] B. S. Pudliner, V. R. Pandharipande, J. Carlson, S. C. Pieper, and R. B. Wiringa, Quantum Monte Carlo calculations of nuclei with A $<$= 7, Phys. Rev. C 56, 1720 (1997), arXiv:nucl-th/9705009 .
  * Gandolfi  _et al._ [2020] S. Gandolfi, D. Lonardoni, A. Lovato, and M. Piarulli, Atomic nuclei from quantum Monte Carlo calculations with chiral EFT interactions, Front. in Phys. 8, 117 (2020), arXiv:2001.01374 [nucl-th] .

##  Appendix A Optimization

From the Rayleigh-Ritz principle, we obtain a lower bound on the expectation
value of the Hamiltonian
$\matrixelement{\Psi}{H}{\Psi}/\innerproduct{\Psi}{\Psi}\equiv E[\Psi]\geq
E_{0}$. We use the energy expectation value with respect to our variational
Ansatz as a cost function to gauge the proximity of the variational state to
the ground state of the Hamiltonian. Formally the ground state is given by:

$\Psi_{0}=\operatorname*{arg\,min}\limits_{\Psi}E[\Psi]\,.$ (15)

where $H\ket{\Psi_{0}}=E_{0}\ket{\Psi_{0}}$.

Since the exact expected value of the Hamiltonian requires us to analytically
solve a high-dimensional integral, which in general is not feasible, we resort
to Monte Carlo sampling and integration. The energy expectation is evaluated
as the average over a set of local energy
$E_{\mathrm{loc}}(\mathbf{X})=\matrixelement{\mathbf{X}}{H}{\Psi}/\bra{\mathbf{X}}\ket{\Psi}$,
where the samples $\mathbf{X}$ are obtained from the probability distribution
$|\Psi(\mathbf{X})|^{2}/\bra{\Psi}\ket{\Psi}$ using the Metropolis-Hastings
algorithm.

To update the variational parameters in our Ansatz such that it exhibits
progressively lower energy expectation values, we apply the stochastic
reconfiguration (SR) algorithm [31], which can be shown to be equivalent to
imaginary-time evolution on the variational manifold and is related to the
Natural Gradient descent method. The update rule for the variational
parameters is given by $\delta\boldsymbol{\theta}=-\eta S^{-1}\mathbf{F}$,
where $\mathbf{F}$ is the force vector consisting of the log-derivatives of
the energy w.r.t. the variational parameters $\boldsymbol{\theta}$ given by

$\displaystyle
F_{i}=2\left(\frac{\bra{\partial_{i}\Psi}H\ket{\Psi}}{\innerproduct{\Psi}{\Psi}}-E[\Psi]\frac{\innerproduct{\partial_{i}\Psi}{\Psi}}{\innerproduct{\Psi}{\Psi}}\right),$
(16)

and the so-called quantum geometric tensor $S$ is

$\displaystyle
S_{ij}=\frac{\bra{\partial_{i}\Psi}\ket{\partial_{j}\Psi}}{\innerproduct{\Psi}{\Psi}}-\frac{\bra{\partial_{i}\Psi}\ket{\Psi}\bra{\Psi}\ket{\partial_{j}\Psi}}{\innerproduct{\Psi}{\Psi}\innerproduct{\Psi}{\Psi}},$
(17)

and $\eta$ is the learning rate [59]. To regularize the computation of
$S^{-1}$, we apply a small diagonal shift such that we actually compute
$(S+\epsilon\mathbb{I})^{-1}$. For all experiments, we used
$\epsilon=10^{-4}$. The learning rate was chosen dependent on the density as
$\eta=\\{0.05,0.05,0.05,0.1,0.1,0.5,1,2.5\\}$ for
$r_{s}=\\{1,2,5,10,20,50,100,110\\}$.

##  Appendix B Message-Passing Neural Networks

In the following, we provide an overview on the general structure of message-
passing neural networks [36, 60] (MPNNs) and relate it to the MP-NQS. First
the relevant data has to be encoded into a graph structure, where the nodes of
the graph $\mathbf{g}_{i}^{(t)}=[\mathbf{x}_{i}^{(0)},\mathbf{h}_{i}^{(t)}]$
describes single-body information and the edges between nodes
$\mathbf{g}_{ij}^{(t)}=[\mathbf{x}_{ij}^{(0)},\mathbf{h}_{ij}^{(t)}]$ contain
relational information about the connected nodes. The graph
$(\mathbf{g}_{i}^{(t)},\mathbf{g}_{ij}^{(t)})$ then gets successively updated
over $T$ iterations, using a two-step message-passing procedure. In the first
step, the messages are typically constructed for all the nodes of the graphs,
but not the edges:

$\displaystyle\mathbf{m}_{i}^{(t+1)}=\sum_{j\in\mathcal{N}(i)}\mathbf{M}^{(t)}(\mathbf{g}_{i}^{(t)},\mathbf{g}_{j}^{(t)},\mathbf{g}_{ij}^{(t)})$
(18)

where $\mathbf{M}^{(t)}$ is a learnable function and the sum runs over all
nodes in the neighborhood of node $i$, denoted by $\mathcal{N}(i)$. The second
step then consists of updating the graph nodes, using the constructed
messages:

$\displaystyle\mathbf{g}_{i}^{(t+1)}=\mathbf{U}^{(t)}(\mathbf{g}_{i}^{(t)},\mathbf{m}_{i}^{(t+1)})$
(19)

where $\mathbf{U}^{(t)}$ is again a learnable function. In the MP-NQS, the
message construction is done via the particle attention mechanism.
Additionally the explicit dependence on the nodes of the graph is removed and
an update of the graph edges is added:

$\displaystyle\mathbf{m}_{ij}^{(t+1)}$
$\displaystyle=\mathbf{M}^{(t)}(\mathbf{g}_{ij}^{(t)})$ (20)
$\displaystyle\mathbf{m}_{i}^{(t+1)}$ $\displaystyle=\sum_{j\neq
i}\mathbf{m}_{ij}^{(t+1)}$ (21)

where $\mathcal{N}(i)$ is chosen to be all other particles, to retain
permutation equivariance. The nodes and edges of the graph are then updated as
follows:

$\displaystyle\mathbf{g}_{i}^{(t+1)}$
$\displaystyle=\mathbf{U}_{1}(\mathbf{g}_{i}^{(t)},\mathbf{m}_{i}^{(t+1)})$
(22) $\displaystyle\mathbf{g}_{ij}^{(t+1)}$
$\displaystyle=\mathbf{U}_{2}(\mathbf{g}_{ij}^{(t)},\mathbf{m}_{ij}^{(t+1)})$
(23)

where we introduced another two learnable functions
$\mathbf{U}_{1},\mathbf{U}_{2}$.

In traditional MPNN applications, the intermediate and final graphs’ nodes are
processed, in a so-called readout phase, to predict a scalar quantity, e.g.
the energy of an atom in a specific molecular configuration:

$\displaystyle E_{i}=\sum_{t}R^{(t)}(\mathbf{g}_{i}^{(t)})$ (24)

where $R^{(t)}$ is again a learnable function. In the MP-NQS, we replace this
step with a transformation of the final graph’s nodes to the physical
configuration space of the particles:

$\displaystyle\mathbf{{}\delta
r}_{i}=\mathbf{W}\cdot\mathbf{g}_{i}^{(T)}\in\mathbb{R}^{d}$ (25)

##  Appendix C Attention Mechanism

In this section, we summarize the original self-attention mechanism introduced
in [43] and provide a comparison to the particle attention mechanism
introduced in this work. The attention mechanism builds on the query/key/value
concept, notably recognized in search algorithms. Given a query vector (e.g.
search phrase), it gets compared to a set of key vectors associated to
candidates of a database (e.g. words on websites). The comparison of the query
and key provides a weight determining the relevance of a corresponding value
(e.g. website) given the query. Given (single-body) input data
$\\{\mathbf{x}_{1},...,\mathbf{x}_{N}\\},\quad\mathbf{x}_{i}\in\mathbb{R}^{D}$,
keys/queries/values are formally computed using learnable weight matrices
$W_{i}^{(Q)}\in\mathbb{R}^{D\times D_{k}}$, $W_{i}^{(K)}\in\mathbb{R}^{D\times
D_{k}}$, $W_{i}^{(V)}\in\mathbb{R}^{D\times D_{v}}$:

$\displaystyle\mathbf{Q}_{i}=W_{i}^{(Q)}\mathbf{x}_{i},\quad\mathbf{K}_{i}=W_{i}^{(K)}\mathbf{x}_{i},\quad\mathbf{V}_{i}=W_{i}^{(V)}\mathbf{x}_{i}$
(26)

To obtain permutation-invariant self-attention from the above, the weight
matrices cannot depend on the index $i$ but are the same independent of the
input data:

$\displaystyle\mathbf{Q}_{i}=W^{(Q)}\mathbf{x}_{i},\quad\mathbf{K}_{i}=W^{(K)}\mathbf{x}_{i},\quad\mathbf{V}_{i}=W^{(V)}\mathbf{x}_{i}$
(27)

Comparing query and key vectors is commonly performed using the (scaled) dot-
product along their feature dimension (with dimension $D_{k}$), such that the
self-attention map can be written as:

$\displaystyle\mathrm{\mathbf{Att}}(\mathbf{Q},\mathbf{K},\mathbf{V})=\mathrm{softmax}(\frac{\mathbf{Q}^{T}\mathbf{K}}{\sqrt{d}})\mathbf{V}\in\mathbb{R}^{N\times
D_{v}}$ (28)

where $\mathbf{Q},\mathbf{K},\mathbf{V}$ are matrices of stacked queries, keys
and values, respectively
($\mathbf{Q}=(\mathbf{Q}_{1},...,\mathbf{Q}_{N})\in\mathbb{R}^{N\times
D_{k}}$). The above can also be written as a weighted sum of value vectors,
where the weights are given by the inner product between queries and keys.

In the MP-NQS we develop and extend the idea of self-attention to two-body
input data
$\\{\mathbf{x}_{11},...,\mathbf{x}_{1N},...\mathbf{x}_{NN}\\},\quad\mathbf{x}_{ij}\in\mathbb{R}^{D}$.
As in self-attention, we utilize query/key weight matrices to compute
query/key pairs:

$\displaystyle\mathbf{Q}_{ij}$
$\displaystyle=W^{(Q)}\mathbf{x}_{ij}\in\mathbb{R}^{D_{k}}$ (29)
$\displaystyle\mathbf{K}_{ij}$
$\displaystyle=W^{(K)}\mathbf{x}_{ij}\in\mathbb{R}^{D_{k}}$ (30)

The comparison between queries and keys is again performed by taking their
inner product but importantly, along the _particle dimension_ rather than the
feature dimension. Additionally, we replace the softmax-activation with a
GELU-activation:

$\displaystyle\boldsymbol{\omega}_{ij}$
$\displaystyle=\mathrm{GELU}\left(\sum_{l}\mathbf{Q}_{il}\mathbf{K}_{lj}\right).$
(32)

where we introduced the weights $\boldsymbol{\omega}_{ij}$. Instead of using a
linear map for the values (as is done for the queries and keys), we use a full
MLP to compute values, $\mathbf{V}_{ij}=\boldsymbol{\phi}(\mathbf{x}_{ij})$
from the input data. We obtain the following attention mechanism, dubbed
particle attention:

$\displaystyle\mathrm{\mathbf{PAtt}}(\mathbf{Q},\mathbf{K},\mathbf{V})=\boldsymbol{\omega}_{ij}\odot\mathbf{V}_{ij}\in\mathbb{R}^{N\times
N\times D_{v}}$ (33)

where $\odot$ denotes element-wise multiplication along the feature dimension.
Using the above as messages for the MPNN, (21), results in permutation-
equivariant, weighted sums of edge features:

$\displaystyle\mathbf{m}_{ij}^{(t)}$
$\displaystyle=\mathrm{\mathbf{PAtt}}(\mathbf{Q},\mathbf{K},\mathbf{V})=\boldsymbol{\omega}_{ij}\odot\mathbf{V}_{ij}$
(34) $\displaystyle\mathbf{m}_{i}^{(t)}$ $\displaystyle=\sum_{j\neq
i}\boldsymbol{\omega}_{ij}\odot\mathbf{V}_{ij}$ (35)

##  Appendix D MP-NQS

We show in Figure 4, for a system of $N=14$ particles at the highest densities
$r_{s}=1,2,5$, that increasing the number of backflow iterations (and
therefore the number of variational parameters) systematically improves the
accuracy of our Ansatz. At $r_{s}=1$ we score slightly worse than WAPNet for a
single backflow iteration, while reaching the same accuracy for three backflow
iterations. Similarly, for $r_{s}=2$ we obtain a slightly higher energy than
WAPNet for a single backflow iteration, but we surpass WAPNet’s performance
after two backflow iterations. For $r_{s}=5$ we obtain more precise results
than WAPNet for a single backflow iteration. Upon comparison to FermiNet, our
results demonstrate that we consistently outperform the unrestricted, single
Slater determinant version of FermiNet at these densities. Interestingly, for
lower densities ($r_{s}=5$), our Ansatz even outperforms the unrestricted
FermiNet with 16 determinants with a single message-passing iteration.

Figure 4:  Energy obtained with our variational Ansatz for $N=14$ particles at
the densities $r_{s}=1,2,5$ as a function of the number of message-passing
backflow iterations. We compare to WAPNet and FermiNet with a single Slater
determinant as well as FermiNet with 16 determinants.

##  Appendix E Comparison to other NQS architectures

The most thoroughly investigated NQS architectures for continuous space are
FermiNet and PauliNet [18, 19]. Both have been shown to provide good
approximations to the ground-state energies of molecular systems. FermiNet has
additionally been applied to the HEG system with very accurate results [23].
Although FermiNet shows slightly better results on molecular systems than
PauliNet, differences in the architectures allow PauliNet to reach comparable
energies with considerably less parameters than FermiNet. Here we compare our
MP-NQS to these two Ansaetze and highlight similarities and differences.

Both FermiNet and PauliNet construct the variational wave function Ansatz by
using the antisymmetry of the Slater determinant. PauliNet uses physically
motivated orbital functions as input to the Slater determinant and alters
these by an orbital-aware multiplicative symmetric factor. Additionally, an
overall Jastrow factor is used to increase the expressiveness of the Ansatz.
FermiNet imposes no physical knowledge on the orbitals (up to an overall
envelope function) and does not need an overall Jastrow factor. To construct
the multiplicative factor in PauliNet and the orbitals in FermiNet, the number
of parameters scales linearly with the system size. In addition, both Ansaetze
construct multiple Slater determinants to increase their accuracy, thereby
multiplying the number of parameters by a constant factor. Both Ansaetze treat
the different spin species by utilizing distinct parameter sets for the
different possible spin projections.

Similar to PauliNet, we use physically motivated orbitals which are evaluated
at backflowed coordinates. These coordinates are constructed by using a novel
and powerful MPNN that leverages the expressiveness of the attention
mechanism. There exists yet another version of FermiNet, called PsiFormer,
that uses the original self-attention mechanism introduced in [43], but has
not been applied to extended systems. While PsiFormer applies the attention
mechanism along the _feature dimension_ of the single-particle stream, we find
substantially improved results by applying attention along the _particle
dimension_ instead. This approach introduces additional correlations between
the particles rather than between the features of the edge-states of our
graphs. As PauliNet, we use a multiplicative factor to alter the orbitals even
further, which allows us to omit the overall Jastrow factor that is present in
PauliNet. By inputting orbital information to the multiplicative factor
directly, we can keep the number of parameters independent of the number of
particles. Notably, our Ansatz does not require multiple Slater determinants.

As mentioned above, FermiNet has been applied to study the HEG in continuous
space. The orbitals are constructed from scratch with plane-wave envelope
functions. In contrast to our approach, FermiNet uses large networks (512
width) to process the positions of the simulated particles and considerably
smaller networks (32 width) to act on the distance vectors between the
particles. The usage of single-particle coordinates breaks the translation
invariance of the system. In our MP-NQS, we only use the distance information
between the particles and ignore the single-particle information to stay
translation invariant. We only use small networks (32 width) to construct the
messages $\mathbf{m}_{ij}^{(t+1)}$ (Eq. 21) that are used to update the nodes
and edges of the particle graph. Overall, we use around 19000 parameters in
our Ansatz which consists of a single Slater determinant. PauliNet usually
uses 70000-100000 parameters and up to 36 Slater determinants, while FermiNet
uses on the order of a million parameters with up to 32 Slater determinants.
The fact that we have considerably fewer parameters allows us to use SR to
optimize our variational Ansatz (see Appendix A). Both FermiNet and PauliNet
use $\mathcal{O}(10^{5})$ optimization steps, but we can reduce this number to
less than $\mathcal{O}(10^{3})$ iterations. To illustrate this point, we show
the optimization curve for $N=14$ particles at $r_{s}=5$ in Figure 5.

Figure 5:  Moving average of the optimization curve (energy per particle in
Hartree averaged over 100 optimization iterations) for $N=14$ particles at
$r_{s}=5$ obtained with a learning rate $\eta=0.05$ and a diagonal shift
$\epsilon=10^{-4}$ for SR. We compare to results from WAP-net [22] and
FermiNet [23] to our MP-NQS. Note that the physically motivation orbitals used
in our Ansatz makes us start close to the ground-state energy with our
optimization.

##  Appendix F Diffusion Monte Carlo

The diffusion Monte Carlo (DMC) method projects the variational states
$|\Psi_{V}\rangle$ in imaginary time to filter out its ground-state component

$|\Psi_{0}\rangle\propto\lim_{\tau\rightarrow\infty}e^{-H\tau}|\Psi_{V}\rangle\,.$
(36)

The variational state is taken to be of the Slater-Jastrow form
$\Psi_{V}=e^{J(\mathbf{X})}S_{\uparrow}S_{\downarrow}$, where $S_{\uparrow}$
and $S_{\downarrow}$ are the Slater determinants of single-electron states for
the spin-up and spin-down species, respectively. The Jastrow factor captures
correlation between electrons, and hence is a function of electron-electron
separation coordinates $\mathbf{r}_{ij}=\\{x_{ij},y_{ij},z_{ij}\\}$

$J(\mathbf{X})=\sum_{\begin{subarray}{c}i<j\\\
\sigma,\tau\in\\{\uparrow,\downarrow\\}\end{subarray}}J_{\sigma\tau}(\mathbf{r}_{ij})\,.\\\
$ (37)

The Pauli principle requires the functional form of $J_{\sigma\tau}$ must be
even under exchange of particles $i$ and $j$. In addition, $J_{\sigma\tau}$
must satisfy the periodic-boundary conditions at the edge of the simulation
box. To fulfill the above requirements, we utilize the parametrization
introduced in Ref. [61]

$J_{\sigma,\tau}(\mathbf{r}_{ij})=\sum_{n=1}^{N_{j}}c_{n,\sigma\tau}[j(x_{ij})^{2}+j(y_{ij})^{2}+j(z_{ij})^{2}]^{n/2}\,,$
(38)

where $j(x)=|x|[1-2(|x|/L)^{3}]$. By taking
$c_{1,\uparrow\uparrow}=c_{1,\downarrow\downarrow}=\frac{1}{4}$ and
$c_{1,\uparrow\downarrow}=c_{1,\downarrow\uparrow}=\frac{1}{2}$, the Kato cusp
conditions at particle coalescence is also automatically fulfilled. The
parameters $c_{n>1,\uparrow\uparrow}$ are determined by minimizing the energy
of the systems by using the “linear” optimization scheme [30]. We find that
using more than $N_{J}=6$ variational parameters does not improve the energy
noticeably.

The imaginary-time diffusion is broken into many small steps $\delta\tau$. The
Trotter-Suzuki decomposition can be applied to the short-time propagator as
$e^{-H\delta\tau}\simeq e^{-V\delta\tau/2}e^{-T\delta\tau}e^{-V\delta\tau/2}$,
where $T$ and $V$ are the kinetic and potential energy operators,
respectively. At each imaginary-time step, we use the free propagator

$\displaystyle G_{0}(\mathbf{X}^{\prime},\mathbf{X})$
$\displaystyle\equiv\langle\mathbf{X}^{\prime}|e^{-T\delta\tau}|\mathbf{X}\rangle$
$\displaystyle=\Big{[}\sqrt{\frac{m}{2\pi\hbar^{2}\delta\tau}}\Big{]}^{3N}\exp[-\frac{(\mathbf{R}^{\prime}-\mathbf{R})^{2}}{2\hbar^{2}\
\delta\tau/m}\Big{]}$ (39)

to sample the new coordinates $\mathbf{X}^{\prime}$ of all particles. As
routinely done in nuclear-physics applications [62, 63], to remove the linear
terms coming from the exponential of Eq. (39), we use two mirror samples
$\mathbf{X}^{\prime}=\mathbf{X}\pm\delta\mathbf{X}$, and we consider the
corresponding importance-sampled weights

$w_{\pm}=\frac{\Psi_{V}(\mathbf{X}\pm\delta\mathbf{X})}{\Psi_{V}(\mathbf{X})}e^{-[V(\mathbf{X}\pm\delta\mathbf{X})+V(\mathbf{X})]\delta\tau/2}\,.$
(40)

Only one of the two samples is kept in the propagation according to a heat-
bath sampling among the two normalized weights $w_{\pm}/(\sum_{\pm}w_{\pm})$
and the average weight $\sum_{\pm}w_{\pm}/2$ is assigned to the propagated
configuration.

The fermion-sign problem is controlled by employing the fixed-node
approximation, which amounts to evaluating the weights of Eq. (40) with the
replacement

$\frac{\Psi_{V}(\mathbf{X}\pm\delta\mathbf{X})}{\Psi_{V}(\mathbf{X})}\to\Re{\frac{\Psi_{V}(\mathbf{X}\pm\delta\mathbf{X})}{\Psi_{V}(\mathbf{X})}}\,.$
(41)

Note that if the real part of the above ratio is negative, the weight of the
configuration is set to zero only after computing the average weight.

##  Appendix G Results

In Tables 1, 2, 4 and 5 we show the ground-state energies of the unpolarized
and fully polarized HEG for different system sizes $N=\\{14,19,27,54\\}$ and
densities $r_{s}=\\{1,2,5,10,20,30,50,70,90,100\\}$.

N | $r_{s}$ | MP-NQS | WAP [22] | FermiNet [23] | FCI∗/DCD∗∗ [34]
---|---|---|---|---|---
14 | 1 | $0.568967(6)$ | $0.568965(1)$ | $0.568904(1)$ | $0.56861(1)^{\ast}$
| 2 | $-0.008391(1)$ | $-0.0083310(3)$ | $-0.008427(1)$ | $-0.00868(2)^{\ast}$
| 5 | $-0.0798544(4)$ | $-0.0798360(1)$ | $-0.079821(1)$ | $-0.08002(2)^{\ast}$
| 10 | $-0.0552126(6)$ | $-0.05520380(3)$ | N/A | $-0.05509^{\ast\ast}$
| 20 | $-0.0324553(2)$ | $-0.0324434(1)$ | N/A | $-0.03201^{\ast\ast}$
| 50 | $-0.01462631(6)$ | $-0.01462211(4)$ | N/A | $-0.01384^{\ast\ast}$
| 100 | $-0.00773018(3)$ | $-0.007729980(2)$ | N/A | N/A
Table 1: Total energy per particle in Hartree for unpolarized system of $N=14$ particles. WAPNet and FermiNet are alternative NQS architectures optimized via VMC. We include FCI and DCD results as benchmarks from quantum chemistry. N | $r_{s}$ | MP-NQS | MP-NQS (Gauss) | LiNet [27] | FN-DMC | BF-DMC [52, 53] | FCI∗/DCD∗∗ [34]
---|---|---|---|---|---|---|---
54 | 1 | $0.52973(1)$ | N/A | $0.530019(1)$ | $0.53094(2)$ | $0.52989(4)$ | $0.52973(3)^{\ast}$
| 2 | $-0.014046(8)$ | $-0.01390(1)$ | $-0.013840(1)$ | $-0.01326(2)$ | $-0.013966(2)$ | $-0.01379^{\ast\ast}$
| 5 | $-0.079090(2)$ | $-0.079064(4)$ | $-0.0788354(2)$ | $-0.07867(1)$ | $-0.079036(3)$ | $-0.07837^{\ast\ast}$
| 10 | $-0.054448(1)$ | $-0.054445(1)$ | $-0.0542785(1)$ | $-0.054269(8)$ | $-0.054443(2)$ | $-0.05322^{\ast\ast}$
| 20 | $-0.0320524(5)$ | $-0.0320480(6)$ | $-0.0316886(1)$ | $-0.031976(8)$ | $-0.032047(2)$ | $-0.03113^{\ast\ast}$
| 50 | $-0.0145015(1)$ | $-0.0144999(1)$ | N/A | $-0.01387(2)$ | $-0.0144877(1)$ | $-0.01281^{\ast\ast}$
| 100 | $-0.0076793(1)$ | $-0.00769203(5)$ | N/A | $-0.007674(3)$ | N/A | N/A
Table 2: Total energy per particle in Hartree for the unpolarized system of $N=54$ particles. FN-DMC results are obtained using the method in Appendix F. We include FCI and DCD results as benchmarks from quantum chemistry. N | $r_{s}$ | MP-NQS | MP-NQS (Gauss) | FN-DMC
---|---|---|---|---
128 | 110 | $-0.0070203(1)(1)$ | $-0.0071121(2)$ | $-0.007009518(1)$
| 200 | $-0.0040023(2)$ | $-0.0040615(1)$ | $-0.0039923(7)$
Table 3: Total energy per particle in Hartree for the unpolarized system of $N=128$ particles. FN-DMC results are obtained using the method in Appendix F. N | $r_{s}$ | MP-NQS | WAP [22]
---|---|---|---
19 | 1 | $1.046262(2)$ | $1.046241(3)$
| 2 | $0.096307(1)$ | $0.096303(1)$
| 5 | $-0.0672489(3)$ | $-0.06725105(3)$
| 10 | $-0.0528624(2)$ | $-0.05286035(1)$
| 20 | $-0.0320101(1)$ | $-0.0320082(1)$
| 50 | $-0.01456685(3)$ | $-0.01456571(2)$
| 100 | $-0.00771832(2)$ | $-0.00771362(2)$
Table 4: Total energy per particle in Hartree for the polarized system of
$N=19$ particles. We compare to results obtained using WAPNet.

N | $r_{s}$ | MP-NQS | FN-DMC
---|---|---|---
27 | 1 | $1.051771(4)$ | $1.05200(3)$
| 2 | $0.098460(2)$ | $0.09866(2)$
| 5 | $-0.0665523(6)$ | $-0.066429(4)$
| 10 | $-0.0526015(2)$ | $-0.052540(6)$
| 30 | $-0.0228184(1)$ | $-0.0228080(8)$
| 50 | $-0.01457976(5)$ | $-0.0145752(6)$
| 70 | $-0.01074986(3)$ | $-0.010749(2)$
| 90 | $-0.00853156(2)$ | $-0.0085314(5)$
Table 5: Total energy per particle in Hartree for the polarized system of
$N=27$ particles. FN-DMC results are obtained using the method in Appendix F.