Neural incomplete factorization: learning preconditioners for the conjugate gradient method

TMLR Paper2645 Authors

08 May 2024 (modified: 21 Jun 2024)Under review for TMLREveryoneRevisionsBibTeXCC BY-SA 4.0
Abstract: The convergence of the conjugate gradient method to solve large-scale and sparse linear equation systems depends on the conditioning of the system matrix, which can be improved by preconditioning. In this paper, we develop a computationally efficient data-driven approach to accelerate the generation of effective preconditioners. We, therefore, replace the typically hand-engineered preconditioners by the output of graph neural networks. Optimizing the condition number of the linear system directly is computationally infeasible. Instead, our method generates an incomplete factorization of the matrix and is, therefore, referred to as neural incomplete factorization (NeuralIF). For efficient training, we utilize a stochastic approximation of the Frobenius loss which only requires matrix-vector multiplications. At the core of our method is a novel message-passing block, inspired by sparse matrix theory, that aligns with the objective of finding a sparse factorization of the matrix. We evaluate our proposed method on both synthetic problem instances and on problems arising from the discretization of the Poisson equation on varying domains. Our experiments show that by utilizing data-driven preconditioners within the conjugate gradient method we are able to speed up the convergence of the iterative procedure.
Submission Length: Regular submission (no more than 12 pages of main content)
Changes Since Last Submission: Based on the reviews we updated the following sections: - we corrected the paragraph and included that Algorithm 1 shows the preconditioned version of the conjugate gradient method - we clarified the connection of the used message passing scheme with the Coates and Königs graph representation - we added the time required for the training of the GNN to the result section
Assigned Action Editor: ~Vikas_Sindhwani1
Submission Number: 2645
Loading