# Simple Fermionic backflow states via a systematically improvable tensor
decomposition

Massimo Bortone<EMAIL_ADDRESS>Department of Physics, King’s
College London, Strand, London WC2R 2LS, United Kingdom  Yannic Rath
Department of Physics, King’s College London, Strand, London WC2R 2LS, United
Kingdom  National Physical Laboratory, Teddington, TW11 0LW, United Kingdom
George H. Booth<EMAIL_ADDRESS>Department of Physics, King’s College
London, Strand, London WC2R 2LS, United Kingdom

###### Abstract

We present an effective ansatz for the wave function of correlated electrons
that brings closer the fields of machine learning parameterizations and tensor
rank decompositions. We consider a CANDECOMP/PARAFAC (CP) tensor factorization
of a general backflow transformation in second quantization for a simple,
compact and systematically improvable Fermionic state. This directly encodes
$N$-body correlations without the ordering dependence of other tensor
decompositions. We consider and explicitly demonstrate various controllable
truncations, in the rank and range of the backflow correlations or magnitude
of local energy contributions, in order to systematically affect scaling
reductions to $\mathcal{O}[N^{3-4}]$. Benchmarking against small Fermi-Hubbard
and chemical systems reveals an improvement over other NQS-like models, while
extending towards larger strongly correlated ab initio systems demonstrates
competitive accuracy with more established DMRG techniques on ab initio 2D
hydrogenic lattices with realistic long-range Coulomb interactions.

##  I Introduction

Computational methods which can solve the physics of strongly correlated
electrons play an important role in the study of molecular and condensed
matter systems, where common perturbative or empirical density functional
approaches fail. In these systems, the interactions between electrons in some
or all of their degrees of freedom contend with the kinetic energy of the
electrons, leading to competition between localization and delocalization of
the electronic structure, and the emergence of many remarkable properties and
low-energy phases. Electronic structure poses a particular challenge amongst
the broader umbrella of quantum many-body problems [1] due to both charge and
spin degrees of freedom, as well as the requirement for antisymmetry which
significantly complicates the form of the solution. However, these features
are crucial in the understanding of the emergent physical behaviour in many
technologically relevant advanced materials, from high-temperature
superconductors to catalytic transition metal complexes.

A recent trend has emerged in the use of systematically improvable
parameterized ansätze for quantum states, which hold the promise of an exact
limit, providing confidence and the ability to internally validate results.
Naturally, the universal approximators devised in the field of machine
learning (ML) have inspired ansätze for this purpose, leading to the
development of Neural Quantum States (NQS) [2] with a wide range of network
architectures, from restricted Boltzmann machines [3, 4], to networks with
autoregressive [5], recurrent [6] or transformer architectures [7] of
differing depths and widths. Concurrently, approaches based on kernel methods
have also been considered [8, 9]. All these parameterized ansätze can in
principle approximate the complex functional dependencies between the
probability amplitudes of the electronic configurations, and have proven
capable of obtaining accurate results with minimal user intervention over a
variety of systems. Crucially, as long as the states can be sampled
efficiently and the Hamiltonian of the system remains $k$-local, then the
techniques of variational Monte Carlo (VMC) can stochastically optimize the
parameters of the model to minimize the energy of the system,
$E_{\theta}=\braket{\Psi_{\theta}}{\hat{H}}{\Psi_{\theta}}$, which then upper-
bounds the true ground state energy, $E_{0}$.

It should be stressed that the use of systematically improvable ansätze in
electronic structure is certainly not a new phenomenon with the emergence of
NQS. One of the most successful variational methods relies on tensor network
states, which provide an improvable tensor factorization of the many-body
amplitudes. For a one-dimensional network, the efficient contraction of these
amplitudes has led to the prominence of Matrix Product State (MPS)
descriptions of correlated systems. The single (hyper)parameter controlling
the expressivity of the model is the ‘bond dimension’, which can be quasi-
continuously enlarged to describe higher levels of entanglement towards a
complete model. Importantly, the simple structure of these states also allows
for new probes and insights into the emergent many-body physics of the model,
and is able to characterize entanglement measures and structures [10, 11].
Insights into the nature of the correlations and the entanglement can be
harder to quantify for NQS, where the diversity of different architectures and
model parameters can also cloud a clear path towards practical improvability
for the states, while it can also be unclear how to precisely design optimal
parameterizations.

In this work we return to tensor factorizations for a novel wave function
parameterization, yet within the framework of NQS. There has been much
research to indicate that the prodigious flexibility of complex NQS
architectures is not being fully exploited for many correlated problems, due
to the challenges in their optimization and initialization within the VMC
framework [12, 13]. Many simple parameterizations have performed as accurately
as more complex forms, and a premium is placed on the compactness of the
ansatz for ease of practical optimization alongside the overall flexibility.
The simpler form for these models can also potentially provide tools for
easier interpretability and improvability of the many-body physics, and open
avenues to alternative optimization strategies. An example of this is the
Gaussian process state (GPS), which was originally motivated via Bayesian
regression as a systematically improvable kernel model with a single model
parameter controlling the expressibility. It was shown to achieve similar
quality results to NQS, while often being more compact and open to novel
insights [14, 15, 9]. Further development exposed a duality of the GPS wave
function model to an exponential of a CANDECOMP/PARAFAC (CP) tensor-rank
decomposition of the wave function amplitudes in second quantization [8]. This
represented an interesting simplification of the model, and suggests further
developments to bring together the domains of NQS and tensor decompositions
for systematically improvable descriptions of correlated states.

This is the topic of this work in application to general Fermionic systems,
which have proven a particular challenge for NQS methods to date. Arguably the
most success for ab initio electronic structure with NQS has come in a first
quantized real-space formulation, where permutation-invariant multielectron
orbitals are parameterized as a neural network, before their product is
explicitly antisymmetrized as a determinant (or short linear combination
thereof) [16, 17]. However, in this work we consider a fixed basis and second
quantization. While this introduces a (necessarily incomplete) basis set
approximation, it also allows for more flexibility in the choice of model
(permutational invariance and antisymmetry are automatically enforced). This
also allows for problems to be defined in a finite and discrete space for the
stochastic sampling where additional approximations can be devised and
chemical insights from atomic orbital correlators are easily accessible.
Furthermore, this formulation allows for a straightforward treatment of core
electrons and direct comparison to established quantum chemical methods, as
well as natural application to multi-resolution and quantum embedding
methodologies [18, 19, 20]. The basis set approximation, in common with
traditional quantum chemical methods, is also much studied with a number of
approaches available which can substantially ameliorate it [21, 22, 23, 24].

Direct application of NQS-like ansätze in second quantization have often
struggled to clearly extend beyond state-of-the-art quantum chemistry, such as
coupled-cluster methods (CCSD) or exact diagonalization (FCI), with results
often restricted to small molecules and/or minimal basis sets [5, 4, 25, 26].
The commutation relations of second quantized operators enforced by a
necessarily unphysical choice of ordering of the degrees of freedom can induce
highly non-local and high-rank parity flips to the probability amplitudes. In
principle, these long-range structures can be described by NQS, but in
practice are very difficult to model and to appropriately optimize within VMC
frameworks, which have mainly been developed for quantum spin systems [4]. As
such, finding better Fermion to spin (qubit) mappings to reduce the rank or
range of these non-local parity changes is an active area of research [27, 28,
29, 30]. Alternatively, NQS-like states can be multiplied by an explicitly
antisymmetric state (e.g. Slater determinant, Pfaffian or antisymmetrized
geminal power) that will subsume much of the impact of these parity flips, at
the cost of potentially limiting the rigorous systematic improvability of the
resulting state. Nevertheless, this approach in combination with symmetry-
breaking and restoration has achieved impressive results in Fermionic models
[3, 8, 15].

Parallel to these developments, backflow transformations have been
parameterized via neural networks as an alternative approach to describe
Fermionic correlations in strongly interacting systems [31, 32]. This approach
modifies the single-electron functions of a Slater determinant (or other
antisymmetric function) to depend parametrically on many (potentially all $N$)
electron coordinates in a configuration-dependent way. A highly related
approach modifies the Slater determinant by coupling the physical degrees of
freedom to a set of configuration-dependent auxiliary or ‘hidden’ fermions,
which can similarly be parameterized as a neural network [33, 34]. The
parameterization of these backflow-type states has undergone a similar
development to other classes of variational wave functions, starting initially
from physically-motivated few-body parameterizations [35, 36, 37, 38], to a
more general ML architecture which allows (in principle) for systematic
improvability to exactness where each orbital and electron can arbitrarily
change based on all other electronic positions. These configuration-dependent
orbitals in backflow states have been defined by a number of different ML
architectures, both in real-space and discrete Fock space models, with and
without an additional Jastrow factor in the parameterization [16, 34, 39, 40].
They were shown to be effective in describing the ground states of Fermi-
Hubbard models [41, 33, 42], homogeneous electron gases [39], ultra-cold Fermi
gases [43] and (primarily in first quantization) ab initio molecular systems
[16, 44, 17, 45], achieving energies comparable or surpassing those from
Diffusion MC, as well as high accuracy coupled-cluster quantum chemical
methods.

Here, we consider a particularly simple CP tensor rank decomposition for these
configuration-dependent backflow orbitals, which allows for a straightforward
yet systematically improvable form for the overall state [46, 47]. We show in
Sec. II how this introduces explicit many-body correlated physics whilst
retaining antisymmetry of the state. Two previous approaches to combine tensor
decompositions with backflow parameterizations have been considered in the
literature; a backflow-inspired extension of the MPS ansatz [48] to build non-
local entanglement beyond the native MPS tensor ordering constraints for spin
systems, as well as a fixed tensor factorization of Fermionic backflow form
[42]. The latter study is more relevant to this work, but the particular
choice of tensor factorization did not admit a parameter by which the results
could be systematically improved, instead relying on a Lanczos step to improve
results. In contrast, the CP decomposition of the backflow parameterization
introduced in this work overcomes these issues, providing a simple and
improvable form which is invariant to orbital ordering.

We also develop a practical approach for ab initio systems to truncate the
length scale of the backflow correlations, providing a further compression of
the model with little loss of accuracy. We apply this variational ansatz to
find the ground state of (doped) Fermi-Hubbard models in Sec. III.1,
significantly outperforming a comparable neural network backflow
parameterization, as well as considering a challenging ab initio molecular
electronic structure problem in Sec. III.2. In common with other studies, we
find increasing the sampling of the VMC optimization important to improve
results, indicating that despite the simple form of the state it is still
challenging to optimize to the expressibility limit of the ansatz [4, 45].
Nevertheless, we still achieve a comparable level of accuracy of more complex
parameterizations. In Sec. III.3, we consider a $6\times 6$ 2D lattice of
hydrogen atoms as a step towards extended systems, with the CP backflow state
comparing favourably to state-of-the-art DMRG calculations and significantly
beyond the scope of exact approaches. Finally, in Sec. IV we discuss the
computational scaling of the method and approaches to reduce this as an
outlook towards larger systems and widespread application.

##  II Backflow determinants via CP tensor-rank decomposition

The wave function for a system of $N$ interacting electrons can be defined in
first quantization by assigning a unique label to each electron, and
introducing real-space $\mathbf{r}_{\alpha}\in\mathbb{R}^{3}$ and spin
$\sigma_{\alpha}\in\\{\uparrow,\downarrow\\}$ coordinates, so that
$\Psi(\mathbf{x})=\Psi(\mathbf{x}_{1},\dots,\mathbf{x}_{N})$, with
$\mathbf{x}_{\alpha}=(\mathbf{r}_{\alpha},\sigma_{\alpha})$. The simplest wave
function that satisfies the required antisymmetry is a Slater determinant of
$N$ single-particle spin-orbitals, $\phi_{i}(\mathbf{x}_{\alpha})$:

$\displaystyle\Phi_{0}(\mathbf{x}_{1},\dots,\mathbf{x}_{N})$
$\displaystyle=\frac{1}{\sqrt{N!}}\left\lvert\begin{matrix}\phi_{i}(\mathbf{x}_{1})&\phi_{j}(\mathbf{x}_{1})&\cdots&\phi_{k}(\mathbf{x}_{1})\\\
\phi_{i}(\mathbf{x}_{2})&\phi_{j}(\mathbf{x}_{2})&\cdots&\phi_{k}(\mathbf{x}_{2})\\\
\vdots&\vdots&\ddots&\vdots\\\
\phi_{i}(\mathbf{x}_{N})&\phi_{j}(\mathbf{x}_{N})&\cdots&\phi_{k}(\mathbf{x}_{N})\end{matrix}\right\rvert,$
(1)
$\displaystyle=\mathcal{A}[\phi_{i}(\mathbf{x}_{1})\phi_{j}(\mathbf{x}_{2})\dots\phi_{k}(\mathbf{x}_{N})],$
(2)

where $\mathcal{A}$ antisymmetrizes and normalizes the subsequent product of
orbitals with respect to exchange of their arguments. We can consider these
single-particle (molecular) orbitals as linear combinations of an underlying
basis (e.g. atomic orbitals, AOs) $\chi_{\mu}(\mathbf{r})$, as

$\phi_{i}(\mathbf{r})=\sum_{\mu=1}^{L}\varphi_{\mu i}\chi_{\mu}(\mathbf{r}),$
(3)

where $L$ is the size of this basis and $\varphi_{\mu i}$ are the coefficients
of the linear combination.

The key-idea of backflow ansätze is to extend the Slater determinant by
generalizing the single-particle orbitals to functions with non-linear
parametric dependencies on all electron coordinates. Historically this meant
transforming the electron coordinates $\mathbf{r}_{\alpha}$ with a new set of
coordinates
$\mathbf{r}^{bf}_{\alpha}=\mathbf{r}_{\alpha}+\sum_{\beta\neq\alpha}\eta(|\mathbf{r}_{\beta}-\mathbf{r}_{\alpha}|)(\mathbf{r}_{\beta}-\mathbf{r}_{\alpha})$,
where the function $\eta(\mathbf{r})$ describes the effective displacement of
the $\alpha$ electron due to the instantaneous position of the other electrons
[49]. This configurational-dependence on all other electron positions can also
be directly encoded into the variational parameters of a linear expansion of
single-particle orbitals, as first introduced by Ref. [31] for lattice models,
yielding a new set of backflow orbitals
$\phi^{bf}_{i}\left(\mathbf{r}_{\alpha};\\{\mathbf{r}_{/\alpha}\\}\right)$. In
an effort to improve the systematic description of these configuration-
dependent backflow orbitals, recent work has proposed to model
$\phi^{bf}_{i}\left(\mathbf{r}_{\alpha};\\{\mathbf{r}_{/\alpha}\\}\right)$
using neural networks [41, 16, 44], ensuring that the electron labels in
$\\{\mathbf{r}_{/\alpha}\\}$ are invariant under permutation to retain overall
antisymmetry.

Within a second quantization representation, the permutational invariance of
electrons and antisymmetry of the state is automatically ensured by the action
and commutation relations of the second quantized operators, independent of
the ansatz chosen. A Slater determinant can thus be obtained from the vacuum
state $\ket{0}$ by creating $N$ electrons in the corresponding single-particle
orbitals as

$\ket{\Phi_{0}}=\prod_{i=1}^{N}\hat{c}_{i}^{\dagger}\ket{0}=\prod_{i=1}^{N}\left(\sum_{\mu=1}^{L}\varphi_{\mu
i}\hat{c}_{\mu}^{\dagger}\right)\ket{0},$ (4)

where $\hat{c}_{\mu}^{\dagger}$ ($\hat{c}_{\mu}$) is now the operator that
creates (annihilates) an electron in the $\mu$-th basis state. To model
electron correlation, Eq. 4 can now be straightforwardly extended via analogy
to the backflow transformations by including in each orbital a parametric
dependence on the full instantaneous orbital occupation vector,
$\mathbf{n}=(n_{1},\dots,n_{L})$, where $n_{\mu}$ indexes instantaneous
occupancy of the four Fock states of spin-$\frac{1}{2}$ fermions in the chosen
orthonormal representation of degree of freedom $\mu$. This modifies the
creation operator of orbital $i$ to be

$\hat{c}^{\dagger}_{i}(\mathbf{n})=\sum_{\mu=1}^{L}\varphi_{\mu
i;\mathbf{n}}\hat{c}^{\dagger}_{\mu},$ (5)

resulting in an exact model, as each orbital can vary independently according
to the instantaneous occupation over the full state. However, it is of limited
use as it is an over-parameterization of the full state, with an exponential
number of variables. We therefore consider a specific tensor-rank
decomposition, the Canonical Decomposition (CANDECOMP) or Parallel Factor
(PARAFAC) decomposition (CPD)[46, 47]. This allows for a systematic and
improvable decomposition of this tensor for each orbital into a polynomial and
low-rank form that is independent of the choice of ordering of the degrees of
freedom defining the occupation vector, $\mathbf{n}$. The CP decomposition
factorizes the occupation number vector over all states of Eq. 5 into a sum of
$M$ tensor products, with each term in the product depending on each degree of
freedom in the full occupation number vector, as

$\varphi^{\mathrm{CPD}}_{\mu
i;\mathbf{n}}=\sum_{m=1}^{M}\prod_{\nu=1}^{L}\epsilon_{\mu i;n_{\nu}\nu m}.$
(6)

We now have a polynomially complex tensor of variational parameters for each
orbital, $\epsilon_{\mu i;n_{\nu}\nu m}$, which encodes the correlation-driven
modifications to orbital $i$ for the specific occupied degree of freedom
$\mu$, based on the fact that state $\nu$ has a local occupation of $n_{\nu}$.
$M$ represents an improvable parameter describing the systematic coupling of
the occupations across all possible occupation strings, providing an
increasingly flexible description of higher-rank correlations in the state
towards exactness. We denote this single parameter controlling the flexibility
of the model as its ‘support dimension’, by analogy with the CP decomposition
within Gaussian process states and kernel model definitions of quantum states
[14, 8, 9]. This CP decomposition splits the $L$-dimensional indices
indicating the $\mathbf{n}$-dependence of the orbital into a sum of products
of rank-3 tensors, depending on each orbital and its occupation. Since this is
a simple product rather than matrix product, there is no change in the
flexibility of these backflow orbitals with the ordering of the degrees of
freedom, ensuring that there should be no explicit dependence on this choice
(as found in tensor network states) or dimensionality of the system.

The proposed ‘CPD’ backflow wave function is obtained by replacing the
orbitals of the Slater determinant in Eq. 4 by those of Eq. 6, giving an
explicitly antisymmetric state where all orbitals depend on the instantaneous
occupation of all degrees of freedom,

$\Psi^{\mathrm{CPD}}=\sum_{\mathbf{n}}\Psi^{\mathrm{CPD}}(\mathbf{n})|\mathbf{n}\rangle,$
(7)

with

$\Psi^{\mathrm{CPD}}(\mathbf{n})=\mathcal{A}[\varphi^{\mathrm{CPD}}_{\mu_{1}i;\mathbf{n}}\varphi^{\mathrm{CPD}}_{\mu_{2}j;\mathbf{n}}\dots\varphi^{\mathrm{CPD}}_{\mu_{N}k;\mathbf{n}}],$
(8)

where the antisymmetrizer acts with respect to the $N$ occupied orbitals of
the configuration $\mathbf{n}$, given by $\mu_{1},\mu_{2}\dots\mu_{N}$. This
model can be evaluated naively via building a matrix and computing a
determinant in $\mathcal{O}[N^{2}ML+N^{3}]$ cost, with each orbital evaluated
according to Eq. 6. However, for low-rank changes to $\mathbf{n}$ where only
$\mathcal{O}[1]$ orbital occupations change, a fast updating scheme can be
devised to reduce the scaling in the matrix build by a factor of $L$. The
update for each orbital in Eq. 6 can be found in $\mathcal{O}[M]$ time by
dividing out contributions from the previous occupations and multiplying by
the new occupations, analogous to the approach in Ref. [15]. Since all
configurational updates in VMC can be formulated in this way, the evaluation
of configurational amplitudes of this CPD state can be reduced to
$\mathcal{O}[N^{2}M+N^{3}]$.

The total number of variational parameters in this state (which in this work
are all real) is therefore $\mathcal{O}[4L^{2}NM]$, where $L$ is the size of
the underlying basis, $N$ is the number of electrons, and $M$ the ‘support
dimension’ controlling the flexibility of the model. This scaling in terms of
the evaluation of the model and number of parameters allows the standard
techniques of VMC to be used for its sampling, optimization and extraction of
observables. We note that extending this CPD form to an explicit
antisymmetrization of geminal two-particle states within a Pfaffian or
antisymmetrized geminal power rather than single-particle orbitals of a
determinant would also be possible in this framework and will be explored in
the future [43, 50].

Unless otherwise indicated, in this work, we conserve a definite spin-
polarization quantum number for each of the $N$ orbitals labelled
$i,j,k,\dots$ in the product in Eq. 8, in which case the overall state must
conserve $\hat{S}_{z}$ symmetry. This is ensured by only allowing spin-orbital
degrees of freedom with the same spin-polarization to be included in the
expansion coefficients of the orbital (i.e. the $\mu$ labels in Eq. 6). This
allows the state for each $\mathbf{n}$ to factorize into a product of spin-up
and spin-down determinants (which are allowed to independently optimize,
analogous to an ‘unrestricted’ single determinant). An alternative (which is
considered in Sec. III.1) is to form a ‘generalized’ determinant by allowing
spin-polarization to mix in each orbital definition, formally breaking
$\hat{S}_{z}$ symmetry in the state. This symmetry is nevertheless restored
via the sampling of configurations with definite $\hat{S}_{z}$ in the Markov
chain during the VMC procedure. This $\hat{S}_{z}$ symmetry-breaking and
projective restoration can improve results by allowing further flexibility in
the state, but increases the cost in the evaluation of the determinant
defining the amplitude by a factor of eight, and doubles the number of
parameters (as $\mu$ labels spin-orbitals, not spatial orbitals). Importantly,
regardless of whether this spin symmetry is broken or not in the orbital
definition, the backflow correlations act both for same-spin and opposite-spin
correlations, with the orbital dependence in Eq. 6 running over the spin-full
occupations of all other degrees of freedom, $n_{\nu}$, ensuring that spin-
dependent correlated physics is captured.

Finally, we note that, although the functional form of the configuration-
dependent orbitals of Eq. 6 is linear in $M$, it does not reduce to an
uncorrelated determinant in the limit of $M=1$. Correlated physics such as
that captured via Gutzwiller or Jastrow correlators are included even in this
limit, since the dependence between the instantaneous occupation of sites
$\mu$ and $\nu$ can be independently addressed in a product form. Indeed, full
non-trivial $N$-body correlations are included even at $M=1$, as the
exponentially large sum of products of these orbitals formed from the
determinant in Eq. 8 builds in an exponential sum of these $N$-fold products
of variational parameters for each orbital. This results in a highly-
expressive state even for very low $M$, which is systematically improvable to
exactness as $M$ is increased.

###  II.1 Initialization

A practical bottleneck in working with parameterized quantum states with many
variational parameters can often be their reliable stochastic optimization.
This can be highly sensitive to the initialization of the state, and the
simple functional form of the CPD backflow orbitals in Eq. 6 allows for a
straightforward and effective initialization of the variational parameters.
Specifically, the tensor $\epsilon_{\mu i;n_{\nu}\nu m}$ can be initialized to
ensure that the CPD wave function exactly spans a given single determinant
such as that found from a prior mean-field solution. For most practical cases,
this provides a good starting point for the VMC optimization.

In this work, we initialize from a restricted Hartree–Fock state, extracting
the molecular orbital coefficients $\varphi^{\textrm{HF}}_{\mu i}$ in the
basis in which the state is to be sampled. The CPD variational parameter
tensor can then be initialized as follows:

$\epsilon_{\mu i;n_{\nu}\nu
m}=\mathcal{N}(0,\sigma)+\begin{cases}\varphi^{\mathrm{HF}}_{\mu i}&\text{if
}m=1\text{ and }\nu=1,\\\ 1&\text{if }m=1\text{ and }\nu>1\\\ 0&\text{if
}m>1\text{ and }\nu\geq 1\\\ \end{cases}$ (9)

where $\mathcal{N}(0,\sigma)$ is a random number drawn from a normal
distribution with standard deviation $\sigma$. This small amount of random
noise is optional, but is added to the initialization in case the Hartree–Fock
solution is too close to a local minimum of the optimization surface.

###  II.2 Backflow truncation via exchange cutoff

While the CPD backflow state only has a polynomial number of parameters, the
$\mathcal{O}[4L^{2}NM]$ scaling is still significantly higher than the native
non-backflow (e.g. GPS) state, and there is are significant benefits in
attempting to reduce this further with a controllable compromise on the
flexibility of the state. Largely redundant parameters in VMC add to
statistical noise without improving accuracy and can be particularly
deleterious in the optimization of the state [51, 52]. In particular, the
scaling with respect to the underlying basis size ($L$) is quadratic in the
CPD state, and in this section we motivate a physical and black-box truncation
of this scaling to further improve the overall performance of the state and
enable access to larger systems.

We do this by restricting the number of degrees of freedom that the backflow
parameterization considers for each $\mu$-indexed site, reducing it from $L$
to a new parameter $K$. This can be motivated as a range-truncation of the
backflow correlations, as has also been considered in other truncated
expansions [32]. If application of this methodology was purely to local
lattice models, then strictly truncating by a distance criteria would likely
be sufficient to capture the dominant correlations. However, we intend the
methodology to be applied equally across lattice models and ab initio systems
and therefore seek an alternative proxy to define the choice of entangled
orbital subspace in which these backflow correlations are defined for each
degree of freedom. This is because an ab initio basis will necessarily be
extended in space and perhaps not even able to be uniquely associated with an
atomic centre. Additionally, the inclusion of the long-range Coulomb
interaction in these systems does not necessarily favour purely distance-based
criteria. We therefore take inspiration from ab initio formulations of the
Density Matrix Renormalization Group (DMRG), where heuristics for the
entanglement between two orbitals are necessary in order to find an
approximately optimal ordering of the extended orbitals for an effective MPS
ansatz. While there are a number of options in the literature, it has been
found that the importance of one orbital in describing the dominant
correlations with another can be reasonably quantified by the magnitude of the
exchange integral between them [53], as

$\mathcal{K}_{\mu\nu}=\int\int
d\mathbf{r}_{1}d\mathbf{r}_{2}\chi_{\mu}^{*}(\mathbf{r}_{1})\chi_{\nu}^{*}(\mathbf{r}_{1})\frac{1}{|\mathbf{r}_{1}-\mathbf{r}_{2}|}\chi_{\mu}(\mathbf{r}_{2})\chi_{\nu}(\mathbf{r}_{2}).$
(10)

This exchange-based metric should decay exponentially between localized
orbitals, tending towards a flexible locality based truncation in the limit of
fully local orbitals, while including the full range of the Coulomb
interaction in the kernel. More rigorous definitions of entanglement between
orbitals such as their mutual information (pair entanglement entropy) [54]
could also be used, but require an initial correlated level of theory on which
to build these metrics. Since we initialize the CPD backflow molecular
orbitals from a Hartree-Fock calculation, the exchange matrix
$\mathcal{K}_{\mu\nu}$ is readily available for no additional cost. The set of
$K$ most entangled orbitals for each orbital $\chi_{\mu}(\mathbf{r})$
according to this metric are selected, defining an $L\times K$ lookup table
which maps to the relevant orbital indices $x_{\mu\nu}\in\\{1,\dots,L\\}$. The
choice of orbitals in the CPD decomposition of Eq. 6 are therefore restricted
as

$\varphi_{\mu i;\mathbf{n}}=\sum_{m=1}^{M}\prod_{\nu=1}^{K}\epsilon_{\mu
i;n_{x_{\mu\nu}}\nu m},$ (11)

thus reducing the number of variational parameters to $\mathcal{O}(LKMN)$ and
formally linear with the size of the system, assuming that $K$ is sufficiently
large to capture the range of correlations around each degree of freedom. As
$K$ tends to $L$, the state returns to the original definition (albeit with an
inconsequential reordering of sites in the backflow) giving the full
flexibility of backflow correlations.

(a) $d=1.0$ Å

(b) $d=2.0$ Å

Figure 1: Electron density of the orbital subspace chosen via the exchange
truncation in the range and rank of explicit backflow correlations. We
consider the $K=5$ subspace found for a central atom in a $6\times 6$ lattice
of hydrogen atoms in a STO-6G basis at both compressed and extended
geometries.

To illustrate the action of this exchange cutoff heuristic, in Figure 1 we
consider the electron density of the $K=5$ most entangled orbitals about a
specific atom for a $6\times 6$ square grid of ab initio hydrogen atoms in a
Boys localized basis [55] at two different interatomic distances, $d$. As
expected for a local representation, the exchange cutoff reduces to selection
of the nearest-neighbor atomic-localized orbitals around the chosen hydrogen
atom in both geometries considered, providing a black-box metric to select the
backflow subspace of correlations for each orbital via exploitation of
locality of these correlations. We note again that the product structure of
the CPD ansatz will build longer-ranged and higher-rank correlations outside
the chosen subset implicitly, albeit no longer explicitly for each orbital
independently.

##  III Results

In all results below we initialize the CPD backflow state from the restricted
Hartree–Fock solution as outlined in Eq. 9, with a noise scale value
$\sigma=0.01$. We optimize parameters using the Stochastic Reconfiguration
(SR) method [56], and when the number of parameters is larger than that of the
samples, we take advantage of the recently introduced kernel formulation from
Ref. [57] to improve the computational cost of the optimization, as outlined
further in Section IV. On the Fermi-Hubbard model, we found that a SR
optimizer with RMSProp momentum, as introduced in Ref. [58], outperforms
standard SR, and we therefore use this optimizer for the results presented in
Section III.1. The final energies presented in the results are computed as
averages over 50 independent energy evaluations with the final optimized
parameters and a large sample size ($2^{16}$ for the Fermi-Hubbard model and
the water molecule, and $2^{14}$ for the $6\times 6$ hydrogen lattice). Error
bars are computed as the standard deviation of these independent energy
evaluations.

The VMC calculations are implemented in the NetKet package [59, 60], which we
interface with our own plugin module, GPSKet for the required custom
functionality. For ab initio systems, Hartree–Fock orbital coefficients and
Hamiltonians are supplied from PySCF [61, 62].

###  III.1 Fermi-Hubbard model

While the main ambition of this work is to apply the newly developed CPD
backflow ansatz to ab initio systems, we first consider a small Fermi-Hubbard
model on a 2D square lattice as a prototypical system for strongly correlated
electrons, where comparison to exact results and neural-network parameterized
backflow states from the literature are both available. The Hamiltonian for
this system is defined as

$\hat{H}=-t\sum_{\langle
i,j\rangle,\sigma}\hat{c}_{i,\sigma}^{\dagger}\hat{c}_{j,\sigma}+U\sum_{i}\hat{n}_{i,\uparrow}\hat{n}_{i,\downarrow},$
(12)

where $\hat{c}_{i,\sigma}^{\dagger}$ ($\hat{c}_{i,\sigma}$) is the operator
that creates (annihilates) a fermion with spin $\sigma$ on site $i$,
$\hat{n}_{i,\sigma}=\hat{c}_{i,\sigma}^{\dagger}\hat{c}_{i,\sigma}$ is the
number operator, $t$ is the hopping amplitude, and $U$ is the on-site
interaction strength. We apply the CPD backflow state, allowing for spin-
polarization breaking and restoration of the orbitals as described in Sec. II,
in the strong interaction regime at $U/t=8$ on a $4\times 4$ lattice with
periodic boundary conditions, at half-filling ($n=N/L=1.0$) and in the hole
doped case ($n=0.875$). This hole-doped case is of particular interest as the
point at which superconductivity and striped orders strongly compete and is
much debated in the literature to date [63, 64]. We compare our results with
those obtained by backflow ansätze based on neural networks (NNB) with similar
numbers of parameters (${\sim}35,000$) taken from Ref. [41] as well as exact
diagonalization (ED) [65].

Figure 2: Percentage relative energy error for the ground state of the
$4\times 4$ square Fermi-Hubbard model at $U/t=8$ compared to exact
diagonalization results [65] at a hole-doped filling of $n=0.875$ (upper
panel) and half-filling (lower panel). CPD backflow results ($\Psi^{CPD}$) are
shows as a function of the configurational sample size in the optimization of
the parameters, for two different model complexities of $M=1$ and $M=2$.
Neural network backflow results are taken from Ref. [41].

In Fig. 2 we show the percentage relative energy error compared to exact
diagonalization for the CPD state of this system, plotted against the number
of samples used in the Markov chain for each update of the parameters in the
SR steps. Our results significantly improve upon the comparable published
neural-network backflow results for this system, even when these are
extrapolated with respect to the complexity of the network architecture in the
NNB ansatz. We find percentage relative errors as low as $0.5\%$ for the doped
case and $0.1\%$ for the half-filled case, which is competitive and within the
scatter of other state-of-the-art techniques in the literature for this
correlation regime [66], albeit with this system too small to be compared in
the thermodynamic limit.

We also show the variational improvability as the support dimension $M$ of the
CPD decomposition is increased from $M=1$ to $M=2$, with a systematic lowering
of all energies found, leading to a maximum of $65,536$ parameters.
Nevertheless, we unfortunately find that it is still generally more
advantageous to increase the number of configurational samples in the Markov
chain than to formally increase the flexibility of the state by increasing
$M$. This is due to noise in the estimates of the expectation values required
for the optimization of the CPD parameters, which amongst other things affects
the inversion of the sampled quantum geometric tensor. This indicates that we
cannot be confident of a complete optimization to the global minimum of this
state, despite the simple parameterization of the CPD form, with the
optimization still limited more by noise in the samples than flexibility in
the model, as found in many other studies of comparable states. However, we do
find a reliable a systematic improvement in the optimized state as this number
of samples is increased, and a high level of accuracy overall for this
correlated state.

###  III.2 Water molecule

We now consider ab initio molecular systems, which are described in second
quantization by an electronic Hamiltonian of the form

$\hat{H}=\sum_{ij,\sigma}h^{(1)}_{ij}\hat{c}_{i,\sigma}^{\dagger}\hat{c}_{j,\sigma}+\frac{1}{2}\sum_{ijkl,\sigma\tau}h^{(2)}_{ijkl}\hat{c}_{i,\sigma}^{\dagger}\hat{c}_{j,\tau}^{\dagger}\hat{c}_{l,\tau}\hat{c}_{k,\sigma},$
(13)

where the sums run over the degrees of freedom in the system and $\sigma,\tau$
are binary spin variables. The $h^{(1)}_{ij}$ matrix elements describe the
kinetic energy operator and interaction with the external potential in these
degrees of freedom, while the $h^{(2)}_{ijkl}$ terms model the Coulomb
interaction between particles. Compared to Fermi-Hubbard models, the
computational complexity of these Hamiltonians is significantly increased by
the $N^{2}(2L-N)^{2}$ scaling of the connected configurations required in
evaluating the local energy (compared to $\mathcal{O}[N]$ terms in Hubbard and
other lattice models). Since the evaluation of the CPD wave function model at
each configuration is $\mathcal{O}[N^{2}M+N^{3}]$ with the fast update (see
Sec. II) this constrains the number of configurational samples that can be
afforded.

For our initial benchmark system, we consider the water molecule in the 6-31G
basis set at the equilibrium geometry used in Ref. [4]. While this seems an
unassuming system from an electron correlation perspective, it has emerged as
somewhat of a benchmark system in the Quantum Monte Carlo (QMC) community,
where it has been studied extensively using a variety of ansätze [51, 67, 68].
Recently developed NQS architectures have struggled to reach state-of-the-art
accuracy for this molecule, despite it still being of a size where exact
diagonalization is possible. Part of the issue with this comes from the fact
that the weakly correlated physics and compact nature of the molecule mean
that it is hard to define an appropriate representation for the basis which
can enable efficient, faithful and representative sampling of the state with
few configurational samples.

Minimizing the number of samples relies on finding a representation which can
be faithfully approximated by a small stochastic selection of configurations,
necessitating an orbital representation of the basis in which the wave
function amplitudes are as flat as possible throughout the Hilbert space. This
maximizes the acceptance rates of the Metropolis-Hastings Markov chain growth,
and ensures that as small a sample as possible can represent the wave function
distribution. Canonical bases of mean-field (e.g. Hartree–Fock) theories are
therefore particularly poorly suited, as they are (away from very strong
correlation) dominated by the configuration of a single Slater determinant.
These bases have been found for NQS with restricted Boltzmann machine
architectures to obtain relatively large correlation energy errors ($\approx
5-10\%$) despite scaling up to $10^{6}$ configurational samples [4]. The
development of autoregressive NQS models has been able to improve upon this by
allowing a direct sampling algorithm of unique configurations that is not
constrained by the limitations of the Metropolis-Hastings algorithm [5].
However, these models still require large sample sizes and have only been
benchmarked in a STO-3G minimal basis set for this system [5, 26, 25]. Rather
than changing the sampling algorithm, in a previous work, we considered the
effect of different orbital representations for the configurational basis
[15]. Following this, we consider orthogonal Foster-Boys orbitals for the
configurations, localized over all degrees of freedom to minimize the physical
spread of the resulting orbitals [55].

Figure 3: Relative ground state correlation energy error (compared to exact
diagonalization) for CPD backflow ansätze ($\Psi^{CPD}$) on the water molecule
(6-31G basis, equilibrium geometry as specified in Ref. [4]) as the number of
configurational samples in each Markov chain is increased. Two support
dimensions corresponding to $M=1$ and $M=4$ are shown. Gaussian process state
(GPS) and restricted Boltzmann machine (RBM) comparison energies are taken
from Refs. [15] and [4], respectively, while the CCSD energy is calculated
with PySCF [61, 62].

The results of the CPD backflow state for different number of configurational
samples in this system are shown in Fig. 3. Once again, we find that
increasing the number of samples can affect the accuracy more significantly
than increasing the complexity of the model, showing results going from $M=1$
($6.7k$ parameters) to $M=4$ ($27k$ parameters). Nevertheless, we outperform
the accuracy of the native RBM state ($728$ parameters) with up to $10^{6}$
samples [4], as well as the GPS multiplied by a symmetry-broken Pfaffian state
($585$ parameters) with about $10^{4}$ samples [15], approaching the accuracy
of coupled-cluster with singles and doubles level of theory which excels for
weakly correlated equilibrium molecular systems such as this. While we still
don’t quite reach this level (and the significant ‘chemical accuracy’ hurdle –
albeit defined with respect to the finite basis set energy), to the best of
our knowledge, these results represent the state-of-the-art for an NQS-like
variational ansatz for this system, reaching a relative correlation energy
error of $\approx 2\%$ with support dimension $M=4$ and using a comparatively
small number of $\mathcal{O}[10^{4}]$ samples.

###  III.3 Towards hydrogen materials

Extending the CPD backflow ansatz beyond benchmark studies and comparison to
exact results, we consider a two-dimensional ab initio lattice of hydrogen
atoms as a step towards combining strong correlation, long-range interactions
and extended systems. These hydrogenic systems have been studied by a variety
of methods in the recent years given their simple specification and challenge
of realistic interactions, whilst maintaining a close connection to the Fermi-
Hubbard model [69, 70, 71, 72, 73, 15, 74]. In particular, different
correlation regimes can be probed by simply changing the interatomic distance
of the lattice, similar to tuning the interaction strength in the Fermi-
Hubbard model. However, crucially these hydrogen lattices require the accurate
treatment of realistic long-range Coulomb interactions and their effects,
which are not present in the Fermi-Hubbard model. Accurately capturing the
ground state of these systems for different interatomic distance is thus a
challenging task for most quantum chemistry methods, as it requires a flexible
and expressive model with a treatment of long-range interactions and high-
energy scattering physics that gives rise to states of significantly different
character.

Figure 4: Ground state energy per atom of a $6\times 6$ square hydrogen
lattice in a STO-6G basis for increasing lattice constants. Shown (top) are
RHF (orange) and UCCSD energies (violet), while the horizontal grey line
indicates the exact energy of the fully dissociated limit in this basis. The
CPD backflow ansatz ($\Psi^{CPD}$) with support dimension $M=1$ (solid black
line) is compared to a GPS ansatz augmented with a Slater determinant
($\Psi^{GPS}\times|\Phi\rangle$) with support dimension $M=72$ [15] (dotted
black line), as well as energies obtained with DMRG (red dots and line) using
a bond dimension of $M=1024$. A further CPD backflow ansatz with support
dimension $M=1$ and a truncation in the backflow subspace to $K=5$ orbitals is
shown as a dashed black line. The backflow CPD (GPS augmented by Slater
determinant) VMC ansätze were optimized with $4096$ ($10,000$) samples, while
DMRG results were obtained using the block2 package [75, 76]. The bottom panel
shows energy difference relative to DMRG of the CPD backflow ansätze and the
GPS times Slater determinant ansatz.

In Figure 4, we report the ground state energy per atom obtained from the CPD
backflow ansätze (with and without the rank and range truncation of Sec. II.2
for the backflow transform to $K=5$ sites) for a $6\times 6$ hydrogen lattice
in a minimal basis (STO-6G). We consider both compact (lower effective $U/t$)
and extended (higher effective $U/t$) lattice structures by varying the
interatomic distances all the way to essentially dissociated non-interacting
hydrogen atoms. We compare our results to energies obtained with restricted
Hartree-Fock (RHF) and unrestricted coupled-cluster with single and double
excitations (UCCSD), as well as an efficient ab initio implementation of
density matrix renormalization group (DMRG) going up to bond dimension of
$1024$ in a fully spin-adapted basis implemented in the block2 package [75,
76]. For an additional comparison between contrasting approaches to Fermionic
variational wave functions, we also include the results obtained from a GPS
ansatz with support dimension $M=72$ acting as a Jastrow in front of a co-
optimized Slater determinant [15], to compare the CPD backflow to this
approach. We optimize the CPD backflow and GPS multiplied by Slater
determinant ansätze in a Boys localized basis for the orbitals, whereas for
the DMRG results we rely on a split-localized basis, in which occupied and
virtual orbitals are localized separately. We found this choice to give the
most consistent results for DMRG across the range of geometries studied.

The RHF and UCCSD description of this equation of state qualitatively breaks
down quite early in this stretching coordinate, with UCCSD failing to converge
beyond $1.5~{}\text{\AA}$. Furthermore, the UCCSD exhibits quantitative error
of $\sim 2~{}\text{mE}_{h}$ per atom even around equilibrium geometries,
confirming that substantial correlation effects are present even in this
regime. As another point of reference, the fully dissociated limit can be
computed via exact diagonalization, where the assumption of simple energy
extensivity from a single atom can be applied. In this limit the energy is
$\approx 0.03~{}\text{E}_{h}$ above the analytic result for the hydrogen atom
due to the basis set incompleteness error, which nevertheless will exhibit a
large degree of cancellation for energy differences along this changing
geometry. The DMRG provides the best variational comparison for this system,
with (apart from $2.5~{}\text{\AA}$) the CPD and GPS results being within
$2~{}\text{mE}_{h}$ per atom of this value.

The CPD backflow ansatz manages to quantitatively capture the features of the
expected potential energy surface, reaching the correct dissociation limit at
large interatomic distances, and showing an overall smooth transition from
weak to strong correlation regimes. We can directly compare different
systematically improvable variational ansätze (DMRG, GPS multiplied by a
Slater determinant and the CPD backflow), all of which are competitive and
variationally optimal at different points in the changing physics of this
system. Around the equilibrium of this system, the CPD backflow and DMRG
states are almost identical and variationally optimal amongst the comparison.
In the intermediate regime ($1.5\leq d<2.0~{}\text{\AA}$), the GPS ansatz
augmented with a Slater determinant provides the best variational energies,
despite (or perhaps because of) the smaller number of parameters ($\approx
12k$ vs. $\approx 187k$ for the CPD backflow ansatz without truncation and
$\approx 26k$ for the one with). An outlier appears to be the nearly
dissociated limit of $2.5~{}\text{\AA}$ interatomic distance, where the DMRG
energy appears erroneously high. This could be due to a particularly large
impact on the one-dimensional MPS topology used, the choice of basis for the
orbitals or the DMRG sweep getting stuck in a local minimum. Nevertheless, the
other variational ansatz largely agree at this point.

Comparing the CPD backflow curves with and without the backflow truncation, we
find (as expected) that the $K=5$ results are all variationally higher than
the parent CPD backflow. This truncation has a very small effect on the
energies at larger interatomic distances, but becomes more significant around
the equilibrium distance and mildly stretched geometries where it reaches a
maximum error of $2~{}\text{mE}_{h}$ per atom. This is expected as the range
of the correlations in the compressed lattice will extend further than the
stretched limit. Nonetheless, even with this restriction the ansatz is able to
reach the coupled-cluster level of accuracy around equilibrium, and to
outperform it on stretched geometries, with a reduction in the number of
parameters compared to the parent model by more than a factor of seven. This
validates the exchange cutoff as a practical parameter reduction scheme for
the CPD backflow ansatz, suggesting benefits in the study of larger systems,
and potentially allowing for an increase in the support dimension of the
model.

Method |  $D_{e}$ (eV) |  $r_{e}$ (Å) | $w_{e}~{}(\text{cm}^{-1})$ | $w_{e}x_{e}(\text{cm}^{-1})$
---|---|---|---|---
UCCSD | 0.163 | 1.225 | 2177.988 | 9.007
DMRG | 0.031 | 1.223 | 1900.038 | 36.409
$\Psi^{\textrm{CPD}}$ | 0.029 | 1.219 | 1928.083 | 39.672
$\Psi^{\textrm{CPD}}(K=5)$ | 0.028 | 1.229 | 1865.462 | 38.586
$\Psi^{\textrm{GPS}}\times\ket{\Phi}$ | 0.029 | 1.237 | 1748.939 | 33.247
Table 1: Dissociation energy ($D_{e}$), equilibrium bond length ($r_{e}$),
harmonic ($\omega_{e}$) and anharmonic ($\omega_{e}\chi_{e}$) frequencies in
wavenumbers for the $6\times 6$ hydrogen lattice obtained by fitting a Morse
potential to data from different quantum chemistry and variational methods
described in this work.

To further compare the physical properties of the potential energy surface of
this lattice as described by the different levels of theory, we fit a simple
Morse potential at different interatomic distances ($r$), given by

$V(r)=D_{e}\left(1-e^{-a(r-r_{e})}\right)^{2}+u,$ (14)

where $D_{e}$ and $a$ control the depth and width of the well, $r_{e}$ is the
equilibrium bond length, and $u$ is the energy offset. Although the Morse
potential is generally used for diatomic molecules, the symmetric stretching
coordinate of this system is nevertheless well modelled by this form. The
differences in the quantum chemistry and VMC methods used to obtain the
potential energy data are reflected in the variations of dissociation energy
($D_{e}$), equilibrium bond length ($r_{e}$), harmonic vibrational frequency
($\omega_{e}$), and anharmonicity constant ($\omega_{e}\chi_{e}$) presented in
Table 1.

UCCSD is expected to accurately describe the correlated physics near
equilibrium geometries, however the rapid divergence after this point renders
even the harmonic vibrational frequencies unreliable. In contrast, DMRG, which
handles both the strong and weak correlations on a consistent level, presents
a more accurate dissociation energy ($D_{e}=0.031~{}\text{eV}$) and harmonic
vibrational frequency ($\omega_{e}=1900.038~{}\text{cm}^{-1}$), while yielding
anharmonicity of the vibrational motion of the atomic lattice describing the
beyond-parabolic nature of the binding as
$\omega_{e}\chi_{e}=36.409~{}\text{cm}^{-1}$. The values obtained from CPD
backflow ansätze with and without truncation closely track those from DMRG. On
the other hand, while agreeing on the dissociation energy, the GPS multiplied
by Slater determinant ansatz stands out amongst the variational methods with a
marginally softer bond, with a larger equilibrium lattice parameter
($r_{e}=1.237~{}\text{\AA}$) and the lowest harmonic vibrational frequency
($\omega_{e}=1748.939~{}\text{cm}^{-1}$). Overall, the methods agree on an
equilibrium bond length around $1.22~{}\text{\AA}$ and a dissociation energy
of $0.03~{}\text{eV}$ (except UCCSD), with variations in the harmonic and
anharmonic wavenumbers.

These variations highlight the strengths and limitations of each method in
modelling the potential energy surfaces and vibrational properties of hydrogen
materials. Taking all these results into consideration, the CPD backflow
ansatz emerges as a competitive method for the study of strong electron
correlation, providing a variational description of the ground state of a two-
dimensional lattice of hydrogen atoms that is in good agreement with other
state-of-the-art methods, while being able to capture strong correlations and
anharmonic effects in the system in a low-energy basis.

####  III.3.1 Spin-spin correlations

The local atomic basis framework of the CPD backflow ansatz allows for the
straightforward computation of atom-resolved expectation values for further
insights into the electronic structure. In the context of hydrogen materials,
local spin-spin correlation functions are of particular interest, as they can
provide insights into the nature of the ground state of the system, and the
emergence of magnetic order. By analogy with Hubbard models, we would expect
some anti-ferromagnetic order to emerge in the electronic structure of this
system, with this order decaying algebraically in the thermodynamic limit.
However, in the presence of long-range interactions this behaviour is far from
confirmed in two-dimensions. While admittedly far from this thermodynamic
limit, we consider the two-point spin-spin correlation function $C(r)$ between
the centre of the $6\times 6$ hydrogen lattice, and atoms at a distance $r$
from the centre. We can define this function via instantaneous (equal-time)
spin-spin correlators
$\left\langle\hat{S}^{z}_{\vec{r}_{a}}\hat{S}^{z}_{\vec{r}_{b}}\right\rangle$
between two atoms as

$C(r)=\frac{1}{N_{bulk}}\sum_{\vec{r}_{a}\in\text{bulk}}\sum_{|\vec{r}_{a}-\vec{r}_{b}|=r}\left\langle\hat{S}^{z}_{\vec{r}_{a}}\hat{S}^{z}_{\vec{r}_{b}}\right\rangle,$
(15)

where $\vec{r}_{a}$ and $\vec{r}_{b}$ are the positions of atoms $a$ and $b$,
and $N_{bulk}$ is the number of equivalent atoms in the centre that we average
over (four). We use the atom-centered atomic orbitals themselves as natural
projectors for the spin operators of each atom. More details about the
calculation of the instantaneous spin-spin correlation function can be found
in Appendix A.

Figure 5: Two-point instantaneous spin-spin correlation function for the
ground state of the $6\times 6$ hydrogen lattice between atoms at two
different interatomic distances $d=1.2$ Å and $d=3.0$ Å. Correlation function
is given as a function of normalized radial interatomic distance between atoms
in the correlator, showing nearest, next-nearest and next-next-nearest
magnetic correlations. The correlators were computed from the optimized CPD
backflow ansatz with $M=1$.

In Figure 5, we compare the radial spin-spin correlation function for the
ground state approximation obtained with the CPD backflow ansatz at near-
equilibrium interatomic distance $d=1.2$ Å, and at a large stretching of
$d=3.0$ Å, normalized for the changing inter-atomic distances. We find the
emergence of the short-range anti-ferromagnetic order in the material, as
anticipated by analogy with Hubbard models. The magnitude of this
antiferromagnetic order increases with increasing interatomic separation,
again keeping with anticipated Hubbard behaviour of increasing $U/t$ values.
However, this order is very short ranged, with the spin in the extended
lattice not directly affecting a lattice site beyond its nearest neighbours.
At more compressed geometries, this order does extend beyond this to the outer
atoms in the lattice (next-next-nearest-neighbours), due to the shorter
distance in real space, but the overall magnitude of these magnetic
correlations is reduced.

##  IV Scaling

As illustrated in the results above, the CPD backflow ansatz performs well on
small Fermionic systems, but further developments for scaling to significantly
larger systems are still required for this to become a clearly competitive
method for the wider electronic structure community. Simplifying the scaling
to assume a general growth of both the basis and electron number such that
$N\sim L$ and assuming $M$ and $K$ are independent of system size, the
parameters grow with system size as $N_{P}\sim\mathcal{O}[N^{3}]$ for the full
ansatz, with the subspace truncation of the backflow reducing this
asymptotically to $N_{P}\sim\mathcal{O}[N^{2}]$ (see Sec. II.2). The fast
updating of backflow orbitals also enables the evaluation of the wave function
log-amplitudes to be performed in $\mathcal{O}[MN^{2}+N^{3}]$ (regardless of
whether a backflow truncation is applied). However, we find in practice the
determinant evaluation has a significantly smaller prefactor than the
construction of the orbitals, so that the dominant scaling is rather
$\mathcal{O}[N^{2}]$ for small to medium-sized systems. Given that the number
of terms in general second quantized ab initio hamiltonians scales as
$\mathcal{O}[N^{4}]$, the resulting scaling of the local energy evaluation is
then $\mathcal{O}[N^{7}]$ for the CPD backflow state in the asymptotic limit,
or $\mathcal{O}[N^{6}]$ for small to medium-sized systems. While this should
be competitive with accurate quantum chemical methods such as coupled-cluster,
it is clear that the prefactor is significantly larger.

Rather than just the local energy evaluation, we should also consider the
computational scaling for the update of the parameters. For larger numbers of
parameters (such as the $\sim 187,000$ of the hydrogen lattice above), their
update used to be the main bottleneck for VMC large-scale ansätze when using
the original SR algorithm [56]. For a model with $N_{P}$ parameters, SR would
scale as $\mathcal{O}(N_{P}^{3})$, since it involves inverting the
$N_{P}\times N_{P}$ quantum geometric tensor matrix. This is not the case with
the kernel formulation of SR recently introduced in Ref. [57], which takes
advantage of a simple linear algebra identity to reduce the dimension of the
matrix that is inverted in the SR algorithm from $N_{P}$ to $N_{S}$, i.e. the
number of samples used during the optimization. When $N_{P}\gg N_{S}$, as in
large-scale models, the scaling of the parameter update becomes
$\mathcal{O}(N_{S}^{2}N_{P}+N_{S}^{3})$, i.e. linear in the number of
parameters. Thus, the evaluation of the local energy remains the computational
bottleneck of the algorithm in the case of ab initio systems.

We show this scaling explicitly in Fig. 6, where we measure the mean runtime
for a full VMC parameter update step for the CPD backflow ansatz, including
the Markov chain sampling of a fixed number of configurations, evaluation of
the local energy, and the subsequent parameter update. By increasing the
number of hydrogen atoms in a chain with fixed equilibrium inter-atomic
distances ($d=1.68$ Å) in a STO-6G basis up to 90 atoms we can extract a
realistic asymptotic scaling of the approach. We set the support dimension of
the ansatz to $M=1$ and choose a sample size of $N_{S}=128$, in order to fit
the data in memory even for the largest system sizes. For each system size, we
let the VMC algorithm run for 50 iterations on a single Nvidia A100 GPU with
40GB of memory. Extracting the scaling from the large system limit gives a
scaling of $\mathcal{O}[N^{6.5}]$, which is only evident for this system when
we reach $>40$ atoms.

Figure 6: Mean VMC step runtime for the CPD backflow ansatz as a function of
the number of atoms in a hydrogen chain with fixed inter-atomic distances
($1.68$ Å) in a STO-6G basis. The support dimension is $M=1$ and the sample
size is $N_{S}=128$. The blue points are obtained with the full Hamiltonian,
while the green and orange points are obtained with the local energy evaluated
with a Hamiltonian pruned with a threshold of $10^{-9}$ and $10^{-5}E_{h}$.
The dashed lines show the observed scaling for the largest system sizes.

As discussed in the context of the GPS model in Ref. [15] and building on
other works in this area [77, 78, 79, 69], this scaling for ab initio systems
can be reduced further by truncating the number of terms in the sum over
connected configurations at each evaluation of the local energy. This
truncation is performed on the magnitude of the Hamiltonian matrix element
connecting the configurations. By presorting the electron repulsion integrals
between the degrees of freedom, this truncation can be implemented without
having to consider the entire set of Hamiltonian matrix elements for each
evaluation of the local energy. Formally, the exponentially decreasing overlap
between the orbitals in the sampled space should reduce the number of
connected determinants which contribute to the local energy asymptotically to
$\mathcal{O}[N^{2}]$ rather than $\mathcal{O}[N^{4}]$ – a scaling which then
matches the scaling of the local energy evaluation in a first quantized
perspective. This results in a practical scaling of the CPD ansatz of
$\sim\mathcal{O}[N^{4-6}]$. However, since this method comes with a certain
overhead in terms of data structures, the lower-bounds on this scaling only
materialize after a certain crossover system size, which depends on the system
and on the truncation threshold on the Hamiltonian matrix elements.

Figure 6 also shows the analogous results including this energetic threshold
of $10^{-5}$ and $10^{-9}E_{h}$, with this tighter threshold expected to incur
negligible change in the sampled energy for a given state. The results show
that a practical crossover point, after which the pruning of Hamiltonian
elements below a small threshold yields a speed-up, is reached already around
system sizes as low as 20-30 electrons for this ansatz, noting that a one-
dimensional chain is advantageous in terms of affecting an advantage from this
approach. However, for large system sizes the speed-up is more than an order
of magnitude and provides the expected asymptotic quadratic improvement,
giving an overall scaling of $\mathcal{O}[N^{3-4}]$ up to $\sim 100$
electrons, a scaling competitive with hybrid Density Functional Theory (DFT)
techniques. Overall, the combination of this ansatz with the various scaling
reductions outlined represent a real potential towards a second quantized,
systematically improvable, VMC algorithm with a practical and competitive
$\mathcal{O}(N^{4})$ scaling for medium to large ab initio systems. Clearly,
further developments for prefactor reductions are key to take advantage of
this improved scaling and access these system sizes.

##  V Conclusions and outlook

In this work, we introduce a general and simple ansatz suitable for ab initio
fermions, based on a systematically improvable tensor rank decomposition of a
general backflow form. This systematically builds configuration-dependent
orbitals of a single antisymmetric Slater determinant in second quantization,
directly encoding non-trivial $N$-body electron-electron correlations with a
parameter scaling of $\mathcal{O}[N^{2-3}]$. We have shown that the ansatz can
achieve competitive accuracy on small Fermionic systems, such as the Fermi-
Hubbard model and the water molecule, and that it can be used to model larger
strongly correlated lattices of ab initio hydrogen atoms with an accuracy
comparable to state-of-the-art DMRG techniques. Finally, we have discussed the
scalability of the ansatz and shown that we can affect various reductions in a
practical fashion to demonstrate $\mathcal{O}(N^{4})$ scaling on medium to
large ab initio systems. We are working on further improvements in the
accuracy and efficiency of the ansatz, as well as taking advantage of the
benefits of working in a second quantized formalism to integrate with
multiscale methods and quantum embedding methodologies to provide a practical
route in the modelling of truly extended systems within this CPD backflow
framework [80, 80].

## Code Availability

The code for this project is implemented as part of the freely available
GPSKet plugin for NetKet [59, 60]. It is made available, together with
configurations files to reproduce the figures in the paper, at
https://github.com/BoothGroup/GPSKet/tree/master/scripts/cpd-backflow.

###### Acknowledgements.

The authors gratefully acknowledge support from the Air Force Office of
Scientific Research under award number FA8655-22-1-7011. We are grateful to
the UK Materials and Molecular Modelling Hub for computational resources,
which is partially funded by EPSRC (EP/P020194/1 and EP/T022213/1).
Furthermore, we acknowledge the use of the high performance computing
environment CREATE at King’s College London [81].

## References

  * Choo et al. [2019] Kenny Choo, Titus Neupert, and Giuseppe Carleo.  Two-dimensional frustrated ${}J{}_{1}\text{\ensuremath{-}}{}J{}_{2}$ model studied with neural network quantum states.  _Physical Review B_ , 100(12):125124, September 2019.  doi: 10.1103/PhysRevB.100.125124.
  * Carleo and Troyer [2017] Giuseppe Carleo and Matthias Troyer.  Solving the quantum many-body problem with artificial neural networks.  _Science_ , 355(6325):602–606, February 2017\.  doi: 10.1126/science.aag2302.
  * Nomura et al. [2017] Yusuke Nomura, Andrew S. Darmawan, Youhei Yamaji, and Masatoshi Imada.  Restricted Boltzmann machine learning for solving strongly correlated quantum systems.  _Physical Review B_ , 96(20):205152, November 2017\.  doi: 10.1103/PhysRevB.96.205152.
  * Choo et al. [2020] Kenny Choo, Antonio Mezzacapo, and Giuseppe Carleo.  Fermionic neural-network states for ab-initio electronic structure.  _Nature Communications_ , 11(1):2368, May 2020\.  ISSN 2041-1723.  doi: 10.1038/s41467-020-15724-9.
  * Barrett et al. [2022] Thomas D. Barrett, Aleksei Malyshev, and A. I. Lvovsky.  Autoregressive neural-network wavefunctions for ab initio quantum chemistry.  _Nature Machine Intelligence_ , 4(4):351–358, April 2022.  ISSN 2522-5839.  doi: 10.1038/s42256-022-00461-z.
  * Hibat-Allah et al. [2020] Mohamed Hibat-Allah, Martin Ganahl, Lauren E. Hayward, Roger G. Melko, and Juan Carrasquilla.  Recurrent neural network wave functions.  _Physical Review Research_ , 2(2):023358, June 2020.  doi: 10.1103/PhysRevResearch.2.023358.
  * Sprague and Czischek [2024] Kyle Sprague and Stefanie Czischek.  Variational Monte Carlo with large patched transformers.  _Communications Physics_ , 7(1):1–11, March 2024\.  ISSN 2399-3650.  doi: 10.1038/s42005-024-01584-y.
  * Rath and Booth [2022] Yannic Rath and George H. Booth.  Quantum Gaussian process state: A kernel-inspired state with quantum support data.  _Physical Review Research_ , 4(2):023126, May 2022\.  doi: 10.1103/PhysRevResearch.4.023126.
  * Giuliani et al. [2023] Clemens Giuliani, Filippo Vicentini, Riccardo Rossi, and Giuseppe Carleo.  Learning ground states of gapped quantum Hamiltonians with Kernel Methods.  _Quantum_ , 7:1096, August 2023.  doi: 10.22331/q-2023-08-29-1096.
  * Orús [2019] Román Orús.  Tensor networks for complex quantum systems.  _Nature Reviews Physics_ , 1(9):538–550, September 2019.  ISSN 2522-5820.  doi: 10.1038/s42254-019-0086-7.
  * Eisert et al. [2010] J. Eisert, M. Cramer, and M. B. Plenio.  Area laws for the entanglement entropy.  _Reviews of Modern Physics_ , 82(1):277–306, February 2010.  doi: 10.1103/RevModPhys.82.277.
  * Bukov et al. [2021] Marin Bukov, Markus Schmitt, and Maxime Dupont.  Learning the ground state of a non-stoquastic quantum Hamiltonian in a rugged neural network landscape.  _SciPost Physics_ , 10(6):147, June 2021.  ISSN 2542-4653.  doi: 10.21468/SciPostPhys.10.6.147.
  * Inack et al. [2022] Estelle M. Inack, Stewart Morawetz, and Roger G. Melko.  Neural Annealing and Visualization of Autoregressive Neural Networks in the Newman–Moore Model.  _Condensed Matter_ , 7(2):38, June 2022.  ISSN 2410-3896.  doi: 10.3390/condmat7020038.
  * Glielmo et al. [2020] Aldo Glielmo, Yannic Rath, Gábor Csányi, Alessandro De Vita, and George H. Booth.  Gaussian Process States: A Data-Driven Representation of Quantum Many-Body Physics.  _Physical Review X_ , 10(4):041026, November 2020\.  doi: 10.1103/PhysRevX.10.041026.
  * Rath and Booth [2023] Yannic Rath and George H. Booth.  Framework for efficient ab initio electronic structure with Gaussian Process States.  _Physical Review B_ , 107(20):205119, May 2023\.  doi: 10.1103/PhysRevB.107.205119.
  * Pfau et al. [2020] David Pfau, James S. Spencer, Alexander G. D. G. Matthews, and W. M. C. Foulkes.  Ab initio solution of the many-electron Schr\”odinger equation with deep neural networks.  _Physical Review Research_ , 2(3):033429, September 2020.  doi: 10.1103/PhysRevResearch.2.033429.
  * Hermann et al. [2020] Jan Hermann, Zeno Schätzle, and Frank Noé.  Deep-neural-network solution of the electronic Schrödinger equation.  _Nature Chemistry_ , 12(10):891–897, October 2020\.  ISSN 1755-4349.  doi: 10.1038/s41557-020-0544-y.
  * Nusspickel and Booth [2021] Max Nusspickel and George H. Booth.  Systematic improvability in quantum embedding for real materials.  _arXiv:2107.04916 [cond-mat, physics:physics]_ , July 2021.
  * Sun and Chan [2016] Qiming Sun and Garnet Kin-Lic Chan.  Quantum Embedding Theories.  _Accounts of Chemical Research_ , 49(12):2705–2712, December 2016.  ISSN 0001-4842.  doi: 10.1021/acs.accounts.6b00356.
  * Kotliar et al. [2006] G. Kotliar, S. Y. Savrasov, K. Haule, V. S. Oudovenko, O. Parcollet, and C. A. Marianetti.  Electronic structure calculations with dynamical mean-field theory.  _Reviews of Modern Physics_ , 78(3):865–951, August 2006.  doi: 10.1103/RevModPhys.78.865.
  * de Lara-Castells et al. [2001] M. P. de Lara-Castells, R. V. Krems, A. A. Buchachenko, G. Delgado-Barrio, and P. Villarreal.  Complete basis set extrapolation limit for electronic structure calculations: Energetic and nonenergetic properties of HeBr and HeBr2 van der Waals dimers.  _The Journal of Chemical Physics_ , 115(22):10438–10449, December 2001.  ISSN 0021-9606.  doi: 10.1063/1.1415078.
  * Grüneis et al. [2017] Andreas Grüneis, So Hirata, Yu-ya Ohnishi, and Seiichiro Ten-no.  Perspective: Explicitly correlated electronic structure theory for complex systems.  _The Journal of Chemical Physics_ , 146(8):080901, February 2017.  ISSN 0021-9606.  doi: 10.1063/1.4976974.
  * Booth et al. [2012] George H. Booth, Deidre Cleland, Ali Alavi, and David P. Tew.  An explicitly correlated approach to basis set incompleteness in full configuration interaction quantum Monte Carlo.  _The Journal of Chemical Physics_ , 137(16):164112, October 2012.  ISSN 0021-9606.  doi: 10.1063/1.4762445.
  * Yanai and Shiozaki [2012] Takeshi Yanai and Toru Shiozaki.  Canonical transcorrelated theory with projected Slater-type geminals.  _The Journal of Chemical Physics_ , 136(8):084107, February 2012.  ISSN 0021-9606.  doi: 10.1063/1.3688225.
  * Wu et al. [2023] Yangjun Wu, Chu Guo, Yi Fan, Pengyu Zhou, and Honghui Shang.  NNQS-Transformer: An Efficient and Scalable Neural Network Quantum States Approach for Ab initio Quantum Chemistry, July 2023.
  * Zhao et al. [2023] Tianchen Zhao, James Stokes, and Shravan Veerapaneni.  Scalable neural quantum states architecture for quantum chemistry.  _Machine Learning: Science and Technology_ , 4(2):025034, June 2023.  ISSN 2632-2153.  doi: 10.1088/2632-2153/acdb2f.
  * Bravyi and Kitaev [2002] Sergey B. Bravyi and Alexei Yu. Kitaev.  Fermionic Quantum Computation.  _Annals of Physics_ , 298(1):210–226, May 2002\.  ISSN 0003-4916.  doi: 10.1006/aphy.2002.6254.
  * Nys and Carleo [2022] Jannes Nys and Giuseppe Carleo.  Variational solutions to fermion-to-qubit mappings in two spatial dimensions.  _Quantum_ , 6:833, October 2022.  doi: 10.22331/q-2022-10-13-833.
  * Nys and Carleo [2023] Jannes Nys and Giuseppe Carleo.  Quantum circuits for solving local fermion-to-qubit mappings.  _Quantum_ , 7:930, February 2023.  doi: 10.22331/q-2023-02-21-930.
  * Derby et al. [2021] Charles Derby, Joel Klassen, Johannes Bausch, and Toby Cubitt.  Compact fermion to qubit mappings.  _Physical Review B_ , 104(3):035118, July 2021\.  doi: 10.1103/PhysRevB.104.035118.
  * Tocchio et al. [2008] Luca F. Tocchio, Federico Becca, Alberto Parola, and Sandro Sorella.  Role of backflow correlations for the nonmagnetic phase of the t-t’ Hubbard model.  _Physical Review B_ , 78(4):041101, July 2008\.  doi: 10.1103/PhysRevB.78.041101.
  * Tocchio et al. [2011] Luca F. Tocchio, Federico Becca, and Claudius Gros.  Backflow correlations in the Hubbard model: An efficient tool for the study of the metal-insulator transition and the large-$U$ limit.  _Physical Review B_ , 83(19):195138, May 2011\.  doi: 10.1103/PhysRevB.83.195138.
  * Moreno et al. [2022] Javier Robledo Moreno, Giuseppe Carleo, Antoine Georges, and James Stokes.  Fermionic wave functions from neural-network constrained hidden states.  _Proceedings of the National Academy of Sciences_ , 119(32):e2122059119, August 2022.  doi: 10.1073/pnas.2122059119.
  * Liu and Clark [2023] Zejun Liu and Bryan K. Clark.  A Unifying View of Fermionic Neural Network Quantum States: From Neural Network Backflow to Hidden Fermion Determinant States, November 2023.
  * Feynman and Cohen [1956] R. P. Feynman and Michael Cohen.  Energy Spectrum of the Excitations in Liquid Helium.  _Physical Review_ , 102(5):1189–1204, June 1956\.  doi: 10.1103/PhysRev.102.1189.
  * Wigner and Seitz [1934] E. Wigner and F. Seitz.  On the Constitution of Metallic Sodium. II.  _Physical Review_ , 46(6):509–524, September 1934\.  doi: 10.1103/PhysRev.46.509.
  * Kwon et al. [1993] Yongkyung Kwon, D. M. Ceperley, and Richard M. Martin.  Effects of three-body and backflow correlations in the two-dimensional electron gas.  _Physical Review B_ , 48(16):12037–12046, October 1993.  doi: 10.1103/PhysRevB.48.12037.
  * Kwon et al. [1998] Yongkyung Kwon, D. M. Ceperley, and Richard M. Martin.  Effects of backflow correlation in the three-dimensional electron gas: Quantum Monte Carlo study.  _Physical Review B_ , 58(11):6800–6806, September 1998.  doi: 10.1103/PhysRevB.58.6800.
  * Pescia et al. [2024] Gabriel Pescia, Jannes Nys, Jane Kim, Alessandro Lovato, and Giuseppe Carleo.  Message-passing neural quantum states for the homogeneous electron gas.  _Physical Review B_ , 110(3):035108, July 2024\.  doi: 10.1103/PhysRevB.110.035108.
  * Romero et al. [2024] Imelda Romero, Jannes Nys, and Giuseppe Carleo.  Spectroscopy of two-dimensional interacting lattice electrons using symmetry-aware neural backflow transformations, June 2024.
  * Luo and Clark [2019] Di Luo and Bryan K. Clark.  Backflow Transformations via Neural Networks for Quantum Many-Body Wave Functions.  _Physical Review Letters_ , 122(22):226401, June 2019.  doi: 10.1103/PhysRevLett.122.226401.
  * Zhou et al. [2024] Yu-Tong Zhou, Zheng-Wei Zhou, and Xiao Liang.  Solving Fermi-Hubbard-type Models by Tensor Representations of Backflow Corrections.  _Physical Review B_ , 109(24):245107, June 2024\.  ISSN 2469-9950, 2469-9969.  doi: 10.1103/PhysRevB.109.245107.
  * Kim et al. [2024] Jane Kim, Gabriel Pescia, Bryce Fore, Jannes Nys, Giuseppe Carleo, Stefano Gandolfi, Morten Hjorth-Jensen, and Alessandro Lovato.  Neural-network quantum states for ultra-cold Fermi gases.  _Communications Physics_ , 7(1):1–12, May 2024\.  ISSN 2399-3650.  doi: 10.1038/s42005-024-01613-w.
  * von Glehn et al. [2022] Ingrid von Glehn, James S. Spencer, and David Pfau.  A Self-Attention Ansatz for Ab-initio Quantum Chemistry, November 2022.
  * Liu and Clark [2024] An-Jun Liu and Bryan K. Clark.  Neural network backflow for ab-initio quantum chemistry, March 2024.
  * Kiers [2000] Henk A. L. Kiers.  Towards a standardized notation and terminology in multiway analysis.  _Journal of Chemometrics_ , 14(3):105–122, 2000\.  ISSN 1099-128X.  doi: 10.1002/1099-128X(200005/06)14:3¡105::AID-CEM582¿3.0.CO;2-I.
  * Kolda and Bader [2009] Tamara G. Kolda and Brett W. Bader.  Tensor Decompositions and Applications.  _SIAM Review_ , 51(3):455–500, August 2009.  ISSN 0036-1445.  doi: 10.1137/07070111X.
  * Lami et al. [2022] Guglielmo Lami, Giuseppe Carleo, and Mario Collura.  Matrix product states with backflow correlations.  _Physical Review B_ , 106(8):L081111, August 2022\.  doi: 10.1103/PhysRevB.106.L081111.
  * 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.  _Physical Review E_ , 74(6):066701, December 2006\.  doi: 10.1103/PhysRevE.74.066701.
  * Lou et al. [2024] Wan Tong Lou, Halvard Sutterud, Gino Cassella, W. M. C. Foulkes, Johannes Knolle, David Pfau, and James S. Spencer.  Neural Wave Functions for Superfluids.  _Physical Review X_ , 14(2):021030, May 2024.  doi: 10.1103/PhysRevX.14.021030.
  * Casula et al. [2004] Michele Casula, Claudio Attaccalite, and Sandro Sorella.  Correlated geminal wave function for molecules: An efficient resonating valence bond approach.  _The Journal of Chemical Physics_ , 121(15):7110–7126, October 2004.  ISSN 0021-9606.  doi: 10.1063/1.1794632.
  * Park and Kastoryano [2020] Chae-Yeun Park and Michael J. Kastoryano.  Geometry of learning neural quantum states.  _Physical Review Research_ , 2(2):023232, May 2020\.  doi: 10.1103/PhysRevResearch.2.023232.
  * Olivares-Amaya et al. [2015] Roberto Olivares-Amaya, Weifeng Hu, Naoki Nakatani, Sandeep Sharma, Jun Yang, and Garnet Kin-Lic Chan.  The ab-initio density matrix renormalization group in practice.  _The Journal of Chemical Physics_ , 142(3):034102, January 2015.  ISSN 0021-9606.  doi: 10.1063/1.4905329.
  * Rissler et al. [2006] Jörg Rissler, Reinhard M. Noack, and Steven R. White.  Measuring orbital interaction using quantum information theory.  _Chemical Physics_ , 323(2):519–531, April 2006\.  ISSN 0301-0104.  doi: 10.1016/j.chemphys.2005.10.018.
  * Foster and Boys [1960] J. M. Foster and S. F. Boys.  Canonical Configurational Interaction Procedure.  _Reviews of Modern Physics_ , 32(2):300–302, April 1960.  doi: 10.1103/RevModPhys.32.300.
  * Sorella [2001] Sandro Sorella.  Generalized Lanczos algorithm for variational quantum Monte Carlo.  _Physical Review B_ , 64(2):024512, June 2001\.  doi: 10.1103/PhysRevB.64.024512.
  * Rende et al. [2023] Riccardo Rende, Luciano Loris Viteritti, Lorenzo Bardone, Federico Becca, and Sebastian Goldt.  A simple linear algebra identity to optimize Large-Scale Neural Network Quantum States, October 2023.
  * Lovato et al. [2022] Alessandro Lovato, Corey Adams, Giuseppe Carleo, and Noemi Rocco.  Hidden-nucleons neural-network quantum states for the nuclear many-body problem.  _Physical Review Research_ , 4(4):043178, December 2022.  doi: 10.1103/PhysRevResearch.4.043178.
  * Carleo et al. [2019] Giuseppe Carleo, Kenny Choo, Damian Hofmann, James E. T. Smith, Tom Westerhout, Fabien Alet, Emily J. Davis, Stavros Efthymiou, Ivan Glasser, Sheng-Hsuan Lin, Marta Mauri, Guglielmo Mazzola, Christian B. Mendl, Evert van Nieuwenburg, Ossian O’Reilly, Hugo Théveniaut, Giacomo Torlai, Filippo Vicentini, and Alexander Wietek.  NetKet: A machine learning toolkit for many-body quantum systems.  _SoftwareX_ , 10:100311, July 2019.  ISSN 2352-7110.  doi: 10.1016/j.softx.2019.100311.
  * Vicentini et al. [2022] Filippo Vicentini, Damian Hofmann, Attila Szabó, Dian Wu, Christopher Roth, Clemens Giuliani, Gabriel Pescia, Jannes Nys, Vladimir Vargas-Calderón, Nikita Astrakhantsev, and Giuseppe Carleo.  NetKet 3: Machine Learning Toolbox for Many-Body Quantum Systems.  _SciPost Physics Codebases_ , page 007, August 2022.  ISSN 2949-804X.  doi: 10.21468/SciPostPhysCodeb.7.
  * Sun et al. [2018] Qiming Sun, Timothy C. Berkelbach, Nick S. Blunt, George H. Booth, Sheng Guo, Zhendong Li, Junzi Liu, James D. McClain, Elvira R. Sayfutyarova, Sandeep Sharma, Sebastian Wouters, and Garnet Kin-Lic Chan.  PySCF: The Python-based simulations of chemistry framework.  _WIREs Computational Molecular Science_ , 8(1):e1340, 2018.  ISSN 1759-0884.  doi: 10.1002/wcms.1340.
  * Sun et al. [2020] Qiming Sun, Xing Zhang, Samragni Banerjee, Peng Bao, Marc Barbry, Nick S. Blunt, Nikolay A. Bogdanov, George H. Booth, Jia Chen, Zhi-Hao Cui, Janus J. Eriksen, Yang Gao, Sheng Guo, Jan Hermann, Matthew R. Hermes, Kevin Koh, Peter Koval, Susi Lehtola, Zhendong Li, Junzi Liu, Narbe Mardirossian, James D. McClain, Mario Motta, Bastien Mussard, Hung Q. Pham, Artem Pulkin, Wirawan Purwanto, Paul J. Robinson, Enrico Ronca, Elvira R. Sayfutyarova, Maximilian Scheurer, Henry F. Schurkus, James E. T. Smith, Chong Sun, Shi-Ning Sun, Shiv Upadhyay, Lucas K. Wagner, Xiao Wang, Alec White, James Daniel Whitfield, Mark J. Williamson, Sebastian Wouters, Jun Yang, Jason M. Yu, Tianyu Zhu, Timothy C. Berkelbach, Sandeep Sharma, Alexander Yu. Sokolov, and Garnet Kin-Lic Chan.  Recent developments in the PySCF program package.  _The Journal of Chemical Physics_ , 153(2):024109, July 2020.  ISSN 0021-9606.  doi: 10.1063/5.0006074.
  * Zheng et al. [2017] Bo-Xiao Zheng, Chia-Min Chung, Philippe Corboz, Georg Ehlers, Ming-Pu Qin, Reinhard M. Noack, Hao Shi, Steven R. White, Shiwei Zhang, and Garnet Kin-Lic Chan.  Stripe order in the underdoped region of the two-dimensional Hubbard model.  _Science_ , 358(6367):1155–1160, December 2017\.  doi: 10.1126/science.aam7127.
  * Sorella [2023] Sandro Sorella.  Systematically improvable mean-field variational ansatz for strongly correlated systems: Application to the Hubbard model.  _Physical Review B_ , 107(11):115133, March 2023\.  doi: 10.1103/PhysRevB.107.115133.
  * Dagotto et al. [1992] E. Dagotto, A. Moreo, F. Ortolani, D. Poilblanc, and J. Riera.  Static and dynamical properties of doped Hubbard clusters.  _Physical Review B_ , 45(18):10741–10760, May 1992.  doi: 10.1103/PhysRevB.45.10741.
  * Simons Collaboration on the Many-Electron Problem et al. [2015] Simons Collaboration on the Many-Electron Problem, J. P. F. LeBlanc, Andrey E. Antipov, Federico Becca, Ireneusz W. Bulik, Garnet Kin-Lic Chan, Chia-Min Chung, Youjin Deng, Michel Ferrero, Thomas M. Henderson, Carlos A. Jiménez-Hoyos, E. Kozik, Xuan-Wen Liu, Andrew J. Millis, N. V. Prokof’ev, Mingpu Qin, Gustavo E. Scuseria, Hao Shi, B. V. Svistunov, Luca F. Tocchio, I. S. Tupitsyn, Steven R. White, Shiwei Zhang, Bo-Xiao Zheng, Zhenyue Zhu, and Emanuel Gull.  Solutions of the Two-Dimensional Hubbard Model: Benchmarks and Results from a Wide Range of Numerical Algorithms.  _Physical Review X_ , 5(4):041041, December 2015\.  doi: 10.1103/PhysRevX.5.041041.
  * Gurtubay and Needs [2007] I. G. Gurtubay and R. J. Needs.  Dissociation energy of the water dimer from quantum Monte Carlo calculations.  _The Journal of Chemical Physics_ , 127(12):124306, September 2007.  ISSN 0021-9606.  doi: 10.1063/1.2770711.
  * Clark et al. [2011] Bryan K. Clark, Miguel A. Morales, Jeremy McMinis, Jeongnim Kim, and Gustavo E. Scuseria.  Computing the energy of a water molecule using multideterminants: A simple, efficient algorithm.  _The Journal of Chemical Physics_ , 135(24):244105, December 2011.  ISSN 0021-9606.  doi: 10.1063/1.3665391.
  * Hachmann et al. [2006] Johannes Hachmann, Wim Cardoen, and Garnet Kin-Lic Chan.  Multireference correlation in long molecules with the quadratic scaling density matrix renormalization group.  _The Journal of Chemical Physics_ , 125(14):144101, October 2006.  ISSN 0021-9606.  doi: 10.1063/1.2345196.
  * Simons Collaboration on the Many-Electron Problem et al. [2017] Simons Collaboration on the Many-Electron Problem, Mario Motta, David M. Ceperley, Garnet Kin-Lic Chan, John A. Gomez, Emanuel Gull, Sheng Guo, Carlos A. Jiménez-Hoyos, Tran Nguyen Lan, Jia Li, Fengjie Ma, Andrew J. Millis, Nikolay V. Prokof’ev, Ushnish Ray, Gustavo E. Scuseria, Sandro Sorella, Edwin M. Stoudenmire, Qiming Sun, Igor S. Tupitsyn, Steven R. White, Dominika Zgid, and Shiwei Zhang.  Towards the Solution of the Many-Electron Problem in Real Materials: Equation of State of the Hydrogen Chain with State-of-the-Art Many-Body Methods.  _Physical Review X_ , 7(3):031059, September 2017\.  doi: 10.1103/PhysRevX.7.031059.
  * Stella et al. [2011] Lorenzo Stella, Claudio Attaccalite, Sandro Sorella, and Angel Rubio.  Strong electronic correlation in the hydrogen chain: A variational Monte Carlo study.  _Physical Review B_ , 84(24):245117, December 2011\.  doi: 10.1103/PhysRevB.84.245117.
  * Sinitskiy et al. [2010] Anton V. Sinitskiy, Loren Greenman, and David A. Mazziotti.  Strong correlation in hydrogen chains and lattices using the variational two-electron reduced density matrix method.  _The Journal of Chemical Physics_ , 133(1):014104, July 2010.  ISSN 0021-9606.  doi: 10.1063/1.3459059.
  * Tsuchimochi and Scuseria [2009] Takashi Tsuchimochi and Gustavo E. Scuseria.  Strong correlations via constrained-pairing mean-field theory.  _The Journal of Chemical Physics_ , 131(12):121102, September 2009.  ISSN 0021-9606.  doi: 10.1063/1.3237029.
  * Motta et al. [2020] Mario Motta, Claudio Genovese, Fengjie Ma, Zhi-Hao Cui, Randy Sawaya, Garnet Kin-Lic Chan, Natalia Chepiga, Phillip Helms, Carlos Jiménez-Hoyos, Andrew J. Millis, Ushnish Ray, Enrico Ronca, Hao Shi, Sandro Sorella, Edwin M. Stoudenmire, Steven R. White, and Shiwei Zhang.  Ground-State Properties of the Hydrogen Chain: Dimerization, Insulator-to-Metal Transition, and Magnetic Phases.  _Physical Review X_ , 10(3):031058, September 2020\.  doi: 10.1103/PhysRevX.10.031058.
  * Zhai et al. [2023] Huanchen Zhai, Henrik R. Larsson, Seunghoon Lee, Zhi-Hao Cui, Tianyu Zhu, Chong Sun, Linqing Peng, Ruojing Peng, Ke Liao, Johannes Tölle, Junjie Yang, Shuoxue Li, and Garnet Kin-Lic Chan.  Block2: A comprehensive open source framework to develop and apply state-of-the-art DMRG algorithms in electronic structure and beyond.  _The Journal of Chemical Physics_ , 159(23):234801, December 2023.  ISSN 0021-9606.  doi: 10.1063/5.0180424.
  * Sharma and Chan [2012] Sandeep Sharma and Garnet Kin-Lic Chan.  Spin-adapted density matrix renormalization group algorithms for quantum chemistry.  _The Journal of Chemical Physics_ , 136(12):124121, March 2012.  ISSN 0021-9606.  doi: 10.1063/1.3695642.
  * Mahajan and Sharma [2020] Ankit Mahajan and Sandeep Sharma.  Efficient local energy evaluation for multi-Slater wave functions in orbital space quantum Monte Carlo.  _The Journal of Chemical Physics_ , 153(19):194108, November 2020.  ISSN 0021-9606.  doi: 10.1063/5.0025055.
  * Wei and Neuscamman [2018] Haochuan Wei and Eric Neuscamman.  Reduced scaling Hilbert space variational Monte Carlo.  _The Journal of Chemical Physics_ , 149(18):184106, November 2018.  ISSN 0021-9606.  doi: 10.1063/1.5047207.
  * Sabzevari and Sharma [2018] Iliya Sabzevari and Sandeep Sharma.  Improved Speed and Scaling in Orbital Space Variational Monte Carlo.  _Journal of Chemical Theory and Computation_ , 14(12):6276–6286, December 2018.  ISSN 1549-9618.  doi: 10.1021/acs.jctc.8b00780.
  * Nusspickel et al. [2023] Max Nusspickel, Basil Ibrahim, and George H. Booth.  Effective Reconstruction of Expectation Values from Ab Initio Quantum Embedding.  _Journal of Chemical Theory and Computation_ , 19(10):2769–2791, May 2023.  ISSN 1549-9618.  doi: 10.1021/acs.jctc.2c01063.
  * King’s College London e-Research team [2022] King’s College London e-Research team.  King’s Computational Research, Engineering and Technology Environment (CREATE).  https://doi.org/10.18742/rnvf-m076, 2022.
  * Boyn and Mazziotti [2021] Jan-Niklas Boyn and David A. Mazziotti.  Accurate singlet–triplet gaps in biradicals via the spin averaged anti-Hermitian contracted Schrödinger equation.  _The Journal of Chemical Physics_ , 154(13):134103, April 2021.  ISSN 0021-9606.  doi: 10.1063/5.0045007.
  * Yun et al. [2023] Sujun Yun, Werner Dobrautz, Hongjun Luo, Vamshi Katukuri, Niklas Liebermann, and Ali Alavi.  Ferromagnetic domains in the large-$U$ Hubbard model with a few holes: A full configuration interaction quantum Monte Carlo study.  _Physical Review B_ , 107(6):064405, February 2023\.  doi: 10.1103/PhysRevB.107.064405.

##  Appendix A Spin-spin Correlations

For a pair of atoms $a$ and $b$, we define the instantaneous spin-spin
correlation function as

$\begin{split}\left\langle\hat{S}_{\vec{r}_{a}}^{z}\hat{S}_{\vec{r}_{b}}^{z}\right\rangle&=\frac{1}{4}\sum_{ijkl}P_{ij}^{a}P_{kl}^{b}\left(\Gamma_{ijkl}^{\alpha\alpha}-\Gamma_{ijkl}^{\alpha\beta}-\Gamma_{ijkl}^{\beta\alpha}+\Gamma_{ijkl}^{\beta\beta}\right)\\\
&+\frac{1}{4}\sum_{ijk}P_{ik}^{a}P_{jk}^{b}\left(\gamma_{ij}^{\alpha}+\gamma_{ij}^{\beta}\right),\end{split}$
(16)

where $P^{a}$ and $P^{b}$ are projectors onto the orbital subspaces of atom
$a$ and $b$, the indices $i,j,k,l$ are orbital labels, and
$\gamma_{ij}^{\sigma}$ and $\Gamma_{ijkl}^{\sigma\tau}$ are spinned one- and
two-body reduced density matrices (RDMs), respectively.

During the evaluation of the local energy, the spin-free one- and two-body
RDMs, $\gamma_{ij}$ and $\Gamma_{ijkl}$, are computed, so that it pays off to
express the previous equation in terms of these quantities. This can be done
using the following identities [82, 83]:

$\displaystyle\gamma_{ij}^{\alpha}$
$\displaystyle=\gamma_{ij}^{\beta}=\gamma_{ij}$ (17)
$\displaystyle\Gamma_{ijkl}^{\alpha\alpha}$
$\displaystyle=\Gamma_{ijkl}^{\beta\beta}=\frac{1}{6}\left(\Gamma_{ijkl}-\Gamma_{kjil}\right)$
(18) $\displaystyle\Gamma_{ijkl}^{\alpha\beta}$
$\displaystyle=\Gamma_{ijkl}^{\beta\alpha}=\frac{1}{6}\left(2\Gamma_{ijkl}+\Gamma_{kjil}\right).$
(19)

Then, Eq. 16 can be written as

$\begin{split}\left\langle\hat{S}_{\vec{r}_{a}}^{z}\hat{S}_{\vec{r}_{b}}^{z}\right\rangle&=-\frac{1}{2}\sum_{ijkl}P_{ij}^{a}P_{kl}^{b}\left(\frac{\Gamma_{ijkl}}{6}+\frac{\Gamma_{kjil}}{3}\right)\\\
&+\frac{1}{4}\sum_{ijk}P_{ik}^{a}P_{jk}^{b}\left(\gamma_{ij}\right).\end{split}$
(20)

Since the $6\times 6$ hydrogen lattice results of Section III.3 are obtained
in a minimal basis of localized orthogonal orbitals, $P^{a}$ projects exactly
onto one local orbital used in the sampling basis of the VMC and can thus be
written as

$P_{ij}^{a}=\begin{cases}1&\text{if }i=a\text{ and }j=a,\\\
0&\text{otherwise}.\end{cases}$ (21)

For a given a variational state $\ket{\psi}$, the one- and two-body RDMs are
estimated via Monte Carlo estimation of local operators

$\displaystyle\gamma_{ij}^{loc}$
$\displaystyle=\frac{\braket{\mathbf{n}}{\hat{c}_{i}^{\dagger}\hat{c}_{j}}{\psi}}{\braket{\mathbf{n}}{\psi}},$
(22) $\displaystyle\Gamma_{ijkl}^{loc}$
$\displaystyle=\frac{\braket{\mathbf{n}}{\hat{c}_{i}^{\dagger}\hat{c}_{j}^{\dagger}\hat{c}_{k}\hat{c}_{l}}{\psi}}{\braket{\mathbf{n}}{\psi}},$
(23)

Evaluating these local operators involves finding all the connected Fock
configurations $\mathbf{n}^{\prime}$ for a given $\mathbf{n}$ sampled from
$\ket{\psi}$ according to the corresponding Born distribution
$|\psi(\mathbf{n})|^{2}$. In the case of the one-body RDM, connected
configurations are obtained by moving one electron from orbital $j$ to orbital
$i$, while for the two-body RDM, they are obtained by a two-electron move. The
Fermionic commutation relations then imply that a sign is picked up for every
electron that is passed during the move, within the pre-established orbital
ordering. This can be accounted for by computing a parity factor
$P_{\mathbf{n}\mathbf{n}^{\prime}}$ for each pair of configurations
$\mathbf{n}$ and $\mathbf{n}^{\prime}$.