Research map / Optimization / Optimizers
Learning Sparse Approximate Inverse Preconditioners for Conjugate Gradient Solvers on GPUs
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes a GNN that outputs a sparse approximate inverse (SPAI) preconditioner for CG on GPUs, trained with a scale-invariant SAI loss that needs only the matrix A (no solution vectors).
Problem
Preconditioning conjugate gradient for sparse SPD linear systems on GPUs: incomplete Cholesky requires triangular solves that parallelize poorly, and learned IC-based GNN preconditioners need solution vectors for training and suffer from long-range dependencies.
Approach
Learn a sparse approximate inverse G (M^-1 = GG^T + eps I) with a GNN whose sparsity matches A, so each CG step needs only two SpMVs; train with a label-free loss ||(AM^-1)/||A|| - I||_F^2 estimated stochastically, which is invariant to the scale of A.
Key results
- On GPU, the proposed method achieves the lowest total time across four test cases (e.g., Heat 36 ms vs 77 ms for Diag, a 113% relative promotion), per Table 1.
- Compared with Neural PCG and Neural IF datasets, the method is fastest on GPU in-distribution (26 ms vs 29 ms Diag on Neural PCG data; 132 ms vs 166 ms Diag on Neural IF data).
- The SAI loss gives 1122 average iterations on the synthetic dataset versus 2109.8 for the L2 loss, per the loss ablation.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use the learned SPAI preconditioner instead of Diagonal | GPU-preconditioned CG · Heat Equation | total solve time (ms): 77 -> 36 | Helped |
| use the learned SPAI preconditioner instead of Diagonal | GPU-preconditioned CG · Poisson Equation | total solve time (ms): 45 -> 26 | Helped |
| use the learned SPAI preconditioner instead of Diagonal | GPU-preconditioned CG · Hyperelasticity | total solve time (ms): 86 -> 51 | Helped |
| use the learned SPAI preconditioner instead of Diagonal | GPU-preconditioned CG · Synthetic System | total solve time (ms): 445 -> 253 | Helped |
| use SAI loss instead of L2 or cosine similarity loss (ablation) | learned SPAI preconditioner · Heat, Poisson, Hyperelasticity, and Synthetic | average CG iterations: Heat: 195.6 / 207.3 / 197.4; Poisson: 134.1 / 133.4 / 128.8; Hyperelasticity: 185.4 / 182.5 / 175.7; Synthetic: 2109.8 / | Mixed |
| replace the mean matrix norm in SAI loss with the Frobenius norm (ablation) | SAI loss on Heat · Heat | average PCG iterations: 197 -> 222 | Hurt |
| replace the mean matrix norm in SAI loss with the L1 norm (ablation) | SAI loss on Heat · Heat | average PCG iterations: 197 -> 231 | Hurt |
| set epsilon to 3×10^-5 instead of the reported default 3×10^-3 (ablation) | learned SPAI preconditioner · Heat validation | CG iterations: 197 -> 222 | Hurt |
| set epsilon to 3×10^-4 instead of 3×10^-3 (ablation) | learned SPAI preconditioner · Heat validation | CG iterations: 197 -> 208 | Hurt |
| set epsilon to 3×10^-2 instead of 3×10^-3 (ablation) | learned SPAI preconditioner · Heat validation | CG iterations: 197 -> 205 | Hurt |
| set epsilon to 3×10^-1 (ablation) | learned SPAI preconditioner · Heat validation | training outcome | Hurt |
Surprising findings
- On CPU, the proposed method is not the fastest; Diagonal preconditioning wins on CPU for Heat, Poisson, Hyperelasticity and Synthetic cases.
- Standalone AMG beats the proposed method on GPU for the Heat and Poisson problems in the AMG comparison appendix.
Limitations
- Sparsity pattern of G is fixed to that of A; no dynamic dropping or two-hop connections.
- Scope limited to SPD systems and CG; extension to GMRES or multigrid not evaluated.
- Limited to single-GPU memory; no multi-GPU scaling.
Experimental setup
Datasets: TetWild PDE-derived FEM hyperelasticity beam twist synthetic random SPD matrices
Benchmarks: Heat equation FEM Poisson equation FEM Hyperelasticity Synthetic system Neural PCG dataset Neural IF dataset
Metrics: total solve time (ms) CG iterations condition number Kaporin condition number construction time
Baselines: Diagonal Incomplete Cholesky AINV/SPAI AMG Neural PCG Neural IF
Builds on: Neural PCG Neural IF SPAI MeshGraphNets
Closest papers by meaning
- Graph Neural Preconditioners for Iterative Solutions of Sparse Linear Systems (ICLR 2025 · similarity 0.90)
- Neural operators meet conjugate gradients: The FCG-NO method for efficient PDE solving (ICML 2024 · similarity 0.88)
- Error Feedback Can Accurately Compress Preconditioners (ICML 2024 · similarity 0.87)
- DOGE-Train: Discrete Optimization on GPU with End-to-End Training (AAAI 2024 · similarity 0.86)
- Schwarz–Schur Involution: Lightspeed Differentiable Sparse Linear Solvers (ICML 2025 · similarity 0.85)
- SymMaP: Improving Computational Efficiency in Linear Solvers through Symbolic Preconditioning (NeurIPS 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.