Graph Neural Preconditioners for Iterative Solutions of Sparse Linear Systems

## Abstract

Abstract Preconditioning is at the heart of iterative solutions of large, sparse linear systems of equations in scientific disciplines. Several algebraic approaches, which access no information beyond the matrix itself, are widely studied and used, but ill-conditioned matrices remain very challenging. We take a machine learning approach and propose using graph neural networks as a general-purpose preconditioner. They show attractive performance for many problems and can be used when the mainstream preconditioners perform poorly. Empirical evaluation on over 800 matrices suggests that the construction time of these graph neural preconditioners (GNPs) is more predictable and can be much shorter than that of other widely used ones, such as ILU and AMG, while the execution time is faster than using a Krylov method as the preconditioner, such as in inner-outer GMRES. GNPs have a strong potential for solving large-scale, challenging algebraic problems arising from not only partial differential equations, but also economics, statistics, graph, and optimization, to name a few.

## 1 Introduction

Iterative methods are commonly used to solve large, sparse linear systems of equations $\mathbf{A}\mathbf{x}=\mathbf{b}$. These methods typically build a Krylov subspace, onto which the original system is projected, such that an approximate solution within the subspace is extracted. For example, GMRES [], one of the most popularly used methods in practice, builds an orthonormal basis $\{\mathbf{v}_{1},\mathbf{v}_{2},\ldots,\mathbf{v}_{m}\}$ of the $m$-dimensional Krylov subspace by using the Arnoldi process and defines the approximate solution $\mathbf{x}_{m}\in\mathbf{x}_{0}+\operatorname{span}(\{\mathbf{v}_{1},\mathbf{v
}_{2},\ldots,\mathbf{v}_{m}\})$, such that the residual norm $\|\mathbf{b}-\mathbf{A}\mathbf{x}_{m}\|_{2}$ is minimized.

The effectiveness of Krylov methods heavily depends on the conditioning of the matrix $\mathbf{A}$. Hence, designing a good preconditioner is crucial in practice. In some cases (e.g., solving partial differential equations, PDEs), the problem structure provides additional information that aids the development of a preconditioner applicable to this problem or a class of similar problems. For example, multigrid preconditioners are particularly effective for Poisson-like problems with a mesh. In other cases, little information is known beyond the matrix itself. An example is the SuiteSparse matrix collection [], which contains thousands of sparse matrices for benchmarking numerical linear algebra algorithms. In these cases, a general-purpose (also called “algebraic”) preconditioner is desirable; examples include ILU [], approximate inverse [], and algebraic multigrid (AMG) []. However, it remains very challenging to design one that performs well universally.

In this work, we propose to use neural networks as a general-purpose preconditioner. Specifically, we consider preconditioned Krylov methods for solving problems in the form $\mathbf{A}\mathbf{M}\mathbf{u}=\mathbf{b}$, where $\mathbf{M}\approx\mathbf{A}^{-1}$ is the preconditioner, $\mathbf{u}$ is the new unknown, and $\mathbf{x}=\mathbf{M}\mathbf{u}$ is the recovered solution. We design a graph neural network (GNN) [, ] to assume the role of $\mathbf{M}$ and develop a training method to learn $\mathbf{M}$ from select $(\mathbf{b},\mathbf{x})$ pairs. This proposal is inspired by the universal approximation property of neural networks [] and encouraged by their widespread success in artificial intelligence [].

Because a neural network is, in nature, nonlinear, our preconditioner $\mathbf{M}$ is not a linear operator anymore. We can no longer build a Krylov subspace for $\mathbf{A}\mathbf{M}$ when a Krylov subspace is defined for only linear operators. Hence, we focus on flexible variants of the preconditioned Krylov methods instead. Specifically, we use flexible GMRES (FGMRES) [], which considers $\mathbf{M}$ to be different in every Arnoldi step. In this framework, all what matters is the subspace from which the approximate solution is extracted. A nonlinear operator $\mathbf{M}$ can also be used to build this subspace and hence the neural preconditioner is applicable to FGMRES.

We use a *graph* neural network because of the natural connection between a sparse matrix and the graph adjacency matrix, similar to how AMG interprets the coefficient matrix with a graph. Many GNN architectures access the graph structure through $\mathbf{A}$-multiplications and they become polynomials of $\mathbf{A}$ when the nonlinearity is omitted [, ]. This observation bridges the connection between a GNN as an operator and the polynomial approximation of $\mathbf{A}^{-1}$, the latter of which is an essential ingredient in the convergence theory of Krylov methods. Interestingly, a side benefit resulting from the access pattern of $\mathbf{A}$ is that the GNN is applicable to not only sparse matrices, but also structured matrices that admit fast $\mathbf{A}$-multiplications (such as hierarchical matrices [, , ], Toeplitz matrices [, ], unassembled finite-element matrices [], and Gauss-Newton matrices []).

There are a few advantages of the proposed graph neural preconditioner (GNP). It performs comparatively much better for ill-conditioned problems, by learning the matrix inverse from data, mitigating the restrictive modeling of the sparsity pattern (as in ILU and approximate inverse), the insufficient quality of polynomial approximation (as in using Krylov methods as the preconditioner), and the challenge of smoothing over a non-geometric mesh (as in AMG). Additionally, empirical evaluation suggests that the construction time of GNP is more predictable, inherited from the predictability of neural network training, than that of ILU and AMG; and the execution time of GNP is much shorter than GMRES as the preconditioner, which may be bottlenecked by the orthogonalization in the Arnoldi process.

### 1.1 Contributions

Our work has a few technical contributions. First, we offer a convergence analysis for FGMRES. Although this method is well known and used in practice, little theoretical work was done, in part because the tools for analyzing Krylov subspaces are not readily applicable (as the subspace is no longer Krylov and we have no isomorphism with the space of polynomials). Instead, our analysis is a posteriori, based on the computed subspace.

Second, we propose an effective approach to training the neural preconditioner; in particular, training data generation. While it is straightforward to train the neural network in an online-learning manner—through preparing streaming, random input-output pairs in the form of $(\mathbf{b},\mathbf{x})$—it leaves open the question regarding what randomness best suits the purpose. We consider the bottom eigen-subspace of $\mathbf{A}$ when defining the sampling distribution and show the effectiveness of this approach.

Third, we develop a scale-equivariant GNN as the preconditioner. Because of the way that training data are generated, the GNN inputs $\mathbf{b}$ can have vast different scales; meanwhile, the ground truth operator $\mathbf{A}^{-1}$ is equivariant to the scaling of the inputs. Hence, to facilitate training, we design the GNN by imposing an inductive bias that obeys scale-equivariance.

Fourth, we adopt a novel evaluation protocol for general-purpose preconditioners. The common practice evaluates a preconditioner on a limited number of problems; in many occasions, these problems come from a certain type of PDEs or an application, which constitutes only a portion of the use cases. To broadly evaluate a new kind of general-purpose preconditioner and understand its strengths and weaknesses, we propose to test it on a large number of matrices commonly used by the community. To this end, we use the SuiteSparse collection and perform evaluation on a substantial portion of it (*all* non-SPD, real matrices falling inside a size interval; over 800 from 50 application areas in this work). To streamline the evaluation, we define evaluation criteria, limit the tuning of hyperparameters, and study statistics over the distribution of matrices.

### 1.2 Related Work

It is important to put the proposed GNP in context. The idea of using a neural network to construct a preconditioner emerged recently. The relevant approaches learn the nonzero entries of the incomplete LU/Cholesky factors [, ] or of the approximate inverse []. In these approaches, the neural network does not directly approximate the mapping from $\mathbf{b}$ to $\mathbf{x}$ like ours does; rather, it is used to complete the nonzeros of the predefined sparsity structure. A drawback of these approaches is that the predefined sparsity imposes a nonzero lower bound on the approximation of $\mathbf{A}^{-1}$, which (and whose factors) are generally not sparse. Moreover, note that although the neural network can be trained on a PDE problem (by varying the coefficients, grid sizes, or boundary conditions) and generalizes well to the same problem, it is unlikely that a single neural network can work well for all problems and matrices in life.

Approaches wherein the neural network directly approximates the matrix inverse are more akin to physics-informed neural networks (PINNs) [] and neural operators (NOs) [, , , , ]. However, the learning of PINNs requires a PDE (i.e., the physics) whose residual forms a partial training loss, whereas NOs consider infinite-dimensional function spaces, with PDEs being the primary application. In our case of a general-purpose preconditioner, there is not an underlying PDE; and not every matrix problem relates to function spaces with additional properties to exploit (e.g., spatial coordinates, smoothness, and decay of the Green’s function). For example, one may be interested in finding the commute times between a node and all other nodes in a graph; this problem can be solved by solving a linear system with respect the graph Laplacian matrix.
The only information we exploit is the matrix itself.

## 2 Method

Let $\mathbf{A}\in\mathbb{R}^{n\times n}$. From now on, $\mathbf{M}$ is an operator $\mathbb{R}^{n}\to\mathbb{R}^{n}$ and we write $\mathbf{M}(\mathbf{v})$ to mean applying the operator on $\mathbf{v}\in\mathbb{R}^{n}$. This notation includes the special case $\mathbf{M}(\mathbf{v})=\mathbf{M}\mathbf{v}$ when $\mathbf{M}$ is a matrix.

### 2.1 Flexible GMRES

A standard preconditioned Krylov solver solves the linear system $\mathbf{A}\mathbf{M}\mathbf{u}=\mathbf{b}$ by viewing $\mathbf{A}\mathbf{M}$ as the new matrix and building a Krylov subspace, from which an approximate solution $\mathbf{u}_{m}$ is extracted and $\mathbf{x}_{m}$ is recovered from $\mathbf{M}\mathbf{u}_{m}$. It is important to note that a subspace, by definition, is *linear*. However, applying a nonlinear operator $\mathbf{M}$ on the linear vector space does not result in a vector space spanned by the mapped basis. Hence, convergence theory of $\mathbf{x}_{m}=\mathbf{M}(\mathbf{u}_{m})$ is broken.

We resort to flexible variants of Krylov solvers [, , ] and focus on flexible GMRES for simplicity. FGMRES was designed such that the matrix $\mathbf{M}$ can change in each Arnoldi step; we extend it to a fixed, but nonlinear, operator $\mathbf{M}$.

Algorithm 1 in Section A summarizes the solver. The algorithm assumes a restart length $m$. Starting from an initial guess $\mathbf{x}_{0}$ with residual vector $\mathbf{r}_{0}$ and 2-norm $\beta$, FGMRES runs $m$ Arnoldi steps, resulting in the relation

$$ $\mathbf{A}\mathbf{Z}_{m}=\mathbf{V}_{m}\mathbf{H}_{m}+h_{m+1,m}\mathbf{v}_{m+1 }\mathbf{e}_{m}^{\top}=\mathbf{V}_{m+1}\overline{\mathbf{H}}_{m},$ (1) $$

where $\mathbf{V}_{m}=[\mathbf{v}_{1},\ldots,\mathbf{v}_{m}]\in\mathbb{R}^{n\times m}$ is orthonormal, $\mathbf{Z}_{m}=[\mathbf{z}_{1},\ldots,\mathbf{z}_{m}]\in\mathbb{R}^{n\times n}$ contains $m$ columns $\mathbf{z}_{j}=\mathbf{M}(\mathbf{v}_{j})$, $\mathbf{H}_{m}=[h_{ij}]\in\mathbb{R}^{m\times m}$ is upper Hessenberg, and $\overline{\mathbf{H}}_{m}=[h_{ij}]\in\mathbb{R}^{(m+1)\times m}$ extends $\mathbf{H}_{m}$ with an additional row at the bottom. The approximate solution $\mathbf{x}_{m}$ is computed in the form $\mathbf{x}_{m}\in\mathbf{x}_{0}+\operatorname{span}(\{\mathbf{z}_{1},\ldots,
\mathbf{z}_{m}\})=\{\mathbf{x}_{0}+\mathbf{Z}_{m}\mathbf{y}\}$, such that the residual norm $\|\mathbf{b}-\mathbf{A}\mathbf{x}_{m}\|_{2}=\|\beta\mathbf{e}_{1}-\overline{
\mathbf{H}}_{m}\mathbf{y}\|_{2}$ is minimized. If the residual norm is sufficiently small, the solution is claimed; otherwise, let $\mathbf{x}_{m}$ be the initial guess of the next Arnoldi cycle and repeat.

We provide a convergence analysis below, which is probably novel. All proofs are in Section B.

###### Theorem 1 .

Assume that FGMRES is run without restart and without breakdown. We use the tilde notation to denote the counterpart quantities when FGMRES is run with a fixed preconditioner matrix $\widetilde{\mathbf{M}}$. For any $\widetilde{\mathbf{M}}$ such that $\mathbf{A}\widetilde{\mathbf{M}}$ can be diagonalized, as in $\mathbf{X}^{-1}\mathbf{A}\widetilde{\mathbf{M}}\mathbf{X}=\bm{\Lambda}=
\operatorname{diag}(\lambda_{1},\ldots,\lambda_{n})$, the residual $\mathbf{r}_{m}=\mathbf{b}-\mathbf{A}\mathbf{x}_{m}$ satisfies

$$ $\|\mathbf{r}_{m}\|_{2}\leq\kappa_{2}(\mathbf{X})\epsilon^{(m)}(\bm{\Lambda})\| \mathbf{r}_{0}\|_{2}+\|\mathbf{Q}_{m}\mathbf{Q}_{m}^{\top}-\widetilde{\mathbf{ Q}}_{m}\widetilde{\mathbf{Q}}_{m}^{\top}\|_{2}\|\mathbf{r}_{0}\|_{2},$ (2) $$

where $\kappa_{2}$ denotes the 2-norm condition number, $\mathbf{Q}_{m}$ (resp. $\widetilde{\mathbf{Q}}_{m}$) is the thin-Q factor of $\mathbf{A}\mathbf{Z}_{m}$ (resp. $\mathbf{A}\widetilde{\mathbf{Z}}_{m}$), $\mathbb{P}_{m}$ denotes the space of degree-$m$ polynomials, and

$$ $\epsilon^{(m)}(\bm{\Lambda})=\min_{p\in\mathbb{P}_{m},\,p(0)=1}\max_{i=1, \ldots,n}|p(\lambda_{i})|.$ $$

Let us consider the two terms on the right-hand side of (2). The first term $\kappa_{2}(\mathbf{X})\epsilon^{(m)}(\bm{\Lambda})\|\mathbf{r}_{0}\|_{2}$ is the standard convergence result of GMRES on the matrix $\mathbf{A}\widetilde{\mathbf{M}}$. The factor $\epsilon^{(m)}(\bm{\Lambda})$ is approximately exponential in $m$, when the eigenvalues $\lambda_{i}$ are located in an ellipse which excludes the origin; see [, Corollary 6.33]. Furthermore, when the eigenvalues are closer to each other than to the origin, the base of the exponential is close to zero, resulting in rapid convergence. Meanwhile, when $\mathbf{A}\widetilde{\mathbf{M}}$ is close to normal, $\kappa_{2}(\mathbf{X})$ is close to 1.

What $\mathbf{M}$ affects is the second term $\|\mathbf{Q}_{m}\mathbf{Q}_{m}^{\top}-\widetilde{\mathbf{Q}}_{m}\widetilde{
\mathbf{Q}}_{m}^{\top}\|_{2}\|\mathbf{r}_{0}\|_{2}$. This term does not show convergence on the surface, but one can always find an $\widetilde{\mathbf{M}}$ such that it vanishes, assuming no breakdown. When $\mathbf{M}$ is close to $\mathbf{A}^{-1}$, $\mathbf{A}\mathbf{Z}_{m}$ is a small perturbation of $\mathbf{V}_{m}$ by definition. Then, (1) suggests that $\overline{\mathbf{H}}_{m}$ is close to the identity matrix for all $m$ and so is $\mathbf{H}_{n}$.

###### Corollary 2 .

Assume that FGMRES is run without restart and without breakdown. On completion, let $\mathbf{H}_{n}$ be diagonalizable, as in $\mathbf{Y}^{-1}\mathbf{H}_{n}\mathbf{Y}=\bm{\Sigma}=\operatorname{diag}(\sigma
_{1},\ldots,\sigma_{n})$. Then, the residual norm satisfies

$$ $\|\mathbf{r}_{m}\|_{2}\leq\kappa_{2}(\mathbf{Y})\epsilon^{(m)}(\bm{\Sigma})\| \mathbf{r}_{0}\|_{2}\qquad\text{for all $m$}.$ (3) $$

See Section C for an approximately exponential convergence of $\|\mathbf{r}_{m}\|_{2}$ due to $\epsilon^{(m)}(\bm{\Sigma})$.

###### Theorem 1 .

###### Corollary 2 .

### 2.2 Training Data Generation

We generate a set of $(\mathbf{b},\mathbf{x})$ pairs to train $\mathbf{M}:\mathbf{b}\to\mathbf{x}$. Unlike neural operators where the data generation is costly (e.g., requiring solving PDEs), which limits the training set size, for linear systems, we can sample $\mathbf{x}$ from a certain distribution and obtain $\mathbf{b}=\mathbf{A}\mathbf{x}$ with negligible cost, which creates a streaming training set of $(\mathbf{b},\mathbf{x})$ pairs without size limit.(^1^11Neural operators learn the mapping from the boundary condition to the solution. Strictly speaking, while $\mathbf{x}$ is the solution, $\mathbf{b}$ is not the boundary condition.)

There is no free lunch. One may sample $\mathbf{x}\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{n})$, which leads to $\mathbf{b}\sim\mathcal{N}(\mathbf{0},\mathbf{A}\mathbf{A}^{\top})$. This distribution is skewed toward the dominant eigen-subspace of $\mathbf{A}\mathbf{A}^{\top}$. Hence, using samples from it for training may result in a poor performance of the preconditioner, when applied to inputs lying close to the bottom eigen-subspace. One may alternatively want the training data $\mathbf{b}\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{n})$, which covers uniformly all possibilities of the preconditioner input on a sphere. However, in this case, $\mathbf{x}\sim\mathcal{N}(\mathbf{0},\mathbf{A}^{-1}\mathbf{A}^{-\top})$ is also skewed in the spectrum, causing difficulty in network training.

We resort to the Arnoldi process to obtain an approximation of $\mathbf{A}^{-1}\mathbf{A}^{-\top}$, which strikes a balance. Specifically, we run the Arnoldi process in $m$ steps without a preconditioner, which results in a simplification of (1):

$$ $\mathbf{A}\mathbf{V}_{m}=\mathbf{V}_{m}\mathbf{H}_{m}+h_{m+1,m}\mathbf{v}_{m+1 }\mathbf{e}_{m}^{\top}=\mathbf{V}_{m+1}\overline{\mathbf{H}}_{m}.$ (4) $$

Let the singular value decomposition of $\overline{\mathbf{H}}_{m}$ be $\mathbf{W}_{m}\mathbf{S}_{m}\mathbf{Z}_{m}^{\top}$, where $\mathbf{W}_{m}\in\mathbb{R}^{(m+1)\times m}$ and $\mathbf{Z}_{m}\in\mathbb{R}^{m\times m}$ are orthonormal and $\mathbf{S}_{m}\in\mathbb{R}^{m\times m}$ is diagonal. We define

$$ $\mathbf{x}=\mathbf{V}_{m}\mathbf{Z}_{m}\mathbf{S}_{m}^{-1}\bm{\epsilon},\qquad \bm{\epsilon}\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{m}).$ (5) $$

Then, with simple algebra we obtain that

$$ $\displaystyle\mathbf{x}$ $\displaystyle\in\operatorname{range}(\mathbf{V}_{m}),$ $\displaystyle\qquad\mathbf{x}$ $\displaystyle\sim\mathcal{N}(\mathbf{0},\bm{\Sigma}_{m}^{\mathbf{x}}),$ $\displaystyle\qquad\bm{\Sigma}_{m}^{\mathbf{x}}$ $\displaystyle=(\mathbf{V}_{m}\overline{\mathbf{H}}_{m}^{+})(\mathbf{V}_{m} \overline{\mathbf{H}}_{m}^{+})^{\top},$ $\displaystyle\mathbf{b}$ $\displaystyle\in\operatorname{range}(\mathbf{V}_{m+1}),$ $\displaystyle\qquad\mathbf{b}$ $\displaystyle\sim\mathcal{N}(\mathbf{0},\bm{\Sigma}_{m}^{\mathbf{b}}),$ $\displaystyle\qquad\bm{\Sigma}_{m}^{\mathbf{b}}$ $\displaystyle=(\mathbf{V}_{m+1}\mathbf{W}_{m})(\mathbf{V}_{m+1}\mathbf{W}_{m}) ^{\top}.$ $$

The covariance of $\mathbf{b}$, $\bm{\Sigma}_{m}^{\mathbf{b}}$, is a projector, because both $\mathbf{V}_{m+1}$ and $\mathbf{W}_{m}$ are orthonormal. As $m$ grows, this covariance tends to the identity $\mathbf{I}_{n}$. By the theory of the Arnoldi process [], $\mathbf{V}_{m+1}$ contains better and better approximations of the extreme eigenvectors as $m$ increases. Hence, $\mathcal{N}(\mathbf{0},\bm{\Sigma}_{m}^{\mathbf{b}})$ has a high density close to the bottom eigen-subspace of $\mathbf{A}$. Of course, because $\bm{\Sigma}_{m}^{\mathbf{b}}$ is low-rank, the distribution has zero density along many directions. Therefore, we sample $\mathbf{x}$ from both $\mathcal{N}(\mathbf{0},\bm{\Sigma}_{m}^{\mathbf{x}})$ and $\mathcal{N}(\mathbf{0},\mathbf{I}_{n})$ to form each training batch, to cover more important subspace directions.

We note that the use of Gaussian distributions brings in the convenience to generally analyze the input/output distributions. It is possible that other distributions are more effective if knowledge of $\mathbf{x}$ exists and can be exploited.

### 2.3 Scale-Equivariant Graph Neural Preconditioner

We now define the neural network $\mathbf{M}$ that approximates $\mathbf{A}^{-1}$. Naturally, a (sparse) matrix $\mathbf{A}=[a_{ij}]$ admits a graph interpretation, similar to how AMG interprets the coefficient matrix with a graph. Treating $\mathbf{A}$ as the graph adjacency matrix allows us to use GNNs [, , , ] to parameterize the preconditioner. These neural networks perform graph convolutions in a neighborhood of each node, reducing the quadratic pairwise cost by ignoring nodes faraway from the neighborhood.

Our neural network is based on the graph convolutional network (GCN) []. A GCN layer is defined as $\text{GCONV}(\mathbf{X})=\operatorname{ReLU}(\widehat{\mathbf{A}}\mathbf{X}
\mathbf{W})$, where $\widehat{\mathbf{A}}\in\mathbb{R}^{n\times n}$ is some normalization of $\mathbf{A}$, $\mathbf{X}\in\mathbb{R}^{n\times d_{\text{in}}}$ is input data, and $\mathbf{W}\in\mathbb{R}^{d_{\text{in}}\times d_{\text{out}}}$ is a learnable parameter.

Figure: Figure 1: Our GNN architecture for the preconditioner $\mathbf{M}$. The nonlinearity $\sigma$ is $\operatorname{ReLU}$. The component $s$ denotes the scale-equivariant operator (8).
Refer to caption: x1.png

Our neural network, as illustrated in Figure 1, enhances from GCN in a few manners:

- 1.
We normalize $\mathbf{A}$ differently. Specifically, we define
$\widehat{\mathbf{A}}=\mathbf{A}/\gamma\quad\text{where}\quad\textstyle\gamma=
\min\Big{\{}\max_{i}\big{\{}\sum_{j}|a|_{ij}\big{\}},\max_{j}\big{\{}\sum_{i}|
a|_{ij}\big{\}}\Big{\}}.$
(6)
The factor $\gamma$ is an upper bound of the spectral radius of $\mathbf{A}$ based on the Gershgorin circle theorem []. This avoids division-by-zero that may occur in the standard GCN normalization, because the nonzeros of $\mathbf{A}$ can be both positive and negative.
- 2.
We add a residual connection to GCONV, to allow stacking a deeper GNN without suffering the smoothing problem []. Specifically, we define a Res-GCONV layer to be
$\text{Res-GCONV}(\mathbf{X})=\operatorname{ReLU}(\mathbf{X}\mathbf{U}+\widehat
{\mathbf{A}}\mathbf{X}\mathbf{W}).$
(7)
Moreover, we let the parameters $\mathbf{U}$ and $\mathbf{W}$ be square matrices of dimension $d\times d$, such that the input/output dimensions across layers do not change. We stack $L$ such layers.
- 3.
We use an MLP encoder (resp. MLP decoder) before (resp. after) the $L$ Res-GCONV layers to lift (resp. project) the dimensions.

Additionally, the main novelty of the architecture is scale-equivariance. Note that the operator to approximate, $\mathbf{A}^{-1}$, is scale-equivariant, but a general neural network is not guaranteed so.
Because the input $\mathbf{b}$ may have very different scales (due to the sampling approach considered in the preceding subsection), to facilitate training, we design a parameter-free scaling $s$ and back-scaling $s^{-1}$:

$$ $\textstyle s(\cdot)=\frac{\sqrt{n}}{\tau}\cdot\quad\text{and}\quad s^{-1}( \cdot)=\frac{\tau}{\sqrt{n}}\,\cdot\,,\quad\text{where }\tau=\|\mathbf{b}\|_{2}.$ (8) $$

We apply $s$ at the beginning and $s^{-1}$ at the end of the neural network. Effectively, we restrict the input space of the neural network from the full $\mathbb{R}^{n}$ to a more controllable space—the sphere $\sqrt{n}\mathbb{S}^{n-1}$. Clearly, $s$ guarantees scale-equivariance; i.e., $\mathbf{M}(\alpha\mathbf{b})=\alpha\mathbf{M}(\mathbf{b})$ for any scalar $\alpha\neq 0$.

## 3 Evaluation Methodology and Experiment Setting

Problems. Our evaluation strives to cover a large number of problems, which come from as diverse application areas as possible. To this end, we turn to the SuiteSparse matrix collection
[https://sparse.tamu.edu](https://sparse.tamu.edu),
which is a widely used benchmark in numerical linear algebra. We select *all* square, real-valued, and non-SPD
matrices whose number of rows falls between 1K and 100K and whose number of nonzeros is fewer than 2M. This selection results in 867 matrices from 50 application areas. Some applications are involved with PDEs (such as computational fluid dynamics), while others come from graph, optimization, economics, and statistics problems.

To facilitate solving the linear systems, we prescale each $\mathbf{A}$ by $\gamma$ defined in (6). This scaling is a form of normalization, which avoids the overall entries of the matrix being exceedingly small or large. All experiments assume the ground truth solution $\mathbf{x}=\mathbf{1}$. FGMRES starts with $\mathbf{x}_{0}=\mathbf{0}$.

Compared methods. We compare GNP with three widely used general-purpose preconditioners—ILU, AMG, and GMRES (using GMRES to precondition GMRES is also called inner-outer GMRES)—and one simple-to-implement but weak preconditioner—Jacobi. The endeavor of preconditioning over 800 matrices from different domains limits the choices, but these preconditioners are handy with Python software support.

ILU comes from scipy.sparse.linalg.spilu; it in turn comes from SuperLU [, ], which implements the thresholded version of ILU []. We use the default drop tolerance and fill factor without tuning.
AMG comes from PyAMG []. Specifically, we use pyamg.blackbox.solver().aspreconditioner as the preconditioner, with the configuration of the blackbox solver computed from pyamg.blackbox.solver_configuration.
GMRES is self-implemented. For (the inner) GMRES, we use 10 iterations and stop it when it reaches a relative residual norm tolerance of 1e-6. For (the outer) FGMRES, the restart cycle is 10.

Note that sophisticated preconditioners require tremendous efforts to implement robustly. Hence, we prioritize ones that are more robust and that come with continual supports. For example, we also experiment with AmgX [] but find that while being faster, it throws more errors than does PyAMG. See Section D.3 for details.

Stopping criteria. The stopping criteria of all linear system solutions are rtol = 1e-8 and maxiters = 100.
For time comparisons, we additionally run experiments by enabling timeout and disabling maxiters. We set the timeout to be the maximum solution time among all preconditioners when using maxiters = 100 as the stopping criterion.

Evaluation metrics. Evaluating the solution qualities and comparing different methods for a large number of problems require programmable metrics, but defining these metrics is harder than it appears to be. Comparing the number of iterations alone is insufficient, because the per-iteration time of each method is different. However, to compare time, one waits until the solution reaches rtol, which is impossible to set if one wants *all* solutions to reach this tolerance. Thus, we can use maxiters and timeout to terminate a solution. However, one solution may reach maxiters early (in time) and attain a poor accuracy, while another solution may reach timeout very late but attain a good accuracy. Which one is better is debatable.

Hence, we propose two novel metrics, depending on the use of the stopping criteria. The first one, *area under the relative residual norm curve with respect to iterations*, is defined as

$$ $\text{Iter-AUC}=\sum_{i=0}^{\text{iters}}\log_{10}r_{i}-\log_{10}\texttt{rtol} ,\quad r_{i}=\|\mathbf{b}-\mathbf{A}\mathbf{x}_{i}\|_{2}/\|\mathbf{b}\|_{2},$ $$

where iters is the actual number of iterations when FGMRES stops and $r_{i}$ is the relative residual norm at iteration $i$. We take logarithm because residual plots are typically done in the log scale. This metric compares methods based on the iteration count, taking into account the history (e.g., convergence speed at different stages) rather than the final error alone. This metric is used when the stopping criteria are rtol and maxiters.

The second metric, *area under the relative residual norm curve with respect to time*, is defined as

$$ $\text{Time-AUC}=\int_{0}^{T}[\log_{10}r(t)-\log_{10}\texttt{rtol}]\,dt\approx \sum_{i=1}^{\text{iters}}[\log_{10}r_{i}-\log_{10}\texttt{rtol}](t_{i}-t_{i-1}),$ $$

where $T$ is the elapsed time when FGMRES stops, $r(t)$ is the relative residual norm at time $t$, and $t_{i}$ is the elapsed time at iteration $i$. This metric compares methods based on the solution time, taking into account the history of the errors. It is used when the stopping criteria are rtol and timeout.

Hyperparameters. To feasibly train over 800 GNNs, we use the same set of hyperparameters for each of them without tuning. There is a strong potential that the GNN performance can be greatly improved with careful tuning. Our purpose is to understand the general behavior of GNP, particularly its strengths and weaknesses. We use $L=8$ Res-GCONV layers, set the layer input/output dimension to 16, and use 2-layer MLPs with hidden dimension 32 for lifting/projection. We use Adam [] as the optimizer, set the learning rate to 1e-3, and train for 2000 steps with a batch size of 16. We apply neither dropouts nor weight decays. Because data are sampled from the same distribution, we use the model at the best training epoch as the preconditioner, without resorting to a validation set.

We use the $\ell_{1}$ residual norm $\|\mathbf{A}\mathbf{M}(\mathbf{b})-\mathbf{A}\mathbf{x}\|_{1}$ as the training loss. We use $m=40$ Arnoldi steps when sampling the $(\mathbf{x},\mathbf{b})$ pairs according to (5). Among the 16 pairs in a batch, 8 pairs follow (5) and 8 pairs follow $\mathbf{x}\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{n})$.

Compute environment. Our experiments are conducted on a machine with one Tesla V100(16GB) GPU, 96 Intel Xeon 2.40GHz cores, and 386GB main memory. All code is implemented in Python with Pytorch. All computations (training and inference) use the GPU whenever supported.

## 4 Results

We organize the experiment results by key questions of interest.

How does GNP perform compared with traditional preconditioners?
We consider three factors: convergence speed, runtime speed, and preconditioner construction cost.

Figure: Figure 2: Percentage of problems on which each preconditioner performs the best.
Refer to caption: extracted/5897006/figs/figs_precond_order_pie_chart.png

Figure: Figure 3: Preconditioner construction time and solution time (using maxiters to stop). The construction time of Jacobi is negligible and not shown. GMRES does not require construction.
Refer to caption: extracted/5897006/figs/figs_time_vs_size.png

In Figure 2 we show the percentage of problems on which each preconditioner performs the best, with respect to iteration counts and solution time (see the Iter-AUC and Time-AUC metrics defined in the preceding section). GNP performs the best for a substantial portion of the problems in both metrics. In comparison, ILU and AMG perform the best for more problems, while Jacobi, GMRES (as a preconditioner), and no preconditioner perform the best for fewer problems in the time metric.

That ILU and AMG are the best for more problems does not compromise the competitiveness of GNP. For one reason, they are less robust and sometimes bare a significantly higher cost in construction, as will be revealed later. Another reason is that GNP can be used when other preconditioners are less satisfactory. In Table 1, we summarize the distribution of the residual-norm ratio between the second best preconditioner and GNP, when GNP is the best. The ratio means a factor of $x$ decrease in the residual norm if GNP is used in place of other preconditioners. Under the Iter-AUC metric, a factor of 6780x decrease can be seen at the 75th percentile, and in the best occasion, more than ten orders of magnitude decrease is seen. Similarly under the Time-AUC metric. These significant reductions manifest the usefulness of GNP.

In Figure 3 and Figure 7 of Section D.1, we plot the preconditioner construction time and solution time for each matrix and each preconditioner. Here, the solution is terminated by maxiters rather than timeout (such that the times are different across matrices). Two observations can be made. First, the training time of GNP is nearly proportional to the matrix size, while the construction time of ILU and AMG is hardly predictable. Even though for many of the problems, the time is only a fraction of that of GNP, there exist quite a few cases where the time is significantly longer. In the worst case, the construction of the AMG preconditioner is more than an order of magnitude more costly than that of GNP.

Second, there is a clear gap between the construction and solution time for GNP, but no such gap exists for ILU and AMG. One is tempted to compare the overall time, but one preconditioner can be used for any number of right-hand sides, and hence different conclusions may be reached regarding time depending on this number. When it is large, the solution time dominates, in which case one sees that in most of the cases GNP is faster ILU, AMG, and GMRES.

On what problems does GNP perform the best?
We expand the granularity of Figure 2 into Figure 4, overlaid with the distributions of the matrices regarding the size, the condition number, and the application area, respectively. One finding is that GNP is particularly useful for ill-conditioned matrices (i.e., those with a condition number $\geq 10^{16}$). ILU is less competitive in these problems, possibly because it focuses on the sparsity structure, whose connection with the spectrum is less clear. Another finding is that GNP is particularly useful in three application areas: chemical process simulation problems, economic problems, and eigenvalue/model reduction problems. These areas constitute many problems in the dataset; both the number and the proportion of them in which GNP performs the best are substantially high.

On the other hand, the matrix size does not affect much the GNP performance. One sees this more clearly when consulting Figure 8 in Section D.2, which is a “proportion” version of Figure 4: the proportion of problems on which GNP performs the best is relatively flat across matrix sizes. Similarly, symmetry does not appear to play a distinguishing role either: GNP performs the best on 7.6% of the symmetric matrices and 17.5% of the nonsymmetric matrices.

Figure: Figure 4: Breakdown of best preconditioners with respect to matrix sizes, condition numbers, and application areas. Only the application areas with the top number of problems are shown. The last bar in the middle plot is for condition number $\geq 10^{16}$.
Refer to caption: extracted/5897006/figs/figs_precond_size_cond_kind_stacked_bar_cnt.png

**Table 2: Failures of preconditioners (count and proportion).**
|  | GNP | ILU | AMG | Jacobi | GMRES as precond |
| --- | --- | --- | --- | --- | --- |
| Construction failure | 0 (0.00%) | 348 (40.14%) | 62 (7.15%) | N/A | N/A |
| Solution failure | 1 (0.12%) | 61 (07.04%) | 5 (0.58%) | 53 (6.11%) | 2 (0.23%) |

How robust is GNP?
We summarize in Table 2 the number of failures for each preconditioner. GNP and GMRES are highly robust. We see that neural network training does not cause any troubles, which is a practical advantage. In contrast, ILU fails for nearly half of the problems, AMG fails for nearly 8%, and Jacobi fails for over 6%. According to the error log, the common failures of ILU construction are that “(f)actor is exactly singular” and that “matrix is singular … in file ilu_dpivotL.c”, while the common failures of AMG are “array … contain(s) infs or NaNs”. Meanwhile, solution failures occur when the residual norm tracked by the QR factorization of the upper Hessenberg $\overline{\mathbf{H}}_{m}$ fails to match the actual residual norm. While ILU and AMG perform the best for more problems than does GNP, safeguarding their robustness is challenging, rendering GNP more attractive.

What does the convergence history look like?
In Figure 5, we show a few examples when GNP performs the best. Note that in this case, ILU fails for most of the problems, in either the construction or the solution phase. The examples indicate that the convergence of GNP either tracks that of other preconditioners with a similar rate (Simon/venkat25), or becomes significantly faster. Notably, on VanVelzen/std1_Jac3, GNP uses only four iterations to reach 1e-8, while other preconditioners take 100 iterations but still cannot decrease the relative residual norm to below 1e-2.

Figure: Figure 5: Example convergence histories.
Refer to caption: extracted/5897006/figs/figs_example_history.png

We additionally plot the convergence curves with respect to time in Figure 9 in Section D.4. An important finding to note is that GMRES as a preconditioner generally takes much longer time to run than GNP. This is because GMRES is bottlenecked by the orthogonalization process, shadowing the trickier cost comparison between more matrix-vector multiplications in GMRES and fewer matrix-matrix multiplications in GNP.

Are the proposed training-data generation and the scale-equivariance design necessary?
In Figure 6, we compare the proposed designs versus alternatives, for both the training loss and the relative residual norm. In the comparison of training data generation, using $\mathbf{x}\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{n})$ leads to lower losses, while using $\mathbf{b}\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{n})$ leads to higher losses. This is expected, because the former expects similar outputs from the neural network, which is easier to train, while the latter expects drastically different outputs, which make the network difficult to train. Using $\mathbf{x}$ from the proposed mixture leads to losses lying in between. This case leads to the best preconditioning performance—more problems having lower residual norms (especially those near rtol). In other words, the proposed training data generation strikes the best balance between easiness of neural network training and goodness of the resulting preconditioner.

In the comparison of using the scaling operator $s$ versus not, we see that the training behaviors barely differ between the two choices, but using scaling admits an advantage that more problems result in lower residual norms. These observations corroborate the choices we make in the neural network design and training data generation.

Figure: Figure 6: Left: comparison of training data generation; right: comparison of scale-equivariance.
Refer to caption: extracted/5897006/figs/figs_neural_net_kde_plot.png

## 5 Discussions and Conclusions

We have presented a GNN approach to preconditioning Krylov solvers. This is the first work in our knowledge to use a neural network as an approximation to the matrix inverse, without exploiting information beyond the matrix itself (i.e., an algebraic preconditioenr). Compared with traditional algebraic preconditioners, this approach is robust with predictable construction costs and is the most effective in many problems. We evaluate the approach on more than 800 matrices from 50 application areas, the widest coverage in the literature.

One common skepticism on neural network approaches for preconditioning is that in general, network training is costly and in our case, the network can be used for only one $\mathbf{A}$. We have shown to the contrary, our training cost sometimes can be much lower than the construction cost of widely used preconditioners such as ILU and AMG, because network training is more predictable while the other preconditioners suffer the manipulation of the irregular nonzero patterns. Moreover, a benefit of neural networks is that compute infrastructures and library supports are in place, which facilitate implementation and offer robustness, as opposed to other algebraic preconditioners whose implementations are challenging and robustness is extremely difficult to achieve. Additionally, traditional algebraic preconditioners are constructed for only one $\mathbf{A}$ as well, not more flexible than ours.

While one may hope that a single neural network can solve many, if not all, linear systems (different $\mathbf{A}$’s), as recent approaches (PINN, NO, or learning the incomplete factors) appear to suggest, it is unrealistic to expect that a single network has the capacity to learn the matrix inverse for all matrices. Existing approaches work on a distribution of linear systems for the *same* problem and generalize under the same problem (e.g., by varying the grid size or PDE coefficients). For GNP to achieve so, we may fine-tune a trained network with a cost lower than training from scratch, or augment the GNN input with extra information (e.g., PDE coefficients). There unlikely exists an approach that remains effective beyond the problem being trained on.

Our work can be extended in many avenues. First, an immediate follow up is the preconditioning of SPD matrices. While a recent work on PDEs [] shows promise on the combined use of NOs and flexible CG [], it is relatively fragile with respect to a large variation of the preconditioner in other problems. We speculate that a split preconditioner can work more robustly and some form of autoencoder networks better serves this purpose.

Second, while GNP is trained for an individual matrix in this work for simplicity, it is possible to extend it to a sequence of evolving matrices, such as in sequential problems or time stepping. Continual learning [], which is concerned with continuously adapting a trained neural network for slowly evolving data distributions, can amortize the initial preconditioner construction cost with more and more linear systems in sequel.

Third, the GPU memory limits the size of the matrices and the complexity of the neural networks that we can experiment with. Future work can explore multi-GPU and/or distributed training of GNNs [, ] for scaling GNP to even larger matrices. The training for large graphs typically uses neighborhood sampling to mitigate the “neighborhood explosion” problem in GNNs [, ]. Some sampling approaches, such as layer-wise sampling [], are tied to sketching the matrix product $\widehat{\mathbf{A}}\mathbf{X}$ inside the graph convolution layer (7) with certain theoretical guarantees [].

Fourth, while we use the same set of hyperparameters for evaluation, nothing prevents problem-specific hyperparameter tuning when one works on an application, specially a challenging one where no general-purpose preconditioners work sufficiently well. We expect that this paper’s results based on a naive hyperparameter setting can be quickly improved by the community. Needless to say, improving the neural network architecture can push the success of GNP much farther.

## Appendix A FGMRES

FGMRES is summarized in Algorithm 1.

Figure: Algorithm 1 FGMRES with $\mathbf{M}$ being a nonlinear operator

## Appendix B Proofs

###### Proof of Theorem 1 .

Write $\|\mathbf{r}_{m}\|_{2}\leq\|\widetilde{\mathbf{r}}_{m}\|_{2}+\|\mathbf{r}_{m}-
\widetilde{\mathbf{r}}_{m}\|_{2}$. It is well known that

$$ $\|\widetilde{\mathbf{r}}_{m}\|_{2}\leq\kappa_{2}(\mathbf{X})\epsilon^{(m)}(\bm {\Lambda})\|\mathbf{r}_{0}\|_{2};$ $$

see, e.g., [, Proposition 6.32]. On the other hand, because $\mathbf{r}_{m}=\mathbf{r}_{0}-\mathbf{A}\mathbf{Z}_{m}\mathbf{y}_{m}$ where $\mathbf{y}_{m}$ minimizes $\|\mathbf{r}_{0}-\mathbf{A}\mathbf{Z}_{m}\mathbf{y}\|_{2}$, we have $\mathbf{r}_{m}=\mathbf{r}_{0}-(\mathbf{A}\mathbf{Z}_{m})(\mathbf{A}\mathbf{Z}_
{m})^{+}\mathbf{r}_{0}=\mathbf{r}_{0}-\mathbf{Q}_{m}\mathbf{Q}_{m}^{\top}
\mathbf{r}_{0}$. Therefore, $\|\mathbf{r}_{m}-\widetilde{\mathbf{r}}_{m}\|_{2}=\|\mathbf{Q}_{m}\mathbf{Q}_{
m}^{\top}\mathbf{r}_{0}-\widetilde{\mathbf{Q}}_{m}\widetilde{\mathbf{Q}}_{m}^{
\top}\mathbf{r}_{0}\|_{2}\leq\|\mathbf{Q}_{m}\mathbf{Q}_{m}^{\top}-\widetilde{
\mathbf{Q}}_{m}\widetilde{\mathbf{Q}}_{m}^{\top}\|_{2}\|\mathbf{r}_{0}\|_{2}$.
∎

###### Proof of Corollary 2 .

By taking $\widetilde{\mathbf{M}}=\mathbf{A}^{-1}\mathbf{V}_{n}\mathbf{H}_{n}\mathbf{V}_{
n}^{\top}$ and noting that $\mathbf{V}_{n}^{\top}=\mathbf{V}_{n}^{-1}$, we have

$$ $(\mathbf{V}_{n}\mathbf{Y})^{-1}\mathbf{A}\widetilde{\mathbf{M}}(\mathbf{V}_{n} \mathbf{Y})=\operatorname{diag}(\sigma_{1},\ldots,\sigma_{n}).$ $$

Then, following Theorem 1, $\|\mathbf{r}_{m}\|_{2}\leq\kappa_{2}(\mathbf{V}_{n}\mathbf{Y})\epsilon^{(m)}(
\bm{\Sigma})\|\mathbf{r}_{0}\|_{2}$. We conclude the proof by noting that $\kappa_{2}(\mathbf{V}_{n}\mathbf{Y})=\kappa_{2}(\mathbf{Y})$.
∎

###### Proof of Theorem 1 .

###### Proof of Corollary 2 .

## Appendix C Approximately Exponential Convergence of ‖ 𝐫 m ‖ 2 subscript norm subscript 𝐫 𝑚 2 \|\mathbf{r}_{m}\|_{2} ∥ bold_r start_POSTSUBSCRIPT italic_m end_POSTSUBSCRIPT ∥ start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT

We could bound the minimax polynomial $\epsilon^{(m)}$ by using Chebyshev polynomials to obtain an (approximately) exponential convergence of the residual norm $\|\mathbf{r}_{m}\|_{2}$ for large $m$.

Assume that all the eigenvalues $\sigma_{i}$ of $\bm{\Sigma}$ are enclosed by an ellipse $E(c,d,a)$ which excludes the origin, where $c$ is the center, $d$ is the focal distance, and $a$ is major semi-axis. We have

$$ $\epsilon^{(m)}(\bm{\Sigma})\leq\min_{p\in\mathbb{P}_{m},\,p(0)=1}\max_{\sigma \in E(c,d,a)}|p(\sigma)|,$ $$

by the fact that the maximum modulus of a complex analytical function is reached on the boundary of the domain. Then, we invoke Eqn (6.119) of [] and obtain

$$ $\epsilon^{(m)}(\bm{\Sigma})\leq\frac{C_{m}(a/d)}{C_{m}(c/d)},$ $$

where $C_{m}$ denotes the (complex) Chebyshev polynomial of degree $m$. When $m$ is large,

$$ $\frac{C_{m}(a/d)}{C_{m}(c/d)}\approx\left(\frac{a+\sqrt{a^{2}-d^{2}}}{c+\sqrt{ c^{2}-d^{2}}}\right)^{m},$ $$

which gives an (approximately) exponential function in $m$. Applying Corollary 2, we conclude that when $m$ is large,

$$ $\|\mathbf{r}_{m}\|_{2}\lessapprox\kappa_{2}(\mathbf{Y})\|\mathbf{r}_{0}\|_{2} \left(\frac{a+\sqrt{a^{2}-d^{2}}}{c+\sqrt{c^{2}-d^{2}}}\right)^{m}.$ $$

## Appendix D Additional Experiment Results

### D.1 Preconditioner Construction Time and Solution Time

In Figure 7, we plot the preconditioner construction time and solution time for each matrix and each preconditioner. This figure is a counterpart of Figure 3 by removing the log scale in the y axis. One sees that the GNP construction time is nearly proportional to the matrix size.

Figure: Figure 7: Preconditioner construction time and solution time (using maxiters to stop). The y axis is in the linear scale.
Refer to caption: extracted/5897006/figs/figs_time_vs_size_linear.png

### D.2 Breakdown of Best Preconditioners With Respect to Matrix Distributions

In Figure 8, we plot the proportion of problems on which each preconditioner performs the best, for different matrix sizes, condition numbers, and application areas. This figure is a counterpart of Figure 4. The main observation is that the proportion of problems on which GNP performs the best is relatively flat across matrix sizes. Other observations follow closely from those of Figure 4; namely, GNP is particularly useful for ill-conditioned problems and some application areas (chemical process simulation problems, economic problems, and eigenvalue/model reduction problems).

Figure: Figure 8: Proportion of problems on which each preconditioner performs the best, for different matrix sizes, condition numbers, and application areas. Only the application areas with the top number of problems are shown. The last bar in the middle plot is for condition number $\geq 10^{16}$. Note that not every matrix has a known condition number.
Refer to caption: extracted/5897006/figs/figs_precond_size_cond_kind_stacked_bar_prop.png

### D.3 Comparison between PyAMG and AmgX

Besides PyAMG, another notable implementation of AMG (which supports using it as a preconditioner and supports GPU) is AmgX [https://github.com/NVIDIA/AMGX](https://github.com/NVIDIA/AMGX). Here, we compare the two implementations in terms of robustness and speed.

Unlike PyAMG, which is a blackbox, AmgX leaves the choice of the AMG method and the parameters to the user. We used the driver examples/amgx_capi.c and the configuration src/configs/FGMRES_AGGREGATION.json, which appears to be the most reasonable for the diverse matrices and application areas at hand. This setting uses aggregation-based AMG rather than classical AMG, which is also consistent with the preference expressed by PyAMG.

Here are some findings.

AmgX is less robust than PyAMG. Table 3 shows that AmgX causes more failures during the construction as well as the solution phase. In fact, there were three cases where AmgX hanged, which hampered automated benchmarking.

**Table 3: Failures of two implementations of AMG.**
|  | PyAMG | AmgX |
| --- | --- | --- |
| Construction failure | 62 (7.15%) | 105 (12.11%) |
| Solution failure | 5 (0.58%) | 36 (4.15%) |

Despite a lack of robustness, the construction of AmgX is much faster than that of PyAMG. For AmgX, the distribution of the preconditioner construction time has a minimum 0.006s, median 0.017s, and maximum 4.895s. In contrast, the distribution for PyAMG has a minimum 0.005s, median 0.286, and maximum 8505.1s. Nevertheless, there are still 16% of the cases where the PyAMG preconditioner is faster to construct than the AmgX preconditioner.

We skip the comparison of the solution time. This is because in AmgX, the Krylov solver and the preconditioner are bundled inside a C complementation. A fair comparison requires isolating the preconditioner from the solver, because our benchmarking framework implements the solver in Python. However, the isolation is beyond scope due to the complexity of the implementation, such as the use of data structures and cuda handling. Nevertheless, we expect that the solution time of AmgX is faster than that of PyAMG, based on the observations made for preconditioner construction.

### D.4 Convergence Histories

In Figure 9, we show the convergence histories of the systems appearing in Figure 5, with respect to both iteration count and time. We also plot the training curves. The short solution time of GNP is notable, especially when compared with the solution time of GMRES as a preconditioner. The training loss generally stays on the level of 1e-2 to 1e-3. For some problems (e.g., VanVelzen_std1_Jac3), training does not seem to progress well, but the training loss is already at a satisfactory point in the first place, and hence the preconditioner can be rather effective. We speculate that this behavior is caused by the favorable architectural inductive bias in the GNN.

Figure: Figure 9: Convergence of the linear system solutions and training history of the preconditioners, for the matrices shown in Figure 5.
Refer to caption: extracted/5897006/figs/GHS_indef_bratu3d.png

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 6
presentation: 4
questions: See weaknesses.
soundness: 4
strengths: * Authors pay their attention to the problem which is linked to using neural networks (in particular, GNN) as the preconditioner. In this sense, the preconditioner is a nonlinear operator and, using a standard preconditioned solver with this preconditioner is not correct. In this case, the convergence theory is broken. They concentrate on flexible GMRES and present an original convergence analysis. 
* Authors propose original scheme of data generation. Suggested sampling is linked with eigen-subspace of the $\mathbf{A}$.
* Of course, the idea of using GNN as the preconditioner is not novel. Authors offer a fresh approach to normalize $\mathbf{A}$. This prevents the potential division-by-zero issue that can arise in the standard normalization of GCN. The major innovation of the architecture is its scale-equivariance, i.e. the input space of the neural network is restricted.
* Proposed architecture is checked on the SuiteSparse matrix collection https://sparse.tamu.edu (square, real-valued 867 matrices (not spd-matrices) from 50 application areas, whose count of non-zero elements is less than 2M. To compare with classical methods such as ILU, AMG, and GMRES author propose two novel metrics.
summary: Authors propose use Graph Neural Network (GNN) as a general-purpose preconditioner. They offer a convergence analysis for flexible GMRES. A new effective training data generation is proposed to training GNN. They develop a scale-equivariant GNN as the preconditioner. A novel evaluation protocol proposed by authors is being for general-purpose preconditioners.
weaknesses: * Authors use $\ell_1$ residual norm as the training loss. Is it possible to use $\ell_2$ norm or others as the training loss?
* Authors compare their architecture with classical methods(GNP, ILU, AMG, Jacobi, and GMREs). It will be good to compare their approach and metrics with the existing general-purpose preconditioners which are used GNN as the preconditioner. See for example, https://proceedings.mlr.press/v202/li23e/li23e.pdf, https://arxiv.org/pdf/2405.15557.
* Are there exist some constraints to use this approach for SPD-matrices?
* Why is not the full GitHub repository made available for review?
* Why do the authors only assume the ground truth solution $\mathbf{x} = 1$?

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. How does the proposed preconditioner handle matrices outside the tested conditions, particularly in real-world applications?

2. Could you elaborate on the training data generation process? How might biases in the training set affect the performance of the GNN?

3. What are the implications of your findings for future research in preconditioning techniques, especially concerning adapting the GNN for a wider array of problems?
soundness: 3
strengths: - **Innovative Approach:** The use of GNNs as preconditioners is a fresh perspective that leverages recent advancements in machine learning.

- **Empirical Validation:** The extensive evaluation on a wide range of matrices provides strong evidence of the proposed method's effectiveness across various problem domains.

- **Robustness:** The proposed method demonstrates high robustness with a low failure rate compared to traditional preconditioners, particularly in handling ill-conditioned problems.

- **Theoretical Contributions:** The paper includes a convergence analysis for the flexible GMRES method, which adds valuable theoretical insights to the practical implementation.
summary: The paper proposes a novel approach using Graph Neural Networks (GNNs) as general-purpose preconditioners for solving large, sparse linear systems. The authors present an empirical evaluation of their method, showing that it outperforms traditional preconditioning techniques like ILU and AMG in terms of construction time and execution efficiency, especially for ill-conditioned matrices. The method is tested on a diverse set of over 800 matrices, suggesting significant advantages in robustness and performance compared to existing methods.
weaknesses: - **Scalability Concerns:** While the paper mentions robustness, it lacks detailed discussions on scalability, particularly regarding the performance of GNNs with significantly larger matrices.

- **Dependence on Training Data:** The effectiveness of the GNN-based preconditioner seems to heavily rely on the quality and diversity of the training data. The paper does not explore the implications of this dependence thoroughly.

- **Limited Theoretical Foundation:** Although the empirical results are compelling, the theoretical grounding could be strengthened, especially concerning the convergence properties of the proposed methods under varied conditions.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. Since neither Iter-AUC nor Time-AUC specifically prioritizes the final residual accuracy, have you considered including a metric that directly measures the final residual? For applications where achieving a specific residual threshold is crucial, this could provide a more balanced evaluation. If not, what are your thoughts on this?
2. Iter-AUC and Time-AUC capture different aspects of convergence (iteration efficiency and time efficiency, respectively). In cases where these metrics might conflict (e.g., one method scores well on Iter-AUC but poorly on Time-AUC), how would you recommend interpreting the results?
soundness: 2
strengths: 1. Introduces a GNN-based preconditioner (GNP) as a novel, general-purpose alternative to traditional methods, approximating the matrix inverse without additional problem-specific information.
2. Offers stable and consistent training times, unlike traditional preconditioners with unpredictable construction times due to matrix irregularities.
3. Contributes a convergence analysis for FGMRES.
4. Tests GNP on over 800 matrices from 50 application areas.
summary: This paper proposes a GNN-based approach to preconditioning Krylov solvers, offering a novel alternative to traditional algebraic preconditioners. Unlike conventional methods, the GNN-based preconditioner (GNP) approximates the matrix inverse without relying on additional information, making it adaptable and robust across a wide range of matrices. GNP demonstrates predictable training costs, good robustness, and, in some cases, faster convergence than traditional preconditioners like ILU and AMG, which are often hampered by unpredictable construction times and structural limitations.
The contributions:
1. Convergence analysis for FGMRES, a widely used but theoretically underexplored method.  
2. New approach to training the neural preconditioner, focusing on effective training data generation by sampling from the bottom eigensubspace of  A to enhance training performance. 
3. Scale-equivariant GNN to serve as the preconditioner, addressing the challenge of varying input scales in the training data by enforcing an inductive bias that maintains scale-equivariance. 
4. Novel evaluation protocol to test the preconditioner broadly, evaluating across over 800 matrices from 50 diverse application areas in the SuiteSparse collection.
weaknesses: Sometimes both Iter-AUC and Time-AUC do not give good "weight" for the final residual accuracy. A method that reaches a lower final residual can be underrepresented if it converges more gradually, while a method with fast early convergence but a higher final residual might appear more favorable.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: - For the training data generation at the beginning of 2.2, is this really meaningful? You cannot sample the whole of $\mathbb{R}^n$ to run the problem forwards as also given possible ill-conditioning of $A$ approximating the inverse is difficult. Is this the reason the authors then switch to Arnoldi? This is the basis for many Krylov solvers and it is not clear what the advantage of the machine learning approach his if we have to do Arnoldi anyways?
- How about the sharpness of the Gershgorin estimates? These are not necessarily sharp.
- Do the authors really observe convergence of GMRES in 10 iterations to $10^{-6}$?
- AMG works better for symmetric problems, for nonsymmetric matrices PyAMG comes with better approaches for nonsymmetric problems such as the AIR approach.
soundness: 3
strengths: - Preconditioning is a relevant problem and discovering new preconditioners for novel applications is resource-demanding
- The paper tackles a large number of test matrices and compares with SOTA black box preconditioners
summary: The authors focus on the task of preconditioning linear systems which remains one of the most important problems in scientific computing and applied maths. The authors use a machine learning approach and propose using graph neural networks as a general-purpose preconditioner.
weaknesses: - The authors state their result as probably novel. Have they not found anything similar in the literature? I would like them to discuss the difference to the work USING FGMRES TO OBTAIN BACKWARD STABILITY IN MIXED PRECISION by Arioli and Duff, which they have done in the revised version.
- I am not a fan of presenting the results as best performer but it could be that the other methods are really close and this method only outperforms slightly and is terrible for the other problems. I would like to see actual  numbers or convergence plots.

Meta-review: This work proposes using graph neural networks (GNNs) as general-purpose preconditioners for solving large, sparse linear systems, particularly for challenging ill-conditioned matrices. Empirical evaluation on 800+ matrices shows that GNN-based preconditioners offer faster and more predictable construction times compared to traditional methods like ILU and AMG, with faster execution times than Krylov-based preconditioners. The approach has strong potential for large-scale problems in various fields, including PDEs, economics, and optimization.

Three reviewers recommended acceptance, emphasizing the substantial contributions of the paper. They highlighted the soundness and novelty of the proposed preconditioning technique, noting its fresh perspective. Another reviewer praised the thoroughness of the experiments, which evaluate performance across 800 matrices from 50 application areas, effectively showcasing the method's capabilities.
However, one reviewer recommended rejection, citing concerns about scalability, reliance on training data, and a limited theoretical foundation.

The AC sides with the majority vote, acknowledging that the paper's strengths significantly outweigh its weaknesses.

Decision: Accept (Poster)