用图神经网络学稀疏逆预条件器，让GPU上解线性方程更快更稳。
Learning Sparse Approximate Inverse Preconditioners for Conjugate Gradient Solvers on GPUs
- 用GNN构造稀疏近似逆预条件器，避免三角求解，适配GPU并行。
- 在3个物理方程数据集和1个合成数据集上提速40%-53%，最快快113%。
- 新设计的尺度不变损失函数，提升泛化能力，适合大规模科学计算。
共轭梯度法（CG）是求解对称正定线性系统Ax=b的常用方法，高效预条件器对快速收敛至关重要。传统预条件器依赖固定算法，难以从数据中优化；现有学习方法多采用图神经网络（GNN）提升性能，但依赖不完全分解，导致三角求解阻碍GPU并行，并引入长程依赖，难被GNN建模。为此，我们提出一种基于GNN的GPU友好型预条件器，专门构建稀疏近似逆（SPAI）预条件器，避免三角求解，每步仅需两次矩阵-向量乘法。矩阵-向量乘法的局部性与GNN的局部传播机制相容。同时，我们设计基于统计的尺度不变损失函数，其特性匹配CG收敛率取决于条件数而非矩阵尺度的性质，显著提升学习效果。在三个偏微分方程衍生数据集和一个合成数据集上的评估表明，该方法优于标准预条件器（对角、IC、传统SPAI）及以往学习型预条件器。在GPU上将求解时间减少40%-53%（最快提升113%），同时获得更优条件数与更强泛化性能。源代码见https://github.com/Adversarr/LearningSparsePreconditioner4GPU。
原文摘要 · Abstract (English)
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
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。