Learned numerical methods for linear algebra and optimization

Learning components of numerical algorithms, such as preconditioners, warm starts and inner solvers, to accelerate them without sacrificing reliability.

Linear algebra is a cornerstone of scientific computing, as well as one of the computational bottlenecks when solving optimization problems or partial differential equations. While reliable general-purpose methods exist, many applications can benefit tremendously from tailored methods that fully exploit the structure of a given problem. This, however, requires expert knowledge, yet even that may not achieve maximum efficiency.

We are pioneering a research direction that uses graph neural networks (GNNs) to learn components of classical numerical algorithms, thereby accelerating them without sacrificing reliability. The key insight is that sparse matrices have a natural graph structure: the sparsity pattern defines edges, and matrix entries define features. GNNs are therefore a natural computational primitive for learning over matrices.

A matrix represented as a bipartite graph

A matrix has a natural graph representation: rows and columns become nodes, and nonzero entries become edges.

Neural incomplete factorization

Our first major result was NeuralIF, a GNN that learns incomplete factorization preconditioners for the conjugate gradient method. At the core is a novel message-passing block, inspired by sparse matrix theory, that aligns with the objective to find a sparse factorization of the matrix. NeuralIF significantly speeds up convergence and reduces total solving time compared to classical preconditioners without fill-ins. In a follow-up, we extended this approach to learn incomplete factorizations for the more general GMRES solver, handling non-symmetric systems.

NeuralIF: message-passing on a matrix graph and convergence comparison

NeuralIF learns a sparse factorization via message passing on the matrix graph (left), dramatically accelerating convergence of the conjugate gradient method (right).

Accelerating optimization solvers

Beyond preconditioning, we are exploring how GNNs can accelerate optimization solvers directly. For quadratic programs (QPs) with linear constraints, we use a GNN to predict a warm-start for an active-set solver, reducing the number of iterations needed to reach the optimal solution. The QP and its constraints are represented as a graph, and the GNN performs node-level prediction to identify the likely active set.

Warm-starting pipeline for active-set QP solvers

A GNN maps the QP graph representation to an active-set prediction, warm-starting the solver.

We have also shown that distributed ADMM iterations have a one-to-one correspondence with GNN message passing. Exploiting this connection, we train a GNN to predict adaptive step sizes and communication weights for ADMM, improving convergence speed and solution quality compared to standard ADMM—without sacrificing convergence guarantees.

The same idea applies beyond graph neural networks. Alternating minimization solves problems with two coupled blocks of variables by optimizing one block at a time, and most of its cost lies in the inner solvers for these subproblems. We replace the inner solvers with neural networks trained to approximate the subproblem solutions. For convex problems with quadratic coupling, a worst-case bound on the optimality gap identifies the per-iteration error that each network is trained to minimize, and a PAC-Bayes argument turns it into a high-probability bound over the distribution of problem instances. On low-dose CT reconstruction, the learned scheme matches the accuracy of the iterative baseline up to 30 times faster.

Earlier work

Our first work in this direction considered nonnegative matrix factorization, where we adopted the König graph representation and framed constrained low-rank factorization as a graph problem solved by interleaving a GNN with ADMM.