Second-order optimisation strategies for neural network quantum states

In the following, we provide a general discussion on optimisation strategies for quantum many-body systems. For concreteness and pedagogical purposes, several aspects of the discussion are illustrated with a simple many-body system [7]. The toy model we consider is a system of $A$ fermions of mass $m$, trapped in a one-dimensional harmonic well [7]. For simplicity, we assume the system to be fully polarised and thus neglect spin effects. We employ Harmonic Oscillator (HO) units and choose a Gaussian-shaped two-body interaction with range $\sigma_{0}$ and strength $V_{0}$, inspired by nuclear-physics applications [7]. The Hamiltonian of the system thus reads

Throughout this work, we use this system as a toy model to benchmark the performance of different energy optimisation schemes. Specifically, we change the number of particles from $A=2$ to $6$ and the interaction strength, $V_{0}$, from $-20$ to $20$. We keep a fixed interaction range, $\sigma_{0}=0.5$, which allows for an exploration of the non-perturbative regime [7].

This system has the advantage to be simple while displaying an interesting phenomenology. Its simplicity allows us to explore many different optimisation scenarios while, its phenomenological diversity supports the fact that our conclusions can not be too specific to the many-body problem we consider. In particular, when varying $V_{0}$ between $-20$ and $20$, the many-body system undergoes a transition from a strongly attractive regime, where spinless fermions behave like a non-interacting bosonic system [36, 37, 38], to a strongly repulsive regime, akin to that of a crystal [39, 40].

As an illustrative example, we show in Fig. 1 the one-body density profile, $n(x)$, as a function of the spatial position for the $A=6$ system. The different panels correspond to different interaction strengths. We show our NQS prediction with a solid line, but provide also two additional benchmarks. Dashed lines correspond to a configuration interaction (CI) result. This employs a full direct diagonalisation of the Hamiltonian within a given configuration space [41]. The associated Hartree-Fock (HF) results are shown with dash-dotted lines.

The central panel (c) shows the non-interacting case, $V_{0}=0$. It displays an overall Gaussian-like profile with a series of $6$ superimposed peaks, which are characteristic of fermion systems. Panels (a) and (b), to the left of the center, show the attractive regime with $V_{0}=-20$ and $V_{0}=-10$, respectively. In this region, the wiggles disappear and the density profile acquires a dumbbell shape. This relatively featureless density profile is reminiscent of a bosonic system [7]. In this regime, the CI results disagree with the NQS prediction because of the truncation of the model space, which is performed for practical numerical reasons. Here, and in the following, we will not show CI results for $V_{0}=-20$ for this reason. In contrast, the HF results agree relatively well with the NQS prediction in the attractive regime.

Panels (d) and (e) correspond to the repulsive regime for values of the interaction strength $V_{0}=10$ and $20$. They suggest a very different behaviour of the density profile in the repulsive regime. The system does not change substantially in average size, but the fermionic oscillations on top of the smooth behaviour become much more pronounced. This is a sign of localisation, or crystallisation, in the system. We note that here the HF density profile differs substantially from the NQS result, although the energy predictions (not shown here) are relatively close to the NQS predictions. For more details on the physics of this model, we refer to our previous work in Ref. [7].

To compute the ground state of the Hamiltonian in Eq. (1) we follow a VMC approach [12] and aim at solving the many-body problem using the Rayleigh-Ritz variational principle,

where $E^{A}_{\text{g.s.}}$ denotes the exact $A$-body ground-state energy of $H$. In VMC, the wave function $\ket{\Psi^{A}}$ is not in general the exact ground state, but rather a trial wave function $\ket{\Psi_{T}}$ that depends on a series of parameters. Using the Rayleigh-Ritz variational principle, the original eigenvalue problem can thus be recast as an optimisation problem in the subspace of trial wave functions.

In practice, three major challenges need to be faced to solve this optimisation problem. First of all, the Hilbert space of $A$-body states in the optimisation search is in principle infinite-dimensional. This difficulty is usually tackled by introducing the much smaller set of trial states $\ket{\Psi_{T}}$ over which the energy is minimised. In doing so, a bias is introduced on the solution, which is usually estimated by varying the size of the set of trial states. A second difficulty comes from the necessity to evaluate the expectation value of the energy over many different trial states. Each estimate consists of an integral of at least dimension $D\times A$, where $D$ denotes the number of space dimensions. To mitigate the cost of evaluating these high-dimensional integrals, one resorts to Monte Carlo estimates, which necessarily introduce a statistical noise on the energy. This can in principle be controlled by using a large enough number of samples in the Monte Carlo estimate. Finally, the third challenge is to find the solution of a global non-linear optimisation problem. Here, the established strategy consists in breaking the problem into a sequence of simpler, local-quadratic optimisation sub-problems, which are solved iteratively. Such iterative optimisers have the advantage of versatility and can be used for many different optimisation problems. However, they often require adjustments to converge in a reasonable time [42].

The third challenge is the main focus of this work, and we will address it in the Section 3. Before diving into it, we provide some details on our NQS strategy which will be useful to frame the global discussion on the optimisation strategy.

Our NQS ansatz is a fully antisymmetric NN, inspired by the implementation of FermiNet [3]. The inputs to our network are the $A$ positions of fermions in the system, $\{x_{i},i=1,\ldots,A\}$, and the output is the many-body wave function $\Psi_{\theta}(x_{1},\ldots,x_{A})$ in the log domain. The state depends on a series of network weights and biases, succinctly summarised by a variable $\theta$ of dimension $N$ [7]. The network is composed of $4$ core components: equivariant layers, generalised Slater matrices (GSMs), Gaussian log-envelope functions, and a summed signed-log determinant function [43]. The equivariant layers have the particularity of sharing a common set of weights and biases among the $A$ input positions. This weight sharing is critical in order to enforce total antisymmetry of the wave function and, as discussed later, will impact the rationale justifying KFAC. For all calculations presented in this work, the number of equivariant layers $L$ is set to $L=2$; their width, $H$, is $H=64$; and we only use one GSM. We stress that the orbitals making a GSM are functions depending on the positions of all the particles, thus departing from a standard Slater matrix. These so-called backflow correlations are key in obtaining a faithful representation of the exact state [7]. In total, our wave function ansatz depends on about $10^{4}$ parameters, which have to be optimised to solve the many-body problem.

The expectation value of the energy, in Eq. (2), is typically formulated as a statistical average over local energies,

where we define $X=(x_{1},\ldots,x_{A})$ as an $A$-dimensional random variable (or walker) distributed according to the Born probability of the many-body wave function, $\absolutevalue{\Psi_{\theta}}^{2}$. For numerical stability, the kinetic energy is computed in the log domain. For a general local interaction, the local energy in a continuous model reads

For our specific model, the interaction terms are

To sample the statistical averages over the Born probability $\absolutevalue{\Psi_{\theta}}^{2}$, we use $N_{W}=4096$ walkers propagated following a Metropolis-Hastings algorithm (MH) [44, 45] with on-the-fly adaptation of the proposal distribution width [30]. In order to neglect any biases due to autocorrelation in our comparison of multiple optimisers, we perform $400$ MH iterations per epoch, to ensure the thermalisation of the Markov chain. The optimisation process requires not only the evaluation of the energy expectation value, $E(\theta)$, but also the gradients with respect to all the parameters of the network $\nabla_{\theta}E(\theta)$. These are evaluated as a statistical average according to

Finally, to provide a reasonable starting point to the energy minimisation algorithms, we drive the NQS through a pre-training phase. This is a supervised learning problem, where we demand that the network reproduce the many-body wave function of the non-interacting system. The goal is to obtain an initial state that is physical and somewhat similar, in terms of spatial extent, to the result after the interaction is switched on. This also helps speed up the simulation in the more demanding cases [43, 7]. Throughout this work, the parameters of the network are randomly initialised before going through $10^{3}$ pre-training epochs using the Adam optimiser. The final pretrained NQS is the starting point for all the optimisation algorithms that we consider. For a given particular system (number of particles $A$ and interaction strength $V_{0}$), we always start from the same exact NQS and set of samples. This guarantees that any differences that arise in the optimisation process are due to the optimisers themselves. We refer the reader to the Supplementary Material and to Refs. [43, 7] for more details on the Monte Carlo sampling and the pre-training used in our calculations.

Solving a global, non-linear optimisation problem is a difficult task. In this work, we consider optimisers that follow the same underlying strategy, sometimes referred to as sequential quadratic programming [42]. In our case, the general problem takes the form,

This problem is divided into a sequence of simpler sub-problems which are solved iteratively. The $n^{\text{th}}$ sub-problem consists of (a) an initial guess on the set of parameters, $\theta_{n}$; (b) a quadratic model, $M_{n}$, approximating the loss function $E(\theta)$ locally; and (c) a trust region, $T_{n}$, specifying where the local model is expected to be a good approximation to the loss function [42]. Formally, solving the $n^{\text{th}}$ sub-problem consists in computing a new energy estimate, $E_{n+1}$, at the updated parameters, $\theta_{n+1}$, namely

Here, $\delta$ represents the update in the parameters and the quadratic model $M_{n}$ takes the form

where $Q$ is a matrix; $L$, a vector; and $C$, a scalar. In practice, the constrained quadratic optimisation sub-problem is replaced by an equivalent unconstrained regularised sub-problem,

where $\lambda_{n}$ is a damping parameter, characterising the size of the equivalent trust region and $R_{n}$ is a symmetric positive semi-definite matrix describing the shape of the trust region. This procedure is often referred to as Tikhonov regularisation [46, 42]. Importantly, not only are the energy and the parameters, $(E_{n},\theta_{n})$, updated at each iteration, but also the local quadratic model and its trust region (or regularisation), $(M_{n}(\delta),T_{n})$.

As an instructive example, the standard gradient descent method can be formulated in this framework. In this case, one solves iteratively the locally constrained sub-problems associated to the model

with trust region

where $\norm{.}_{2}$ denotes the standard Euclidean norm. In this case, the update on the parameter $\delta_{n}$ reads explicitly

When solving this sequence of local constrained sub-problems, one is effectively performing a gradient descent with a learning rate $\alpha$.

Equations (11)-(12) indicate that the standard gradient descent algorithm relies on setting an isotropic constraint on the size of the update parameter, $\delta_{n}$, based on the Euclidean norm. In a general setting, the Euclidean norm has no reason to be an optimal choice. Different gradient descent algorithms have been designed in an attempt to outperform the standard version presented in the previous subsection. The natural gradient descent (NGD) algorithm, which has been advocated to be optimal for a class of optimisation problems, is of particular interest [47, 48]. This algorithm is the starting point of the KFAC optimiser [27].

The NGD was originally defined in a supervised learning setting, where the goal is to learn an empirically sampled distribution $q$ by minimising a loss function. One particular loss function, chosen here for pedagogical reasons, is the cross-entropy, namely

Here, $p_{\theta}$ is a probability distribution, depending on several parameters $\theta$, that is adjusted to minimise the cross-entropy loss. In this case, the manifold of the probability distributions $p_{\theta}$ possesses a natural geometry commonly referred to as information geometry [31, 32]. The information geometry is defined by the Kullback-Leibler (KL) divergence between two probability distributions

Although the KL divergence does not satisfy all the axioms of a distance, it defines locally a metric obtained by Taylor-expanding it at second-order, namely

where $F(\theta)$ is the Fisher information matrix (FIM). The expansion is performed around the "diagonal" point, $p_{\theta}=\ p_{\theta^{\prime}}$. From the expansion, one obtains an explicit expression for the FIM,

where $\theta_{i}$ and $\theta_{j}$ are individual components of the global parameter vector, $\theta$. In this context, the NGD consists simply in performing a gradient descent with a trust region defined by the norm associated to the FIM. For example, $T^{\text{NGD}}_{n}=\set{\delta:\delta^{\mkern-1.5mu\mathsf{T}}F(\theta_{n})\delta\leq r^{2}},$ where $r>0$ denotes the maximum size of an update on $\delta_{n}$. As mentioned earlier, it is more practical to work with an equivalent unconstrained regularised problem. In the case of NGD for VMC, the regularised local model is

where the damping $\lambda_{F}$ is a function of the size of $T^{\text{NGD}}_{n}$. Equivalently, the update minimising $M^{\text{NGD-reg}}_{n}(\delta)$ can be expressed as

where $\alpha$ is a learning rate related to the size of $T^{\text{NGD}}_{n}$.

In practice, NGD has been shown to be efficient for specific kinds of supervised learning problems [47, 48]. Additionally, stochastic reconfiguration, a quantum generalization of NGD, has been observed to perform well in the context of VMC [49, 50, 9]. The exact reasons for the success of the NGD are still under study [29]. Apart from empirical tests, two critical properties of the FIM have been raised for explaining its performance: (i) the FIM is symmetric definite positive and (ii) it provides an approximation to the Hessian of the loss function in the case of supervised learning [29]. Point (i) is important to avoid instabilities in the directions associated to negative eigenvalues of the quadratic term, so that updates in the parameters remain bounded. Point (ii) leads to a re-interpretation of the NGD as a second-order optimisation scheme, where the Hessian of the loss function is approximated by the FIM. This motivates the minimisation of a model

This interpretation leads to the update

which is equivalent to fixing the learning rate to unity, $\alpha=1$, in Eq. (19).

A practical implementation of the NGD algorithm requires computing the update in Eq. (21). In the case of NN models, the number of parameters, $N$, is large ($N\gg 10^{4}$). Computing exactly the update $\delta^{\text{NGD}}_{n}$ at every epoch can be computationally prohibitive. The goal of the KFAC algorithm is to compute an approximation to the update $\delta^{\text{NGD}}_{n}$ with a time complexity of $O\left(N\right)$ while keeping the number of epochs to convergence roughly constant. We note that, for the specific case of NQS, there are alternative proposals to implement this update calculations employing fast Jacobian-vector product algorithms [51, 52, 53].

Consider a feed-forward NN of arbitrary depth. For a given layer $l$, we denote the weight matrix $W^{(l)}$ and, for simplicity, we assume it to be of size $H\times H$ and without bias. The input row-vector $f^{(l-1)}\in\mathbb{R}^{1\times H}$ and the output row-vector $h^{(l)}\in\mathbb{R}^{1\times H}$ are related according to $h^{(l)}=\tanh\left(f^{(l-1)}W^{(l){\mkern-1.5mu\mathsf{T}}}\right)\ ,$ where the activation function $\tanh$ is applied element-wise¹¹1We focus on this activation function because it is the one of choice in our NQS model [7], but any other choice would not change our conclusions. . In this work, we follow the convention used in the PyTorch framework where inputs and outputs of a layer are row vectors.

In the KFAC algorithm originally described in Ref. [27], the FIM is approximated in two ways. First, the matrix elements between parameters of different layers are neglected. The approximated FIM, $\breve{F}(\theta)$, is block-diagonal and reads

where $W^{(l)}_{ij}$ is the parameter associated to the matrix element $(i,j)$ of the $l^{\text{th}}$ layer weight matrix. Second, each layer-dependent block is approximated by a Kronecker product. Using the chain rule, the FIM derivatives can be rewritten as

Introducing the so-called forward activations, $a^{(l-1)}$, and the backward sensitivities, $e^{(l)}$, as the row vectors

we obtain

where $\text{vec}(.)$ and $(.\otimes.)$ denote respectively the standard vectorisation (stacking columns of a matrix) and Kronecker product operations. Using Eq. (25), the $l^{\text{th}}$ block of the FIM can be rewritten in vectorised form as

Finally, the Kronecker-factored approximation on the $l^{\text{th}}$ block of the FIM, $\breve{F}^{(l)}_{\text{KFAC}}(\theta)$, is obtained by exchanging the order of the Kronecker product and the statistical average, i.e.

The resulting block-diagonal KFAC FIM will be denoted as $\breve{F}_{\text{KFAC}}(\theta)$ from now on.

Assuming a unit learning rate, the update in the parameters when using the KFAC approximation thus reads

We employ a capital $\Delta^{\text{KFAC}}_{n}$ to differentiate it from previous parameter updates. The main advantage of the KFAC approximation is that, instead of having to invert a matrix of dimension $(L\times H^{2})\times(L\times H^{2})$, one can invert separately the $2L$ Kronecker factors in Eq. (27), which are much smaller matrices of dimension $H\times H$. The main drawback is that the KFAC approximation is very crude and can only catch the very coarse structure of the exact FIM [27]. For completeness, let us also mention that the two Kronecker factors are averaged over successive epochs, with a maximum exponentially decay value set to $0.9$ [27].

To compensate for the poor quality of the update $\Delta^{\text{KFAC}}_{n}$, a second step was designed in the original KFAC algorithm [27]. The idea is that, whenever $\Delta^{\text{KFAC}}_{n}$ provides a good direction, we would like to take a big step. In contrast, when the direction is poor, we would like to take a small step, to avoid increasing the loss function. In the original KFAC algorithm, this is achieved by computing an optimal rescaling factor $\alpha$ such that the quadratic model $M^{\text{NGD}}_{n}\left(\alpha\Delta^{\text{KFAC}}_{n}\right)$, which uses the exact FIM, is minimised. In addition, a momentum in the previous update $\delta^{\text{KFAC}}_{n-1}$ is usually included. A natural way to include both improvements consists in minimising $M^{\text{NGD}}_{n}\left(\alpha\Delta^{\text{KFAC}}_{n}+\mu\delta^{\text{KFAC}}_{n-1}\right)$ over the two parameters $\alpha$ and $\mu$. Explicitly, the optimal coefficients $\alpha^{*}$ and $\mu^{*}$ are obtained from the linear system

where we momentarily drop the superindex KFAC in all the $\Delta_{n}$ and $\delta_{n}$ updates for simplicity. Using the exact FIM is crucial to help correct cases where the KFAC approximation is too crude. This correction remains relatively cheap, since Eq. (29) only requires computing a couple of matrix-vector calculations.

Last but not least, the original derivation of KFAC describes a particular Tikhonov regularisation in order to stabilise the overall algorithm [27]. Such a regularisation is of critical importance, since it limits how large the update in the parameters can be. If the trust region is too small, the benefits of using a second-order optimiser will be null. If the trust region is too large, instabilities will kick in and ruin any chances of reaching convergence.

The KFAC FIM is used to estimate the direction of the update, while the exact FIM is used to evaluate an optimal rescaling of the direction. Consequently, two distinct trust regions are considered, which reflect the different level of trust one has in the direction, $\Delta^{\text{KFAC}}_{n}$, and the rescaling factor, $\alpha^{*}$. Those are implemented by introducing two damping factors, $\gamma_{n}$ and $\lambda_{n}$, which regularise $\breve{F}_{\text{KFAC}}(\theta_{n})$ and $F(\theta_{n})$, respectively. More explicitly, the regularised matrices are obtained from the substitutions

where $\pi^{(l)}$ is a factor that is automatically adjusted to minimise the cross-term obtained when expanding the Kronecker product [27]. In practice, the damping factors are dynamically adapted every $T$ epochs.

While we try to be as close to the original KFAC algorithm as possible, our numerical implementation for the NQS toy model slightly differ in the dynamically adjustment of the damping factors. The damping factor $\lambda_{n}$ is adapted following a Levenberg-Marquardt rule [54], namely

The reduction ratio,

measures the discrepancy between the actual variation of the loss (numerator) and the anticipated one from the quadratic model (denominator). Unlike in the original KFAC implementation, however, we also account for fluctuation on the energy estimation by introducing a tolerance on increasing the energy, i.e. if $E(\theta)\leq E(\theta+\delta)\leq E(\theta)+\epsilon$, we default to the case where the damping $\lambda_{n}$ is reduced. In all our calculations $\epsilon$ is set to $10\%$ of $E(\theta)$. The damping $\gamma_{n}$ is adapted with a greedy algorithm by testing which damping among $(\gamma_{n},\omega_{2}\ \gamma_{n},\frac{\gamma_{n}}{\omega_{2}})$ leads to the best update. For all the calculations performed in this work, we set $T=5$, $\omega_{1}=({19}/{20})^{T}$, $\omega_{2}=\sqrt{{19}/{20}}^{T}$ and the initial damping values are set to $\lambda_{0}=10^{3}$ and $\gamma_{0}=\sqrt{10^{3}}$. For clarity, we provide a schematic overview of the KFAC algorithm in Fig. 2. For more details, we refer the reader to Ref. [27].

Since its inception, the KFAC optimiser has been tested on several deep autoencoder optimisation problems with standard datasets (MNIST, CURVES and FACES), which are commonly used to benchmark NN optimisation methods [55, 56]. The focus of these models is on supervised learning, with a standard feed-forward NN. Beyond such standard cases, the original KFAC optimiser may need to be adapted. For instance, if convolution layers are a part of the NN architecture, Refs. [28, 57] suggest modifications to the KFAC approximation to the FIM. In the case of VMC with NQS, several attempts have been put forward to exploit the benefits of KFAC [3, 30, 58]. To bypass instabilities encountered when applying the standard KFAC optimiser, several ad hoc adjustments have been implemented to rectify the update on the parameters [3]. In particular, Ref. [3] describes the use of a scheduled learning rate, a constant damping, a norm constraint and the discarding of any momentum in the parameter updates. Apart from the lack of theoretical motivation for such adjustments, it is consistently reported within the literature that heuristics and fine-tuning are required to obtain good performance [3, 30, 58].

In this section, we advocate for a different approach, which avoids any such heuristic adjustment of hyperparameters, while remaining close to the original philosophy of Ref. [27]. We start by providing an overview of the performance associated to our NQS KFAC implementation for the toy model described in Sec. 2, before suggesting improvements to the formulation of the algorithm. We show the evolution of the energy as a function of the epoch in the panels of Fig. 3. From top to bottom, the different rows correspond to increasing number of particles, from $A=2$ in the top to $A=6$ in the bottom row. From left to right, the results scan the interaction strength in a range $V_{0}\in\set{-20,-10,0,10,20}$. Dashed lines provide an indication of the expected results for the HF approximation and for the CI case. The optimisation process in all cases is stopped after $1000$ epochs. Very few minimisations reach this final number and most lead to numerically unstable results much earlier than that. As in previous NQS implementations[3, 30, 58], we confirm that the performance of the KFAC optimiser is indeed fluctuating. In many cases, the energy decreases steadily within the first $50$ to $100$ epochs, only to become extremely erratic after that point. None of the KFAC runs manage to converge to the benchmark values, regardless of the number of particles and of the strength of the interaction. We take this as a confirmation that the KFAC algorithm, in its original formulation, is unreliable for VMC with NQS.

One option to compensate the strong oscillations is to try and improve the rescaling phase of the optimiser. The main argument put forward in favour of using $M^{\text{NGD}}_{n}$ in the rescaling phase is that the FIM can be interpreted as an approximation to the Hessian of the loss function [27, 29]. However, as emphasized in Refs. [27, 29], the validity of this argument relies critically on the assumption that one is trying to solve a supervised learning problem. In this case, the loss function quantifies how far the probability distribution $p_{\theta}$, modelled by the NN, is from the target distribution, $q$. An example of such loss function is the cross-entropy, given in Eq. (14). In VMC, however, one is solving a reinforcement learning problem, which is fundamentally different. In this case, the loss function is the energy $E(\theta)$ which characterises the quality of the NQS independently of the knowledge of the target wave function. Consequently, the underlying rationale justifying the use of the FIM as an approximate Hessian is lost.

In order to restore the validity of KFAC in VMC calculations, we replace the exact FIM used in the rescaling phase by a quasi-Hessian matrix, $H_{Q}(\theta)$, whose role is to provide a better quadratic model than $M^{\text{NGD}}_{n}$, namely

Throughout this work, we will refer to the resulting optimisation algorithm as quasi-Newton KFAC (QN-KFAC), and we employ the superscript QN in the corresponding quantities. For supervised learning problems, the performance of NGD has been connected to the similarity between the FIM and the Generalised-Gauss-Newton matrix [29]. Specifically, this matrix can be obtained by neglecting the Hessian of the output of the NN in the calculation of the Hessian of the loss function [29]. To improve the rescaling phase of the algorithm, we follow a similar path and define the quasi-Hessian $H_{Q}(\theta)$ as the matrix obtained from the Hessian of the energy $E(\theta)$, after neglecting the Hessian of $\ln\absolutevalue{\Psi_{\theta}}$. More specifically, the Hessian of the energy reads

Neglecting the first line, the induced quasi-Hessian reads

A key advantage of using a quasi-Hessian with this structure is that matrix-vector operations can be performed as fast as for the FIM. However, a critical drawback is that the $H_{Q}(\theta)$ is not in general positive semi-definite, which can lead to instabilities if the damping $\lambda_{n}$ is not large enough. In addition, the formula given in Eq. (35) requires to compute third-order derivatives of the wave function. These arise from the second-order derivatives with respect to the input in the kinetic energy term of the local energy, and the derivatives with respect to the parameters. These derivatives should be handled with care in order to avoid numerical instabilities [3]. A schematic overview of the QN-KFAC algorithm is provided in Fig. 2.

We now proceed to compare the VMC energy evolution employing QN-KFAC and KFAC. We stress that we perform systematic calculations with all hyperparameters of the two optimisers kept fixed to the aforementioned values (see also the Supplementary Material). This minimises the differences in the optimisation scheme and provides a general overview in different physical conditions. We observe that QN-KFAC (green lines) leads to a dramatic increase in overall stability. We find that for a small number of particles, $A\leq 3$, QN-KFAC almost always manages to converge to reasonable values. For larger particle numbers, the convergence pattern is more erratic. In some cases, we notice strong oscillations. In others, we observe jumps in energy predictions. The best-behaved convergence patterns are found for $V_{0}=0$ and $V_{0}=-10$. While this improvement is encouraging, we also notice that, even in the best cases, the late stage convergence can be slow. For instance, if we focus on the second column, corresponding to $V_{0}=-10$, we observe that, after a quick convergence during the first $100$ epochs, the QN-KFAC optimiser seems to struggle to converge fully before reaching our stopping criterion of $1000$ epochs. Comparing QN-KFAC with CI in the $A=2$ case (top panel), it is clear that QN-KFAC converged to the exact result. However, for $A=6$ particles and the same interaction strength, in the absence of exact CI benchmarks, it is not clear whether or not the QN-KFAC finished converging.

The QN-KFAC optimiser provides a first improvement over KFAC in our NQS toy model simulations. This was achieved by modifying the quadratic model used in the rescaling phase of the KFAC algorithm. In order to tackle the remaining issues observed when testing QN-KFAC, like the presence of strong oscillations in some cases and the slow convergence pattern in the late stage of the optimisation process, we now proceed to improve on the direction estimation phase of the optimiser.

The original KFAC algorithm was developed for supervised learning problems with standard feed-forward NNs. In the case of a reinforcement learning problem like VMC, the theoretical motivation for using the FIM is lost. This prompted the changes discussed in the previous subsection. In this subsection, we further argue that using an NQS architecture additionally breaks the original rationale motivating the Kronecker factorisation, which lies at the heart of the KFAC algorithm [27].

As mentioned in section 3.3, the direction of the update in the parameters $\delta^{\text{KFAC}}_{n}$ is obtained based on two approximations of the complete FIM: (i) the FIM is assumed to be block-diagonal in the layers and (ii) each block is further approximated by exchanging the Kronecker product and the statistical average, thus replacing Eq. (26) by Eq. (27). The latter approximation is discussed at length in the original derivation of KFAC in Ref. [27]. The validity of Eq. (26), however, is unclear. The chain rule used to derive it relies on the specific NN type and architecture. In the case of a standard feed-forward NN, considered in Ref. [27], Eq. (26) does obviously hold. Equation (26) however breaks down in the case of convolution layers and the original KFAC algorithm was adapted accordingly in Ref. [28].

In the case of NQS for fermionic systems, permutation-equivariant layers also break Eq. (26). Specifically, let us consider the $l^{\text{th}}$ permutation-equivariant layer with a shared weight matrix $W^{(l)}\in\mathbb{R}^{H\times H}$ with input $f^{(l-1)}\in\mathbb{R}^{A\times H}$ and output $h^{(l)}\in\mathbb{R}^{A\times H}$, related through the activation function $h^{(l)}=\tanh\left(f^{(l-1)}W^{(l){\mkern-1.5mu\mathsf{T}}}\right).$ Using the chain rule, the derivatives in the FIM for a fermionic NN now read

In this case, the forward activities $a^{(l-1)}$ and backward sensitivities $e^{(l)}$ are not row vectors, but rather matrices

With this, we recover formally the same result as Eq. (25),

However, because the forward activities and backward sensitivities are now matrices, we no longer find a Kronecker product structure for the partial derivatives, i.e. $\text{vec}\left(e^{(l)}a^{(l-1){\mkern-1.5mu\mathsf{T}}}\right)\neq a^{(l-1)}\otimes e^{(l)}.$ Instead, the $l^{\text{th}}$ block of the FIM in vectorised form reads

where $I_{A}$ is the $A\times A$ identity matrix. In this result, we have used the general formula $\text{vec}\left(BCD\right)=(D^{\mkern-1.5mu\mathsf{T}}\otimes B)\text{vec}\left(C\right)$ [59], which holds for any matrices $B$, $C$ and $D$ with compatible matrix products. Comparing Eq. (26) to Eq. (39), one sees that the weight-sharing introduces an intermediate $A^{2}\times A^{2}$ factor, $\text{vec}\left(I_{A}\right)\text{vec}\left(I_{A}\right)^{\mkern-1.5mu\mathsf{T}}$, which prevents us from recasting the FIM as the average of a Kronecker product.

Interestingly, this analysis suggests that the KFAC rationale remains valid for the $A=1$ case. We expect that it will become less and less valid as the number of particles, $A$, increases. This seems to support our empirical observation that QN-KFAC performs better for small values of $A$. The results obtained with QN-KFAC in Fig. 3 look indeed more stable, and converge faster, for $A=2$ and $3$ than for $A\geq 4$ .

In order to compensate for the incompatibility of the original Kronecker factorisation motivation when employing NQS, we introduce an additional step in QN-KFAC. Our goal is to improve on the direction of the update $\delta^{\text{KFAC}}_{n}$ starting from the one obtained from KFAC. To do so, we use an iterative solver for linear equations. We employ MINRES [60], an algorithm that solves linear systems, $Mx=y,$ for sparse, symmetric matrices, $M$, and output vectors, $y$. The underlying idea of MINRES consists in combining a Lanczos algorithm with an LQ factorisation in order to minimise the squared residual $r(x)^{2}\equiv(Mx-y)^{\mkern-1.5mu\mathsf{T}}(Mx-y).$ In our case, we use MINRES to find an approximate solution to the system

using $\Delta^{\text{KFAC}}_{n}$ as an initial guess. We recall that $\breve{F}$ is the block-diagonal approximation of the FIM as defined in Eq. (26). For simplicity, we choose to reuse $\lambda_{n}$ to regularise the directional update improvement, instead of $\gamma_{n}$. Since we use the block-diagonal FIM without the KFAC in Eq. (40) we consider it to be reasonable to choose the same trust region used for the rescaling (which uses the exact FIM). Following Ref. [61], we also use the preconditioner $\left(\text{diag}(\breve{F}(\theta_{n}))+\kappa I\right)^{\xi}$ with $\kappa=10^{-2}$ and $\xi=0.75$ to accelerate the convergence of the linear solver. In addition, MINRES is stopped as soon as we achieve a relative residual smaller than $10^{-8}$ or when the number of iteration reaches a maximum number, $N_{\mathrm{MR}}$. We dub the resulting optimiser QN-MR-KFAC. We note that when $N_{\mathrm{MR}}=0$, QN-MR-KFAC reduces to QN-KFAC. For clarity, we provide a schematic overview of the QN-MR-KFAC algorithm in Fig. 2.

We now test the impact of using a refined directional update, $\delta^{M}_{n}$, instead of the bare KFAC estimate, $\delta^{\text{KFAC}}_{n}$. We focus first on cases where QN-KFAC converges reasonably well without violent oscillations. We show the energy as a function of the epoch number for $V_{0}=-10$ with $A=2$ (top panel) and $A=4$ (bottom panel) particles in Fig. 4. Specifically, we illustrate the impact of MINRES by exploring a different set of iterations, $N_{\mathrm{MR}}\in\set{0,1,5,25,50}$.

The results clearly indicate that improving the direction of the parameter update (or increasing $N_{\mathrm{MR}}$) has two key effects. First, the oscillations in the energy estimate during the minimisation process are substantially reduced. Even for $A=2$ (top panel), where the oscillations were relatively small, increasing $N_{\mathrm{MR}}$ leads to substantially reduced oscillations in the energy predictions. For $4$ particles (bottom panel), QN-KFAC, which corresponds to the $N_{\mathrm{MR}}=0$ line, has relatively large oscillations in energy, of order $0.1$ in HO units. When $N_{\mathrm{MR}}$ increases to $5$, the oscillations are reduced by about a factor of $2$. A further increase up to $N_{\mathrm{MR}}=50$ shows minor oscillations, to less than about $0.03$ in HO units. Second, the convergence in the late stages of the optimisation is considerably faster. For $A=2$, QN-KFAC only converges to the CI result at around the $1000^{\text{th}}$ epoch. When the MINRES solver performs more than $25$ iterations per epoch, the optimisation process reaches the CI value in $\approx 120$ iterations. In the $4$-particle case, the improvement is more dramatic. QN-KFAC plateaus, but does not converge, at the $1000^{\text{th}}$ epoch. Increasing $N_{\mathrm{MR}}$, a more transparent convergence pattern appears and, for $N_{\mathrm{MR}}=50$, about $400$ epochs are enough to reach the minimum. We stress that the characteristically slow approach to the minimum was a major drawback of the QN-KFAC proposal.

Overall, we observe that improving the quality of the direction of $\delta_{n}$ substantially impact the performances. Comparing the $A=2$ and $A=4$ cases, we see that the number of MINRES iterations necessary to reach optimal convergence is larger as $A$ increases. This supports our previous theoretical analysis expecting KFAC to be less and less well motivated as $A$ increases. In practice, we find that setting $N_{\mathrm{MR}}=50$ is sufficient for any case where a stable convergence is reached.

A more systematic analysis of the importance of the improvement in the direction of the updates is provided in Fig. 5. Here, we show the same results as in Fig. 3, but comparing QN-KFAC to QN-MR-KFAC with $N_{\text{MR}}=50$. The results confirm that QN-MR-KFAC greatly improves on both the oscillations of the energy values and the late phase of convergence. In the vast majority of cases, QN-MR-KFAC performs better than QN-KFAC in terms of accuracy and speed of convergence. Concerning the stability of the optimisation process, the improvements are more marginal. The attractive and non-interacting cases (columns to the left of, and including, the center) only show some erratic peaks in energy in a few, selected cases. In contrast, the repulsive cases are still prone to important oscillations as soon as $A\geq 4$.

Finally, we notice that in several cases (for example when $A=2$, $V_{0}=10$ and $A=5$, $V_{0}=-10$) QN-MR-KFAC shows instabilities even after having converged. We interpret this behaviour in terms of the dynamically adjusted damping. Once convergence is reached, the trust region keeps increasing. One may eventually hit a remote spurious point in the local quadratic model which has a value that is lower than the true ground-state energy and jump there. The associated numerical instability takes a few epochs to self-adjust. As stated before, this is one of the drawbacks of using a quasi-Hessian, $H_{Q}(\theta)$, which is not positive semi-definite. These pathological cases may in principle be cured by introducing a minimum value on the damping, $\lambda_{n}$. More generally, most of the fluctuations encountered so far can be tackled by being more restrictive on the trust regions. However, in doing so, we also restrict the size of the updates and slow down the overall convergence of the optimiser. We also introduce the risk of an artificial convergence, in the sense that the energy may converge only because the trust region has shrunk to a very narrow point. We leave further analysis of trust regions and their dynamical adjustments for future work.

Despite the multiple improvements to the original KFAC algorithm which we designed to take into account the particularities of VMC and NQS, our finest optimiser, QN-MR-KFAC, is still subject to violent fluctuations in the energy. As discussed in Sec. 3.5, this could be tackled by further refining the trust regions. Doing so, however, may limit the original promise of performance or require fine-tuning of the hyperparameters, as in past attempts to adapt KFAC to VMC with NQS [3, 30]. These attempts, including ours, rely on the use of the FIM, $F$, to estimate the direction of the updates $\Delta_{n}$. The use of the FIM was originally motivated by its connection to the Hessian of the loss function in a supervised learning setting. We have already argued that standard VMC-NQS does not fall in the category of supervised learning problems, and hence using the FIM (or its corresponding approximations) is no longer justified. The substantial improvement observed in Fig. 3, where the quasi-Hessian $H_{Q}$ is used in the rescaling phase of the optimiser instead of the exact FIM $F$, further supports this idea. We hypothesise that, the failure of KFAC are not only due to the approximations it makes to NGD, but also more fundamentally to the failure of NGD itself to provide a good optimiser for VMC.

To test this hypothesis, we perform calculations with an optimiser that is identical to QN-MR-KFAC with $N_{\mathrm{MR}}=50$, except that the exact FIM, $F$, as opposed to the quasi-Hessian, $H_{Q}$, is used in the rescaling phase. This is closer in spirit to Hessian-free optimisers [61]. In addition, to avoid any contamination from KFAC, we replace the initial guess in the MINRES solver by the previous update, rescaled by a prefactor $\zeta=0.95$, i.e. the starting point of the MINRES solver is now $\zeta\times\delta_{n-1}$ instead of $\Delta^{\text{KFAC}}_{n}$. This follows the default recommendation of Ref. [61] for Hessian-free optimisers. All other hyperparameters are fixed to the same value used in our previous QN-MR-KFAC calculations. Although we do not calculate exactly the natural gradient update, $\delta^{\text{NGD}}_{n}$ given in Eq. (21), we deem this to be sufficiently close to the NGD approach. In consequence, we will refer to this optimiser as NGD. A schematic overview of our implementation of NGD is provided in Fig. 2.

We present the results of this NGD optimiser in Fig. 6. The panels are organised in the same fashion as Fig. 3. In all the cases considered, NGD fails to converge smoothly to a final value. We find strong oscillations, akin to those observed for KFAC in Fig. 3. In some cases, particularly in the attractive regime (columns to the left of the center), the optimiser seems to get stuck around the HF solution. Even then, violent oscillations arise in the energy minimisation process. The generic failure of the NGD optimiser across different numbers of particles and interaction strengths further supports our hypothesis that NGD is fundamentally unfit for VMC. As discussed in Sec. 3.2, NGD can be seen as a simple gradient descent where the definition of trust regions are based on information geometry, i.e., on the KL divergence (15). We interpret the failure of NGD as a consequence of the unsuitability of information geometry for VMC.

This raises the question: which geometry should one use for VMC? Ideally, we would like to use a geometry leading to (i) a local metric that is positive semi-definite, for stability purposes; and (ii) a metric informed by the loss function that we are trying to minimise, here the energy $E(\theta)$. While both conditions are satisfied when considering information geometry for supervised learning problems, condition (ii) is violated in the case of VMC simulations. In order to restore condition (ii), we now introduce the concept of decision geometry.

Our starting point consists in reformulating VMC in a game-theory framework. In this paradigm, the standard information geometry is generalised by the concept of decision geometry [62, 63, 34, 32].

Similarly to how the KL divergence, Eq. (15), is related to the cross-entropy loss, Eq. (14), one can define a divergence function associated to any loss function $\mathcal{L}(\theta)$ of the type

so long as $\mathcal{S}(X,p_{\theta})$ is a so-called proper scoring rule [35]. A scoring rule $\mathcal{S}(x,p)$ is any integrable function that takes as inputs a probability distribution, $p$, and a sample, $x$, and returns a real valued quantity. From any scoring rule, one defines the associated expected score $S(p,q)$ for any two probability distributions $p$ and $q$ by

A scoring rule is said to be proper if and only if, for any two probability distributions $p$ and $q$,

In this framework, several typical quantities of information theory can be generalised. For example, the generalised entropy and the generalised cross-entropy are simply $S(p,p)$ and $S(p,q)$, respectively. Their associated divergence function, $D_{S}$, is defined by

The new geometry, associated to $D_{S}$, is referred as a decision geometry [63, 34] because of its relation with statistical decision problems [64]. Similarly to the KL divergence, $D_{S}$ defines a metric locally. This is obtained by Taylor-expanding at second-order, namely

where $G_{S}(\theta)$ is the decision metric associated to $S$. Let us emphasize that, by construction, $G_{S}(\theta)$ is symmetric positive semi-definite.

This symmetric and positive definite decision metric provides a natural starting point for optimisation schemes that are not tied to the FIM. In this context, we define a novel optimiser, the decisional gradient descent (DGD) which consists simply in performing a gradient descent with trust regions defined by the norm associated to the decision metric, $T^{\text{DGD}}_{n}=\set{\delta:\delta^{\mkern-1.5mu\mathsf{T}}G_{S}(\theta_{n})\delta\leq r^{2}}\ .$ Here, $r>0$ denotes the maximum size on the update $\delta_{n}$. Just like for NGD, it is more practical to work with an equivalent unconstrained regularised problem. This leads to the definition of an update

where $\alpha$ is again a learning rate that can be related to the size of $T^{\text{DGD}}_{n}$.

If we choose the scoring rule $S^{\text{sup}}(X,p_{\theta})\equiv-\ln{p_{\theta}(X)}$, the decision geometry framework is such that one recovers exactly the standard cross-entropy, Eq. (14), and the KL divergence, Eq. (15). In other words, DGD reduces to NGD. To go beyond NGD, we move away from the standard logarithmic scoring rule and choose instead an alternative one based on the local energy, namely