Title: Learning Sparse Approximate Inverse Preconditioners for Conjugate Gradient Solvers on GPUs
Venue: NeurIPS 2025
Accepted: yes
Keywords: Graph Neural Networks

Preconditioning

GPU Computing

Iterative Solvers

Scientific Computing
Abstract: The conjugate gradient solver (CG) is a prevalent method for solving symmetric and positive definite linear systems $\mathbf{Ax} = \mathbf{b}$, where effective preconditioners are crucial for fast convergence. Traditional preconditioners rely on prescribed algorithms to offer rigorous theoretical guarantees, while limiting their ability to exploit optimization from data. Existing learning-based methods often utilize Graph Neural Networks (GNNs) to improve the performance and speed up the construction. However, their reliance on incomplete factorization leads to significant challenges: the associated triangular solve hinders GPU parallelization in practice, and introduces long-range dependencies which are difficult for GNNs to model. To address these issues, we propose a learning-based method to generate GPU-friendly preconditioners, particularly using GNNs to construct Sparse Approximate Inverse (SPAI) preconditioners, which avoids triangular solves and requires only two matrix-vector products at each CG step. The locality of matrix-vector product is compatible with the local propagation mechanism of GNNs. The flexibility of GNNs also allows our approach to be applied in a wide range of scenarios. Furthermore, we introduce a statistics-based scale-invariant loss function. Its design matches CG's property that the convergence rate depends on the condition number, rather than the absolute scale of $\mathbf{A}$, leading to improved performance of the learned preconditioner. Evaluations on three PDE-derived datasets and one synthetic dataset demonstrate that our method outperforms standard preconditioners (Diagonal, IC, and traditional SPAI) and previous learning-based preconditioners on GPUs. We reduce solution time on GPUs by 40%-53% (68%-113% faster), along with better condition numbers and superior generalization performance.


Peer reviews:
comments: SUMMARY:
A new type of sparse approximate inverse is presented, based on GNNs, which incorporates a scaled minimization of the inverse in the Frobenius norm.  A stochastic estimator is used to reduce the cost of evaluating in this norm.  The results shows both fewer iterations in PCG and lower times in general for a suite of problems ranging from the heat equation to hyperelasticity to random (SPD) matrices.

STRENGTHS AND WEAKNESSES:
The paper has several strengths.  It clearly motivates issues with previous GNN-based solvers and classical SPAI solvers.  The paper is also concise, very quickly summarizing the details of the method, the model, and the experimental setup.  Overall the results are quite good as well.  The appendix results on AMG are a good baseline (and address a predictable question).  Of course AMG will out perform for the Poisson problem here, but the reference point is very helpful.  The paper also compares with prior methods, another helpful data point.

One weakness is the size of the system.  Which appear to be limited to fairly small sizes at this point.  Also the train/test splitting of 4:1 raise the question of whether this is a practical solver in practice.

QUESTIONS:
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?  Is performance (of the SpMV) better on the GPU due to density?

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?

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 4
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-06-30

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a GPU-friendly preprocessing technique based on a graph neural network (GNN) that learns a Sparse Approximate Inverse. Most existing approaches rely on Cholesky factorization, whose triangular solves become bottlenecks on GPUs when used with the Conjugate Gradient (CG) method. To eliminate this bottleneck, the authors directly construct an approximation $A^{-1} \\approx GG^\\top$ with a GNN and use it as the preconditioner, reducing the it to sparse matrix–vector multiplications that run efficiently on GPUs. They also devise a scale-invariant loss function for training the GNN, which further improves performance. Numerical experiments on several PDE problems and on synthetic data demonstrate that the proposed method outperforms.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n- The manuscript is well organized and clearly written, making it very easy to follow.\n- Algorithm design that takes hardware characteristics into account is important; given the rapid performance gains of modern GPUs, GPU-oriented preconditioners such as the one proposed here are valuable.\n- The paper identifies triangular solves in factorization-based preconditioners as a key bottleneck and offers a simple yet effective way to avoid them.\n- The numerical experiments are comprehensive and thoughtfully designed; the purpose of each experiment is clear. Reporting both iteration counts and wall-clock times gives a full picture of the algorithm’s behavior. The per-stage timing analysis in Table 3 is particularly informative.\n\n## Weaknesses\nNo major shortcomings were found; if anything:\n- Figures are generally small and hard to read.\n- The method offers no special advantage in CPU computations (though this is arguably outside the scope of the work).\n\nQUESTIONS:\n1. 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?\n2. If the proposed method does generalize better than competing approaches, what is the underlying reason? Could the locality of the computations be a factor?\n3. The parameter $\\varepsilon$ in $GG^\\top + \\varepsilon I$ is a hyper-parameter. How sensitive is the method’s performance to its value?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses the challenge of accelerating the Conjugate Gradient (CG) solver for symmetric positive definite linear systems on GPUs. The paper proposes a method that use GNNs to construct Sparse Approximate Inverse (SPAI) preconditioners. This method is suitable for GPU because required only matrix-vector products which are highly parallelizable. The results demonstrate that the proposed method outperforms standard traditional and previous learning-based preconditioners on GPUs.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n* The paper is well written. I appreciated the smooth introduction on the preconditioning of linear systems, as well as the mathematical rigor throughout the paper.\n* The method is tested on three different PDE-derived problems (Heat, Poisson, Hyperelasticity) and a synthetic dataset.\n* The proposed method is compared against a comprehensive set of baselines, including standard traditional methods (Diag, IC, AINV) and recent learning-based approaches ([1], [2]).\n* The author provides results for both CPU and GPU across different datasets.\n* The author provides a comparison of the proposed loss function $L_{\\text{SAI}}$ to the cosine similarity loss $L_{\\text{CS}}$ and scale-variant loss $L_{2}$.\n\n**Weaknesses**\n* This loss function $L_{\\text{SAI}}$ was previously introduced in article [1], with a detailed investigation and theoretical justification provided in [3].\n* In Table 4, the GPU performance for the learning-based baselines [1, 2] is marked as \"N/A\" (or \"/\"). The authors justify this by stating that a component they use IC is slow on a GPU. While this justification is reasonable and supports their main idea, it also means there is no direct GPU-to-GPU comparison against other learning-based baselines. Consequently, the GPU speedup is demonstrated only against traditional methods.\n* The authors consider heat equation with $a(x)$ (randomly sampled on the mesh, with values ranging from $1 \\cdot 10^{-4}$ to $5 \\cdot 10^{-4}$). The proposed method seems to be mostly tested on very smooth equations.\n\n\n[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.\n\n[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.\n\n[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.\n\nQUESTIONS:\n* The **SAI loss** is a key contribution. It uses a non-standard definition for the matrix norm: $\\Vert A \\Vert = mean_{A_{ij} \\not= 0} \\vert A_{ij}\\vert$.\n\t* Could authors explain the motivation for choosing this specific heuristic instead of more conventional matrix norms (e.g., the Frobenius, spectral, or L1 norm)?\n\t* How does the performance of this specific matrix form compare to that of conventional matrix norms, such as the Frobenius, spectral, or L1 norm?\n* Could authors give the results of GPU-to-GPU comparison against other learning-based baselines?\n* Could authors  consider heat equation with $a(x)$ with high contrast values? \n* Could authors give results  of the method's performance on non-smooth PDE fields?\n\nLIMITATIONS:\nThe authors discussed the method limitations in the conclusion.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-03"}

comments: SUMMARY:
This paper aims to solve the linear system with sparse weights via the proposed
precondition conjugate gradient (CG) method. The linear system focuses on the graph
field applications, and this precondition CG method utilizes a learnable graph neural
network’s (GNN’s) output as the preconditioner.

STRENGTHS AND WEAKNESSES:
#### Strengths

- Problem setup is clear.
- The author designed a learning-based presconditioner that can utilize GPU for the CG computing.

#### Weaknesses

- The approach of using GNN to generate apreconditioner may not be very general. For each type of applications, it seems we must train a GNN, a process that is time consuming. For example, for the four types of applications considered in Section 4, we must rain a GNN for each of them. Further, for the same type of problems, the GNN training is uder some fixed settings (e.g., 9,147 meshes in line 210). The effectiveness of the proposed methods may depend on these settings.

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

- It is useful to discuss the training time for GNN models.
- In line 86, SPAI stands for Sparse Approximate Inverse. 
  However, in line 237, the authors switch to use AINV.

QUESTIONS:
See "weaknesses"

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 2
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-03