Learning Sparse Approximate Inverse Preconditioners for Conjugate Gradient Solvers on GPUs

We are grateful to Reviewer GmV6 for the positive feedback and for asking insightful questions that help us strengthen our claims. The figures will be clarified in the final version to ensure readability.

Concern 1 & 2: The authors claim strong generalization, and Table 5 indeed shows good performance on out-of-distribution problems. However, is this superiority also observed relative to other learning-based methods? For example, do the baselines in Table 4 fail to generalize under the conditions of Table 5?

If the proposed method does generalize better than competing approaches, what is the underlying reason? Could the locality of the computations be a factor?

We appreciate this excellent question. We compare the generalization capability of each learning-based methods on the same dataset shown in Table (4) following Table 5's OOD settings. Total time (ms) and total iterations k (in parentheses) are listed in the tables below.

vs. Neural PCG [1]: we decrease the density of mesh in the heat problem to an unseen value.

vs. Neural IF [2]: we increase the sparsity of the synthesized matrix from 1e-3 to 2e-3.

• In distribution: our method is 96% faster compared with Neural PCG[1] and 687% faster compared with Neural IF[2];
• Out of distribution: our method is 60% faster compared with Neural PCG[1] and 627% faster compared with Neural IF[2].

We strongly concur with your intuition that locality is fundamental to our method's strong generalization. Our preconditioner inherently captures local, transferable physical patterns from the graph structure. This property enables robustness across diverse scales and geometries, unlike global methods (e.g., IC-based) that tend to overfit to training data and struggle with distribution shifts.

Concern 3: The parameter $\varepsilon$ in $GG^\top + \varepsilon I$ is a hyper-parameter. How sensitive is the method’s performance to its value?

Thank you for asking about this important hyper-parameter. We conducted an ablation study to evaluate the sensitivity to $\varepsilon$, confirming that our method is robust and does not require careful per-instance tuning. We trained different models from scratch using various $\varepsilon$ values on the "Heat" dataset. The validation performance (in iterations) confirms that the model is not sensitive to the precise choice of $\varepsilon$ within a reasonable range.

Performance is stable and efficient for $\varepsilon$ between 3e-4 and 3e-2. Too large a value (1e-1) may over-regularize, causing the preconditioner to deviate from $A^{-1}$, which leads to performance degradation. This analysis robustly supports our claim that $\varepsilon$ generally does not require instance-specific fine-tuning within a practical range.

We thank Reviewer umir for the encouraging review and for the questions that allow us to elaborate on the core strengths of our method.

Concern 1: The paper identifies locality in this SPAI approach as being well-suited to the GPU, arguing two hops in $G G^T$. Could the discussion be expanded?

Our method's suitability for GPUs stems from two synergistic aspects:

1. Computational Locality: The preconditioner application ($M^{-1}r=GG^T r$) relies on sparse matrix-vector multiplication (SpMV), a highly parallelizable operation that is heavily optimized in modern GPU libraries.
2. Learning Locality: More fundamentally, the approximation $A^{-1} \approx GG^T$ is itself local (see Equation (3) in our paper). This aligns perfectly with the local message-passing mechanism of GNNs, allowing the model to efficiently and effectively learn these structural patterns.

Concern 2: What happens for very large matrices? Is the main limitation using a single GPU or are there other limitations? For example if training on sizes from 400 -- 32k, can the model be used on matrices of size 1M?

This is an excellent question about the method's scalability.

• Inference Stage: The computational complexity during inference scales near-linearly with the number of non-zero entries (nnz) under fixed GNN architecture (e.g., layers/hidden dimensions). This suggests that applying the trained model to million-scale sparse matrices is potentially feasible, provided that both the model and matrix data can fit within GPU memory. Our largest validated case with matrices up to ~100k ("Heat-Large" in Table 5) shows stable performance without retraining. However, we acknowledge that for extremely large or structurally complex matrices, deeper message-passing or adaptive architectures might be needed — which could affect scalability.
• Training Stage: The primary limitation is in training, where GPU memory for holding very large graphs becomes a bottleneck. This is a general challenge in the GNN field, not specific to our method. To contextualize scale: recent GNN works handle graphs with 50k–2M nodes [1], 14k–1.5M nodes [2], and 89k–2.4M nodes [3] (the number of nodes corresponds to the size of matrix in this work). We believe training scalability could potentially be addressed through established GNN techniques (e.g., graph sampling, distributed training), though this is beyond our current scope and would require dedicated validation.

[1] Wei-Lin Chiang, Xuanqing Liu, Si Si, Yang Li, Samy Bengio, and Cho-Jui Hsieh. 2019. Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional Networks. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining (KDD '19).

[2] Hanqing Zeng, Hongkuan Zhou, Ajitesh Srivastava, Rajgopal Kannan and Viktor Prasanna. 2020. GraphSAINT: Graph Sampling Based Inductive Learning Method. International Conference on Learning Representations

[3] Rui Xue, Haoyu Han, MohamadAli Torkamani, Jian Pei, and Xiaorui Liu. 2023. LazyGNN: large-scale graph neural networks via lazy propagation. In Proceedings of the 40th International Conference on Machine Learning (ICML'23)

We sincerely thank Reviewer JerM for the detailed feedback and for raising several critical questions. Your comments have helped us significantly improve the clarity and completeness of our work.

Concern 1: This loss function was previously introduced in article [1], with a detailed investigation and theoretical justification provided in [3].

We would like to clarify the fundamental difference between our work and references [1, 3]. Those methods aim to learn a Cholesky decomposition of the original matrix ($A \approx LL^T$). In contrast, our goal is to directly learn a decomposition of the inverse matrix ($A^{-1} \approx GG^T$). This is a fundamentally different objective, which necessitates our novel SAI loss function designed specifically for this purpose. We will emphasize this distinction more clearly in the revised manuscript.

Concern 2: The SAI loss is a key contribution. It uses a non-standard definition for the matrix norm.

• Could authors explain the motivation for choosing this specific heuristic instead of more conventional matrix norms?
• How does the performance of this specific matrix form compare to that of conventional matrix norms?

We chose this specific norm ($\Vert A \Vert = \mathrm{mean}_{A_{ij} \not= 0} \vert A_{ij}\vert$) for three main reasons:

1. Dimension-Agnostic: Unlike the Frobenius norm (where $\|I_n\|_F = \sqrt{n}$), our norm is insensitive to matrix size, which is crucial for training on datasets with varying matrix sizes.
2. Robustness: Compared to $L_1$ norms ($\|A\|_1 = \max_{1 \leq j \leq n} \sum_{i=1}^{m} |a_{ij}|$), which sum absolute values and can be dominated by sparse but large entries (common in ill-conditioned matrices), our averaging approach is less sensitive to individual entries, leading to more robust performance.
3. Simplicity: The spectral norm, defined as the largest singular value of $A$, is costly to compute, while our norm is a simple and inexpensive measure of average nonzero magnitude.

To quantitatively validate our choice, we conducted an ablation study on the “Heat” dataset. We trained our model with different loss norms and compared the average PCG iterations required for convergence. The results clearly demonstrate the superiority of our proposed norm.

Concern 3: Could authors give the results of GPU-to-GPU comparison against other learning-based baselines?

We thank the reviewer for this suggestion. We conducted the requested GPU-to-GPU comparison on the respective datasets of the baselines. The results clearly demonstrate the significant computational advantage of our method.

• vs. Neural PCG [1]:
  • Neural PCG: 51 ms (108 iterations)
  • Ours: 26 ms (102 iterations) (1.96x speedup)
• vs. Neural IF [2]:
  • Neural IF: 1040 ms (354 iterations)
  • Ours: 132 ms (456 iterations) (7.88x speedup)

Concern 4: Could authors consider heat equation with high contrast values? Could authors give results of the method's performance on non-smooth PDE fields?

Thank you for this insightful suggestion. We designed a heat problem with piecewise constant density values to demonstrate performance on non-smooth, high-contrast PDE problems. We set the higher density to 1e-2 and lower density to 1e-4 (100× contrast ratio), while keeping other parameters and meshes unchanged. We used the model trained on the “Heat” dataset described in the paper, and directly applied the preconditioner on this high-contrast dataset.

Notably, our method achieves the fastest solution time, substantially outperforming all baselines despite the discontinuous coefficients. Although Incomplete Cholesky requires fewer iterations, its higher per-iteration computational cost results in slower overall performance. These results validate that our method effectively handles high-contrast, non-smooth PDE fields without requiring specialized modifications.

[1] Li, Y., Chen, P. Y., Du, T., & Matusik, W. (2023, July). Learning preconditioners for conjugate gradient PDE solvers. In International Conference on Machine Learning (pp. 19425-19439). PMLR.

[2] Häusner, P., Öktem, O., & Sjölund, J. (2023). Neural incomplete factorization: learning preconditioners for the conjugate gradient method. arXiv preprint arXiv:2305.16368.

[3] Trifonov, V., Rudikov, A., Iliev, O., Laevsky, Y. M., Oseledets, I., & Muravleva, E. (2024). Learning from linear algebra: A graph neural network approach to preconditioner design for conjugate gradient solvers. arXiv preprint arXiv:2405.15557.

We thank Reviewer Rbu2 for the constructive feedback and positive evaluation of our work.

Concern 1: The approach of using GNN to generate a preconditioner may not be very general. For each type of application, it seems we must train a GNN, a process that is time consuming. Further, for the same type of problems, the GNN training is under some fixed settings. The effectiveness of the proposed methods may depend on these settings.

It is useful to discuss the training time for GNN models.

Our model, once trained offline, can be directly applied to solve the same PDE under various boundary conditions, physical parameters, and meshes. This "train once, use many times" paradigm offers significant generality.

Compared to traditional methods like AMG, which often require expert-level manual parameter tuning for each new problem instance, our GNN-based approach automates this process. The GNN learns to generate efficient preconditioners directly from data, representing a more general paradigm.

Regarding the training time, we provide the following details for our one-time offline training process:

• Heat & Poisson: ~24 hours on a single A100 GPU.
• Elasticity: ~10 hours on a single RTX 3060 GPU.
• Synthetic: ~12.2 hours on a single A100 GPU.

Concern 2: On GPU, while the proposed method is the fastest, the simplest diagonal preconditioner works quite well. So in the practical use, will people switch to the proposed method?

We appreciate the reviewer raising this practical point. We fully acknowledge that diagonal preconditioners (e.g., Jacobi) offer significant advantages in simplicity, minimal implementation overhead, and efficient GPU parallelization, making them a strong and widely adopted baseline for many real-world applications—especially on smaller or less complex problems.

However, for large-scale and highly challenging problems, our method provides substantial practical gains. It achieves a 2x speedup over the diagonal preconditioner on standard benchmarks (e.g., Table 3). Crucially, this gap widens with problem complexity: on the “Heat-Large” dataset (Table 5), our method saves ~100ms per solve, and in the extreme heat equation experiment (density=1e-5, as detailed in our response to Reviewer JerM), it solves the system in 29 ms versus 56 ms for the diagonal preconditioner. In contexts where milliseconds matter—such as real-time simulation or billion-scale systems -- this reduction translates to meaningful efficiency improvements.

Therefore, while diagonal preconditioners remain ideal for lightweight tasks, our method offers a valuable alternative for demanding scenarios where computational speed is critical, justifying its adoption in such cases.

Concern 3: In line 86, SPAI stands for Sparse Approximate Inverse. However, in line 237, the authors switch to use AINV.

Thank you for pointing this out. Sparse Approximate Inverse (SPAI) is a class of methods, aiming at using a sparse matrix to approximate the $A^{-1}$. AINV is a specific, widely-used implementation of SPAI that uses a factorization to guarantee the positive definiteness of the preconditioner, i.e., $A^{-1}\approx G G^T$, where $G$ is sparse [1]. We will clarify this distinction and add the appropriate citation in the final version of the paper.

[1] J. Scott and M. Tůma, Algorithms for Sparse Linear Systems. in Nečas Center Series. Cham: Springer International Publishing, 2023. doi: 10.1007/978-3-031-25820-6.