Graph Neural Preconditioners for Iterative Solutions of Sparse Linear Systems

Thank you for your feedback. Please find point-to-point replies below. We also updated the paper with blue text. We will be happy to clarify further if you have additional questions.

RE: 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.

Thank you for bringing up the paper by Arioli and Duff (2009). We could comment on the nature of the convergence analysis of this paper and that of ours. Arioli and Duff (2009) are concerned with backward error analysis, which is typically in the following form:

There exists $\hat{k}$ such that for all $k \ge \hat{k}$, $\|\mathbf{r}_k\| \le c \varepsilon + O(\varepsilon^2)$. Here, $\varepsilon$ is the machine precision, signaling convergence, and $c$ may depend on $n$, $k$, $\|\mathbf{b}\|+\|\mathbf{A}\|\|\mathbf{x}_k\|$, and other terms.

Our analysis, on the other hand, is in the following form:

There exists $\hat{k}$ such that for all $k \ge \hat{k}$, $\|\mathbf{r}_k\| \lessapprox a \cdot b^k$, where $b<1$ and $a$ is independent of $k$.

The main difference is that our result spells out the convergence rate, while Arioli and Duff (2009) state convergence without a rate. Another difference is that the analysis of Arioli and Duff (2009) assumes finite arithmetics, while our analysis assumes infinite precision.

We have updated the paper with this discussion and removed the novelty claim to avoid confusion, Our novelty is the result but not the act of performing the analysis.

RE: 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.

In Figure 5, we showed the convergence plots for four matrices. In addition, we updated the paper by adding 138 convergence plots at the end (see Appendix Section E.6). These matrices come from application areas identified by Figure 4 to favor GNP: chemical process simulation problems, economic problems, and eigenvalue/model reduction problems. On many occasions, when GNP performs the best, it outperforms the second-best substantially. It can even decrease the relative residual norm under rtol while other preconditioners barely solve the system.

Thank you for your feedback. Please find point-to-point replies below. We also updated the paper with blue text. We will be happy to clarify further if you have additional questions.

RE: 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.

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?

Thank you for the thoughtful suggestion. In the updated paper, we summarized the results based on the suggested residual metric in Figure 7 (see Appendix Section E.1). Because for each problem, we solve the linear system twice, once using rtol and maxiters as the stopping criteria; and the other time using rtol and timeout. Hence, there are two charts in the figure, depending on how the iterations are terminated.

It is interesting to compare the two charts, as well as to compare them with Figure 2, plotted under the Iter-AUC and Time-AUC metrics. A few observations follow. First, across the two termination scenarios, observations are similar. Second, compared to the use of the Iter-AUC and Time-AUC metrics, GNP is the best for a similar portion of problems under the residual metric. Third, the portion of problems that ILU and GMRES (as a preconditioner) perform the best increases under the residual metric, while the portion that AMG performs the best decreases. Finally, Jacobi and no preconditioner remain less competitive.

RE: 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?

Indeed, not a single metric paints the full picture. The iteration metric is concerned with the convergence speed (more so from a theoretical standpoint), but if the per-iteration costs are different across methods, it does not properly reflect the time-to-solution progress. This happens in many studies of numerical algorithms, not only in the comparison of two preconditioners but also in the comparison of first-order and second-order optimizers, for example.

Practitioners may care more about time. However, the time metric has its limitations, because in order to measure time, one waits either until the solution reaches a preset residual tolerance, or until the number of iterations max out. For the former, it is impractical to set a universal tolerance for hundreds of problems from diverse domains. For the latter, one solution may max out the iterations early with poor accuracy, while another solution may reach timeout very late but attain good accuracy. Which one is better is arguable. Hence, the time metric is tricky to use.

An easier case is that one evaluates methods on only one (or a few) problem(s). In this case, it is possible to set a problem-specific residual tolerance. Then, one prioritizes the time metric over the iteration metric.

Thank you again for your comments. Please let us know if there are other concerns.

Dear authors, thank you for your detailed responses, espesially for adding residual metric and analyzing it. Now my comments are addressed. However, I believe my assessment is already correct.

Thank you for your feedback. Please find point-to-point replies below. We also updated the paper with blue text. We will be happy to clarify further if you have additional questions.

RE: Scalability Concerns: While the paper mentions robustness, it lacks detailed discussions on scalability, particularly regarding the performance of GNNs with significantly larger matrices.

The computational cost of GNP is dominated by matrix-matrix multiplications, where the left matrix is either $\mathbf{A}$ or a tall one that has $n$ rows. This forward analysis holds for training and inference, because the back propagation also uses these matrix-matrix multiplications. Hence, the computational cost of GNP scales as $O(nnz(\mathbf{A}))$, the same as that of the Krylov solver.

Practical considerations, on the other hand, are more sensitive to the constants hidden inside the big-O notation. For example, one needs to store many (but still a constant number of) vectors of length $n$, due to the batch size, the hidden dimensions of the neural network, and the automatic differentiation. Denote by $c$ this constant number; then, the GPU faces a storage pressure of $nnz(\mathbf{A}) + cn$. Once this amount is beyond the memory capacity of the GPU, one either performs training (and even inference) by using CPU only, offloads some data to the CPU and transfers them back to the GPU on demand, or undertakes more laborious engineering by distributing $\mathbf{A}$ and other data across multiple GPUs.

We have included this discussion in Appendix D of the updated paper.

RE: 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.

In Figure 6, we compared three approaches of generating the training data, including two naive choices ($\mathbf{x}$ is drawn from standard normal and $\mathbf{b}$ is drawn from standard normal) and our proposed design (denoted by "x ~ mix"). We compare their performance in terms of training quality and preconditioner quality. In terms of training quality, the figure suggests that using x ~ normal leads to lower losses, while using b ∼ normal 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 is harder to train. Using x from the proposed mixture leads to losses lying in between.

More important is the preconditioner quality. Our design leads to the best preconditioning performance. In other words, our design strikes the best balance between the easiness of neural network training and goodness of the resulting preconditioner.

RE: 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.

In Corollary 2 and Appendix Section C, we established the exponential convergence of the residual norm $\| \mathbf{r}_m \|_2$ with respect to the number of iterations $m$, namely

This convergence is based on mild assumptions (i.e., FGMRES does not incur breakdown and the resulting upper Hessenberg matrix can be diagonalized).

RE: How does the proposed preconditioner handle matrices outside the tested conditions, particularly in real-world applications?

The proposed preconditioner is learned separately for each matrix. The learning is based on a set of $(\mathbf{b},\mathbf{x})$ data pairs sampled from a distribution. The actual tested condition, $\mathbf{x}=\mathbf{1}$, does not come from this distribution. Hence, our evaluation can test out-of-distribution generalization.

For a new problem, one would train a preconditioner for it by sampling fresh $(\mathbf{b},\mathbf{x})$ data pairs consistent with the problem dimension.

Our tested problems come from SuiteSparse, which collects real-world matrices from numerous applications. For example, the HB collection, from which 30 matrices we tested, "stem from actual applications and exhibit numerical pathologies that arise in practice." [1] Testing 867 matrices from 50 applications in the experiments allows us to draw observations based on a large population. This is way more comprehensive than existing preconditioner papers that focus on one or two problems only.

Reference:

[1] I. S. Duff, Roger G. Grimes, and John G. Lewis. Sparse matrix test problems. ACM Trans. Math. Softw. 15(1), 1--14. 1989.

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

The training data is generated in a streaming fashion, meaning that it is not a finite dataset where a data point will be reused in each epoch. Each batch contains 16 $(\mathbf{b},\mathbf{x})$ pairs, where 8 of them have $\mathbf{x}$ generated from the standard normal and the rest of them have $\mathbf{x}$ generated from $\mathbf{V}_m\mathbf{Z}_m\mathbf{S}_m^{-1}\mathbf{\epsilon}$, with $\mathbf{\epsilon}$ being standard normal, $\mathbf{Z}_m$ and $\mathbf{S}_m$ from the SVD of $\overline{\mathbf{H}}_m$, and $\overline{\mathbf{H}}_m$ and $\mathbf{V}_m$ from the Arnoldi process. The Arnoldi process and the SVD are computed only once before all the batches. The vector $\mathbf{b}$ is computed as $\mathbf{A}\mathbf{x}$.

We are unsure what you meant by "biases in the training set". We do not introduce bias in the training data. Happy to elaborate further if you could clarify.

RE: What are the implications of your findings for future research in preconditioning techniques, especially concerning adapting the GNN for a wider array of problems?

Our findings reassert the folklore wisdom that there is not a one-size-fits-all preconditioner, especially when one moves beyond solving PDEs and enters the territories of data science and machine learning (for example, some of our test matrices come from statistical problems, graph problems, and optimization problems). Traditionally, practitioners choose the preconditioner for a particular problem by experience; trials and errors are an integral part of their work. Our preconditioner offers an additional choice to them. Our experiments suggest that a substantial portion of problems benefit from this choice.

We consider benchmarking to be an innovative contribution of this work, as there is rarely a preconditioner paper testing the proposed technique on a large number of problems and painting a full picture of its applicability. A benefit of large-scale benchmarking is that there is little room for problem-specific tuning. Hence, we offer default parameters, which a practitioner can try first. From there, the practitioner can tune the parameters and even the neural architecture for a particular problem.

For future research, we believe that split preconditioners, sequential fine-tuning, and distributed implementation are the imminent follow-up work, as laid out in the final section of the paper. These directions extend this work to problems that are SPD, that are time-varying, and that are exa-scale, respectively. The empirical success of our findings should also encourage the exploration of universal approximation of GNNs for matrix function (e.g., matrix inverse) times a vector.

Thank you again for your comments. Please let us know if the rebuttal and the update of the paper clear your concerns.

Thank you for your feedback. Please find point-to-point replies below. We also updated the paper with blue text. We will be happy to clarify further if you have additional questions.

RE: Authors use $\ell_1$ residual norm as the training loss. Is it possible to use $\ell_2$ norm or others as the training loss?

We mostly follow the general practice of robust regression, as the $\ell_1$ loss is less susceptible to outliers. Using other loss functions, such as the $\ell_2$ loss, is possible. In preliminary experiments, we did not find a substantial difference between the use of $\ell_1$ versus $\ell_2$. We speculate that this is because the signal $\mathbf{A}\mathbf{x}$ follows a Gaussian distribution and thus rarely incurs outliers. We have clarified the rationale of using $\ell_1$ in the updated paper.

RE: 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.

Existing GNN approaches for learning general-purpose preconditioners include the cited works by Li et al. (2023), Häusner et al. (2023), Bånkestad et al. (2024), and your suggested reference Trifonov et al. (2024). They are inapplicable to our scenario for two reasons. First, they learn a neural network that generates a preconditioner for SPD matrices (because, e.g., they learn the incomplete Cholesky factor), whereas in our work, we are concerned with non-SPD matrices only (please see more discussions in the next response item regarding the separate treatment of the SPD case). Second, the cited works learn the neural network from a distribution of matrices parameterized by the given problem (e.g., a parameterized PDE, or a PDE with different boundary conditions or mesh discretizations). In our case, we are given an algebraic matrix $\mathbf{A}$ without the problem parameterization.

We have included your suggested reference in the updated paper.

RE: Are there exist some constraints to use this approach for SPD-matrices?

Both (F)GMRES and the proposed GNP preconditioner apply to all matrices, including SPD ones. However, there are theoretical and practical advantages to using the alternative CG solver for SPD matrices, because in this case, the Arnoldi process is equivalent to the Lanczos process, which results in a tridiagonal matrix and gets rid of orthogonalization in practice. In this case, an SPD preconditioner works better with CG. Hence, we discuss in the final section that an extension of this work is to develop GNNs and training techniques that lead to a self-adjoint operator (the operator version of SPD preconditioner).

RE: Why is not the full GitHub repository made available for review?

Thank you for your interest in the code. Indeed, we plan to organize the code, package the preconditioner as a python module, and release the benchmarking framework when this paper is accepted for publication. By releasing the code, we hope to encourage the field to adopt the preconditioner and build their evaluation on the benchmarking framework.

RE: Why do the authors only assume the ground truth solution $\mathbf{x}=\mathbf{1}$?

Setting $\mathbf{x}=\mathbf{1}$ is a common practice in testing linear solvers and preconditioners, as is used in, for example, the original FGMRES paper. We deliberately avoid setting $\mathbf{x}$ to be Gaussian random, because this (partially) matches the training data generation, resulting in the inability to test out-of-distribution cases. We have updated the paper with this rationale.

Thank you again for your comments. Please let us know if there are other concerns.