Learning Sparse Approximate Inverse Preconditioners for Conjugate Gradient Solvers on GPUs

2025-10-31

AI Summary

Traditional incomplete factorization preconditioners for symmetric positive definite (SPD) linear systems exhibit poor parallelism and long-range dependencies on GPUs, hindering efficient conjugate gradient (CG) acceleration. Method: We propose the first GPU-oriented, graph neural network (GNN)-driven sparse approximate inverse (SPAI) preconditioner learning framework. It eliminates triangular solves—requiring only two sparse matrix-vector (SpMV) operations—and models local matrix structural propagation via GNNs. A statistics-driven, scale-invariant loss function ensures strong correlation between preconditioner quality and system condition number. The model is end-to-end trainable and generalizes across problems. Results: Evaluated on four benchmark SPD datasets, our method reduces GPU CG solve time by 40–53% (achieving 68–113% speedup) over state-of-the-art traditional and learned preconditioners. It significantly lowers condition numbers and markedly improves convergence rate and robustness.

Technology Category

• Machine Learning: Graph-based Machine Learning
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Web Mining and Content Analysis: Large pretrained models with web data

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

Problem

• Developing GPU-friendly sparse approximate inverse preconditioners for conjugate gradient solvers
• Addressing triangular solve limitations in learning-based preconditioners for GPU parallelization
• Improving convergence through scale-invariant loss functions matching CG properties

Innovation

• Uses GNNs to construct Sparse Approximate Inverse preconditioners
• Employs statistics-based scale-invariant loss function
• Avoids triangular solves for GPU parallelization efficiency