Learning incomplete factorization preconditioners for GMRES

Venue: NLDL_org_2025_Conference
Published: 2024-09-06 18:20:31

Abstract: In this paper, we develop a data-driven approach to generate incomplete LU factorizations of large-scale sparse matrices. The learned approximate factorization is utilized as a preconditioner for the corresponding linear equation system in the GMRES method. Incomplete factorization methods are one of the most commonly applied algebraic preconditioners for sparse linear equation systems and are able to speedup  the convergence of Krylov subspace methods. However, they are sensitive to hyper-parameters and might suffer from numerical breakdown or lead to slow convergence when not properly applied. We replace the typically hand-engineered algorithms with a graph neural network-based approach that is trained against data to predict an approximate factorization. This allows us to learn preconditioners tailored for a specific problem distribution. We analyze and empirically evaluate different loss functions to train the learned preconditioners and show their effectiveness to decrease the number of GMRES iterations and improve the spectral properties on our synthetic dataset.

Main paper: 
 background: 
 We start by providing a brief overview of the GM-RES algorithm, emphasizing the critical role of preconditioners in accelerating convergence. Then, we describe graph neural networks which we use to parameterize the mapping to learn data-driven preconditioners.We focus on GMRES [1] , a popular and widely adopted iterative Krylov subspace method to solve general linear equation systems of the form Ax = b where A is a n × n non-singular matrix denoted by the set GL(n, R). The core idea of Krylov subspace methods is to iteratively refine the solution to the problem starting from an initial guess x 0 . At each iteration, the approximate solution is updated by minimizing the Euclidean norm of the residual within the corresponding Krylov subspace:EQUATIONwhere the k-th Krylov subspace K k (A, r 0 ) = span{r 0 , Ar 0 , A 2 r 0 . . . , A k-1 r 0 } of A is generated by the initial residual r 0 = b -Ax 0 [2] . The new approximation x k is computed by extending the previous solution x k-1 within an increasingly larger Krylov subspace. This subspace progressively offers more directions to approximate the solution by capturing more of the matrix A's action on the residual, improving the conditioning and accuracy of the minimization problem over iterations. The efficiency arises because x k is found by solving a smaller, reduced system in the subspace rather than the original larger system. The GMRES algorithm is guaranteed to find the exact solution to the linear equation system in at most n steps.Each iteration consists of two main steps. In the first step of the algorithm, an orthonormal basis for the subspace K k is computed. This can be achieved, for example, using Arnoldi's algorithm which is based on the Gram-Schmidt procedure. In the second step, an unconstrained least squares problem is solved via the QR factorization of the subspace basis obtained in the first step. Based on this, the solution of Equation ( 1) can be constructed efficiently.The detailed algorithm is shown in Appendix A.Preconditioning Arguably, the most important 149 design choice for any Krylov subspace method is 150 the choice of preconditioner [3] . The goal of precon-151 ditioning is to replace the original linear equation 152 system with a new preconditioned system that ex-153 hibits a better clustering of singular values, which 154 in turn leads to faster convergence of the iterative 155 scheme [12] . For a non-singular and easy-to-invert 156 matrix P , the linear equation system AP -1 y = b is 157 solved instead of the original problem. The solution 158 to the original problem is then given by x = P -1 y. 1 159 Constructing a preconditioner involves trading off 160 the time required to compute the preconditioner 161 itself and the resulting speed up in the iterative 162 scheme [3] . Since the original system A is typically 163 sparse, we restrict the preconditioning matrix P to 164 be subject to sparsity constraints as well.165 Algebraic preconditioners do not assume any addi-166 tional problem information and compute the precon-167 ditioning matrix solely based on the structure and 168 values of the matrix A in contrast to methods tak-169 ing into account the underlying problem domain. In 170 the simplest case, the Jacobi preconditioner approx-171 imates P using a diagonal matrix. More advanced 172 methods such as incomplete factorizations or the 173 Gauss-Seidel method compute approximate factor-174 izations of the original system A which implicitly 175 give rise to preconditioners as they allow efficient in-176 versions [3, 13] . Most commonly, the approximation 177 utilizes triangular factorizations since efficient inver-178 sion in at most O(n 2 ) operations can be achieved and 179 sparsity can be exploited [14] . However, It is also pos-180 sible to approximate the matrix P -1 directly [15] . 181Graph neural networks (GNNs) are a recently pop-183 ularized family of neural network architectures de-184 signed to process data represented on unstructured 185 grids or graphs [16] . A graph is a tuple G = (V, E) 186 consisting of vertices (or nodes) V and a bivariate 187 relation of edges E ⊆ V × V. With slight abuse of 188 notation, we refer to the edge features of the edge 189 e = (i, j) ∈ E as e ij ∈ R m and node features of node 190 i ∈ V as n i ∈ R k . The message-passing framework 191 updates the initial edge and node features -rep-192 resenting the input data -in each layer l as follows: EQUATIONEQUATION 
network architecture: 
 The goal of this paper is to learn a mapping f θ , parameterized by a graph neural network, that takes a non-singular matrix A ∈ GL(n, R) as input and outputs the corresponding preconditioner P for the equation system. In order to ensure that the preconditioning system is easily invertible -such that we can apply it within the GMRES algorithm This constraint is directly encoded in the network architecture as we discuss next.Network parameterization To parameterize the function f θ that outputs the preconditioner, we exploit the strong connection of matrices and linear algebra with graph neural networks [18, 19] .We treat the system A as the adjacency matrix of a corresponding graph that we use directly within the message-passing framework. This transformation is known in the literature as Coates graph representation is depicted in Figure 1 [20] . Similar to Häusner et al. [5] and Tang et al. [21] , we use node features describing the structural properties of the corresponding rows and columns of the matrix. We Additionally to the non-zero elements in the matrix, the graph has been modified by adding edges for all missing diagonal elements, corresponding to self-loops.use the non-zero elements of A as the corresponding 252 edge features of the graph.Compared to previous approaches which only out-254 put a single lower triangular factor, our model out-255 puts two separate matrices L and U . Häusner et 256 al. [5] incorporate positional encodings by modifying 257 the Coates graph representation, allowing them to 258 output the lower triangular factor as the message 259 passing is only executed over this set of edges. In 260 contrast, we add positional encodings directly to the 261 input data by adding an edge feature to each directed 262 edge in the graph. This feature indicates whether 263 the final output embedding of the corresponding 264 edge belongs to the upper-or lower-triangular part 265 (±1 respectively) without changing the underlying 266 graph structure used for message passing.Invertibility We require the predicted factors L 268 and U of the preconditioner obtained from the for-269 ward pass through the model to be invertible in 270 order to obtain a valid preconditioner for the GM-271 RES method during inference. Therefore, we need 272 to enforce non-zero elements on each diagonal.Since the input matrix A is not guaranteed to 274 have non-zero diagonal elements, we add the corre-275 sponding edges to the graph if necessary before the 276 message passing. This is visualized for the matrix 277 element A 33 in Figure 1 .278 However, it is still possible that the output of the 279 GNN results in a singular matrix as the final edge 280 representations, obtained from message passing, can 281 result in zero diagonal elements. Through additional 282 measures, we avoid this in both factors and enforce 283 non-singularity. Inspired by the classical LU factor-284 ization, we enforce unit diagonal elements for the U 285 factor. For the diagonal elements in the lower trian-286 gular factor L we are using the activation function 287 ζ ϵ (x) =      x, if |x| > ϵ, ϵ, if 0 ≤ x ≤ ϵ, -ϵ, if -ϵ ≤ x < 0,313 σ max (AP -1 ) ≤ ϵ -1 ∥A -P ∥ F + 1.Here, ϵ is the hyper-parameter of the preconditioner 315 chosen in Equation ( 3). Further, we can estimate 316 the Frobenius norm using Hutchinson's trace 317 estimator [23] which allows us to express the loss to 318 minimize the upper bound on the largest singular 319 value of the system in the following form:EQUATION)321where w is a standard normal distributed vector.This loss function has been previously applied by 323Häusner et al. [5] as an unbiased estimator of the 324 squared Frobenius norm distance.Lemma 2 The smallest singular value of AP -1 is 326 lower bounded by the following inequalities:327 σ min (AP -1 ) ≥ ∥P A -1 ∥ -1 F ≥ (∥P A -1 -I∥ F +1) -1 .Based on this lemma, we obtain the following loss function2 by using a similar approximation as in the previous step and optimizing over the inverse:EQUATIONThe drawback of this loss is that it requires computing A -1 w during training. In other words, we need to solve many potentially ill-conditioned linear systems in order to train the learned preconditioner.To maintain computational efficiency we can make use of a supervised dataset instead, consisting of tuples (A i , x i , b i ) such that A i x i = b i . This requires us to solve each training problem only once which can be computed beforehand in an offline fashion. Further, this allows us to include information about the distribution of the right-hand side of the systems in the training [6] . We replace the standard distributed samples w in Equation ( 6) with the available training samples -to avoid solving problem instances during training -leading to a biased approximation of the previous loss function:EQUATION)This loss function has been previously introduced by Li et al. [6] as an inductive bias introduced by the dataset for the Froebnius norm. Recently, Trifonov et al. [7] derived this loss function and the unbiased variant from equation ( 6) based on reweighing the result from Lemma 1 using the inverse matrix as weights. They conjecture that optimizing this loss leads to minimizing the low-frequency components in the preconditioned equation system resulting in larger small singular values.By combining the two previous loss functions from Equation ( 6) and the unregularized version of Equation (5), we introduce a novel combined loss with the goal to further improve the conditioning by taking into account both ends of the spectrum. The combined loss function -and its approximation using a supervised dataset similar to Equation (7) -areEQUATION)where α is a hyper-parameter chosen to balance the two individual loss terms. During the training, we are minimizing the empirical risk over the corresponding unsupervised or supervised dataset to learn the parameters:EQUATION)The proofs for the lemmas and the connection of the loss functions to previous methods from the literature can be found in Appendix C. Table 1 . Average performance of classical and data-driven preconditioners with different loss functions evaluated on 10 test samples. We show the conditioning of the right-preconditioned system AP -1 in the first three columns indicating the smallest (σmin) and largest (σmax) singular values as well as the condition number (κ). For each preconditioner we compute different Frobenius norm loss functions proportional to the bounds in Lemma 1 and Lemma 2. Further, we compare the convergence of the preconditioner in the GMRES algorithm in terms of the total time (including both the time to compute the preconditioner and solve the systems) and the number of iterations. For the unpreconditioned method, we compute the loss for P = I. Note that, the problem at hand is unbalanced as the spectrum of the original system is already skewed towards small singular values. The performance in terms of GMRES iterations of the different loss functions also depends on the initial distribution 416 of singular values. In Figure 3 we show this distri-417 bution for different preconditioners compared to the 418 original distribution. We compare the singular value 419 distributions and discuss the results in more detail 420 in Appendix D.1.Method σ min ↑ σ max ↓ κ ↓ ∥P -A∥ F ↓ ∥P A -1 -I∥ F ↓ Time ↓Comparison with data-driven preconditioner 422We additionally compare our method with previ-423 ously developed data-driven preconditioners for in-424 complete Cholesky factorizations (Learned IC) with 425 an exp-activation function for the diagonal elements 426 to ensure invertability. We use the loss function 427 L max to train the method, similar to previous ap-428 proaches [5] , and the same hyper-parameters as our 429 model for a fair comparison.From the results we can see that the loss and 431 conditioning is worse than our proposed method. 432 The method is able to reduce the number of iter-433 ations compared to the unpreconditioned GMRES 434 solver. However, the GMRES performance is signif-435 icantly worse than the proposed LU factorization 436 based techniques. This is, however, not very surpris-437 ing as the LU factorization offers more flexibility to 438 approximate the true matrix inverse. Comparison with classical preconditioners Additionally, we compare in Table 1 the performance of the learned preconditioner with the classical Jacobi and incomplete LU method without fill-ins, as well as the case where no preconditioner is applied.In general, the results show the vital importance of selecting an appropriate preconditioner. For the problem considered, ILU and the learned preconditioner perform nearly on par in terms of required iterations to solve the problem instances. We can see that even though the conditioning in terms of the largest and smallest singular value of the ILU method is by far inferior to the other preconditioners, it still manages to substantially reduce the number of iterations for convergence. This shows that the spectrum edges, while important, are not the only factor affecting the method's efficiency. It underscores the importance of considering the entire spectrum of the preconditioned system when analyzing the method's convergence. 
discussion & conclusion: 
 In this paper, we learn incomplete factorizations of sparse matrices directly from data. When applied as a preconditioner our approach leads to a significant reduction in terms of iterations compared to the unpreconditioned GMRES. Combining data-driven methods with existing algorithms has a huge potential since it allows obtaining the best results from both worlds: guarantees about the convergence is obtained from classical analysis of the algorithms while the data-driven components allow us to learn directly from data to improve upon the existing algorithms.Limitations Our approach assumes having access to a distribution of linear equation systems that share similarities between the problem instances. However, in practice, it is not always trivial to obtain such a distribution of problems. Further, the initial cost of training the network needs to be amortized over many future problems that need to be solved. However, given the initial sample problems, new problems can be solved faster in an online fashion which is important for time-critical applications.In our experiments we only train and evaluate our method on a single dataset of very limited size which is inspired by numerical computations. Applying the method to more diverse problem classes with different initial distribution of singular values is required in the future to obtain more general results about the performance. 
: 
 ficient matrix inversion. To decrease the storage and computational time not all elements of the factorization are computed for the preconditioner leading to a sparse but incomplete factorization. The computation of the preconditioner is time-consuming, difficult to parallelize, and can suffer from numerical instabilities [3, 4] . In this work, we replace the classical computation of the incomplete LU factors with a graph neural network (GNN). The GNN learns to predict the corresponding incomplete factors of the matrix directly by training against data. The main contributions of our paper are the following:• We design a GNN model and train it to output a non-singular incomplete LU factorization.• We theoretically analyze existing loss functions from the literature and showcase their connection to large and small singular values.• Building on these insights, we derive a novel loss function for training learned preconditioners that combines the benefits of prior approaches.In numerical experiments, we showcase the effectiveness of our model in reducing GMRES iterations in combination with the different loss functions and validate the theoretical results on a synthetic dataset.Most similar to our work, Häusner et al. [5] and Li et al. [6] both learn an incomplete factorization for the conjugate gradient method. Trifonov et al. [7] extend this to combine data-driven and classical algorithms by correcting the output of the incomplete Cholesky factorization with a learned component. However, all of these methods assume the matrix to be symmetric and positive definite and utilize the incomplete Cholesky factorization as the underlying preconditioning technique. In contrast, rather focus on the more general incomplete LU factorization, which does not require positive definiteness or symmetry, within the GMRES algorithm similar to the recently proposed GraphPAN [8] .In previous work, Chen [9] instead proposes directly estimating the inverse of the matrix using a non-linear neural network as a preconditioner, requiring the use of the flexible GMRES method, which has a more complex convergence behavior. However, the proposed method requires retraining for each new problem. In contrast, we propose leveraging the similarity within a class of linear equation systems to first train the model offline and then generate preconditioners for new problems with a negligible computational overhead during inference time.For the class of sparse approximate inverse preconditioner, Bånkestad et al. [10] learn the sparsity pattern of the approximate inverse matrix for which the preconditioner can be computed efficiently. Li et al. [11] instead propose to learn auto-encoder-based generative models to generate preconditioners. In contrast, this paper only focuses on learning factorized preconditioners that allow easy inversion. 
 
 Appendix: The GNN is designed so that the learned preconditioner matches the sparsity pattern of the input matrix A. While dynamically learning the sparsity pattern and matrix elements could improve preconditioner quality, the large search space and combinatorial complexity make this challenging.Our proposed loss function accounts for the spec-495 trum edges in the preconditioned system. However, 496 beyond the edges and condition number, the distribu-497 tion of singular values also affects the convergence 498 of Krylov subspace methods [24] . Designing loss 499 functions that consider this can further improve the 500 performance of learned preconditioners. It is also 501 possible to combine our data-driven LU factorization 502 with classical ILU methods along the lines of the 503 method proposed by Trifonov et al. [7] In lines 21-23 of Algorithm A.2, the situation where the element h k+1,k in the upper Hessenberg matrix H k+1,k becomes zero is addressed. If this condition is met, the loop breaks, indicating that the algorithm has prematurely found an invariant Krylov subspace. When this occurs, the (k + 1)th column of V k+1 does not exist because the last row of H k+1,k is zero. As a result, Equation (10) simplifies toEQUATIONAV k = V k H k .This indicates that the Krylov subspace has reached its full dimension, and no further vectors can be generated. Consequently, the GMRES algorithm can terminate early, as the exact solution lies within the current subspace. This premature convergence is referred to as "lucky" because it implies that the solution has been found in less than n iterations. Also, note that in practice, it is not necessary to run the full Arnoldi algorithm in each iteration but the basis vectors found in the previous iterations can be reused and only the most recent orthonormalization step needs to be executed. w i+1 = AP -1 v i 14:▷ Gram-Schmidt ortho-normalization 15:for j = 1 to i do 16: ▷ Exact solution found (lucky breakdown)EQUATION23: break 24: return V k+1 , H k+1,kPreconditioning The preconditioning matrix P (or its inverse) does not need to be stored explicitly. Instead, it is sufficient to have access to an implicit representation of the matrix which can be applied to a vector v to evaluate the product P -1 v as in line 22 of Algorithm A.1 and Algorithm A.2 respectively. In the case of preconditioner given by triangular factorizations such as ILU and our learned approach, the matrix is obtained by solving the system LU v = r which can be achieved in O(n 2 ) operations using the forward-backward substitution and can be further accelerated by exploiting the sparse structure of the matrices. The solution to solving these systems is equivalent to explicitly inverting the matrices and computing the matrix-vector product U -1 L -1 r which forms the preconditioning matrix as P -1 = (LU ) -1 . However, solving the triangular systems is more efficient in practice as the generally non-sparse inverse matrices are not required to be stored.Here, we provide additional details for the implementation of the dataset, our learned preconditioner, and the baseline preconditioners used in the experiments.Our goal with the provided dataset is not solving a real-world problem. Rather, we focus on a highly ill-conditioned synthetic dataset which allows us to systematically test the different loss functions derived. For the problem scale used, direct methods are much superior compared to the shown results.of the most fundamental problems in numerical computational science [25] . The general form of the Poisson equation is given by:EQUATION)EQUATIONwhere f (x) is the source function and u(x) is the unknown function to be solved for.In our study, we generate the matrix A using Noisy data However, we are not specifically interested in spd matrices. Therefore, to obtain a distribution A over a large number of Poisson equation PDE problems that can be efficiently solved using GMRES, we perturb all the non-zero entries with standard Gaussian noise, which yields arbitrary matrices not necessarily spd. Let us denote the perturbed matrix asA ′ = (a ′ ij ).We can represent this mathematically by:743 a ′ ij = a ij + X ij if a ij ̸ = 0, 0 if a ij = 0, 744where X ij ∼ N (0, 1) are i.i.d. random variables.By introducing randomness into A, we are ef-746 fectively modeling a scenario where the system's 747 properties are not perfectly regular, which could 748 represent, for example, some form of physical or 749 numerical irregularity or perturbation in the grid 750 or medium. During testing, we use a deterministic 751 right-hand side b as described previously. However, 752 for the supervised training dataset we sample the 753 vectors b to be normally distributed in order to 754 avoid overfitting as this would lead to a significant 755 decrease in model performance.Problem instances All problem instances gen-757 erated come from the same distribution but using 758 non-overlapping random seeds we ensure that there 759 is no data leakage between the training and test data. 760 We create 200 problem instances for training and 10 761 instances for validation and testing respectively.All problems are of size n = 2 500 with approx-763 imately 50 000 -150 000 non-zero elements. The 764 condition number κ of the unpreconditioned system 765 is in the range between 5 000 to 50 000. 766 Thus, even though the problem parameters only 767 vary slightly, the resulting conditioning and poten-768 tial performance of the different loss functions can 769 be very different between the individual problem 770 instances.We implement the neural network based on the 773 pytorch-geometric framework which offers direct 774 support for computing the node features and util-775 ity functions such as adding remaining self-loops 776 [26] . During inference, we utilize the numml sparse 777 matrix package to efficiently compute the forward-778 backward substitution [27] . The pseudocode for the 779 forward pass of the learned preconditioner is shown 780 in Algorithm B.1.Inputs and transformations As described in 782 Section 3, we transform the original linear equation 783 system matrix A into a graph G using the Coates 784 graph representation. We use the non-zero elements 785 of the matrix as edge features and add a second edge 786 feature as a positional encoding of the output. For 787 the node features, we utilize the same set of 8 input 788 features as previously applied by Häusner et al. [5] 789 describing the structural properties of the matrix: 790• the node degree deg(v)• the maximum degree of neighboring nodes e l+1 ij = ψ θ (e l ij , n l i , n l j )12:for i ∈ {1, . . . , n} do 13:▷ Aggregate edge features per node 14: ▷ Compute updated node features 16:m l+1 i = j∈N (i) e l+1n l+1 i = ϕ θ (n k i , m l+1 i ) 17:if not final layer in the network then 18:▷ Add skip connections 19:for (i, j) ∈ E do 20:e l+1 ij ← [e l+1 ij , a ij ] ⊤ 21: ▷ Form the lower triangular matrix L = (l ij ): (l ij ) ← e (L)ij , if (i, j) ∈ E and i > j, 0, otherwise.22: ▷ Form the upper triangular matrix U = (u ij ):(u ij ) ← e (L)ij , if (i, j) ∈ E and i < j, 0, otherwise.23: ▷ Ensure invertibility 24: for i ∈ {1, . . . , n} do 25:l ii ← ζ(e (L)ii )26:u ii ← 1 27: return (L, U ) the objective becomes numerically more difficult as 847 the two terms are in conflict with each other. We 848 choose the hyper-parameter α = 1 /7 based on the 849 validation set performance to balance out the two 850 loss terms of the combined loss.The GMRES algorithm and the Arnoldi method, 853 given in Algorithm A.1 and Algorithm A.2, respec-854 tively, are implemented in pytorch directly.Classical Preconditioner The Jacobi precon-856 ditioner can be implemented and applied very 857 efficiently using pytorch sparse data structures for 858 diagonal matrices. For the ILU(0) preconditioner, 859 we use the ILU++ implementation which is a 860 high-performance implementation of the algorithm 861 in C++ [28, 29] . 897 σ max (AP -1 ) = ∥AP -1 ∥ 2 = ∥(A -P + P )P -1 ∥ 2 = ∥(A -P )P -1 + I∥ 2 ≤ ∥(A -P )P -1 ∥ 2 + ∥I∥ 2 ≤ ∥A -P ∥ 2 ∥P -1 ∥ 2 + 1.Now, we exploit the structure of the learned preconditioner to obtain an upper bound for ∥P -1 ∥ 2 . Specifically, we use two key observations. First, the matrix U has ones on the diagonal. Second, the absolute values of the diagonal elements on the matrix L are bounded from below by ϵ, a property that arises from the specific choice of activation function in Equation (3).These properties, when combined with the fact that the singular values of a triangular matrix coincide with the absolute value of its diagonal elements, allow us to derive the following upper bound on the spectral norm of the preconditioner P :∥P -1 ∥ 2 = ∥U -1 L -1 ∥ 2 ≤ ∥U -1 ∥ 2 ∥L -1 ∥ 2 = 1 σ min (U ) • 1 σ min (L) ≤ ϵ -1 .Combining this bound on ∥P -1 ∥ 2 with the previous inequality, we obtain:σ max (AP -1 ) ≤ ϵ -1 ∥A -P ∥ 2 + 1.Proof of Lemma 2 In the second lemma, we show how to obtain a lower bound on the smallest singular value of the preconditioned system and it is derived as follows:σ min (AP -1 ) = 1 σ max (P A -1 ) = 1 ∥P A -1 ∥ 2 ≥ 1 ∥P A -1 ∥ F .Previous learned preconditioner approaches in the literature [6, 7] used the following additional steps to obtain a weaker bound on the smallest singular value of the system: The reason for this is that optimizing first approximation directly leads to degenerate solutions as the minimum is attained when P ≈ 0 which would avoid having small singular values but the large singular value of the preconditioned system would be unbounded. Therefore, one can see the additional term involving I as a regularizer enforcing P to be non-degenerate.∥P A -1 ∥ 2 =∥P A -1 ∥ 2 -1 + 1 =∥P A -1 ∥ 2 -∥I∥ 2 + 1 ≤ ∥P A -1 ∥ 2 -∥I∥ 2 + 1 ≤∥P A -1 -I∥ 2 + 1 ≤∥P A -1 -I∥ F + 1.Frobenius norm approximation Equivalently to optimizing over the Frobenius norm, we can optimize the preconditioner over the squared Frobenius norm instead as we are only interested in the best parameters, not the best objective value. This has the additional advantage that the loss function is fully differentiable. Further, we can approximate the squared Frobenius norm using Hutchinson's trace estimator which can be easily verified [23] . For a random vector w that satisfies E [ww T ] = I we can approximate the squared Froebnius norm using the following transformations:∥M ∥ 2 F = trace(M T M ) = trace(M T M E [ww T ]) = E [trace(M T M ww T )] = E [trace(w T M T M w)] = E [w T M T M w] ≈ w T M T M w = ∥M w∥ 2 2 .In practice, w is typically chosen to follow either a standard normal distribution or a Rademacher distribution.While it is possible to use multiple samples to obtain a more accurate approximation of the Frobenius norm, we follow previous approaches by estimating the norm using a single random vector to limit computational resources to create the dataset.Here, we present additional results for the singular values, choice of activation function and robustness with respect to the chosen random seed.Here, we provide additional results for the baseline and learned preconditioners. In methods, we can see that using the loss from Equation (6) leads to a distribution far away from zero but the singular values are not very clustered. The best clustering is achieved using the loss corresponding to large singular values L max . However, the spectrum does not show a fast decay towards small singular values. In the spectrum of the combined loss, we can see the combined effect of the two approaches: the spectrum is clustered around a single value and the singular values are away from zero. These two properties of the spectrum make the preconditioned system easier to solve leading to better performance observed in the experiments presented in Table 1 . We can see that with a larger ϵ parameter, both the small singular values of the preconditioned system as well as the large singular values decrease. The best performance is achieved when balancing the two different ends of the singular value spectrum.Finally, we can see that all models trained with the combined loss outperform the previous models trained on only one loss function showing the robustness of our proposed method.Noteworthy, in the NeuralPCG method proposed by Li et al. [6] the loss function L min is used leading to better results than observed in our experiments.However, a crucial difference between this method and our approach is the fact that the diagonal elements are not learned in the NeuralPCG preconditioner. Instead, the square root of the original elements of the matrix A is used. Since the matrices considered are positive definite, the diagonal Further, it is easy to integrate the method devel-1056 oped here with data-driven methods for incomplete 1057 Cholesky methods [5] . This can be achieved by 1058 adding the ϵ parameter directly to the diagonal ele-1059 ments in order to ensure a lower bound on the values. 1060Here, right preconditioning is used, but other formulations use left-or split-preconditioning instead often leading to similar but not equivalent results.By convention, we always minimize the loss function. 


Decision: The reviewers are all convinced of accepting this paper.

Reviews: review: title: title: Review of "Learning incomplete factorization preconditioners for GMRES"
paper_summary: paper_summary: This paper proposes to learn incomplete LU factorisations of sparse matrices through graph neural network data driven methods to be used a preconditioner for linear equations of GMRES. The paper highlights the flaws in traditional handcrafted methods that require significant time and strong heuristics to obtain a fast and accurate solutions for equation systems. This works proposes that the data driven approaches aim to learn tailored systems thus improving speed when adapting to different systems. The work is trained and evaluated on a synthetic dataset comprising problems arising from the discretization of the Poisson equation. Here results are presented for different losses and predconditioners. The results demonstrate improved performance both in terms of computational efficiency and task performance.
main_review: main_review: The authors address my weaknesses and answer questions, given my already positive outlook I maintain my score.
strength_weakness: strength_weakness: Major Comments:

1. The clear weakness of this work lies in the experimental analysis. The use of one problem set, and synthetic data limits the understanding of method generalisation given there is an assumption that is not validated that the test data lies significantly outside the training distribution as to be a fair evaluation of performance. I understand the nature of this problem limits possible evaluation scenarios, however, further empirical analyses would significantly strengthen the work.
2. The proposed combined loss seemingly performs worse than individual losses but does achieve this result faster. However, the concern here is that the GMRES method with no preconditioner is also performing at a similar standard, yet much slower. Can the authors quantify more clearly how much benefit they obtain from their method given pre-training and performance gains not being drastic.
3. No comparisons made against existing works. Although, your related works clearly state there are other data driven methods to  generate preconditioners these are not empirically analysed. How do I know your method is better than these alternatives?
4. What is the distinct purpose of employing a graph neural network? Correct me if I’m mistaken, however, there seems to be no advantage or reason to employ a graph neural network over a standard ANN or CNN approach? While you employ a GNN to enforce the scarcity, the same effect can be achieved with other NN approaches and slightly modified constraints.

Minor Comments:

1. Your introduction could perhaps better explain the challenges and limitations of existing approaches to better define the problem statement. 
2. The caption for Table 1 could be more informative, it is initially hard to interpret without careful reading of the main text. More effective caption could help the reader better interpret what metrics are being used to evaluate performance and what methods are being analysed. Also I assume time is in seconds? This needs clarification.
3. Given your method is primarily comparing against iterations and time, and as mentioned in your limitations, it would be beneficial to make more clear the compute time for pre-training and a cost (or at least approximate) of the time taken with pre-training in mind. 
4. The background of graph neural networks could perhaps be omitted to the appendix given the audience of the conference, and the resulting space used to explain prior misconceptions and misunderstandings in further details.

Questions: 
1. For the limits that you define as bounds of the training case to only explore he edges of the equation system. I understand that investigating the whole system is not tractable and your lemma provide the bounds, however, were any further analyses performed on cases outside of these bounds?
2. How does the impact of epsilon in the activation function affect performance of the method? How was this value found?
title: title: Interesting data-driven approach to lower-upper factorisation of large-scale sparse matrices using graph neural networks
paper_summary: paper_summary: This manuscript introduces a data-driven approach to perform lower-upper (LU) factorisation of large-scale sparse matrices. The starting point is the generalised minimal residual method (GMRES) algorithm, which is one of the most popular methods for solving the task of LU factorisation in the case of large-scale spare matrices. This task is typically computationally demanding, and iterative methods are necessary. These iterative methods are highly dependent on the preconditioning methods, which are hyperparameter sensitive and can have difficulties converging. Therefore, a graph neural network (GNN)-based approach is proposed to learn the preconditioner instead. By precomputing matrix factorisation to create a supervised dataset, the GNN can learn to perform the task and solve new factorisation in less iterations. The writing is clear and the results are encouraging and well presented.
main_review: main_review: I think the authors have done a good job with the rebuttal. The addition of another data-driven approach is welcome, and it is also interesting to see the effect of different hyperparameters. Given my initial positive response, this has been reinforced and I keep my original recommendation.
strength_weakness: strength_weakness: 1. Other deep learning-based baselines could improve the impression of the results. The comparison with the classical preconditions is suitable and welcome, but it would also be useful to see how alternative deep learning-based approaches could work in this problem setting. For instance, the work of Chen [1] would be impractical due to the requirement of retraining for each problem, but would at least give an indication to what the performance of a different deep learning-based approach could be. I do not expect this baseline to be implemented in an update version as I believe it would be too much work for this iteration. But a discussion on what deep learning-based baseline that would be most suitable to compare to in future works would be beneficial.

2. The evaluation could be made more robust. Training on 200 samples and testing on 10 is reasonable, but generating a larger dataset would be useful to ensure that the performance estimates are reliable. It would also be interesting to see the variation in performance between different training runs to shed light on the stability of the proposed methodology.

3. The introduction could be made more friendly towards reader that are unfamiliar with the field of LU factorisation. The introduction starts with "The GMRES algorithm" without defining the acronym and without a reference. Similarly, LU factorisation is also presented without being defined. While both of these are well-known methods, it would increase the clarity of the writing if they were written out on their first appearance.

[1] J. Chen. “Graph Neural Preconditioners for Iterative Solutions of Sparse Linear Systems”. In: arXiv preprint arXiv:2406.00809 (2024)