Graph Neural Preconditioners for Iterative Solutions of Sparse Linear Systems

Official Comment by Authors (1/2)

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.

Official Comment by Authors (1/2)

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 $$ \| \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. $$ This convergence is based on mild assumptions (i.e., FGMRES does not incur breakdown and the resulting upper Hessenberg matrix can be diagonalized).

Official Comment by Authors (2/2)

**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.