A simple linear algebra identity to optimize Large-Scale Neural Network
  Quantum States
Neural-network architectures have been increasingly used to represent quantum
many-body wave functions. These networks require a large number of variational
parameters and are challenging to optimize using traditional methods, as
gradient descent. Stochastic Reconfiguration (SR) has been effective with a
limited number of parameters, but becomes impractical beyond a few thousand
parameters. Here, we leverage a simple linear algebra identity to show that SR
can be employed even in the deep learning scenario. We demonstrate the
effectiveness of our method by optimizing a Deep Transformer architecture with
$3 \times 10^5$ parameters, achieving state-of-the-art ground-state energy in
the $J_1$-$J_2$ Heisenberg model at $J_2/J_1=0.5$ on the $10\times10$ square
lattice, a challenging benchmark in highly-frustrated magnetism. This work
marks a significant step forward in the scalability and efficiency of SR for
Neural-Network Quantum States, making them a promising method to investigate
unknown quantum phases of matter, where other methods struggle.
Riccardo Rende, Luciano Loris Viteritti, Lorenzo Bardone, Federico Becca, Sebastian Goldt
# A simple linear algebra identity to optimize

###### Abstract

Neural network architectures have been increasingly used to represent quantum many-body wave functions. In this context, a large number of parameters is necessary, causing a crisis in traditional optimization methods. The Stochastic Reconfiguration (SR) approach has been widely used in the past for cases with a limited number of parameters \(P\), since it involves the inversion of a \(P\times P\) matrix and becomes computationally infeasible when \(P\) exceeds a few thousands. This is the major limitation in the context of deep learning, where the number of parameters significantly surpasses the number of samples \(M\) used for stochastic estimates (\(P\gg M\)). Here, we show that SR can be applied exactly by leveraging a simple linear algebra identity to reduce the problem to the inversion of a manageable \(M\times M\) matrix. We demonstrate the effectiveness of our method by optimizing a Deep Transformer architecture featuring approximately 300,000 parameters without any assumption on the sign structure of the ground state. Our approach achieves state-of-the-art ground-state energy on the \(J_{1}\)-\(J_{2}\) Heisenberg model in the challenging point \(J_{2}/J_{1}=0.5\) on the \(10\times 10\) square lattice, a demanding benchmark problem in the field of highly-frustrated magnetism. This work marks a significant step forward in the scalability and efficiency of SR for Neural Network Quantum States (NNQS), making them a promising method to investigate unknown phases of matter of quantum systems, where other methods struggle.

+
Footnote †: These authors contributed equally.

_Introduction._ Deep learning is crucial in many research fields, with neural networks being the key to achieve impressive results. Well-known examples include using Deep Convolutional Neural Networks for image recognition [1; 2] and Deep Transformers for language-related tasks [3; 4]. The success of deep networks comes from two ingredients: architectures with a large number of parameters (often in the billions), which allow for great flexibility, and training these networks on large amounts of data. However, to successfully train these large models in practice, one needs to navigate the complicated and highly non-convex landscape associated with this extensive parameter space.

The most used methods rely on stochastic gradient descent (SGD), where the gradient of the loss function is estimated from a randomly selected subset of the training data. Over the years, variations of traditional SGD, such as Adam [5] or AdamW [6], have proven highly effective. These modifications enable more efficient optimization processes and lead to improved accuracy. In the late 1990s, Amari and collaborators [7; 8] suggested to use the knowledge of the geometric structure of the parameter space to adjust the gradient direction for non-convex landscapes, defining the concept of Natural Gradients. In the same years, Sorella [9; 10] proposed a similar method, now known as Stochastic Reconfiguration (SR), to enhance optimization of variational functions in quantum many-body systems. Importantly, this latter approach consistently outperforms other methods as SGD or Adam, leading to significantly lower variational energies. The main idea of SR is to exploit information about the curvature of the loss landscape, thus improving the convergence speed in landscapes which are steep in some directions and very shallow in others [11]. For physical inspired wave functions (e.g., Jastrow-Slater [12] or Gutzwiller-projected states [13]) the original SR formulation is a highly efficient method since there are few variational parameters, typically from \(O(10)\) to \(O(10^{2})\).

Over the past few years, neural networks have been extensively used as powerful variational _Ansatze_ for studying interacting spin models [14], and the number of parameters have increased significantly. Starting of simple Restricted Boltzmann Machines (RBMs) [14; 15; 16; 17; 18; 19], more complicated architectures as Convolutional Neural Networks (CNNs) [20; 21; 22] and Recurrent Neural Networks (RNNs) [23; 24; 25; 26] have been introduced to handle challenging systems. In particular, Deep CNNs [27; 28; 29; 30] have proven to be highly accurate for two-dimensional models, undoing methods as Density-Matrix Renormalization Group (DMRG) approaches [31] and Gutzwiller-projected states [13]. These deep learning models have great performances when the number of parameters is large. However, a significant bottleneck arises when employing the original formulation of SR for optimization, as it is based on the inversion of a matrix of size \(P\times P\), where \(P\) denotes the number of parameters. Consequently, this approach becomes computationally infeasible as the parameter count exceeds \(O(10^{4})\), primarily due to the constraints imposed by the limited memory capacity of current-generation GPUs.

Recently, Ref. Chen and Heyl [29] made a step forward in the optimization procedure by introducing an alternative method, dubbed MinSR, to train NNQS wave
functions. MinSR does not require inverting the original \(P\times P\) matrix but instead a much smaller \(M\times M\) one, where \(M\) is the number of configurations used to estimate the SR matrix. This is convenient in the deep learning setup where \(P\gg M\). Most importantly, this procedure avoids allocating the \(P\times P\) matrix, reducing the memory cost to \(P\times M\). However, this formulation is obtained as a minimization of the Fubini-Study distance with an _ad hoc_ constraint. In this work, we first use a simple relation from linear algebra to show, in a transparent way, that SR can be rewritten exactly in a form which involves inverting a small \(M\times M\) matrix and that only a standard regularization of the SR matrix is required. Then, we exploit our technique to optimize a Deep Vision Transformer (Deep ViT) model, which has demonstrated exceptional accuracy in describing quantum spin systems in one and two spatial dimensions [38; 39; 40; 41]. Using almost \(3\times 10^{5}\) variational parameters, we are able to achieve state-of-the-art ground-state energy on the most paradigmatic example of quantum many-body spin model, the \(J_{1}\)-\(J_{2}\) Heisenberg model on square lattice:

\[\hat{H}=J_{1}\sum_{\langle i,j\rangle}\hat{\mathbf{S}}_{i}\cdot\hat{\mathbf{S}}_{j}+J_{ 2}\sum_{\langle\langle i,j\rangle\rangle}\hat{\mathbf{S}}_{i}\cdot\hat{\mathbf{S}}_{j} \tag{1}\]

where \(\hat{\mathbf{S}}_{i}=(S_{i}^{x},S_{i}^{y},S_{i}^{z})\) is the \(S=1/2\) spin operator at site \(i\) and \(J_{1}\) and \(J_{2}\) are nearest- and next-nearest-neighbour antiferromagnetic couplings, respectively. Its ground-state properties have been the subject of many studies over the years, often with conflicting results [31; 37; 13]. Still, in the field of quantum many-body systems, this is widely recognized as the benchmark model for validating new approaches. Here, we will focus on the particularly challenging case with \(J_{2}/J_{1}=0.5\) on the \(10\times 10\) cluster, where there are no exact solutions.

Within variational methods, one of the main difficulties comes from the fact that the sign structure of the the ground state is not known for \(J_{2}/J_{1}>0\). Indeed, the Marshall sign rule [42] gives the correct signs (on each cluster) only when \(J_{2}=0\), while it largely fails close to \(J_{2}/J_{1}=0.5\). In this respect, it is important to mention that some of the previous attempts to define variational wave functions included the Marshall sign rule as a first approximation of the correct sign structure, thus taking advantage of a prior probability (_Marshall prior_). By contrast, within the present approach, we do not need to use any prior knowledge of the sign structure, thus defining a very general and flexible variational _Ansatz_.

In the following, we first show the alternative SR formulation, then discuss Deep Transformer architecture, recently introduced by some of us [38; 39] as a variational state, and finally present our results obtained by combining the two techniques on the \(J_{1}\)-\(J_{2}\) Heisenberg model.

_Stochastic Reconfiguration._ Finding the ground state of a quantum system with the variational principle involves minimizing the variational energy \(E(\theta)=\left\langle\Psi_{\theta}|\hat{H}|\Psi_{\theta}\right\rangle/\left \langle\Psi_{\theta}|\Psi_{\theta}\right\rangle\), where \(|\Psi_{\theta}\rangle\) is a variational state parametrized through a vector \(\theta\) of \(P\)_real_ parameters (in case of complex parameters, we can treat their real and imaginary parts separately [43]). For a system of \(N\)\(1/2\)-spins, \(|\Psi_{\theta}\rangle\) can be expanded in the computational basis \(|\Psi_{\theta}\rangle=\sum_{\{\sigma\}}\Psi_{\theta}(\sigma)\left|\sigma\right\rangle\), where \(\Psi_{\theta}(\sigma)=\left\langle\sigma|\Psi_{\theta}\right\rangle\) is a map from spin configurations of the Hilbert space, \(\{\ |\sigma\rangle=|\sigma_{1}^{z},\sigma_{2}^{z},\cdots,\sigma_{N}^{z}\rangle\,\ \sigma_{i}^{z}=\pm 1\}\), to complex numbers. In a gradient-based optimization approach, the fundamental ingredient is the evaluation of the gradient of the loss, which in this case is the variational energy, with respect to the parameters \(\theta_{\alpha}\) for \(\alpha=1,\ldots,P\). This

\begin{table}
\begin{tabular}{l c c c c} \hline \hline
**Energy per site** & **Wave function** & **\# parameters** & **Marshall prior** & **Reference** & **Year** \\ \hline -0.48941(1) & NNQS & 893994 & Not available & [32] & 2023 \\ -0.494757(12) & CNN & Not available & No & [22] & 2020 \\ -0.4947359(1) & Shallow CNN & 11009 & Not available & [21] & 2018 \\ -0.49516(1) & Deep CNN & 7676 & Yes & [20] & 2019 \\ -0.495502(1) & PEPS + Deep CNN & 3531 & No & [33] & 2021 \\ -0.495530 & DMRG & 8192 SU(2) states & No & [31] & 2014 \\ -0.495627(6) & aCNN & 6538 & Yes & [34] & 2023 \\ -0.49575(3) & RBM-fermionic & 2000 & Yes & [15] & 2019 \\ -0.49586(4) & CNN & 10952 & Yes & [35] & 2023 \\ -0.4968(4) & RBM (\(p=1\)) & Not available & Yes & [36] & 2022 \\ -0.49717(1) & Deep CNN & 106529 & Yes & [28] & 2022 \\ -0.497437(7) & GCNN & Not available & No & [27] & 2021 \\ -0.497468(1) & Deep CNN & 421953 & Yes & [30] & 2022 \\ -0.4975490(2) & VMC (\(p=2\)) & 5 & Yes & [13] & 2013 \\ -0.497627(1) & Deep CNN & 146320 & Yes & [29] & 2023 \\ -0.497629(1) & RBM+PP & 5200 & Yes & [37] & 2021 \\ **-0.497634(1)** & **Deep ViT** & **267720** & **No** & **Present work** & **2023** \\ \hline \end{tabular}
\end{table}
Table 1: Ground-state energy on the 10\(\times\)10 square lattice at \(J_{2}/J_{1}=0.5\).
gradient can be expressed as a correlation function [43]

\[F_{\alpha}=-\frac{\partial E}{\partial\theta_{\alpha}}=-2\Re\left[\langle(\hat{H} -\langle\hat{H}\rangle)(\hat{O}_{\alpha}-\langle\hat{O}_{\alpha}\rangle)\rangle \right]\, \tag{2}\]

which involves the diagonal operator \(\hat{O}_{\alpha}\) defined as \(O_{\alpha}(\sigma)=\partial\text{Log}[\Psi_{\theta}(\sigma)]/\partial\theta_{\alpha}\). The expectation values \(\langle\dots\rangle\) are computed with respect to the variational state. The SR updates [19; 43; 9] are constructed according to the geometric structure of the landscape:

\[\mathbf{\delta\theta}=\tau\left(S+\lambda\mathbb{I}_{P}\right)^{-1}\mathbf{F}\, \tag{3}\]

where \(\tau\) is the learning rate and \(\lambda\) is the regularization parameter. The matrix \(S\) has shape \(P\times P\) and it is defined in terms of the \(\hat{O}_{\alpha}\) operators [43]

\[S_{\alpha,\beta}=\Re\left[\langle(\hat{O}_{\alpha}-\langle\hat{O}_{\alpha} \rangle)^{\dagger}(\hat{O}_{\beta}-\langle\hat{O}_{\beta}\rangle)\rangle \right]. \tag{4}\]

The Eq. (3) defines the standard formulation of the SR, which involves the inversion of a \(P\times P\) matrix, being the bottleneck of this approach when the number of parameters is larger than \(O(10^{4})\). To address this problem, we start reformulating Eq. (3) in a more convenient way. For a given sample of \(M\) spin configurations \(\{\sigma_{i}\}\) (sampled according to \(|\Psi_{\theta}(\sigma)|^{2}\)), the stochastic estimate of \(F_{\alpha}\) can be obtained as:

\[\bar{F}_{\alpha}=-\Re\left\{\frac{2}{M}\sum_{i=1}^{M}[E_{L}(\sigma_{i})-\bar{ E}_{L}]^{*}[O_{\alpha}(\sigma_{i})-\bar{O}_{\alpha}]\right\}. \tag{5}\]

Here, the local energy is defined as \(E_{L}(\sigma)=\langle\sigma|\hat{H}|\Psi_{\theta}\rangle\left/\,\langle\sigma |\Psi_{\theta}\right.\). We use \(\bar{E}_{L}\) and \(\bar{O}_{\alpha}\) to denote sample means. Throughout this work, we adopt the convention of using latin indices and greek indices to run respectively over configurations and parameters. Equivalently, Eq. (4) can be stochastically estimated as

\[\bar{S}_{\alpha,\beta}=\Re\left[\frac{1}{M}\sum_{i=1}^{M}\left(O_{\alpha i}- \bar{O}_{\alpha}\right)^{*}\left(O_{\beta i}-\bar{O}_{\beta}\right)\right]. \tag{6}\]

To simplify further, we introduce \(\varepsilon_{i}=-2[E_{L}(\sigma_{i})-\bar{E}_{L}]^{*}/\sqrt{M}\) and \(Y_{\alpha i}=(O_{\alpha i}-\bar{O}_{\alpha})/\sqrt{M}\), allowing us to express Eq. (5) in matrix notation as \(\bar{\mathbf{F}}=\Re[Y\mathbf{\varepsilon}]\) and Eq. (4) as \(\bar{S}=\Re[YY^{\dagger}]\). Writing \(Y=Y_{R}+iY_{I}\) we obtain:

\[\bar{S}=Y_{R}Y_{R}^{T}+Y_{I}Y_{I}^{T}=XX^{T} \tag{7}\]

where \(X=\text{Concat}(Y_{R},Y_{I})\in\mathbb{R}^{P\times 2M}\), the concatenation being along the last axis. Furthermore, using \(\mathbf{\varepsilon}=\mathbf{\varepsilon}_{R}+i\mathbf{\varepsilon}_{I}\), the gradient of the energy can be recasted as

\[\bar{\mathbf{F}}=Y_{R}\mathbf{\varepsilon}_{R}-Y_{I}\mathbf{\varepsilon}_{I}=X\mathbf{f}\, \tag{8}\]

with \(\mathbf{f}=\text{Concat}(\mathbf{\varepsilon}_{R},-\mathbf{\varepsilon}_{I})\in\mathbb{R }^{2M}\). Then, the update of the parameters in Eq. (3) can be written as

\[\mathbf{\delta\theta}=\tau(XX^{T}+\lambda\mathbb{I}_{P})^{-1}X\mathbf{f}. \tag{9}\]

This reformulation of the SR updates is a crucial step, which allows the use of a well-known linear algebra identity [44] (see _Supplemental Material_ for a derivation). As a result, Eq. (9) can be rewritten as

\[\mathbf{\delta\theta}=\tau X(X^{T}X+\lambda\mathbb{I}_{2M})^{-1}\mathbf{f}. \tag{10}\]

This derivation is our first result: it shows, in a simple and transparent way, how to exactly perform the SR with the inversion of a \(M\times M\) matrix and without the need to allocate a \(P\times P\) matrix. We emphasise that the last formulation is very useful in the typical deep learning setup, where \(P\gg M\). Employing Eq. (10) instead of Eq. (9) proves to be more efficient both in terms of computational complexity [\(O(4M^{2}P+8M^{3})\) operations instead of \(O(P^{3})\)] and memory usage [requiring only the allocation of \(O(MP)\) numbers instead of \(O(P^{2})\)]. Crucially, we developed a memory-efficient implementation of SR that is optimized for deployment on a multi-node GPU cluster, ensuring scalability and practicality for real-world applications (see _Supplemental Material_). Other methods, based on iterative solvers, require \(O(nMP)\) operations, where \(n\) is the number of steps needed to solve the linear problem in Eq. (3). However, this number increases significantly for ill-conditioned matrices (the matrix \(S\) has a number of zero eigenvalues equal to \(P\)-\(M\)), leading to many non-parallelizable iteration steps and consequently higher computational costs [45].

The variational wave function.The architecture of the variational state employed in this work is based on the Deep Vision Transformer (Deep ViT). In the Deep ViT, the input configuration \(\sigma\) of shape \(L\times L\) is firstly split into \(b\times b\) patches, which are then embedded in a \(d\)-dimensional vector space in order to obtain an input

Figure 1: Optimization of the Deep ViT with batch size \(b=2\), 8 layers, embedding dimension \(d=72\) and \(12\) heads per layer, on the \(J_{1}\)-\(J_{2}\) Heisenberg model at \(J_{2}/J_{1}=0.5\) on the \(10\times 10\) square lattice. The first 200 optimization steps are not shown for better readability. Inset: last 200 optimization steps when recovering sequentially the full translational (green curve), rotational (orange curve) and reflections and parity (red curve) symmetries.
sequence \(\big{(}\mathbf{x}_{1},\ldots,\mathbf{x}_{L^{2}/b^{2}}\big{)}\), where \(L^{2}/b^{2}\) is the total number of patches. A real-valued Deep ViT _transforms_ this sequence into an output sequence \(\big{(}\mathbf{y}_{1},\ldots,\mathbf{y}_{L^{2}/b^{2}}\big{)}\), with \(\mathbf{y}_{i}\in\mathbb{R}^{d}\) being a rich, context aware representation of the spins in \(i\)-th patch. At the end, a complex-valued fully connected output layer, with \(\log[\cosh(\cdot)]\) as activation function, is applied to \(\mathbf{z}=\sum_{i}\mathbf{y}_{i}\) to produce a single number \(\text{Log}[\Psi_{\theta}(\sigma)]\) and predict both the modulus and the phase of the input configuration [39]. We employ _Factored Attention_[46] and two-dimensional relative positional encoding to enforce translational symmetry among patches and only a small summation (\(b^{2}\) terms) is required to restore the one inside the patches. In addition, the other symmetries of the Hamiltonian can be restored with the quantum number projection approach [47, 35]. In particular, we restore rotations, reflections (\(C_{4v}\) point group), and spin parity. As a result, the symmetrized wave function reads:

\[\tilde{\Psi}_{\theta}(\sigma)=\sum_{i=0}^{b^{2}-1}\sum_{j=0}^{7}\left[\Psi_{ \theta}(T_{i}R_{j}\sigma)+\Psi_{\theta}(-T_{i}R_{j}\sigma)\right]\;. \tag{11}\]

In the last equation, \(T_{i}\) and \(R_{j}\) are the translation and the rotation/reflection operators, respectively. Furthermore, due to the \(SU(2)\) spin symmetry of the \(J_{1}\)-\(J_{2}\) Heisenberg model, the total magnetization is also conserved and the Monte Carlo sampling (see _Supplemental Material_) can be limited in the \(S^{z}=0\) sector for the ground state search.

_Results._ Our objective is to approximate the ground state of the \(J_{1}\)-\(J_{2}\) Heisenberg model in the highly frustrated point \(J_{2}/J_{1}=0.5\) on the \(10\times 10\) square lattice. We use the formulation of the SR in Eq. (10) to optimize a variational wave function parametrized through a Deep ViT, as discussed above. The result in Table 1 is achieved with the symmetrized Deep ViT in Eq. (11) using \(b=2\), number of layers \(8\), embedding dimension \(d=72\), number of heads per layer \(12\). This variational state has in total \(267720\) real parameters. Regarding the optimization protocol, we choose the learning rate \(\tau=0.03\) (cosine decay annealing) and the number of samples is fixed to be \(M=6000\). We emphasise that using Eq. (9) to optimize this number of parameters would be infeasible on avaiable GPUs: the memory requirement would be more than \(O(10^{3})\) gigabytes, one order of magnitude bigger than the highest available memory capacity. Instead, with the formulation of Eq. (10), the memory requirement is reduced to \(\approx 26\) gigabytes, which can be easily handled by available GPUs. The simulations took four days on twenty A100 GPUs. Remarkably, as illustrated in Table 1, we are able to obtain the state-of-the-art ground-state energy using an architecture solely based on neural networks, without using any other regularization than the diagonal shift reported in Eq. (10), fixed to \(\lambda=10^{-4}\). We stress that a completely unbiased simulation, without assuming any prior for the sign structure, is performed, in contrast to the other cases where the Marshall sign rule is used to stabilize the optimization [28, 29, 37, 48] (see Table 1). This is an important point since a simple sign prior is not available for the majority of the models (e.g., the Heisenberg model on the triangular or Kagome lattices). We would like also to mention that the definition of a suitable architecture is fundamental to take advantage of having a large number of parameters. Indeed, while a stable simulation with a simple regularization scheme (only defined by a finite value of \(\lambda\)) is possible within the Deep ViT wave function, other architectures require more sophisticated regularizations. For example, to optimize Deep GCNNs it is necessary to add a temperature-dependent term to the loss function [27] or, for Deep CNNs, a process of variance reduction and reweighting [29] helps in escaping local minima. We also point out that physically inspired wave functions, as the Gutzwiller-projected states [13], which give a remarkable result with only a few parameters, are not always equally accurate in other cases.

In Fig. 1 we show a typical Deep ViT optimization on the \(10\times 10\) at \(J_{2}/J_{1}=0.5\). First, we optimize the Transformer having translational invariance among patches (blue curve). Then, starting from the previously optimized parameters, we restore sequentially the full translational invariance (green curve), rotational symmetry (orange curve) and lastly, reflections and spin parity symmetry (red curve). Remarkably, whenever a new symmetry is restored, the energy reliably decreases [47]. We stress that our optimization process, which combines the SR formulation of Eq. (10) with a real valued Deep ViT followed by a complex-valued fully connected output layer [39], is highly stable and insensitive to the initial seed, ensuring consistent results across multiple optimization runs.

_Conclusions._ We have introduced an innovative formulation of the SR that excels in scenarios where the number of parameters (\(P\)) significantly outweighs the number of samples (\(M\)) used for stochastic estimations. Exploiting this approach, we attained the state-of-the-art ground state energy for the \(J_{1}\)-\(J_{2}\) Heisenberg model at \(J_{2}/J_{1}=0.5\), on a \(10\times 10\) square lattice, optimizing a Deep ViT with \(P=267720\) and using only \(M=6000\) samples. It is essential to note that this achievement highlights the remarkable capability of deep learning in performing exceptionally well even with a limited sample sizes relative to the overall parameter count. This also _challenges_ the common belief that large amount of Monte Carlo samples are required to find the solution in the exponentially-large Hilbert space and for precise SR optimizations [30]. This result has important ramifications for investigating the physical properties of challenging quantum many-body problems, where the use of the SR is crucial to obtain accurate results. The use of Neural Network Quantum States with an unprecedented number of parameters can open new opportunities in approximating ground states of quantum spin Hamiltonians, where other methods fail. Additionally, making minor modifications to the equations describing parameter updates within the SR framework [see Eq. (10)] enables us to describe
the unitary time evolution of quantum many-body systems according to the time-dependent variational principle (TDVP) [43; 49; 50]. Extending our approach to address time-dependent problems stands as a promising avenue for future works. Furthermore, this new formulation of the SR can be a key tool for obtaining accurate results also in quantum chemistry, especially for systems that suffer from the sign problem [51]. Typically, in this case, the standard strategy is to take \(M>10\times P\)[43], which may be unnecessary for deep learning based approaches.

###### Acknowledgements.

 We thank R. Favata for critically reading the manuscript and M. Imada for stimulating us with challenging questions. We also acknowledge G. Carleo and A. Chen for useful discussions. The variational quantum Monte Carlo and the Deep ViT architecture were implemented in JAX [52]. The parallel implementation of the Stochastic Reconfiguration was implemented using mpi4jax [53]. R.R. and L.L.V. acknowledge access to the cluster LEONARDO at CINECA through the IsCa5_ViT2d project. S.G. acknowledges co-funding from Next Generation EU, in the context of the National Recovery and Resilience Plan, Investment PE1 - Project FAIR "Future Artificial Intelligence Research". This resource was co-financed by the Next Generation EU [DM 1555 del 11.10.22].
