Learning Sparse Approximate Inverse Preconditioners for Conjugate Gradient Solvers on GPUs
Authors: Zherui Yang
Zhehao Li
Kangbo Lyu
Yixuan Li
Tao Du
Ligang Liu
Venue: NeurIPS 2025 poster ()
Keywords: Graph Neural Networks
Preconditioning
GPU Computing
Iterative Solvers
Scientific Computing
Primary area: machine_learning_for_sciences
TL;DR: Our work introduces a GNN-based method for learning SPAI preconditioners, enabling efficient and robust training across various datasets and outperforming traditional and learning-based approaches in accelerating Conjugate Gradient solvers on GPUs.

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.

Learning Sparse Approximate Inverse Preconditioners for Conjugate Gradient Solvers on GPUs
Zherui Yang' Zhehao Li
Yixuan Li! 'University of Science and Technology of China Tiinguua University 4Shanghai Qi Zhi Institute
Kangbo Lyu
Tao Du
Ligang Liu
Stanford University
Abstract
The conjugate gradient solver (CG) is a prevalent method for solving symmetric and positive definite linear systems Ax b, where effective preconditioners are crucial for fast convergence. Traditional preconditioners rely on prescribed algo- rithms to offer rigorous theoretical guarantees, while limiting their ability to exploit optimization from data. Existing learning-basee methods often utilize Graph Neu- ral 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 intro- duces long-range dependencies which are difficult for GNNs to model. To address these issues, propose a learning- based method to generate GPU-friendly precon- ditioners, particularly using GNNs to construct Sparse Approximate Inverse (SPAT) preconditioners, which avoids triangular solves and requires only 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 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 A, leading to improved performance of the learned preconditioner. Evalua- tions on three PDE-derived datasets and one synthetic dataset demonstrate that our method outperforms standard preconditioners (Diagonal, IC, and traditional SPAI) and previous learning-bast preconditioners on GPUs. We reduce solution time on GPUs by 40%-53% (68%-16 113% faster), along with better condition numbers and superior generalization performance.
Introduction
Solving symmetric positive definite(SPID) sparse linear systems Ax b is essential in scientific computing and has wide applications in science and engineering, while the development of efficient numerical solvers remains challenge. The conjugate gradient (CG) method widely used for its efficiency, yet the convergence rate depends critically preconditioning techniques that effectively transform the original linear system into an easier one to solve at low computational cost.
Traditional preconditioners have long history of development and have achieved considerable success, Although modern GPUs provide massive parallelism, existing preconditioners struggle to leverage it. becoming the computational bottleneck despite their convergence benefits. The diagonal preconditioner offers a simple and GPU-frienll implementation but provides limited convergence improvement. Incomplete Cholesky (IC) demonstrates strong CPU performance and effectiveness in ill-conditioned problems [1] However, its requires two triangular solves at each CG step, which are typically difficult to parallelize 2 3, hindering its performance on GPUs.
CCoresponding author: 1g] iu@ustc..dd...
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Recent learning-basee approaches 4 6 aim to improve traditional approaches using neural networks, particularly factorization-base based methods (e.g., IC) and graph neural networks (GNNs). By substituting the sequential factorization process with GNN evaluation, these methods reduce the computational time required for incomplete factorizations. The trainability of GNNs enables the generation of higher-quality factorizations, thereby accelerating the convergence of iterative solvers. However, these methods also come with certain limitations. Prior works inherit the limitation of factorization-base methods that involve triangular solves, restricting their GPU parallelization efficiency. The triangular solves also pose challenges for GNNs to capture long-range dependencies and global information through the elimination tree, while GNNs aggregate information from local neighborhoods and struggle to effectively model such interactions 77]. Additionally, many existing approaches require computing the solution vector x A b of the original linear system for every A in the dataset to evaluate their loss functions 4
9 Generating datasets containing large matrices becomes increasingly computationally expensive. The varying scales of A further pose difficulties for models in learning to improve the condition number of the preconditioned system.
To address these challenges, we present novel approach for learning GPU-friendly preconditioners using GNNs. Our approach directly approximates the inverse A -1 with a sparse matrix and formulates its construction Sparse Approximate Inverse (SPAI) preconditioner
graph-learning problem, which shares conceptual parallels with the traditional
[10 M. The sparse mratrixvvector product (SpMV) is the only routine required at each CG iteration, allowing our approach to leverage GPU accelcration throughout both the construction and CG solving phases. We argue that the local nature of SpMV naturally complements the local propagation mechanism of GNNs. Our approach avoids imposing restrictive triangular structures on the preconditioner, enabling us to preserve the sparsity pattern of A even use different one. This design results in improved condition numbers compared to previous SPAI and factorization-bann based preconditioners. Furthermore, to eliminate dependence on solution vectors and enable robust training across various datasets, we propose the Scale invariant Aligned Identity Loss (SAI loss). This loss requires only the input matrix A, and its inherent A- scale invariance aligns with the scale-invariant convergence behavior of CG solvers, yielding higher performance of learned preconditioners. Extensive experiments conducted on three -dE-verived datasets and synthetic dataset validate the effectiveness of our approach compared to previous works. The results show that our approach achieves up to 113% speedup compared with standard traditional preconditioning techniques and previous learning-basse preconditioners, Our approach also achieves a better condition number while maintaining good generalization performance.
In summary, we make the following contributions:
1. We propose learning-based approach to generating GPU-friendly preconditioners, leverag- ing the GPU parallelism across both the construction and CG solving stages. Our approach focuses on generating the SPAI preconditioner using GNNs, exploiting the natural alignment between the local computation of SpMV and the local propagation mechanism of GNNs.
2. We introduce the SAI loss, novel statistics- based and A-scale invariant loss function, which reduces datasets' generation costs and enhances the performance of learned preconditioners.
3. Extensive experiments on three PDE- datasets and synthetic dataset demonstrate that our method achieves lower condition numbers and reduced computation time compared to existing approaches, while also maintaining generalizability, scalability, and robustness.
2 Related Works
Traditional Preconditioners Preconditioners M-1 transform the original linear system Ax = into one casier solve M - for iterative solvers at low computational cost. The diagonal preconditioner scales the original system by its diagonal entries [I]]. offering a straightforward and GPU-friendly implementation. However, it provides limited improvement to the convergence rate. Incomplete Cholesky (IC) preconditioner [T] computes sparse and incomplete triangular factors (M LL z A). It offers better spectral approximations but suffers from limited parallelism due to its reliance on triangular solves [12]. Sparse Approximate Inverse (SPAI) preconditioners 13-15 construct a sparse matrix M-1 GGT A-1 under sparsity constraints. At each CG step, SPAI use matrix-vector products instead of triangular solves in IC, enabling parallelism suitable for GPU architectures. However, its construction depends on sequential algorithms and often exhibits suboptimal performance. Alternative techniques, such polynomial [M] algebraic multigrid (AMG)
2

---

[6], and domain decomposition [7 preconditioners, each balance robustness, memory usage, and parallel efficiency differently. While effective in their respective domains, they lack generality and often require careful parameter tuning to ensure performance [8] 19 Machine learning offers promising solution to address this limitation.
Learning-base based Preconditioners Recent research in learning-based preconditioning has pro- gressed across several directions. CNNs (0). GNNs 21 and neural operators 88. which directly approximate the action of matrix inversion (N(A, b) ~ A A b. are utilized as preconditioners in many iterative solvers. Although these approaches can significantly reduce the number of iterations, they involve a full forward pass through the network at each iteration, thereby increasing computa- tional overhead. The auxiliary components neural operators are employed as subspace in iterative solvers 22 24
The usage of neural operators restricts their applicability primarily to PDE-related problems, while neglecting more general problems. Li et al. [4, Häusner al. [5), Trifonov et al.
6 have investigated the integration of machine learning and factorization-based preconditioners, specifically by training GNNs to predict the factor IC preconditioners. However, these methods suffer from the triangular solve due to their sequential computation [T0] and the theoretical challenges for GNNs to model
dependencies along the elimination tree [ Although these methods have achieved performance gains compared with the traditional alternatives, few of them consider the performance on GPUs, while our approach focuses specifically on GPU performance.
Our Approach
3.1
Problem Setup
We consider solving the sparse and SPD linear system Ax = b, where R=xa (or Rubxnb for a blocked sparse matrix with as block sizeD arises from diverse applications, such spatial discretization of PDEs using the finite element method (FEM). The matrix is associated with a graph GA EV,E), defined follows:
1. Vertices V: Represent variables (e.g., mesh nodes in FEM) with features v E Rd that encode geometric or physical properties.
2. Edges & = {(i,j) Aij * 0): Connect interacting variables, with edge features eij representing the corresponding entries or blocks Aij-
To accelerate CG's convergence, preconditioner M-1 is introduced, which transforms the problem into an easier one to solve by modifying the search direction. The residual vector r is replaced with the preconditioned direction
-r at each CG step, shown in Section A good preconditioner M reduces the condition number of the linear system by effectively approximating A and lower condition number typically indicates faster CG convergence. We consider the following preconditioner construction task based on the graph and its associated features:
(GA, (va), (e,,)) Algorithm
M R2xx
where "Algorithm" denotes both traditional approaches (e.g., IC, SPAI, AMG) and neural networks. Traditional approaches construct M through prescribed and heuristic rules based on the graph structure and features, while the learning-based approaches infer the mapping from data.
3.2 Graph Neural Networks for SPAI construction
3.2.1
Sparse Approximate Inverse Preconditioners (SPAI)
To ensure high computational efficiency, particularly on GPUs, we directly employ a sparse approxi- mate inverse M as the preconditioner for (SPAD). Additionally, since the convergence of CG depends on symmetric and positive definite M we propose to factorize the preconditioner as
A 2 M GG
+II,
where the sparse matrix is the output the GNN and e a small positive constant to enforce SPD property, shown in Figure Numerous prior studies [5
14 13
have demonstrated that sparse
In the hyperelasticity problem, there are variables to solve per node, resulting in block size of
3

---

Ilerftive
Groph Message Possing
precond(r):
returg, & G.T r ++ epsilon
Marrix
def cgtA, precond);
precond(r)
XL
Matrix G
Input: Nodo/Edge Features
Edge Oueput: Gij
Preconditioner: M-1 GG
Node ((x))
Overview our approach: By inputting the matrix's nonzero entries Aij and node features a(æ;), the GNN processes these features through message passing, and outputs the entries of Gij- The sparse matrix is assembled and then applied the preconditioned CG solver.
Figure I:
M can effectively reduce the condition number, even when A- iteration, the application of preconditioner s = M is instantiated as two sparse matrix-vector products and vector addition, which is both efficient on GPUs.
1 is dense. As a result, each CG
Previous works [ take an opposite approach to ours. Their approaches output a sparse triangular matrix L, which aim to approximate with LL resulting Incomplete Cholesky (IC) precon- 1 The application of preconditioner LL s involves two triangular solves: forward substitution Ly = and backward substitution L L s = y. Although triangular solves can be efficient CPUs, the inherent sequential nature limits their performance GPUs.
ditioner M-1 = (LL )
Locality SPAI Preconditioner Consider the application of our preconditioner M-'r. Suppose G shares the same sparsity pattern as A, the j-th entry of the output vector is
Sj (M- = er, + Σ Gj Σ
GHIFE.
Ag1
AWFO
Since the non-zero entries A and the entries of vectors correspond to edges and nodes in the graph, Equation (3) reveals that the oulpul on node depends only its two-hop neighborhood, mediated through intermediate nodes I and source nodes k. This makes SPAI fundamentally more compatible with GNN architectures than factorization- based methods.
In contrast, in preconditioner, the forward substitution Ly = r and backward substitution L propagate information along the elimination tree induced by chosen node ordering. The forward substitution aggregates values from ancestor nodes, while backward substitution collects from descendant nodes, resulting in global dependencies. This hierarchical propagation in IC conflicts with the local aggregation mechanism of GNNs over fixed-hop neighborhoods. Furthermore, IC requires explicit directional dependencies (enforcing lower triangular L). which are inconsistent with the undirected/symmerric assumptions in existing GNN-based preconditioners.
3.2.2 GNN Architecture
Our GNN follows encoder-processor-ee procesoor-dcoooee architecture, adapted from Gilmer et al. 25). Given node features V; e phonoe and edge features eij E Rooe the GNN first applies two MLPs E,, E. to encode the node and edge features into their hidden representations x x(0), hh)) Rd.
Fan (v:),
h
F (eir))
(4)
Subsequently, I. message-passing layers applied. Each layer utilizes three MLPs, (t) ((0) (()) which serve as message functions and update functions for node and edge features:
m (1) = Σ
(x (t-1)
eij),
JEN(
(5)
(0) x
x; (1-1) f (m ()),, h h (t-1)
f(()(mm)))))
f(tt 10 ((-1))

---

Finally, decoder MLP D is applied to the edge features of the last layer h(*) to get the block entries:
G D(bl))) E R
(6)
The global sparse matrix G E R constructed by replacing each block Aij with Gty while preserving the original sparsity patter..
3.3 Loss Function
The loss function is crucial for learning effective preconditioners M. While traditional SPAI ap- proaches aim minimize [M A || ||AM I| these formulations suffer from two limitations: (1) Direct evaluation of AM- requires significantly larger memory and computation (2) These losses depend the absolute magnitude of A, while the convergence rate of CG solvers is invariant to the magnitude.
resources for large matrices when both A and M are sparse.
Stochastic Estimation To address the computational bottleneck, we first adopt the stochastic trace estimator 5 26 to approximate the matrix Frobenius norm without explicitly constructing A.M
AM -| tr((Am (AM I)) Ew ||AM w
w
(7)
where w is vector of independent and identically distributed (i.i.d)) random variables drawn from the standard normal distribution 1). This reduces the computation to matrix-vector products AM w, avoiding costly matrix-matril products. The estimator can also be applied to
1 R IM w A A-lw However, this either requires access to A- involves solving = A for each sample in the dataset, making it computationally expensive.
[M A -
w/l..
Scale Invariance The loss function in ( exhibits scale sensitivity because it depends the absolute magnitude of A. To align with the CG solver's scale invariance property, we normalize A by its norm || in the loss function:
I I|? LSAI(A, w) = (IA AM AM-M- I
12
(AA = |A|
AM
W
(8)
This formulation ensures LSAI( (A, M-1 w) CsAI(aA, M-1, w) for any a > O. We refer to this as Scale invariant Aligned Identity (SAT) loss, which explicitly decouples the preconditioner learn- ing from the matrix's absolute scale, while enforcing alignment between the scaled preconditioned matrix and the identity matrix. The condition number k of AM can be estimated as
Tmax (AM-1) mmax (I E) min (AM-1) miin (I E)
+ Tmax E
x(AM-1)
20 max
E
2
maxx E)
where E AM-1 /||A|| is the error matrix and a denotes matrix singular values (see Section B for its detailed proof). This derivation shows that the SAI loss encourages M to focus on improving AM- spectral properties, while being invariant to absolute magnitude of A. Although conventional matrix norms of A like the Frobenius norm are theoretically valid, their scale is often sensitive matrix dimensions and outlier entries. We therefore define more robust and dimension- agnostic norm || = mean /ol Aijl- This setting also ensures that the edge features of the GNN and the loss term remain on a reasonable scale, promoting more stable and efficient optimization across matrices of varying sizes.
Experimental results
We answer the following questions through our experiments: (1) How does our approach compare with traditional and learning-basse approaches? (2) Is our loss more effective than those in prior works? (3) How do our approach generalize to unseen examples? (4) Does our approach yield better condition number compared to previous methods? We describe the experiment setup in Sec.
and provide answers to these questions in Sec. 4.2 to 4.5
5

---

500 time slace
(a) Density plot examples TetWild Dataset for Heat/Poisson Equation test case.
(b) Starting and ending timestep of meshes with differ- ent resolution in the hyperelasticity case.
Figure 2: Examples in our PDE-derived test cases.
4.1 Experiment Setup
Our experiments include three PDE-derived datasets, well as a fully algebraic dataset composed of randomly generated matrices. As shown Figure we use FEM and tetrahedron mesh to discretize the PDEs, yielding the sparse system matrix A. The test cases considered in our study are follows:
1. Heat Equation: alx) u, Au O, where a(x) represents the spatially varying density (randomly sampled on the mesh, with values ranging from le 4 4 to 5e 4).
2. Poisson Equation: - = f with ulp = 9. where D denotes a randomly specified Dirichlet boundary.
3. Hyperelasticity: Mü+ fur (u) = fext. where Ms the mass matrix, and fint. fext represent internal and external forces, respectively. The internal force fint is nonlinear and derived from a stable Neo-Hookean material model [27].
4. Synthetic System: A - PP + el, where P is a random matrix and s = 10 ensures the numerical positive definiteness of A.
For Heat and Poisson problems, we use 9,147 meshes with node counts ranging from 400 32,000 from the TetWild dataset (28). All meshes are normalized to ensure the entire geometry fits within the 1-1 domain. The hyperelasticity simulations involve beam twist scene for 500 timesteps. This experiment shares the same geometry but varies in resolution and topology through remeshing with node counts ranging from 645 to 14,039 (matrix size 3x larger). The synthetic dataset contains 1,000 matrices with a sparsity of approximately 0.12%. with matrix sizes ranging from 10,000 to 20,000, All experiments employ a 4:1 train-test split, with all reported results evaluated on the test set. Detailed problem configurations, their matrix representations and corresponding inputs of GNNs are provided in Appendix C.2
GNN Implementation and Training For all experiments in this work, we fixed the number of message passing steps L to 4, the number of hidden layers in all the MLPs to I, the number of neurons in the hidden layer to 24, and models trained for 500 epochs using batch size of 4 on a single NVIDIA A100 GPU, optimized with AdamW W (29) and an exponentially decaying learning rate scheduler (decay rate = 0.99).
10-4 The GNN has about 24k trainable parameters in total. All
Evaluation Metries The constructed preconditioner M is subsequently applied in CG solve Ax b. Since each CG iteration requires one preconditioner application, the total solving time is governed by:
Troal = T Teonstruct + x (Tapply -Tg))
(10)
where k denotes iteration count, and T Teonstruct- Tapply represents the preconditioner's 's construction time at the beginning and application time per iteration, and Teg accounts for fixed operations per CG iteration (e.g., dot products, matrix-veed vectoo products). The primary objective of preconditioning is to minimize Tooal Tapply and Tronstruct rely on the computational complexity its associated routines (such as SpMV and triangular solve), while preconditioners with better spectral approximation to A -1 or smaller condition number often result in lower k. This establishes fundamental trade-off between solver efficiency and approximation quality in preconditioner design.

---

TrE Titee (710
130
RgHOXD4 Sclae lina
SAGAA
Molic
80U0 2000
1000
Total Tined Irrt
(a) Total time required to converge for rtol 10
(b) Scalability
(c) Convergency
Figure 3: Performance of CG with different preconditioners for the heat problem. Figure (a) compares the average total solve time Trotll and preconditioner's construction time Teonstruct of CG with different preconditioners and devices. Figure (b) illustrates the relationship between matrix size and T Total including its 95% confidence interval, demonstrating the superior scalability of our approach on GPUs. Figure (c) compares the total solve time required to achieve different rtol.
4.2 Preconditioner Performance
We consider three standard traditional preconditioners both GPU and CPU architectures, in-
cluding: (1) Diagonal (Diag), (2) Incomplete Cholesky (IC), and (3) Sparse Approximate Inverse (AINV All the preconditioned CG are implemented in C++ and CUDA with OpenBLAS (30). CuBLAS, cuSPARSE [31], and cusplibrary [32) for their high-performance linear algebra kernels and
preconditioner implementations. We list their implementation details in Appendix C.1
Comparison to Traditional Approaches Table summarizes the time cost Trotal and iteration count ie for the CG method equipped with different preconditioners to converge to a relative tolerance rtol |b Ax||/[||] < 10 The results demonstrate that approach outperforms all baseline approaches and achieves consistent superiority across all categories of solving tasks. Additionally, we also present comparison of our method with AMG in Appendix D.2
Table GPU Benchmark results across different datasets. Total time Teotal (ms) and total iterations k (in parentheses) until relative residual norm is less than rtol = 10 s are listed in the table. Lower values indicate better performance. The lowest value bold and the second lowest underlined. The relative promotion (Rel. Prom.) indicates the ratio of the time saved over the second-bes, one.
Test Case
Diag
IC
AINV
Ours
Rel. Prom. (1)
Heat Equation
77(520)
167(204)
78(330)
36(197)
113%
Poisson Equation
45(320)
101(128)
58(217)
26(128)
73%
Hyperelasticity
86(464)
202(117)
247(266)
51(175)
68%
Synthetic System
445(2775)
1399(1808)
5024(10896)
253((122))
75%
Performance Analysis Across Convergence Thresholds Table 2 summarizes the total solve time Trotll and iteration numbers of CG solvers using different preconditioners up to various convergence thresholds. Although the Diagonal preconditioner has the smallest construction overhead, it results in a limited improvement in convergence rate. Compared with previous preconditioners, our approach achieves the smallest Tootl to the given rtol. We also compare the scalability of our approach with that baseline methods in Figure 3 and our approach still efficient as the matrix size increases.
Runtime Analysis of Routines in CG solvers Table J breaks down the runtime of each routine in CG solvers. Our approach reduces Trntal by balancing three key factors: low setup cost (comparable to Diag) T Teonstructs reasonable iteration counts & (similar IC), and GPU-efficient operations T sppply- This trade-off makes approach particularly competitive in GPU settings. Although IC achieves
'SPAI is specific subclass of AINV preconditioners that focuses constructing a sparse approximate inverse, whereas AINV is the general parent concept of approximate inverse techniques.
7

---

comparable iteration counts k. its computational cost for applying the preconditioner (the triangular solves) is substantially higher than ours on GPUs, resulting in degraded performance.
Table 2:
Comparison between different preconditioners for the heat problem. Total time Tiotal (ms), total iterations (in parentheses), and preconditioner's construction time construuc. (Cons.) are listed in the table. The best value is in bold, and lower value indicates better performance.
CPU
GPU
Stage
Diag
IC
AINV
Ours
Diag
IC
AINV
Ours
Cons.
0.126
8.426
19.308
0.181
0.196
1.866
18.924
0.181
10 10-2
44(309)
57(126)
80(199)
52(124)
43(309)
92(115)
52(199)
22(124)
10
58(383)
72(157)
98(246)
68(154)
53(384)
114(143)
61(246)
27(154)
10 1
71(450)
83(180)
111(282)
79(176)
63(442)
132(164)
67(280)
32(176)
10
83(511)
96(205)
132(328)
90(197)
77(520)
167(204)
78(330)
36(197)
Table 3: Time breakdown of each routine in preconditioned CG solvers on GPUs and CPUs for the heat problem. All timings reported in milliseconds.
CPU
GPU
Method
k:
Tonnstrect
Tapply
Tog.
k
T ronstruct
Tapply
Teg
Diag
511
0.12
0.01
520
0.12
0.01
IC
205
8.54
0.26
205
1.88
0.80
0.16
0.14
AINV
328
19.0
0.17
330
19.0
0.03
Ours
197
0.18
0.29
197
0.18
0.04
Comparison to Learning-based Approaches We compare our method with existing learning- based approaches on the dataset provided by previous works in Table We adopt their open-source implementation with the default settings. Compared to previous learning-based methods, our approach demonstrates better performance particularly on GPUs. The matrix size provided Li al. is sufficiently small that Diagonal preconditioner on CPUs achieves the best performance, while our approach attains the second-best performance. Compared to Häusner et al. [5], our approach achieves lower total solve time (Trota) on GPUs.
Table Comparsion to previous works. Total time Trotal (ms) and total iterations (in parentheses) until relative residual norm is less than rtol 10-8 are listed in the table. Prev. corresponds to the performance of previous works. The lowest value is in bold and the second lowest is underlined.
Device
Test Case
Diag
IC
AINV
Prev.
Ours
Li et 4
12(208)
19(90)
33(108)
26(108)
17(102)
CPU
Häusner et al.
799(970)
1387(438)
3235(753)
1139(354)
1320(456)
Li et al.
29(208)
45(87)
30(108)
51(108)
26(102)
GPU
Häusner et al.
166(970)
440(385)
2456(753)
1040(354)
132(456)
4.3 Generalizability and Robustness
We present the performance of our approach on d-sstiiibiittiinnn test samples in Table For the heat problem, we evaluate at a fixed mesh density of 1e-3 (Heat-Densing and further assess generalization to finer meshes with more than 32k nodes (Heat-Large). For the hyperelasticity problem, we test on the same geometry with finer mesh (22,618 nodes). For the synthetic problem, we use matrices with size ranging from 48k to 96k Wynthetic-Laaree Comparing the relative promotions on each test approach generalizes well to unseen resolutions and physical parameters.

---

Table 5: Test on out-of-distribution data on GPUs. The total time (ms) and total iterations (in parentheses) are reported. The best value is bold, and a lower value indicates better performance. In-distribution (In) and out-of-distribution (Out) relative promotions are recorded.
Rel. Prom. (1)
Test Case
Diag
IC
AINV
Ours
Out
In
DensitD
62(468)
135(175)
67(307)
35(201)
80%
113%
Heat- Large
251(1033)
808(388)
407(740)
154(409)
62%
73%
Hyperelasticity
326(1005)
667(287)
1154(496)
236(359)
72%
68%
Synthetic-Large
485(1575)
1616(781)
19840(5591)
347(559)
39%
75%
4.4 Ablation Study on Loss Functions
7 and the scale variant loss L2 = ||AM- -n--ww ww w|l shown in Table 6 although Le performs similarly on the heat problem, the proposed SAT loss demonstrates better effectiveness on more difficult test cases.
We compare SAT loss with cosine similarity loss Lcs = w AM-1 w ||
Table 6: Comparison of different losses. The average iteration counts are recorded. Since T construct and T apply are shared, lower value indicates a better result and smaller total solving time.
Loss
Heat
Poisson
Hyperelasticity
Synthetic
L2
195.6
134.1
185.4
2109.8
Las
207.3
133.4
182.5
2185.7
LSAI
197.4
128.8
175.7
1122.0
4.5 Condition Number of Preconditioned Matrix
The condition numbers provide theoretical measure of the preconditioner's performance. Besides the standard condition number K, we also compute the Kaporin's condition number [33], which the ratio of the arithmetic mean to the geometric mean of the matrix eigenvalues X;
max, A min; X;
(EN AV/N (A1d2.... ... AN)
x(AM))
KKaporin (AM
(1I)
We evaluate the condition numbers a simplified heat problem single mesh with 3,764 nodes and varying diffusivity coefficients. Figure demonstrates that our method notably reduces the condition number when compared with earlier traditional approaches, revealing its efficacy.
10*
1.5
107
Condition Number
Condbton Numbbe
38a
1.2
10?
(a) K
(b) KKaparin
Figure 4: Condition number distributions. Median, IQR, and outliers are shown. smaller conditioner number indicates better performance of the preconditioner, and the lower bound is 1.

---

5 Conclusion
This work proposes learning-based approach for constructing SPAI preconditioners, aiming accelerating the convergence of CG solvers on GPUs. Our method leverages the natural alignment be- tween the local propagation mechanism of GNNs and the localized structure of SPAI preconditioners. Furthermore, we propose the SAI loss to reduce the computational cost of training data generation and improve the quality of the learned preconditioners. Experimental results demonstrate that our approach delivers consistent speedups over traditional and existing learning-baad -baend methods, with improved robustness, generalization, and compatibility with parallel computing architectures.
Limitations and Future Works Although our approach outperforms existing methods in acceler- ating CG solvers, several limitations remain. First, enforce the sparsity pattern of G to exactly match that of A, while many existing approaches employ dynamic dropping strategies to further limit fill-in or incorporate two-hop connection in G to improve the effectiveness of preconditioner. Second, while our framework could potentially be extended to other Krylov subspace methods such as GMRES integrated into multigrid frameworks (e.g.. as a smoother), this work focuses specifically on SPD systems and the CG solver. Third, our current work is limited by the memory of a single GPU, and promising direction for future work is to scale our approach to multi GPU systems using techniques from distributed GNNs.
Acknowledgments and Disclosure of Funding
This work is supported by the National Key R&D Program of China (2022YFB3303400) and the National Natural Science Foundation of China (62025207). Tao Du acknowledges the research funding support from Tsinghua University and Shanghai Qi Zhi Institute.
References
Chih-Jen Lin and Jorge J. Moré. Incomplete cholesky factorizations with limited memory. SIAM J. Sci. Comput., 21:24 45, 1999, URL https://app www.nticscholee.......... DotpoaoooooooooooooooooroooaaaDa http:/oopusI:::6224::
[2] Weifeng Liu, Ang Li, Jonathan D. Hogg, lain S. Duff, and Brian Vinter. A synchronization-ffee algorithm for parallel sparse triangular solves. In European Conference Parallel Processing, 2016. URL
httpe://api..semntticee cacholaar org/ Drrg/orpusID:: 16983504
[3] Ichitaro Yamazaki, Sivasankaran Rajamanickam, and Nathan D. Ellingwood. Performance portable supernode-based sparse triangular solver for manycore architectures. Proceedings of the 49th International Conference on Parallel Processing, 2020. URL https: http:/ppppsmmmaicccccolllllra org/ Drg/CorpusID |210070204]
[4] Yichen Li, Peter Yichen Chen, Tao Du, and Wojeiech Matusik. Learning preconditioners for conjugate gradient pdc solvers. In International Conference Machine Learning, pages 19425-19439. PMLR, 2023.
[5] Paul Häusner, Ozan Öktem, and Jens Sjölund. Neural incomplete factorization: lcarning preconditioners for the conjugate gradient method. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. URL httppa://oppereeiieee htt.:e//trror/iddoooo..2...0000
[6] Vladislav Trifonov, Alexander Rudikov, Oleg Iliev, Yuri M. Laevsky. Ivan Oseledets, and Ekaterina Muravleva. Learning from linear algebra: A graph neural network approach to preconditioner design for conjugate gradient solvers, 2025. URL ltttps/s:///az arg/abs/2405 15557
[7] Vladislav Trifonov, Ekaterina Muravleva, and Ivan Oseledets. Can message-passinn passing gan approximate triangular factorizations of sparse matrices? arXiv preprint 25v: 2025.
[8] Alexander Rudikov, Vladimir Fanaskov, Ekaterina Muravleva, Yuri M. Laevsky, and Ivan Oseledets. Neural operators meet conjugate gradients: The FCG-NO method for efficient PDE solving. In Ruslan Salakhurdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp, editors. Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 42766 44782. PMLR, 21-27 Jul 2024.
[9] Vladislav Trifonov, Alexander Rudikov, Oleg Iliev, Yuri M Laevsky, Ivan Oseledets, and Ekaterina Muravleva. Learning from linear algebra: A graph neural network approach to preconditioner design for conjugate gradient solvers. arXiv preprint 240.:::4455.5555 2024.
10

---

[10] Jennifer A. Scott and Miroslav Tuma. Algorithms for sparse linear systems. Nečus Center Seriex, 2023.
URL httpppp//app.smmaaiccaaooooooooa www.setiiaaaolllaa......acaa azg//orrpuslD:2584221
[11] Yousef Saad. Iterative methods for sparse linear systems. SIAM, 2003.
[12] Åke Björek et al. Numerical methods in muirix computations, volume 59. Springer, 2015.
[13] Robert Bridson and Wei-Pai Tang. Ordering, anisotropy, and factored sparse approximate inverses. SIAM J. Sci. Comput., 21:867-888 1999. htt.lt/pppp..pppp..aeaaaaaaaeera http://ccccc..cc.ooooooooooooapa 14594893
[14] Michele Benzi. A sparse approximate inverse preconditiione for the conjugate gradient method. SIAM J Sci Comput, 17, 07 1996. doi: 10. 10.11777/11044277977744444
[15] L. Yu. Kolotilina and A. Yu. Yeremin. Factorized sparse approximate inverse preconditionings i. theory. SIAM Journal on Matrix Analysis and Applications. 14(1):45 58, 1993. doi: 10.1137/0614004, URL https://doi. org/ ///g///../0 10. 1137/0614004
[16] Jinchao Xu and Ludmil Zikatanov. Algebraic multigrid methods. Acta Numerica, 26:591-721, 2017.
[17] Barry F Smith. Domain decomposition methods for partial differential equations. In Parallet Numerical Algorithms, pages 225 243. Springer, 1997.
[18] Ali Taghibakhshi, Scott Maclachlan, Luke Olson, and Matthew West. Optimization based algebraic multigrid coarsening using reinforcement learning. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang. and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 12129-12140. Curran Associates, Inc., 2021. URL httptst//pproceeednnggg Paper pdf
[19] Haifeng Zou, Xiaowen Xu, Chen-Song Zhang, and Mo Zeyao. Autoamg(0): An auto-tuned amg method based on deep learning for strong threshold. Communications Computational Physics, 36( 1):0002220
2024. ISSN 1991-7120. doi: https:/doi.orr// http://dbbiorr/10.4220///c..pppp 40..40//cipp.002223-0000000
[20] Kai Weixian Lan, Elias Gueidon, Ayano Kaneda, Julian Panetta, and Joseph Teran. A neural-prnt puturalpeecodiiionne poisson solver for mixed Dirichlet and neumann boundary conditions. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver. Jonathan Scarlett, and Felix Berkenkamp. editors, Proceedings of the 41st International Conference Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 15976-25994 PMLR, 21-27 Jul 2024.
[21] Chen. Graph neural preconditioners for iterative solutions of sparse linear systems. In International Conference on Learning Representations, 2024.
[22] Jian Luo, Jie Wang, Hong Wang, Huanshuo Dong, Zijie Geng, Hanzhu Chen, and Yufei Kuang. Neural krylov iteration for accelerating linear system solving. In Neural Information Processing Systems, 2024. httpss//:://app
http://pp..maatiicccccemmm www.nticschol...........ccc org/ ::tp///rppI:::::P
2763118000
[23] Alena Kopanicáková and George Em Karniadakis. Deeponet based preconditioning strategies for solving parametric linear systems of equations. ArXiv, abs/2401.0200 0224401,2006, 2024. URL https://api
semanticscholar. OrggCorrpuIDD 266756022
[24] Enrui Zhang, Adar Kahana, Eli Turkel, Rishikesh Ranade, Jay Pathak, and George Em Karniadakis. Blending neural operators relaxation methods in pdle numerical solvers. Nature Machine Intelligence, 2022. URL hhttps://aappp semant
org/Corpuslll 251903572
[25] Justin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, and George E. Dahl. Neural message passing for quantum chemistry. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 1263-1272. PMLR, 06-11 Aug 2017.
[26] Michael F Hutchinson. stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications Statisties-Simulationn and Computation, 18(3):1059 1076. 1989.
[27] Theodore Kim and David Eberle. Dynamic deformables: implementation and production practicalities (now with code!). ACM SIGGRAPH 2022 Courses, 2022. URL Corpus ID: 2512578111111
[28] Yixin Hu, Qingnan Zhou, Xifeng Gao, Alec Jacobson, Denis Zorin, and Daniele Panozzo. Tetrahedral meshing in the wild. ACM Trans, Graph. 37(4))60:1 60:14, July 2018. ISSN 0730-0301. doi: 10. 145/ 3197517.3201353. URL http:://tttp./// acm. www./o../111.... 1145 32013553

---

[29] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on org/ Corpua ID: 53592270
Learning Reprexentations, 2017. URL httpp:::/ppp...mmmiicccccaceeeee
[30] Qian Wang, Xianyi Zhang. Yunquan Zhang, and Qing Yi. Augem:
Automatically generate high performance dense linear algebra kemnlss on x86 cpus. 2013 SC International Conference for
High Performance Computing, Networking, Storage and Analysis (SC). pages 1-12. 2013. URL
https://api. www.nti
httg//CrpusD::::0m 566935
[31] Nathan Bell and Michael Garland. Implementing sparse matrix- multiplication throughput- oriented processors. Proceedings of the Conference High Performance Computing Networking, Storage and Analysis, pages 1-11, 11, 2009. URL htll:httpa://am api semanticscholaar wwt./oorgg.....ooooooo 14531936
[32] Steven Dalton, Nathan Bell, Luke Olson, and Michael Garland. Cusp: Generic parallel algorithms for sparse matrix and graph computations, 2014. URLhtttp:/cuuppiirras http:sspppppprrarutrgubbb.b 1o/ Version 0.5.0.
[33] Igor E Kaporin. New convergence results and preconditioning strategics for the conjugate gradient method. Numerical linear algebra with applications, 1(2)17992210 1994.
[34] Richard Barrett, Michael Berry Tony Chan, James Demmel. June Donato, Jack Dongarra, Victor Eijkhout, Roldan Pozo, Charles Romine, and Henk van der Vorst. Templates for the Solution of Linear Systems: Building Blocks for Iterative Methods. Society for Industrial and Applied Mathematics, 1994, doi: 10. 10.1137/1.7781119997
URL https:/:///epub siam. www.aoorg////ooo//eeooeee ott/////abb///....o///o/////////
[35] Wenzel Jakob. nanobind:
tiny
and efficient
c++/python
bindings, 2022.
[36] Gael Guennebaud, Benoît Jacob, et al. Eigen v3. http://eigen.uxxfamilyyoo org. 2010.
[37]
Nathan Bell, Luke N. Olson, Jacob Schroder, and Ben Southworth. PyAMG: Algebraic multigrid solvers in python. Journal of Open Source Software, 8(87):5495, 2023. doi: 00.21105/joss.054955 URL
https://doi. org /1/g//0...1 200...1101////.. 0549)
[38] Maxim Naumov, M. Arsaev, Patrice Castonguay, Jonathan M. Cohen, Julien Demouth, Joc Eaton, Simon K. Layton, N. Markovskiy, István Reguly, Nikolai Sakharnykh, V. Scllappan, and Robert Strzodka. Amgx: library for gpu accelerated algebraic multigrid and preconditioned iterative methods. SIAM Sci. Comput. 37, 2015. URL hattps://aat httpp::/ppi..emmatt semantiiniiic cacholar. warpllarrrrrgeeee org/Cor Corppsrrpss 32694273
[39] Hang Si. Tetgen, delaunay-basee based quality tetrahedral mesh generator. ACM Transactions Mathematical Software (TOMS), 41:1 36, 2015. httppp:://ppp..smmaaiiccc.
org/ CorpusID: 10022108
A Preconditioned Conjugate Gradient Algorithm
We describe the standard PCG algorithm 34
here.
B Proof of the Inequality
The inequality concerns the condition number n(AM-1). The error matrix is defined in the main
text as E = AM-
AM-1/|| || A/||
I. From this definition, we can write:
AM-1
I+E=
(12)
|
The condition number 15 k(AM-1) is given by:
Ki K(AM-1) Tmax (AM-)
(13)
mmin (M-1))
Since for any non-zero scalar c, x((X)) = kx()), we have:
x(AM-)) = AA--1 I: [
O max (I + E)
(14)
|||||
Gmin (I E)
We use Weyl's inequality for singular values. For any square matrices X and Y of the same dimensions, Weyl's inequality states:
(x (X) S F YY),
(15)
12

---

Algorithm Preconditioned Conjugate Gradient
i 0
r ++ b Ax d d M M-1r
// Apply Preconditioner
δnew r'd du t- dnew
while < imax and Snew > Eu do
q Ad
d
X ad
if is divisible by 50 then
b Ax
else
s - M r
// Apply Preconditions
8old t Enew
&new t- r's s
8 Add
d s 3d
end if
i +
end while
where (+(.) is the k-th singular value and Fmax YY is the largest singular value of Y (i.e., its spectral norm (12))
Let X = I (the identity matrix) and Y = E. The singular values of I are all 1, so o (I) = for all k.
Applying Weyl's inequality to + E:
lo [I + E) dx( I 0 max EE)
(16)
| [ E) 1| $ max E))
(17)
This implies that for any singular value 0% [I E):
max EE) S cI + E) 1 mmax (E).
(18)
This inequality holds for both the maximum and minimum singular values of I E:
max [I + E) < max E)
(19)
mmin (T + E) 2 Tmax EE)
(20)
For the lower bound on miin (I+ E) be positive, and thus for I + E be invertible, we require
max (E) < 1. This condition is typically met when the preconditioner M is effective, making
AM /||A|| close to
Under the condition Tiax E)) < 1, we can bound the condition number:
xAM K
max (I + E)
max
E
(21)
Gmin (I + E)
max E
This establishes the first part of the inequality.
For the approximation, if max (E) is small (i.c., mmax E
1), use the Taylor expansion
for (1 x)
+ + + ... for |x| < 1. Let x = Tmax
E). Then:
1 x 1-a = z) x) (1 + x)(1 + x x² + 0(xx)) = 1 2x + 2x + +(x)).
(22)
For << 1, we can approximate this as:
+
maax
20 max (E).
(23)
max
Finally, recall that the spectral norm El is defined as mmax (E). Therefore,
ml. A(A <
0max
1 + max
E + E|/2.
(24)
0max
This completes the proof of the inequality chain presented in the main text.
13

---

C Experiment details
C.1
Implementation of Traditional Baselines
Traditional Baselines We implement all the preconditioned conjugate gradient solvers in C++ code with reduced function call overhead, and use nanobind [35] to generate its Python binding. Even one method can have different implementations on different devices, leading different iteration counts.
Table 7: Implementation of baseline preconditioners
Device
Diag
IC
SPAI
AMG
CPU
Eigen [36
Eigen [36] cuSPARSE
cusplibrary
32
PyAMG [37
GPU
Custom
cusplibrary
32
AMGX [3
All evaluations are performed on an AMD Ryzen 5600 CPU and an NVIDIA GeForce RTX 3060 GPU. The CPU frequency is fixed 4.0 GHz and the OpenMP thread count is set 4. The source code is compiled using GCC 14.2 and CUDA 12.8.
C.2 Dataset Configurations
Heat Equation/Poisson Equation: The equations are discretized using the finite element method. We generate 9, 147 samples on tetrahedral meshes from the TetWild dataset (28], employing spatial discretization. The input node features of the GNNs are the nodes' position and the density.
Hyperelasticity We adopt standard optimization based time integration with stable Reo-Hookean material model 277 v 0.4, and density 1.0. The geometry the beam defined as H x W L x 1 4, with Dirichlet boundary conditions applied to the left and right boundaries. All physical parameters are specified in units. The domain is discretized at varying resolutions using TetGen 39). The input node features are the nodes' position and one-hot vector to indicate whether the node is on the Dirichlet boundary.
Synthetic For each sample, we first construct P with specified sparsity of 3 10-4, enforce the SPD property of A, we compute A PP =I, where -ee-4 is chosen for numerical stability. The resulting matrix A has sparsity of 1.2e-3 approximately. The matrix size of P is randomly selected between 12,000 and 24,000. The input node feature is the i-th row average of matrix A.
D Additional Experiments
D.1 Performance CPUs
Table 8: Benchmark results across different datasets. Total time Tiotl (ms) and total iterations & (in parentheses) until relative residual norm is less than rtol = 10-8 are listed in the table. Lower values indicate better performance. The lowest value is in bold, and the second lowest is underlined.
Device
Test Case
Diag
IC
AINV
Ours
Heat Equation
77(520)
167(204)
78(330)
36(197)
Poisson Equation
45(320)
101(128)
58(217)
26(128)
GPU
Hyperelasticity
86(464)
202(117)
247(266)
51(175)
Synthetic System
445(2775)
1399(1808)
5024(10896)
253(1122)
Heat Equation
83(511)
96(205)
132(328)
90(197)
Poisson Equation
53(331)
68(135)
94(208)
59(128)
CPU
Hyperelasticity
287(464)
299(160)
453(266)
361(175)
Synthetic System
950(2776)
1064(1221)
14377(10906)
1251(1122)

---

D.2 Comparison to AMG
Algebraic Multigrid (AMG) powerful tool for solving large-scale linear systems on both CPU and GPU architectures by exploiting the hierarchical structure of graph nodes. AMG can further serve as a preconditioner for CG solvers. While our method adopts a significantly different approach compared to AMG, we also include a comparison between AMG and our method on our datasets.
Our evaluation employs the default configurations listed in AMGX and PyAMG library for the evaluation across all datasets, while more precise settings could be applied to further improve its performance. As listed Table while effective for Heat/Poisson Equation problems, AMG can degrade the performance of CG solvers compared to other baseline preconditioners in more complex scenarios, due challenges in selecting optimal parameters for smoothers and cycles.
Benchmark Results:
Table 9; Benchmark result across different datasets. Total time Troal (ms) and total iterations & (in parentheses) until relative residual norm less than rtol 10- are listed in the table. We examine two configurations: (I) AMG-preconditiones CG (AMG+CG) and (2) standalone AMG. "/" indicates that the default settings does not to converge within 10 seconds.
CPU
GPU
Case
AMG+CG
AMG
Ours
AMG+CG
AMG
Ours
Heat
47(17)
117(69)
90(197)
20(27)
17(18)
36(197)
Heat-Large
491(22)
1312(103)
1086(409)
41(10)
94(17)
154(409)
Poisson
37(14)
67(42)
59(128)
10(9)
38(30)
26(128)
Hyperelasticity
785(91)
361(175)
286(128)
51(175)
Synthetic
6970(551)
1251(1122)
253(1 122)
D.3
Ablation Study on Matrix Norms
To validate the design choice of the custom matrix norm within our SAI loss, we conduct tar- geted ablation study. discussed in Section 3.3 our proposed mean norm was selected for its dimension dimension-agnosticisnn robustness to outliers, and low computational cost, offering key advantages over conventional norms such as the Frobenius and L1 norms. In this experiment, we evaluate the impact of this choice the Heat dataset. We replace our norm with several standard alternatives while keeping all other model components and hyperparameters identical. The results are summa- rized in Table IO They clearly demonstrate the superiority of our proposed norm. This represents a significant improvement over the Frobenius norm and the LI norm. This quantitative evidence confirms that the specific properties of our chosen norm not merely theoretical advantages but translate directly into tangible performance gains in the optimization process.
Table 10: Ablation study on the choice matrix norm within the SAI loss. We report the average PCG iterations required for convergence on the "Heat" dataset.
Norm Type
Iterations
Frobenius Norm (ll || LI Norm (Il [I))
222
231
Ours
197
D.4 Sensitivity Analysis of the Hyper-parameter
This section addresses the sensitivity of our method's performance the value of the hyper-parameter e in (2). To evaluate this, we conducted an ablation study on the Heat dataset by using different values of e. Iummmarzzes the validation performance (measured in iterations) for the tested values. The results demonstrate that our method is robust, with stable performance across three orders of magnitude (3 x 10-4 4 to 3 10-2), indicating that e requires no careful per-instance tuning. However, an excessively large value of (e.g., x 10-1) leads to training failure. We hypothesize that such strong regularization term over-constrains the preconditioner, causing deviate excessively from the intended structure and thus degrading performance.
15

---

Table 11: Sensitivity analysis of the hyper-parameter e on the Heat dataset. Performance is measured by the number of iterations required on the validation set.
&
10-6 = 10-
=-3 =3. 10
= x 10
= 10-1
222
208
197
205
Training Failure
D.5
Generalization Analysis Against Learning-Basee Baselines
This section provides a detailed analysis to address the generalization capabilities of our method compared to other baaming-based approaches.
Comparison with Neural PCG For this experiment, we follow the OOD setting from Table in the main paper and compare against Neural PCG [ Specifically, problem to an unseen value. TTall presents the total time (ms) and iteration counts (in parentheses) for both in-distribution and OOD settings. Our method not only outperforms Neural PCG in the in-distribution setting but also demonstrates significantly better generalization. While all methods experience performance degradation in the OOD setting, our method remains the most efficient.
decrease the density in the heat
Table 12: OOD generalization comparison with Neural PCG 4
Total time (ms) and total iterations (in parentheses) are reported. The best and second-best results are bolded and underlined, respectively.
Setting
IC
Diagonal
AINV
Neural PCG
Ours
In-Distribution 45(87)
29(208)
30(108)
51(108)
26(102)
cuitof-Distributionn
337(876)
262(1956)
179(1039)
268(743)
167(1022)
Comparison with Neural IF We further test generalization against Neural IF [5] by increasing the sparsity the synthesized matrix from 10 to 10-3 The results are summarized in Table 13 Again, our method exhibits superior performance and robustness. In the more challenging OOD setting, our approach achieves remarkable speedup over Neural IF, confirming strong generalization capability.
Table 13: OOD generalization comparison with Neural IF 5]. The best and second-best results are bolded and underlined, respectively.
Setting
IC
Diagonal
AINV
Neural IF
Ours
In In-Distribution Out-of-Distrbuttion
440(385)
166(970)
2456(753)
1040(354) 1222(375)
132(456) 168(470)
616(414)
207(1040)
3220(804)
16

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The main algorithm and theoretical claims are outlined and analyzed in Section The effectiveness of our method is validated and discussed in Section
g
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The paper discusses the limitations of the proposed method in Section 5 including the fixed sparsity pattern constraint and the focus on SPD systems and CG solvers.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sellsspeciiiationn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
17

---

Answer: [Yes]
Justification: All theorems include clearly stated assumptions and complete proofs in Section Supporting lemmas and prior results are properly referenced.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
provides detailed descriptions of the experimental setup, including model architecture, training procedure, datasets, and evaluation metrics.
Justification: Section 4.1
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
18

---

reproducible, they Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: Yes, the code will be made publicly available upon publication, along with detailed instructions to reproduce all main experimental results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://niis.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/
Cooee//dddeeebmimiia ssi onPolicy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Section 4.1
provides detailed descriptions of the experimental setup, including model architecture, training procedure, datasets, and evaluation metrics.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification: The paper reports error bars for all key experimental results.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
19

---

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g., negative еттоr rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of exccution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Section and Appendix C.1 provides detailed descriptions of our experiment enviromments.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer:
[Yes]
Justification:
We confirm that this research complies with the NeurlPS Code Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
The work is primarily focused improving numerical methods for solving linear systems and does not involve direct human interaction, personal data, or deployment societal contexts. While the method may inspire future research and contribute to computational efficiency in scientific computing, no specific societal impacts are foresecable this stage.
Justification:
Guidelines:
The answer NA means that there is no societal impact of the work performed.
20

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., dentoymenn of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper does not involve the release of models datasets with high risk for misuse. The proposed method is numerical solver technique and does not generate content, process personal data, operate in safety-critical domains without human oversight. Therefore, no specific safeguards are required.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Yes, as the paper uses existing software libraries including cuSPARSE, Eigen, cusplibrary, AMGX, and PyAMG. All libraries are used in compliance with their respective licenses, and proper citations to the corresponding projects and documentation are provided in the manuscript.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
21

---

The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g.. CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperssi thcode com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-pane se-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The code is well-documented, including instructions for setup, training, and evaluation, along with descriptions of dependencies, limitations, and licensing information. Documentation and code will be made open-source upon publication.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing human subjects. It focuses algorithmic development and numerical experiments for solving sparse linear systems, with data collected from human participants.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for rescarch with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
22

---

Answer: [NA]
Justification: The paper does not involve research with human subjects. The work is focused algorithmic development and numerical experiments for sparse linear systems.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: The core methodology of this research does not involve the of large language models (LLMs).
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.LM policy (https: for what should or should not be described.
23


Reviews:
[{"review_id": "900xkwzfBV", "reviewer": "Reviewer_Rbu2", "summary": "This paper aims to solve the linear system with sparse weights via the proposed\nprecondition conjugate gradient (CG) method. The linear system focuses on the graph\nfield applications, and this precondition CG method utilizes a learnable graph neural\nnetwork’s (GNN’s) output as the preconditioner.", "questions": "See \"weaknesses\"", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "#### Strengths\n\n- Problem setup is clear.\n- The author designed a learning-based presconditioner that can utilize GPU for the CG computing.\n\n#### Weaknesses\n\n- 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.\n\n- 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?\n\n- It is useful to discuss the training time for GNN models.\n- In line 86, SPAI stands for Sparse Approximate Inverse. \n  However, in line 237, the authors switch to use AINV.", "was_revised": true, "final_justification": "An interesting direction though the eventual practical use may still have some way to go.", "rebuttal": "We thank Reviewer Rbu2 for the constructive feedback and positive evaluation of our work.\n\n> **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.\n>\n> It is useful to discuss the training time for GNN models.\n\nOur 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.\n\nCompared 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.\n\n\nRegarding the training time, we provide the following details for our one-time offline training process:\n\n- **Heat & Poisson:** ~24 hours on a single A100 GPU.\n- **Elasticity:** ~10 hours on a single RTX 3060 GPU.\n- **Synthetic:** ~12.2 hours on a single A100 GPU.\n\n> **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?\n\nWe 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.\n\nHowever, 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.\n\nTherefore, 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.\n\n> **Concern 3:** In line 86, SPAI stands for Sparse Approximate Inverse. However, in line 237, the authors switch to use AINV.\n\nThank 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.\n\n[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."}, {"review_id": "I9oydggq5u", "reviewer": "Reviewer_JerM", "summary": "This 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.", "questions": "* 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?", "limitations": "The authors discussed the method limitations in the conclusion.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "Thank you for your reply. My concern have been resolved", "rebuttal": "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.\n\n> **Concern 1:** This loss function was previously introduced in article [1], with a detailed investigation and theoretical justification provided in [3].\n\nWe 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.\n\n> **Concern 2:** The SAI loss is a key contribution. It uses a non-standard definition for the matrix norm.\n>\n> - Could authors explain the motivation for choosing this specific heuristic instead of more conventional matrix norms?\n> - How does the performance of this specific matrix form compare to that of conventional matrix norms?\n\nWe chose this specific norm ($\\Vert A \\Vert = \\mathrm{mean}_{A_{ij} \\not= 0} \\vert A_{ij}\\vert$) for three main reasons:\n\n1. **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.\n2. **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.\n3. **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.\n\nTo 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.\n\n| Loss Norm       | Avg. Iterations |\n| :-------------- | :-------------- |\n| Frobenius Norm  | 222             |\n| L1 Norm         | 231             |\n| **Mean (Ours)** | **197**         |\n\n> **Concern 3:** Could authors give the results of GPU-to-GPU comparison against other learning-based baselines?\n\nWe 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.\n\n- **vs. Neural PCG [1]:**\n  - Neural PCG: 51 ms (108 iterations)\n  - Ours: **26 ms** (102 iterations) (1.96x speedup)\n- **vs. Neural IF [2]:**\n  - Neural IF:  1040 ms (354 iterations)\n  - Ours: **132 ms** (456 iterations) (7.88x speedup)\n\n> **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?\n\nThank 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.\n\n| Preconditioner      | Time (ms) | Iterations |\n| ------------------- | --------- | ---------- |\n| Diagonal            | 60        | 460        |\n| Incomplete Cholesky | 121       | **159**    |\n| AINV                | 65        | 292        |\n| **Ours**            | **40**    | 184        |\n\nNotably, 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.\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."}, {"review_id": "9OJtsgWIte", "reviewer": "Reviewer_GmV6", "summary": "This 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.", "questions": "1. 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?", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## 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).", "was_revised": true, "final_justification": "Thank you for your reply. I will keep my score.", "rebuttal": "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.\n\n> **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?\n>\n> If the proposed method does generalize better than competing approaches, what is the underlying reason? Could the locality of the computations be a factor?\n\nWe 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.\n\n**vs. Neural PCG [1]:** we decrease the density of mesh in the heat problem to an unseen value.\n\n|                     | IC       | Diagonal  | AINV      | Neural PCG[1] | Ours          |\n| ------------------- | -------- | --------- | --------- | ------------- | ------------- |\n| In distribution     | 45(87)   | 29(208)   | 30(108)   | 51(108)       | **26**(102)   |\n| Out of distribution | 337(876) | 262(1956) | 179(1039) | 268(743)      | **167**(1022) |\n\n**vs. Neural IF [2]:** we increase the sparsity of the synthesized matrix from 1e-3 to 2e-3.\n\n|                     | IC        | Diagonal  | AINV      | Neural IF[2] | Ours         |\n| ------------------- | --------- | --------- | --------- | ------------ | ------------ |\n| In distribution     | 440(385)  | 166(970)  | 2456(753) | 1040(354)    | **132**(456) |\n| Out-of-distribution | 616 (414) | 207(1040) | 3220(804) | 1222(375)    | **168**(470) |\n\n- In distribution: our method is 96% faster compared with Neural PCG[1] and 687% faster compared with Neural IF[2];\n- Out of distribution: our method is 60% faster compared with Neural PCG[1] and 627% faster compared with Neural IF[2].\n\nWe 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.\n\n> **Concern 3:** The parameter $\\varepsilon$ in $GG^\\top + \\varepsilon I$ is a hyper-parameter. How sensitive is the method’s performance to its value?\n\nThank 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.\n\n| $\\varepsilon$ | 3e-5 | 3e-4 | 3e-3    | 3e-2 | 3e-1 |\n| :-------------- | :--- | :--- | :------ | :--- | :--- |\n| Val. Iterations | 222  | 208  | **197** | 205  | Training Failure |\n\nPerformance 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."}, {"review_id": "0PdJOyMTYr", "reviewer": "Reviewer_umir", "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.", "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?\n\nWhat 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", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "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.\n\nOne 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.", "was_revised": true, "final_justification": "Overall this work advances the field in a meaningful way.  There have been many attempts at advancing solvers and this paper make a clear contribution with the SPAI-based loss.  The results are solid and the paper is clear.  The authors did well with the rebuttals all around and mostly answered the questions.  I rated accept for the clarity and for the clear/meaningful contributions.", "rebuttal": "We thank Reviewer umir for the encouraging review and for the questions that allow us to elaborate on the core strengths of our method.\n\n> **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?\n\nOur method's suitability for GPUs stems from two synergistic aspects:\n\n1. **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.\n2. **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.\n\n> **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?\n\nThis is an excellent question about the method's scalability.\n\n- **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.\n- **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.\n\n[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).\n\n[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\n\n[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)"}]

Author rebuttal:

Meta-review: This paper studies the problem of finding effective preconditioners for solving symmetric positive definite linear systems. The main contribution of the work is a learning-based method that relies on graph neural networks to construct sparse approximate inverse preconditioners, and which is naturally suitable for computation on GPUs. All reviewers agreed that the paper is well written and clearly introduces the problem setting and how their approach fits with existing works, and Reviewers JerM and GmV6 both noted the comprehensive nature of the empirical results comparing to established baselines and reporting both iterations and wall-clock time. Overall, the reviewers were positive about the paper, and I agree with the consensus that it would be a worthy addition to the conference.
Decision: Accept (poster)