RAPNet: Accelerating Algebraic Multigrid with Learned Sparse Corrections
Abstract
The scalable solution of large sparse linear systems is a bottleneck in scientific computing and graph analysis. While algebraic multigrid (AMG) offers optimal linear scaling, its performance is severely constrained by the trade-off between the sparsity and convergence quality of coarse-grid operators. Classical AMG heuristics struggle to balance these objectives, often sacrificing stability or performance for sparsity. We propose RAPNet, a graph neural network (GNN) framework that resolves this trade-off by learning to generate sparse, robust coarse operators directly from the sparse algebraic system. Key to our approach is a level-wise training strategy that enables learning from small subgraphs and generalization to million-node domains, bypassing the bottlenecks of prior neural AMG attempts. RAPNet executes exclusively during the solver setup phase, ensuring that the solve phase retains its favorable computational properties. We show that our method outperforms classical non-Galerkin baselines on diverse PDE discretizations and graph Laplacians, making it particularly effective for multi-query tasks such as eigenproblems, time-dependent simulations, and inverse or design problems.
Lay Summary
Solving massive math equations is a major roadblock in science and computer analysis. While traditional methods are helpful, they force a difficult trade-off between keeping the math simple (using less memory) and solving the problem quickly. We built RAPNet, an artificial intelligence tool that resolves this issue. Using a neural network, RAPNet learns to make highly efficient, lightweight corrections to these traditional mathematical solvers. It only runs once during the initial setup phase, ensuring the actual step-by-step solving process stays incredibly fast. By training our AI on small examples, we found it successfully applies its knowledge to solve massive problems with millions of interconnected parts. Our method beats older techniques across various physical simulations and complex networks, offering a highly effective upgrade for large-scale computer tasks.