Learning Sparse Approximate Inverse Preconditioners for Conjugate Gradient Solvers on GPUs

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 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 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. Source code available at https://github.com/Adversarr/LearningSparsePreconditioner4GPU

• The paper introduces a novel learning-based method that uses GNNs to construct sparse approximate inverse preconditioners for CG solvers, avoiding traditional triangular solves.
• It leverages a scale-invariant loss function and local message-passing to ensure GPU-friendly parallel operations with only two SpMV operations per iteration.
• Experimental results show significant speedups (68–113%) and improved condition numbers on large PDE-derived and synthetic datasets, highlighting scalability and robustness.

Learning Sparse Approximate Inverse Preconditioners for Conjugate Gradient Solvers on GPUs

Introduction and Motivation

The paper addresses the challenge of efficiently solving large sparse symmetric positive definite (SPD) linear systems, a core problem in scientific computing, by accelerating the conjugate gradient (CG) method through improved preconditioning. Traditional preconditioners such as Diagonal, Incomplete Cholesky (IC), and Sparse Approximate Inverse (SPAI) have well-understood theoretical properties but are limited in their ability to exploit data-driven optimization and, crucially, in their compatibility with GPU architectures. Factorization-based preconditioners, while effective on CPUs, suffer from sequential triangular solves that impede parallelization on GPUs. Recent learning-based approaches using Graph Neural Networks (GNNs) have attempted to address these limitations but inherit the same bottlenecks due to their reliance on incomplete factorizations.

This work proposes a learning-based method for constructing GPU-friendly SPAI preconditioners using GNNs, leveraging the natural alignment between the local computation of sparse matrix-vector products (SpMV) and the local message-passing mechanism of GNNs. The approach avoids triangular solves, requires only two SpMV operations per CG iteration, and introduces a scale-invariant loss function tailored to the convergence properties of CG. The method is evaluated on PDE-derived and synthetic datasets, demonstrating significant speedups and improved condition numbers compared to both traditional and previous learning-based preconditioners.

Methodology

SPAI Preconditioner Construction via GNNs

The core idea is to directly approximate $A^{-1}$ with a sparse matrix $M^{-1} = GG^\top + \varepsilon I$, where $G$ is constructed by a GNN operating on the graph representation of $A$. The GNN processes node and edge features derived from the matrix's nonzero entries and associated physical or geometric properties, outputting the entries of $G$ in a block-wise fashion.

Figure 1: Overview of the approach: The GNN ingests the matrix's nonzero entries and node features, performs message passing, and outputs the entries of the sparse matrix $G$, which is assembled and used in the preconditioned CG solver.

The locality of the SPAI preconditioner ensures that each output entry depends only on two-hop neighborhoods, making it fundamentally compatible with GNN architectures. This contrasts with IC preconditioners, where triangular solves induce global dependencies along the elimination tree, which are difficult for GNNs to model and inefficient for GPU parallelization.

GNN Architecture

The GNN follows an encoder-processor-decoder paradigm. Node and edge features are encoded via MLPs, followed by $L$ message-passing layers that update hidden representations using local aggregation. The decoder MLP outputs block entries for $G_{ij}$, which are assembled into the global sparse matrix $G$ while preserving the original sparsity pattern of $A$.

Scale-Invariant Aligned Identity (SAI) Loss

A key innovation is the SAI loss, which is both statistics-based and scale-invariant. It is defined as

$$\mathcal{L}_{\mathrm{SAI}}(A, M^{-1}, \mathbf{w}) = \left\| \left( \frac{1}{\|A\|} A M^{-1} - I \right) \mathbf{w} \right\|_2^2$$

where $\mathbf{w}$ is a random vector and $\|A\|$ is a robust, dimension-agnostic norm (mean of nonzero entries). This loss aligns with the scale-invariant convergence behavior of CG, avoids the need to compute solution vectors for every $A$, and focuses optimization on improving the spectral properties of the preconditioned system.

Experimental Results

Datasets and Setup

Experiments are conducted on three PDE-derived datasets (Heat, Poisson, Hyperelasticity) and a synthetic dataset, with matrix sizes ranging from $10^3$ to $10^5$. All models are trained and evaluated on GPUs, with a fixed GNN architecture and hyperparameters.

Figure 2: Examples of PDE-derived test cases, illustrating mesh density and resolution for the Heat/Poisson and Hyperelasticity problems.

Performance and Scalability

The proposed method consistently outperforms Diagonal, IC, and traditional SPAI preconditioners, as well as previous learning-based approaches, in terms of total solution time and iteration count. On GPUs, speedups of 68–113% are observed, with the method maintaining efficiency as matrix size increases.

Figure 3: Performance of CG with different preconditioners for the heat problem, showing total solve time, scalability with matrix size, and convergence behavior across thresholds.

The method achieves a favorable trade-off between setup cost, iteration count, and GPU-efficient operations, making it particularly competitive for large-scale problems.

Condition Number Analysis

The learned SPAI preconditioner yields lower condition numbers for the preconditioned system compared to all baselines, directly correlating with faster CG convergence.

Figure 4: Condition number distributions for different preconditioners, with the proposed method achieving lower median and variance, indicating improved spectral properties.

Generalization and Robustness

The approach generalizes well to out-of-distribution samples, including finer meshes and matrices with different physical parameters, maintaining significant speedups and low iteration counts. Ablation studies confirm the effectiveness of the SAI loss and the robustness of the method to hyperparameter choices.

Practical and Theoretical Implications

The proposed method demonstrates that learning-based SPAI preconditioners can be effectively constructed via GNNs, fully leveraging GPU parallelism and overcoming the limitations of factorization-based approaches. The scale-invariant loss function aligns optimization with the theoretical properties of CG, enabling robust training across diverse datasets and matrix scales.

Practically, this enables faster and more scalable solvers for large-scale scientific computing tasks, particularly in domains where matrix structure and sparsity are dictated by physical discretizations (e.g., FEM for PDEs). The method is compatible with existing GPU-accelerated linear algebra libraries and can be integrated into broader simulation pipelines.

Theoretically, the work suggests that local message-passing architectures are well-suited for constructing preconditioners that operate via local computations, and that scale-invariant objectives are critical for generalization in numerical linear algebra tasks.

Future Directions

Several avenues for future research are identified:

• Sparsity Pattern Optimization: Allowing dynamic sparsity patterns or two-hop connections in $G$ could further improve preconditioner effectiveness and reduce fill-in.
• Extension to Other Solvers: Adapting the framework to non-SPD systems (e.g., GMRES) or integrating with multigrid methods as smoothers.
• Distributed Training and Inference: Scaling to multi-GPU or distributed environments for even larger problems.
• Automated Hyperparameter Tuning: Further automating norm selection and regularization for diverse application domains.

Conclusion

This work presents a principled, learning-based approach for constructing sparse approximate inverse preconditioners via GNNs, tailored for GPU acceleration and scale-invariant optimization. The method achieves strong empirical results in terms of speed, convergence, and generalization, and provides a foundation for further advances in data-driven numerical linear algebra and scientific computing.