Learning Sparse Approximate Inverse Preconditioners for Conjugate Gradient Solvers on GPUs

AI summary

Overview

• Research area: Machine learning for numerical linear algebra — specifically learned preconditioners for iterative linear solvers, combining graph neural networks with GPU-oriented sparse linear algebra.
• Technical level: Intermediate. The paper assumes familiarity with the conjugate gradient method, preconditioning, and basic graph neural network message passing.
• 1-sentence scope: The paper trains a graph neural network to output a sparse approximate inverse preconditioner so that preconditioned conjugate gradient solvers run faster on GPUs than with Diagonal, Incomplete Cholesky, traditional SPAI/AINV, or prior learned preconditioners.

What This Paper Is About

Conjugate gradient (CG) solvers are the standard tool for solving sparse symmetric positive definite linear systems Ax = b, but they need good preconditioners to converge quickly. Traditional preconditioners are either fast to apply but weak (Diagonal) or strong but poorly parallelizable on GPUs (Incomplete Cholesky, whose triangular solves are sequential). Existing learning-based approaches inherit the Incomplete Cholesky structure, so they inherit the triangular-solve bottleneck and force GNNs to model long-range dependencies along an elimination tree. This paper instead learns a Sparse Approximate Inverse (SPAI) preconditioner directly, which needs only matrix-vector products at each CG step and therefore matches both GPU parallelism and the local propagation of GNNs.

Key Contributions

1. A learning-based SPAI preconditioner. The method uses a GNN to generate a sparse matrix G such that the preconditioner is M⁻¹ = GGᵀ + εI, avoiding triangular solves entirely; each CG step requires only two sparse matrix-vector products and a vector addition. The authors argue this locality aligns with GNN message passing, whereas IC's elimination-tree dependencies are long-range.

2. The Scale invariant Aligned Identity (SAI) loss. A statistics-based loss that depends only on the input matrix A and is invariant to the absolute scale of A (for any α > 0, the loss value is unchanged). This removes the need to precompute solution vectors x = A⁻¹b for every matrix in the dataset, reducing dataset generation cost, and it matches CG's property that convergence depends on the condition number rather than the scale of A.

3. A GPU-focused empirical study. Evaluation on three PDE-derived datasets (Heat, Poisson, Hyperelasticity) and one synthetic dataset shows lower total solve time than Diagonal, IC, AINV, and prior learned preconditioners, together with lower condition numbers and generalization to out-of-distribution problems.

4. Analysis of why the design works. The paper derives the two-hop dependency structure of the SPAI application and gives a bound relating the SAI loss to the condition number of the preconditioned matrix (Appendix B).

Main Findings

• GPU speedups over traditional preconditioners. On the main GPU benchmark (total time in ms with iteration count in parentheses, converged to rtol < 10⁻⁸): Heat — Diag 77 (520), IC 167 (204), AINV 78 (330), Ours 36 (197); Poisson — 45 (320), 101 (128), 58 (217), Ours 26 (128); Hyperelasticity — 86 (464), 202 (117), 247 (266), Ours 51 (175); Synthetic — 445 (2775), 1399 (1808), 5024 (10896), Ours 253 (1122). Relative promotions were 113%, 73%, 68%, and 75% respectively.
• Headline speedup range. The abstract reports reducing solution time on GPUs by 40%–53% (68%–113% faster) versus standard and previous learning-based preconditioners.
• Favorable cost breakdown. The paper's Table 3 decomposes runtime into iteration count k, construction time, preconditioner application time, and other CG time. The method combines a low construction cost (0.18 ms on GPU, comparable to Diagonal's 0.12 ms), iteration counts similar to IC (197 versus 205 on CPU; 197 versus 204 on GPU for the heat problem), and cheap GPU application, whereas IC's application cost is the triangular solve.
• Better than prior learned preconditioners on GPUs. Against the dataset of Li et al. [4]: on GPU, Prev. scored 51 (108) total ms versus Ours 26 (102), with Diag 29 (208), IC 45 (87), AINV 30 (108). Against Häusner et al. [5]: on GPU, Prev. 1040 (354) versus Ours 132 (456), with Diag 166 (970), IC 440 (385), AINV 2456 (753). On CPU, the Li et al. [4] matrices are small enough that Diagonal achieves the best time (12 ms), with the proposed method second-best at 17 ms.
• Generalization to out-of-distribution problems. On harder unseen cases the method still led: Heat-Density — 35 (201) versus Diag 62, IC 135, AINV 67, with out-of-distribution relative promotion 80% and in-distribution 113%; Heat-Large (>32k nodes) — 154 (409) versus 251, 808, 407, promotion 62% (out) and 73% (in); Hyperelasticity with a finer mesh (22,618 nodes) — 236 (359) versus 326, 667, 1154, promotion 72% (out) and 68% (in); Synthetic-Large (48k–96k matrices) — 347 (559) versus 485, 1616, 19840, promotion 39% (out) and 75% (in).
• The SAI loss helps most on difficult problems. In the loss ablation (average iteration counts): Heat — ℒ₂ 195.6, ℒ_CS 207.3, ℒ_SAI 197.4; Poisson — 134.1, 133.4, 128.8; Hyperelasticity — 185.4, 182.5, 175.7; Synthetic — 2109.8, 2185.7, 1122.0. The authors note ℒ₂ performs similarly on the heat problem but SAI is better on the harder cases.
• Lower condition numbers. Evaluated on a simplified heat problem on a single mesh with 3,764 nodes and varying diffusivity coefficients, the method reduces both the standard condition number κ and Kaporin's condition number compared with earlier traditional approaches (reported as distributions with median, IQR, and outliers in Figure 4).
• Scale invariance property. The SAI loss satisfies ℒ_SAI(A, M⁻¹, w) = ℒ_SAI(αA, M⁻¹, w) for any α > 0, and the associated bound is κ(AM⁻¹) ≤ (1 + σ_max(E))/(1 − σ_max(E)) ≈ 1 + 2σ_max(E) = 1 + 2‖E‖₂, where E = AM⁻¹/‖A‖ − I.

Methodology in Plain English

The starting point is that a preconditioned CG step replaces the residual r with s = M⁻¹r. The paper chooses to approximate A⁻¹ directly with a sparse matrix rather than approximating A with a factorization. Because CG needs the preconditioner to be symmetric positive definite, the sparse matrix G produced by the network is combined as M⁻¹ = GGᵀ + εI, with ε a small positive constant (set to 10⁻⁴ in all experiments). Applying the preconditioner then costs two sparse matrix-vector products instead of the two triangular solves that Incomplete Cholesky requires.

The matrix A is treated as a graph: vertices are variables (mesh nodes in the finite element setting) carrying geometric or physical features, and edges connect entries where A_ij is nonzero, with edge features derived from those entries. A GNN with an encoder-processor-decoder structure first encodes node and edge features with MLPs, applies L message-passing layers that also update edge features, and finally decodes each edge feature into a block entry G_ij. The global G is assembled by replacing each nonzero block of A with the corresponding learned block, preserving A's sparsity pattern.

Training uses the SAI loss, which measures how close AM⁻¹ is to the identity after normalizing A by a scale ‖A‖. To avoid ever forming the dense product AM⁻¹, the loss is estimated stochastically using random vectors w drawn from a standard normal distribution, which reduces the computation to matrix-vector products. The norm used is the mean absolute value over nonzero entries of A, chosen because it is robust to matrix dimension and outlier entries. Because the loss needs only A, no ground-truth solution vectors must be computed to build the training set.

The setup is: Heat, Poisson, Hyperelasticity (block size 3, a beam twist scene over 500 timesteps), and a synthetic system A = PPᵀ + εI with ε = 10⁻⁴. Heat and Poisson use 9,147 meshes with node counts from 400 to 32,000 from the TetWild dataset; the hyperelasticity meshes have node counts from 645 to 14,039 with matrices 3× larger; the synthetic set has 1,000 matrices of roughly 0.12% sparsity and sizes from 10,000 to 20,000. All use a 4:1 train-test split. The GNN uses 4 message-passing steps, one hidden layer per MLP, hidden width 24, about 24k trainable parameters, 500 training epochs, batch size 4, AdamW with an exponentially decaying learning rate (decay rate 0.99), on a single NVIDIA A100 GPU. All preconditioned CG implementations use C++ and CUDA with OpenBLAS, cuBLAS, cuSPARSE, and cusplibrary.

Why This Matters

• Impact on research. The work reframes learned preconditioning as a sparsity-pattern-preserving, GPU-first problem rather than a factorization problem. By matching the preconditioner's computation to the GNN's local propagation, it sidesteps the long-range-dependency argument against GNNs on elimination trees, and by removing the need for solution vectors in the loss, it lowers the cost of building training datasets for large matrices.
• Real-world applications.
  • Finite element simulation of heat conduction and Poisson-type problems on irregular tetrahedral meshes.
  • Nonlinear structural and soft-body simulation, specifically the hyperelasticity (Neo-Hookean beam twist) case studied in the paper.
  • Computer graphics and physically based animation, where large sparse systems from implicit time stepping must be solved every frame.
  • General algebraic sparse SPD systems, represented here by the synthetic dataset with sizes up to 96k.
• Industry relevance. The bottleneck the paper targets is practical: on GPUs, Incomplete Cholesky's triangular solves make it slower than a simple diagonal preconditioner despite far better convergence. A learned preconditioner with near-diagonal construction cost and IC-like iteration counts is directly useful for simulation pipelines that already run on GPU hardware from cuBLAS/cuSPARSE-style libraries.

Future Directions

• Sparsity pattern flexibility. The current method forces G to match the sparsity pattern of A exactly. The authors suggest dynamic dropping strategies to limit fill-in, or including two-hop connections in G, could improve preconditioner quality.
• Beyond CG and SPD systems. The framework could potentially be extended to other Krylov subspace methods such as GMRES or used as a smoother inside multigrid, but the paper restricts itself to SPD systems and the CG solver.
• Scaling beyond one GPU. The work is limited by single-GPU memory; distributed GNN techniques are proposed as a route to multi-GPU scaling.
• Open evaluation questions. The provided text includes a comparison to AMG in Appendix D.2 but not its results, and Appendix C.1/C.2 details (implementation specifics and problem configurations) are outside the content shown. The loss-function proof in Appendix B is also truncated in the supplied content.

Target Audience

This paper suits researchers and practitioners working at the intersection of machine learning and numerical simulation: people building learned solvers or preconditioners, GNN researchers interested in scientific computing applications, and GPU/HPC engineers who need fast sparse linear solvers for finite element or physics-based simulation. Readers should already be comfortable with the conjugate gradient method, preconditioning concepts such as condition number, and basic GNN message passing; the paper is not an entry-level introduction to any of these topics.

Authors’ 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