Gauge-equivariant neural networks as preconditioners in lattice QCD
We demonstrate that a state-of-the art multi-grid preconditioner can be
learned efficiently by gauge-equivariant neural networks. We show that the
models require minimal re-training on different gauge configurations of the
same gauge ensemble and to a large extent remain efficient under modest
modifications of ensemble parameters. We also demonstrate that important
paradigms such as communication avoidance are straightforward to implement in
this framework.
Christoph Lehner, Tilo Wettig
# Gauge-equivariant neural networks as preconditioners in lattice QCD

###### Abstract

We demonstrate that a state-of-the art multi-grid preconditioner can be learned efficiently by gauge-equivariant neural networks. We show that the models require minimal re-training on different gauge configurations of the same gauge ensemble and to a large extent remain efficient under modest modifications of ensemble parameters. We also demonstrate that important paradigms such as communication avoidance are straightforward to implement in this framework.

## I Introduction

Our current understanding of nature at the most fundamental level is to a large extent based on quantum field theories. In particle physics, Quantum Chromodynamics (QCD) explains, for example, how the proton is made up of smaller constituents, quarks and gluons. To describe current and future experiments, and to search for physics beyond the Standard Model, we need to be able to solve QCD to high precision. Lattice QCD constitutes a systematically improvable tool to solve QCD in the nonperturbative regime by numerically simulating the theory on a finite space-time lattice. It has evolved over more than four decades and is now of direct phenomenological relevance, see [1] and references therein. It is also very compute-intensive and employs the largest supercomputers worldwide [2]. Therefore much research is focused on improving the algorithms that dominate the run time of these simulations.

The most time-consuming element, both in the generation of gauge-field configurations and in the computation of physical observables, is typically the solution of the Dirac equation in the presence of a given gauge field. For physical values of the light quark masses and large lattice volumes, the condition number of the matrix representing the Dirac operator becomes very large, and consequently very sophisticated methods are required to solve the Dirac equation in a feasible time frame. The current state of the art is to use a suitable preconditioner inside a Krylov subspace solver. The construction of the preconditioner is a complicated problem whose solution requires deep knowledge of the underlying physics. The aim of this paper is to reformulate the problem in the language of gauge-equivariant neural networks and to show that such networks can learn the general paradigms of state-of-the-art preconditioners and efficiently reduce the iteration count of the outer solver. We also provide a flexible implementation interface in the Grid Python Toolkit (GPT) [3] that allows for experimentation and further studies.

We briefly relate this paper to previous work. We will concentrate on multi-grid preconditioners [4; 5; 6; 7; 8; 9; 10] and refer to [11] for an introduction. The idea of learning the elements of multi-grid preconditioners with neural networks has been pursued in a number of earlier publications, see, e.g., [12; 13; 14; 15; 16; 17; 18]. These works differ in the details of their approaches, e.g., the choice of the loss function, the network architecture, and the kind of learning (supervised or unsupervised). The main difference to our work is that we have to address the gauge degrees of freedom. More precisely, our approach must be gauge-equivariant, i.e., the map implemented by the neural network must commute with local gauge transformations [19; 20].

A number of papers have introduced gauge-equivariant neural networks in the context of lattice quantum field theory: Refs. [21; 22; 23] mainly addressed the question of gauge-field sampling in several different theories, while Ref. [24] showed how any gauge-covariant function on the lattice can be approximated by neural networks. Our work builds on and extends these papers.

The structure of this paper is as follows. In Sec. II, we introduce gauge-equivariant layers as the building blocks of the models we study in this work. In Sec. III, we discuss the problem of solving the preconditioned Dirac equation with the Wilson-clover Dirac operator. In Sec. IV, we construct preconditioner models that address the high-mode component of the Dirac operator. In Sec. V, we discuss a model to address the low-mode component of the Dirac operator. In Sec. VI, we combine the specialized models to a multi-grid model that addresses both the low-mode and high-mode components. We conclude in Sec. VII, where we also give an outlook to future work.

## II Gauge-equivariant layers

In this section, we define the building blocks of the gauge-equivariant neural networks considered in this work and explain their properties in detail. We begin with a discussion of the concepts of parallel transport and gauge equivariance.

### Parallel transport and gauge equivariance

We consider a discrete \(d\)-dimensional space-time lattice with \(L_{\mu}\) sites in dimension \(\mu\in\{1,\ldots,d\}\) and \(d\in\mathbb{N}\). The canonical unit vector in dimension \(\mu\) is denoted by \(\hat{\mu}\). The set of all lattice sites shall be \(S=\{(x_{1},\ldots,x_{d})\,|\,x_{\mu}\in\{1,\ldots,L_{\mu}\}\}\). Consider a field \(\varphi:S\to V_{I},x\mapsto\varphi(x)\) with internal vector space \(V_{I}\). The internal vector space
shall be a product of a gauge vector space \(V_{G}=\mathbb{C}^{N}\) and a non-gauge vector space \(V_{G}=\mathbb{C}^{N}\) with \(N,N\in\mathbb{N}\), i.e.,

\[V_{I}=V_{G}\otimes V_{G}\,. \tag{1}\]

We also consider gauge fields \(U_{\mu}:S\to\mathrm{SU}(N),x\mapsto U_{\mu}(x)\) with \(\mathrm{SU}(N)\) acting on \(V_{G}\). The set of fields \(\varphi\) shall be \(\mathcal{F}_{\varphi}\), and the set of fields \(U_{\mu}\) shall be \(\mathcal{F}_{U}\).

We define the parallel-transport operator \(T_{p}:\mathcal{F}_{\varphi}\to\mathcal{F}_{\varphi},\varphi\mapsto T_{p}\varphi\) as

\[T_{p}=H_{p_{n_{p}}}\cdots H_{p_{2}}H_{p_{1}} \tag{2}\]

for a path \(p\) defined as the sequence \(p_{1},\ldots,p_{n_{p}}\) with \(n_{p}\in\mathbb{N}\) and \(p_{i}\in\{\pm 1,\pm 2,\ldots,\pm d\}\). The operator \(H_{p_{i}}:\mathcal{F}_{\varphi}\to\mathcal{F}_{\varphi},\varphi\mapsto H_{p_{ i}}\varphi\) acts on a field according to1

Footnote 1: Note that the operator \(H_{p_{i}}\) does not act on the numerical value \(\varphi(x)\). Rather, it acts on the field \(\varphi\), resulting in the new field \(H_{p_{i}}\varphi\), which is then evaluated at \(x\). Note also that in Eq. (3), the information is transported from \(x-\hat{p}_{i}\) to \(x\).

\[H_{p_{i}}\varphi(x)=U_{p_{i}}^{\dagger}(x-\hat{p}_{i})\varphi(x-\hat{p}_{i}) \tag{3}\]

so as to transport information by a single hop in direction \(\hat{p}_{i}\). Here, we introduced the convention \(\hat{\nu}=-\hat{\mu}\) for \(\nu=-\mu\), and we identify \(U_{-\mu}(x)=U_{\mu}^{\dagger}(x-\hat{\mu})\). Addition and subtraction of coordinate tuples are defined component-wise. Note that a single path \(p\) defines the transport for any site \(x\in S\) to

\[x^{\prime}=x+\sum_{i=1}^{n_{p}}\hat{p}_{i} \tag{4}\]

and may be illustrated using a representative starting point. If \(x^{\prime}=x\), the path is closed. Note that the trivial path \(0\) with \(n_{0}=0\) and \(T_{0}=\mathbb{1}\) is allowed as well.

A field \(\varphi\in\mathcal{F}_{\varphi}\) acquires a phase \(\theta_{\mu}\) when translated by \(L_{\mu}\) in direction \(\hat{\mu}\), i.e.,

\[\varphi(x+L_{\mu}\hat{\mu})=e^{i\theta_{\mu}}\varphi(x) \tag{5}\]

for any coordinate tuple \(x\). A gauge field \(U_{\mu}\in\mathcal{F}_{U}\) is periodic in all dimensions, i.e.,

\[U_{\mu}(x+L_{\nu}\hat{\nu})=U_{\mu}(x) \tag{6}\]

with \(\nu\in\{1,\ldots,d\}\). These equations define \(\varphi(x)\) and \(U_{\mu}(x)\) for all sites \(x\) outside of \(S\).

In Fig. 1, we illustrate the transport from the red starting point along a path \(p\) to the black site. This path corresponds to

\[T_{p}=H_{-1}H_{-2}H_{-1}H_{2}H_{2}\,, \tag{7}\]

where \(\hat{1}\) and \(\hat{2}\) is the horizontal and vertical unit vector, respectively, in Fig. 1.

A gauge transformation is parametrized by a field \(\Omega:S\to\mathrm{SU}(N),x\mapsto\Omega(x)\) that acts on all \(\varphi\in\mathcal{F}_{\varphi}\) and \(U_{\mu}\in\mathcal{F}_{U}\) by

\[\varphi(x) \to\Omega(x)\varphi(x)\,, \tag{8}\] \[U_{\mu}(x) \to\Omega(x)U_{\mu}(x)\Omega^{\dagger}(x+\hat{\mu})\,. \tag{9}\]

It is straightforward to show that under such a gauge transformation we have

\[T_{p}\varphi(x)\to\Omega(x)T_{p}\varphi(x) \tag{10}\]

for any path \(p\), i.e., the parallel-transport operator \(T_{p}\) commutes with gauge transformations, and thus it is a gauge-equivariant operator. For a comprehensive discussion of gauge equivariance we refer to Ref. [20].

### Parallel-transport convolutions

The models discussed in this work will be composed of individual layers that map \(n\) input features \(\varphi_{1},\ldots,\varphi_{n}\in\mathcal{F}_{\varphi}\) to \(m\) output features \(\psi_{1},\ldots,\psi_{m}\in\mathcal{F}_{\varphi}\).

We consider a parallel-transport convolution (PTC) layer defined by2

Footnote 2: Equation (11) is a convolution with kernel \(W\) and input \(\varphi\), whose argument is shifted by \(T_{p}\).

\[\psi_{a}(x)\stackrel{{\mathrm{PTC}}}{{=}}\sum_{b=1}^{n}\sum_{p\in P }W_{a}^{bp}T_{p}\varphi_{b}(x) \tag{11}\]

for \(a=1,\ldots,m\), with a set of paths \(P\) and an endomorphism \(W_{a}^{bp}\in\mathrm{End}(V_{\tilde{G}})\). This extends the definition of Ref. [23] from nearest-neighbor hops to a sum over arbitrary paths. For closed paths \(p\), we recover the case discussed in Ref. [24]. Note that in lattice QCD \(W_{a}^{bp}\) is a \(4\times 4\) spin matrix.

We also consider a local parallel-transport convolution (LPTC) layer defined by

\[\psi_{a}(x)\stackrel{{\mathrm{LPTC}}}{{=}}\sum_{b=1}^{n}\sum_{p\in P }W_{a}^{bp}(x)T_{p}\varphi_{b}(x) \tag{12}\]

Figure 1: The path \(p\) defining a parallel-transport operator \(T_{p}\) can be visualized as a sequence of hops from a starting point (red) to an end point (black).
with \(W_{a}^{bp}:S\rightarrow\mathrm{End}(V_{\tilde{G}}),x\mapsto W_{a}^{bp}(x)\). Such a layer is also gauge equivariant and may be able to better address localized features. In the following we refer to the elements of \(W\) as layer weights.

Since we intend to learn a _linear_ preconditioner in this work, we do not apply an activation function in these layers. The expressivity of a deep network composed of such layers is therefore equivalent to a single layer with a larger set \(P\). Nevertheless, it may be computationally more efficient for a given problem to compose multiple layers with smaller sets \(P\).

In Fig. 2, we provide a graphical representation of a (L)PTC layer with two input features and one output feature and \(P=\{p_{1},p_{2}\}\) with

\[T_{p_{1}}=H_{-1}H_{-2}H_{-1}\,,\qquad T_{p_{2}}=H_{-2}H_{1}\,. \tag{13}\]

### Restriction and prolongation layers

In order to let information propagate efficiently over long distances in terms of sites \(x\in S\), we make use of the multi-grid paradigm [4; 5]. To this end, we consider a coarse grid with lattice sites \(\tilde{S}\) and a coarse field \(\tilde{\varphi}:\tilde{S}\rightarrow\tilde{V}_{I},y\mapsto\tilde{\varphi}(y)\) with coarse internal vector space \(\tilde{V}_{I}\). The set of such fields is denoted by \(\mathcal{F}_{\tilde{\varphi}}\). Note that there are no gauge degrees of freedom in \(\tilde{V}_{I}\).

We define a restriction layer mapping a \(\varphi\in\mathcal{F}_{\varphi}\) to a \(\tilde{\psi}\in\mathcal{F}_{\tilde{\varphi}}\) by

\[\tilde{\psi}(y)\stackrel{{\mathrm{RL}}}{{=}}\sum_{x\in B(y)}W(y,x )\varphi(x) \tag{14}\]

with \(W:\tilde{S}\times S\rightarrow\mathrm{Hom}(V_{I},\tilde{V}_{I})\) and block map \(B:\tilde{S}\rightarrow\mathcal{P}(S)\), where \(\mathcal{P}\) denotes the power set. We also define a corresponding prolongation layer mapping a \(\tilde{\varphi}\in\mathcal{F}_{\tilde{\varphi}}\) to a \(\psi\in\mathcal{F}_{\varphi}\) by

\[\psi(x)\stackrel{{\mathrm{PL}}}{{=}}W(y,x)^{\dagger}\tilde{ \varphi}(y) \tag{15}\]

for \(x\in B(y)\). In practice, we choose \(B\) corresponding to a blocking in all dimensions. The linear maps \(W\) satisfy

\[\sum_{x\in B(y)}W(y,x)W(y,x)^{\dagger}=\mathbb{1}_{\tilde{V}_{I}}\,, \tag{16}\]

where \(\mathbb{1}_{\tilde{V}_{I}}\) is the identity in \(\tilde{V}_{I}\). These layers are straightforward to extend to the case of multiple input and output features.

The linear maps \(W\) can be considered layer weights and are constructed from a list of vectors that are blockwise orthonormal, see Sec. V for details. The restriction and prolongation layers are gauge equivariant if

\[W(y,x)\to W(y,x)\Omega(x)^{\dagger} \tag{17}\]

under a gauge transformation. Note that since \(\tilde{V}_{I}\) does not have gauge degrees of freedom there is no \(\Omega(y)\) on the coarse grid. We provide a graphical representation of the restriction and prolongation layers in Fig. 3.

### Parallel and identity layers

In this work, we consider models that act on a given input feature with multiple layers in parallel. Consider applying a layer \(L_{i}\) to input features \(\varphi_{1},\dots,\varphi_{n}\) mapping to output features \(\psi_{i1},\dots,\psi_{im_{i}}\). For several layers \(L_{1},\dots,L_{\ell}\), we concatenate the output features \(\psi_{11},\dots,\psi_{1m_{1}},\dots,\psi_{\ell 1},\dots,\psi_{\ell m_{\ell}}\). The combination of layers \(L_{1},\dots,L_{\ell}\) being applied in parallel can then be considered to be a single layer that maps features \(\varphi_{1},\dots,\varphi_{n}\) to features \(\psi_{11},\dots,\psi_{1m_{1}},\dots,\psi_{\ell 1},\dots,\psi_{\ell m_{\ell}}\).

We also introduce an identity layer that maps the input features without modification to output features (which implies \(m=n\)). Such a layer is represented graphically by a single dashed arrow pointing from the input features to the output features.

We provide a graphical representation for the case of \(n=1\), \(\ell=2\), and \(m_{1}=m_{2}=1\) in Fig. 4.

Figure 3: Graphical representation of the restriction layer (left) and prolongation layer (right) for a single feature. The layers are represented by the gray square frustums, while the input and output features are represented by the planes.

Figure 2: Graphical representation of a (L)PTC layer with two input features and one output feature. The planes represent the features. The layer is represented by the paths drawn and the dashed arrow.
### Communication avoidance

In practice, the performance of a given model in terms of execution time is crucial. For problem sizes of interest to the lattice QCD community, a single problem will be distributed over multiple compute nodes that are connected by a communication network. It is not uncommon that the time needed to exchange information between nodes exceeds the time each node spends performing floating-point operations. Therefore it is an important paradigm in lattice QCD to investigate approaches that avoid communication between nodes even if it possibly increases the computational effort within a given node [25; 26; 27; 28]. In this work, we also investigate layers which do not communicate between different sub-volumes that would typically be mapped to multiple nodes in an MPI job. We perform such investigations by setting the gauge links \(U_{\mu}\) that connect one such sub-volume to another to zero. For such a modified model, we can then avoid the communication step between nodes altogether.

## III The Wilson Dirac operator

The main objective of this work is to precondition the Dirac equation

\[Du=b \tag{18}\]

with Dirac operator \(D:\mathcal{F}_{\varphi}\to\mathcal{F}_{\varphi}\), source \(b\in\mathcal{F}_{\varphi}\), and solution \(u\in\mathcal{F}_{\varphi}\). It is useful to interpret Eq. (18) as a matrix equation with \(u,b\in\mathbb{C}^{k}\) and invertible complex \(k\times k\) matrix \(D\) with

\[k=L_{1}\cdots L_{d}N\bar{N}\,. \tag{19}\]

We train a model to play the role of an invertible complex \(k\times k\) preconditioner matrix \(M\) in

\[(DM)M^{-1}u=b\,, \tag{20}\]

where we attempt to improve the condition number of \(DM\) compared to \(D\). Ideally, \(DM\) is close to the identity matrix up to a trivial scaling factor. The Dirac matrix transforms as

\[D\to\Omega D\Omega^{\dagger} \tag{21}\]

under a gauge transformation with block-diagonal matrix \(\Omega=\oplus_{x\in\mathcal{S}}\Omega(x)\otimes\mathbb{1}_{V_{G}}\), which motivates the use of gauge-equivariant layers to construct \(M\).

We first consider the Wilson Dirac operator [29]

\[D_{\mathrm{W}} =\frac{1}{2}\sum_{\mu=1}^{4}\gamma_{\mu}(H_{-\mu}-H_{+\mu})+m\] \[-\frac{1}{2}\sum_{\mu=1}^{4}(H_{-\mu}+H_{+\mu}-2) \tag{22}\]

with mass \(m\in\mathbb{R}\) and Euclidean gamma matrices \(\gamma_{1},\ldots,\gamma_{4}\) satisfying the anti-commutation relation \(\gamma_{\mu}\gamma_{\nu}+\gamma_{\nu}\gamma_{\mu}=2\delta_{\mu\nu}\) with Kronecker delta \(\delta_{\mu\nu}\). This operator can be mapped to a single PTC layer with a zero-hop path and eight one-hop paths.

We add a clover term that includes closed paths consisting of four hops using

\[Q_{\mu\nu} =H_{-\mu}H_{-\nu}H_{+\mu}H_{+\nu}+H_{-\nu}H_{+\mu}H_{+\nu}H_{-\mu}\] \[\quad+H_{+\nu}H_{-\mu}H_{-\nu}H_{+\mu}+H_{+\mu}H_{+\nu}H_{-\mu}H_{ -\nu} \tag{23}\]

to obtain the Wilson-clover Dirac operator [30]

\[D_{\mathrm{WC}}=D_{\mathrm{W}}-\frac{c_{\mathrm{sw}}}{4}\sum_{\mu,\nu=1}^{4} \sigma_{\mu\nu}F_{\mu\nu} \tag{24}\]

with \(c_{\mathrm{sw}}\in\mathbb{R}\),

\[F_{\mu\nu}=\frac{1}{8}(Q_{\mu\nu}-Q_{\nu\mu})\,, \tag{25}\]

and

\[\sigma_{\mu\nu}=\frac{1}{2}(\gamma_{\mu}\gamma_{\nu}-\gamma_{\nu}\gamma_{\mu} )\,. \tag{26}\]

The operator \(D_{\mathrm{WC}}\) can also be mapped to a single PTC layer, however, paths up to four hops are needed.

For the numerical experiments presented in the following sections, we use gauge group SU(3) and the \(D_{\mathrm{WC}}\) operator tuned to near criticality, i.e., the mass parameter is chosen such that the real part of the smallest eigenvalue is close to zero. This provides a challenging problem even for the small lattice volume with \(L_{1}=L_{2}=L_{3}=8\) and \(L_{4}=16\) used in this work. We set \(m=-0.6\) and \(c_{\mathrm{sw}}=1\) on a pure Wilson gauge configuration [29] with coupling parameter \(\beta=6.0\). We use periodic boundary conditions also for the fields in \(\mathcal{F}_{\varphi}\), i.e., \(\theta_{\mu}=0\) in Eq. (5). We show the spectrum of \(D_{\mathrm{WC}}\) on a representative single gauge configuration in Fig. 5.

We quantify the improvement achieved using the preconditioner \(M\) by the reductio

Figure 4: Graphical representation of two parallel layers \(L_{1}\) and \(L_{2}\) being applied to a single input feature and mapping to two output features. As before, the features are represented by planes. An identity layer (i.e., a copy operation) is represented by a dashed arrow. In this example, the only nontrivial layer is \(L_{1}\), which includes a single path in (11) or (12).
to solve Eq. (20) to \(10^{-8}\) precision in the preconditioned FGMRES [31]. We quote the iteration count gain defined as the iteration count of the unpreconditioned solve divided by the iteration count of the preconditioned solve.

The methods developed in this work also extend to other Dirac matrices. However, particular challenges exist in some cases. For example, in the case of domain-wall fermions [32; 33] the spectrum encircles the origin [9; 10], which limits the convergence of unpreconditioned solves of \(Du=b\) using Krylov-subspace methods.

## IV High-mode preconditioners

We want to learn a preconditioner \(M\) that approximates \(D^{-1}\). For this purpose it is useful to consider an eigendecomposition of \(D\) and first construct optimal models for the high-mode and low-mode components separately. We study the high-mode component in this section and the low-mode component in Sec. V. We then combine the corresponding models in Sec. VI.

### Model setup and training strategy

The high-mode part of the spectrum of \(D_{\rm WC}\) is related to the short-distance behavior. Therefore we expect a single layer with paths up to one hop to already show a gain in iteration count. We consider a linear model \(M\) mapping a vector \(x\) to \(Mx\). We employ a supervised learning approach and describe a single training step in the following.

We first pick a random vector \(v\) with components drawn from a Gaussian distribution with mean zero and unit standard deviation. We then construct the cost function3

Footnote 3: Note that in Eq. (20) we use \(DM\), while in Eq. (27) we use \(MD\). If \(DM\) is close to the identity, then so is \(MD\), and thus Eq. (27) is a suitable cost function.

\[C=|MD_{\rm WC}v-v|^{2} \tag{27}\]

and its derivatives with respect to the model weights using backpropagation. This corresponds to a batch of a single training tuple \((D_{\rm WC}v,v)\), where the model learns to map the first to the second component. This cost function is dominated by the high modes of \(D_{\rm WC}\) and is therefore similar in spirit to using the spectral radius [12; 14]. Since we use a different random vector at every iteration our training data set is unbounded in size and there is no need to add a regulator. This holds even for LPTC layers with a large number of model weights. We then apply a single iteration of the Adam optimizer [34] with parameters \(\beta_{1}=0.9\), \(\beta_{2}=0.98\), and \(\alpha=10^{-3}\) that gave good performance for the models considered in this work. This process is repeated until the model weights are converged sufficiently.

All layers and optimizers are implemented in the Grid Python Toolkit (GPT) [3], and corresponding code samples are provided in App. A.

### Locality and communication avoidance

In Fig. 6 we compare the performance of single-layer models with a maximum of one hop. They correspond to a version of Fig. 2 with a single input and output feature and nine paths corresponding to

\[T_{0} =\mathbb{1}\,,\] \[T_{1} =H_{1}\,, T_{2} =H_{2}\,,\] \[T_{3} =H_{3}\,, T_{4} =H_{4}\,,\] \[T_{5} =H_{-1}\,, T_{6} =H_{-2}\,,\] \[T_{7} =H_{-3}\,, T_{8} =H_{-4}\,. \tag{28}\]

We also investigate communication-avoiding versions with local volume \(4^{3}\times 8\). We find that the LPTC models do not perform better in terms of iteration count gain than the PTC models. However, the LPTC models require more training compared to the PTC models. The slower convergence is expected due to the much larger number weights in the LPTC models. We find that eliminating communication between sub-volumes, as described in Sec. II.5, only leads to a modest reduction in

Figure 5: Eigenvalues \(\lambda\) of the Wilson-clover Dirac operator with \(m=-0.6\) and \(c_{\rm sw}=1\) on a pure-Wilson-gauge configuration with \(\beta=6\), \(L_{1}=L_{2}=L_{3}=8\), and \(L_{4}=16\). The mass \(m\) is tuned to near criticality for the experiments in this work. We computed the boundaries of the spectrum using the Arnoldi method applied to \((D-\lambda)^{-1}\) for several carefully selected values of \(\lambda\) and filled in the bulk of the spectrum by hand for illustrative purposes.
performance. After translating the iteration count gain to a reduction in time-to-solution, we may therefore find the communication-avoiding models to perform best.

### Multiple hops and deep networks

In Fig. 7, we investigate models with multiple hops either in a single layer or distributed over two layers. We use one-hop layers with paths defined in Eq. (28) as well as a two-hop layer extending this set by all combinations

\[H_{a}H_{b} \tag{29}\]

for \(a,b\in\{-4,-3,-2,-1,1,2,3,4\}\) with \(a\neq-b\). The two-hop layer therefore has 65 distinct paths compared to the 9 paths of the one-hop layer.

The first model that we investigate stacks two one-hop layers with one input and one output feature back-to-back. We denote this model as "2 layers (\(1\to 1\to 1\)), 1 hop." The second model is similar but has two output features in the first layer and correspondingly two input features in the second layer. We denote this model as "2 layers (\(1\to 2\to 1\)), 1 hop." The third model consists of a single two-hop layer as described above.

We find that the second model performs best and gives approximately twice the iteration count gain of the corresponding single-layer models with a maximum of one hop shown in Fig. 6. Since the layers are linear, the two-layer models are not more expressive compared to the single-layer model with two hops. We therefore expect the third model to be able to match the performance of the second model with a sufficiently improved training procedure. It is not surprising that the second model can be trained more efficiently compared to the third model given that it has a smaller number of weights. We conclude that while deep models do not increase expressivity, the computational effort needed to train deep models may be reduced compared to a corresponding shallow model with more paths.

### Transfer learning

In Fig. 8 we investigate how well the one-layer one-hop PTC model of Fig. 6 that was trained on a given gauge configuration with \(\beta=6.0\) and \(m=-0.6\) performs when it is used in the case of (i) a different gauge configuration of the same gauge ensemble, (ii) a gauge configuration of a different ensemble with \(\beta=5.9\), and (iii) the same gauge configuration but with a different mass \(m=-0.55\)

Figure 6: Convergence of the cost function (27) and iteration count gain for one-layer and one-hop high-mode preconditioners. The lattice volume is \(8^{3}\times 16\), and the local volume for the communication-avoiding version is \(4^{3}\times 8\).

Figure 7: Convergence of the cost function (27) and iteration count gain for two-layer and two-hop high-mode preconditioners.
In all cases, we investigate the performance without re-training and after additional re-training steps following the same procedure as for the initial training. We find that the high-mode preconditioner model does not require re-training to efficiently perform in all three cases. Once such a model is trained, it can be used efficiently for different gauge configurations of the same and similar ensembles. We note that the maximum iteration count gain for mass \(m=-0.55\) is significantly reduced. In this case, however, the spectrum is not well tuned to criticality and the initial problem is therefore less challenging. Comparing with Fig. 6, we also observe a modest fluctuation in iteration count gain between different configurations.

## V Low-mode preconditioners

We now turn to the low-mode component in the eigen-decomposition of \(D\). Since the low-mode component corresponds to the long-distance behavior of the Dirac operator \(D\), it is not efficient to use the layers discussed in Sec. IV since a rather deep network composed of such layers would be needed to propagate information over sufficiently long distances. The multi-grid paradigm, however, is ideally suited to address this issue. In this section, we focus solely on the low-mode component and then combine low modes and high modes in Sec. VI.

### Model setup and training strategy

In the multi-grid approach, we define an additional coarser version of the lattice as well as restriction and prolongation operations that map between the fine and coarse lattices. These operations must preserve the low-mode component of \(D\)[35].

To achieve this, we first find vectors \(u_{1},\ldots,u_{s}\) in the near-null space of \(D\), i.e., vectors that satisfy

\[Du_{i}\approx 0 \tag{30}\]

with null vector \(0\) and \(i\in\{1,\ldots,s\}\) for \(s=\dim(\tilde{V}_{I})\). These vectors are then blocked such that one site \(y\in\tilde{S}\) on the coarse lattice corresponds to a set of sites, or block, \(B(y)\subset S\) on the fine lattice. Let us denote such a blocked vector, which lives on the sites \(B(y)\), by \(u_{i}^{y}\). One then defines an inner product within each block \(B(y)\) and orthonormalizes the vectors \(u_{1}^{y},\ldots,u_{s}^{y}\) within each block according to this inner product. The resulting vectors are labeled \(\bar{u}_{1}^{y},\ldots,\bar{u}_{s}^{y}\). The linear map \(W^{\dagger}\) discussed in Sec. II.3 is then defined as

\[W(y,x)^{\dagger}=\sum_{i=1}^{s}\bar{u}_{i}^{y}(x)\hat{e}_{i}^{\dagger} \tag{31}\]

with standard basis \(\hat{e}_{1},\ldots,\hat{e}_{s}\) of \(\tilde{V}_{I}\) and \(x\in B(y)\).

In practice a good approximation of such vectors \(u_{i}\) can be found by applying the FGMRES solver for matrix \(D\) with source vector \(0\) and a random vector as initial guess. This procedure removes high-mode components in \(u_{i}\), leaving a linear combination of low-modes. We follow this approach in the numerical experiments presented in the following. While high precision is not needed, we solve to \(10^{-8}\) precision to avoid an additional tuning step. We use a coarse grid of size \(2^{3}\times 4\) and a list of 12 near-null vectors \(u_{1},\ldots,u_{12}\).

We define a coarse-grid operator

\[\tilde{D}=RD_{\text{WC}}P \tag{32}\]

with restriction matrix \(R\) and prolongation matrix \(P\) that are defined according to Eqs. (14) and (15). We then train a coarse-grid model \(\tilde{M}\) that contains a single LPTC layer with gauge fields \(U_{\mu}=\mathbbm{1}\), \(V_{G}=\mathbbm{C}^{1}\), \(V_{\tilde{G}}=\tilde{V}_{I}\), and use only zero-hop and one-hop paths corresponding to \(\{H_{1},H_{2},H_{3},H_{4},H_{-4}\}\). We omit the \(H_{-1}\), \(H_{-2}\), and \(H_{-3}\) paths since they are redundant on a \(2^{3}\times 4\) coarse grid with periodic boundary conditions. The gauge fields are replaced with the identity since the coarse fields do not have a gauge degree of freedom. We refer to this special case of the LPTC layer as cLPTC in the following.

Figure 8: Convergence of the cost function (27) and iteration count gain for one-layer and one-hop high-mode preconditioners. We re-train the model of Fig. 6 for a different gauge configuration in the same ensemble, for a different value of \(\beta=5.9\), and for a different mass value of \(m=-0.55\). The network performs well in all cases even without re-training.
We follow the training procedure described in Sec. IV.1 but replace the cost function with

\[C=\left|\tilde{M}\tilde{D}v-v\right|^{2}. \tag{33}\]

It is worth noting that one could have considered a different cost function

\[C^{\prime}=|\tilde{M}v-\tilde{D}^{-1}v|^{2} \tag{34}\]

in order to project more strongly on the low modes of \(\tilde{D}\). In this case, however, the training tuples require the somewhat costly inversion of \(\tilde{D}\). We find that the cost function Eq. (33) is sufficient for the purpose of training the coarse-grid model. This point will be revisited when we train a combined multi-grid model in Sec. VI.

Note that the gauge equivariance of the restriction and prolongation layers is guaranteed if every vector \(u_{i}\) is a linear combination of eigenmodes of \(D\) with gauge-invariant coefficients. In our procedure the coefficients are gauge invariant in the statistical average over random initial guess vectors. Furthermore, note that the weights \(W\) of the restriction and prolongation layers could also be learned directly [12; 14]. We leave the systematic study of learning the restriction and prolongation layers, including explicitly gauge-equivariant versions, to future work.

### Results

In Fig. 9, we show the cost function (33) and the iteration count gain for the training of the coarse-grid model \(\tilde{M}\). In this case, we consider the iteration count gain for the inverse of \(\tilde{D}\). We find that a significantly longer training process is needed compared to the high-mode preconditioner models of Sec. IV.

We also investigate using the fully trained model from a given gauge configuration and applying it to a different gauge configuration. We use the same definition of the restriction and prolongation layers on the different gauge configuration to preserve the definition of \(\tilde{D}\). For the same reason we also use the same seeds for the random number generator to generate the initial guess for the fields \(u_{1},\dots,u_{12}\). We find that after a modest amount of re-training the model performs very well on the different gauge configuration. The re-training phase is significantly shorter compared to the initial training phase. We note that the maximum iteration count gain again differs to some degree between configurations.

## VI Multi-Grid Preconditioners

In the previous sections we successfully trained separate models \(M\) to approximate the short-distance and long-distance features of \(D^{-1}\). In this section we combine them to obtain a model that approximates \(D^{-1}\) over a wide range of distances.

### Smoother model setup and training strategy

We first create a version of the short-distance model that accepts a second input feature, which provides an initial guess. This model plays the role of a smoother in the multi-grid paradigm. The initial guess is provided by the long-distance model acting on the coarse grid.

Concretely, we aim to find a sequence of \(u_{k}\) that approximately solve \(Du=b\) such that the equation becomes exact in the \(k\to\infty\) limit. The smoother then maps the tuple \((u_{k},b)\) to \(u_{k+1}\). If we have a high-mode model \(M_{\text{h}}\) that approximates \(D^{-1}\) sufficiently well this can be achieved by the iterative relaxation approach

\[u_{k+1} =(\mathbb{1}-M_{\text{h}}D)u_{k}+M_{\text{h}}b\] \[=u_{k}+M_{\text{h}}(b-Du_{k})\,. \tag{35}\]

This approach is also commonly referred to as defect correction with defect \(b-Du_{k}\).

Since both \(D\) and the high-mode model \(M_{\text{h}}\) can be represented by (L)PTC layers we should be able to train a model \(M_{\text{s}}\) only composed of (L)PTC layers to map \((u_{k},b)\) to a \(u_{k+r}\) for \(r\in\mathbb{N}^{+}\). Such a model has two input

Figure 9: Convergence of the cost function (33) and iteration count gain for one-layer and one-hop low-mode preconditioners. We show both the initial training in blue as well as the performance of the trained model on a different gauge field of the same gauge ensemble in orange. We find that after a moderate amount of re-training, the model performs well on a different gauge configuration.
features and one output feature. We may construct \(M_{\text{s}}\) using \(2r\) (L)PTC layers stacked back-to-back since each iteration of Eq. (35) corresponds to two (L)PTC layers. All but the final layer need two output features.

In order to choose a reasonable value for \(r\), we studied the performance of the final multi-grid preconditioner described below and found that \(r=2\) performed significantly better than \(r=1\). We therefore train the model \(M_{\text{s}}\) for \(r=2\) using the cost function

\[C=|M_{\text{s}}(u_{k},b)-u_{k+r}|^{2} \tag{36}\]

with random vectors \((u_{k},b)\) and \(u_{k+r}\) given by Eq. (35). We use the same optimizer as in Secs. IV and V.

In Fig. 10, we show the training progress. The iteration count gain is obtained by using \(M_{\text{s}}\) with initial guess zero as a preconditioner for \(Du=b\). We use both PTC and LPTC layers with zero-hop and one-hop paths. We expect these models to yield an iteration count gain of approximately twice the iteration count gain of the corresponding high-mode models shown in Fig. 6 because of \(r=2\). We find that this expectation is satisfied by our data. In Fig. 10, we first train the PTC model and then use the model weights as initial values for the LPTC model (using the same value for every site \(x\)). We find no additional benefit by using the LPTC model.

### Multi-grid model setup and training strategy

We are now ready to combine the individual models to a complete multi-grid model \(M\) as shown in Fig. 11. We start by duplicating the input feature. One copy is preserved for the smoother, while the other copy is restricted to the coarse grid, where we apply the coarse-grid model of Sec. V. The result is then prolonged to the fine grid, and both the copy of the initial feature and the result of the coarse-grid model are combined to two input features for the last four layers. These layers are the smoother that we have learned in Sec. VI.1.

We may expect this combined model to work well by using the weights obtained in the training of the respective model components. The model performance may, however, be further improved by continued training of the complete multi-grid model \(M\). For such additional training, we need to modify the cost function of Secs. V and IV such that both the low-mode and high-mode components of \(D\) constrain the model in the training phase. To this end, we use

\[C=|Mb_{h}-u_{h}|^{2}+|Mb_{\ell}-u_{\ell}|^{2} \tag{37}\]

with \(b_{h}=D_{\text{WC}}v_{1}\), \(u_{h}=v_{1}\), \(b_{\ell}=v_{2}\), and \(u_{\ell}=D_{\text{WC}}^{-1}v_{2}\). Here, \(v_{1}\) and \(v_{2}\) are random vectors normalized such that \(|b_{h}|=|b_{\ell}|=1\). We therefore use a batch size of two with one training tuple geared towards the high-mode component and the other training tuple geared towards the low-mode component of \(D_{\text{WC}}\). We can shift the focus of the training between both components by adding a relative weight factor to Eq. (37).

### Results

In Fig. 12, we show the performance of the multi-grid (MG) model with initial weights taken from the trained model components as well as progress achieved by continued training of the combined model \(M\). From the start, the model performs substantially better than the smoother by itself. Continued training of the combined model further improves the iteration count gain to approximately 40. Such continued training converges within the first 20 training steps.

We also study using the multi-grid model trained on one configuration applied to a different gauge configuration of the same gauge ensemble. In Fig. 12, we show that after a brief re-training phase of only 20 training steps, the model performs optimally on the different gauge configuration as well.

Note that for concreteness we only present results for a two-level multi-grid preconditioner in this work. The extension to multiple levels is straightforward. In Fig. 11, one merely has to replace the coarse-grid layer limited by the blue features by the entire model as presented in Fig. 11. By repeating this process \(n\) times, one obtains an \((n+2)\)-level multi-grid preconditioner.

Figure 10: Convergence of the cost function (36) and iteration count gain for four-layer and one-hop smoother. The iteration count gain is studied for the case of zero initial guess. We first train the PTC model and use the result as initial weights for the LPTC model.
Also note that we use a rather small lattice volume of \(8^{3}\times 16\) in this work. In future work, we will investigate multi-grid models in more challenging large-volume simulations, where even larger iteration count gains should be achievable.

## VII Summary and Outlook

In this paper we have initiated a program to use gauge-equivariant neural networks to learn preconditioners in lattice QCD. We introduced a number of building blocks from which suitable models can be constructed: (i) parallel-transport convolution layers that can include arbitrary paths, with either global or local weights, (ii) restriction and prolongation layers that implement the multi-grid paradigm, and (iii) parallel layers that act on a single input feature.

To solve the Dirac equation for the Wilson-clover Dirac operator we have first constructed models that approximate the high-mode and low-mode component of the operator separately. We then combined these models in a two-level multi-grid model, which can be extended straightforwardly to an arbitrary number of levels. In all cases we found that the models reduce the iteration count of the outer solver significantly, e.g., by up to \(O(40)\) in the multi-grid model. We also found that transfer learning works: If we consider another gauge configuration (for the same or a slightly different value of \(\beta\)) or a slightly different quark mass, only a modest amount of re-training (or none at all) is required for the model to perform efficiently again.

We also introduced a communication-avoiding algorithm in which layers do not transfer information between sub-volumes assigned to different MPI processes. In our numerical experiments we found that the performance, i.e., the iteration count gain, of the corresponding model is only slightly reduced. We expect that on large supercomputers, the wall-clock time saved by avoiding

Figure 11: The combined two-level multi-grid model studied in this work. The use of the multi-grid paradigm allows for the efficient transport of information over both short and long-distances. Additional levels can be introduced by recursively replacing the coarse-grid layer (limited by the blue features) by the entire model as presented above.

Figure 12: Convergence of the cost function (37) and iteration count gain for the complete multi-grid model. We use the weights of the individually trained model components as starting point and show further improvement by training the combined model. The model also performs well on a different gauge configuration and quickly converges to optimum performance after a modest amount of re-training.
communication more than compensates for this modest reduction.

There are many interesting directions which we plan to explore in future work. For example, we will attempt to learn the weights \(W\) of the restriction and prolongation layers directly, without computing the near-null vectors explicitly. Also, we will investigate the space of possible models that can be constructed from our building blocks in a more comprehensive manner. Furthermore, we plan to perform benchmarks that measure the cost of (re-) training and applying our models and compare the overall wall-clock time to standard state-of-the-art multi-grid methods. It would also be worthwhile to apply our ideas to Dirac operators whose spectrum encircles the origin, such as in the case of domain-wall fermions. Finally, our finding that very little, if any, re-training is needed between configurations suggests that the present approach could also be beneficial in the generation of gauge-field configurations by Markov chain Monte Carlo.

## Appendix A GPT code listings

In this appendix, we provide Grid Python Toolkit (GPT) [3] code listings to implement the models used in this work. We first import the library and load a gauge field \(U\):

```
1importgptasg
2
3#loadgaugefield
4U=g.load("gauge_field")
5grid=U[0].grid
```

The layer drawn in Fig. 2 corresponds to

```
1#objecttypesforQCD
2ot_i=g.ot_vector_spin_color(4,3)
3ot_w=g.ot_matrix_spin(4)
4
5#twodistinctpaths
6paths=[
7g.path().f(0).f(1).f(0),
8g.path().f(1).b(0)
9]
10
11#defineanabbreviation
12l=g.ml.layer
13
14#definethelayerofFig.2
15fig2=layer.parallel_transport_convolution(
16grid,U,paths,ot_i,ot_w,2,1
17)
```

in the case of lattice QCD. Next, we define restriction and prolongation layers to a coarse grid of size \(4^{4}\) defined using vectors \(\bar{u}_{i}\) as

```
1#definecoarsegrid
2coarsegrid=g.grid([4,4,4,4],g.double)
3
4#load\bar{u}_ivectors
5u_bar=g.load("u_bar")
6
7#createblockingmap
8b=g.block.map(coarse_grid,basis)
9
10#createrestrictionandprolongationlayers
11restrict=l.block.project(b)
12prolong=l.block.promote(b)
```

Note that in the numerical work in this paper, we used a \(2^{3}\times 4\) coarse grid, while we present the \(4^{4}\) case here since it lifts the degeneracy of paths mentioned in Sec. V.

The complete multi-grid preconditioner model of Fig. 11 corresponds to

```
1#defineabbreviations
2lptc=l.local_parallel_transport_convolution
3ptc=l.parallel_transport_convolution
4
5#identiessoncoursegrid
6one=g.complex(coarse_grid)
7one[:]=1
8
9I=[g.copy(one)foriinrange(4)]
10
11
12#coarse-gridvectorspace
13cot_i=g.ot_vector_complex_additive_group(
14len(u_bar)
15cot_w=g.ot_matrix_complex_additive_group(
16len(u_bar)
17}
18
19#consideronlynearest-neighborhops
20paths=[
21g.path().forward(i)
22foriinrange(4)
23]+[
24g.path().backward(i)
25foriinrange(4)
26]
27
28#coarse-gridlayer
29defcoarse_lptc(n_in,n_out):
30returnlptc(
31coarse_grid,I,paths,
32cot_i,cot_w,n_in,n_out
33)
34
35#fine-gridlayer
36define_ptc(n_in,n_out):
37returnptc(
38grid,U,paths,ot_i,
39ot_w,n_in,n_out
40)
41
42#combinedmulti-gridmodel
43model_multi_grid=g.ml.model.sequence(
44l.parallel(
45l.sequence(),
46l.sequence(
47restrict,
48coarse_lptc(1,1),
49prolong
50),
51fine_ptc(2,2),
52fine_ptc(2,2),
53fine_ptc(2,1),