Learning from Linear Algebra: A Graph Neural Network Approach to Preconditioner Design for Conjugate Gradient Solvers

TL;DR

This paper introduces a graph neural network approach that learns to design preconditioners for conjugate gradient methods, outperforming classical and neural methods in reducing system condition numbers for parametric PDEs.

Contribution

It leverages well-established linear algebra preconditioners as a starting point to train GNNs that produce more effective preconditioners, improving convergence in iterative solvers.

Findings

GNN-based preconditioners outperform classical methods in reducing condition numbers.

The approach is particularly effective for parametric partial differential equations.

The learned preconditioners lead to faster convergence in conjugate gradient methods.

Abstract

Large linear systems are ubiquitous in modern computational science and engineering. The main recipe for solving them is the use of Krylov subspace iterative methods with well-designed preconditioners. Recently, GNNs have been shown to be a promising tool for designing preconditioners to reduce the overall computational cost of iterative methods by constructing them more efficiently than with classical linear algebra techniques. Preconditioners designed with these approaches cannot outperform those designed with classical methods in terms of the number of iterations in CG. In our work, we recall well-established preconditioners from linear algebra and use them as a starting point for training the GNN to obtain preconditioners that reduce the condition number of the system more significantly than classical preconditioners. Numerical experiments show that our approach outperforms both classical and neural network-based methods for an important class of parametric partial differential equations. We also provide a heuristic justification for the loss function used and show that preconditioners obtained by learning with this loss function reduce the condition number in a more desirable way for CG.