Abstract:This paper presents a learnable solver tailored to iteratively solve sparse linear systems from discretized partial differential equations (PDEs). Unlike traditional approaches relying on specialized expertise, our solver streamlines the algorithm design process for a class of PDEs through training, which requires only training data of coefficient distributions. The proposed method is anchored by three core principles: (1) a multilevel hierarchy to promote rapid convergence, (2) adherence to linearity concerning the right-hand-side of equations, and (3) weights sharing across different levels to facilitate adaptability to various problem sizes. Built on these foundational principles and considering the similar computation pattern of the convolutional neural network (CNN) as multigrid components, we introduce a network adept at solving linear systems from PDEs with heterogeneous coefficients, discretized on structured grids. Notably, our proposed solver possesses the ability to generalize over right-hand-side terms, PDE coefficients, and grid sizes, thereby ensuring its training is purely offline. To evaluate its effectiveness, we train the solver on convection-diffusion equations featuring heterogeneous diffusion coefficients. The solver exhibits swift convergence to high accuracy over a range of grid sizes, extending from $31 \times 31$ to $4095 \times 4095$. Remarkably, our method outperforms the classical Geometric Multigrid (GMG) solver, demonstrating a speedup of approximately 3 to 8 times. Furthermore, our numerical investigation into the solver's capacity to generalize to untrained coefficient distributions reveals promising outcomes.

What problem does this paper attempt to address?

This paper proposes a method called MGCNN (Learnable Multi-Grid Solver) for solving sparse linear systems based on structured grids that arise from the discretization of partial differential equations (PDEs). Unlike traditional algorithm that depends on specialized knowledge, MGCNN simplifies the algorithm design process for specific PDE categories through training, requiring only coefficient distribution data. The method is based on three core principles: (1) adopting multi-level hierarchy to promote fast convergence; (2) following linearity rules regarding the equation's right-hand side; (3) sharing weights between different levels to accommodate different problem scales. MGCNN utilizes a similar computational pattern of convolutional neural network (CNN) to construct a network that can solve PDE linear systems with heterogeneous coefficients. It generalizes right-hand side terms, PDE coefficients, and grid sizes, enabling offline training. The paper demonstrates its effectiveness through experiments on convection-diffusion equations with heterogeneous diffusion coefficients, showing that MGCNN can achieve high accuracy rapidly on different grid sizes from 31×31 to 4095×4095, and it is 3-8 times faster than traditional geometric multigrid (GMG) solvers. Furthermore, the research also discovers that training the solver on a mixed coefficient distribution can significantly enhance its robustness in the face of unseen distributions, providing a potential approach for improving solver robustness in various scenarios. MGCNN combines the traditional multilevel structure and CNN, aiming to create an efficient iterative solver that does not require expert knowledge and is applicable to large-scale problems.

MGCNN: a learnable multigrid solver for sparse linear systems from PDEs on structured grids

Learning to Optimize Multigrid PDE Solvers

Multilevel CNNs for Parametric PDEs

UGrid: An Efficient-And-Rigorous Neural Multigrid Solver for Linear PDEs

Multigrid-Augmented Deep Learning Preconditioners for the Helmholtz Equation using Compact Implicit Layers

Solving the Discretised Multiphase Flow Equations with Interface Capturing on Structured Grids Using Machine Learning Libraries

Learning a generalized multiscale prolongation operator

A deep learning algorithm to accelerate algebraic multigrid methods in finite element solvers of 3D elliptic PDEs

A fully differentiable GNN-based PDE Solver: With Applications to Poisson and Navier-Stokes Equations

Monte Carlo Neural PDE Solver for Learning PDEs via Probabilistic Representation

Solving multiscale elliptic problems by sparse radial basis function neural networks

Multigrid-augmented deep learning preconditioners for the Helmholtz equation

Multi-Level GNN Preconditioner for Solving Large Scale Problems

Deep Learning-Enhanced Preconditioning for Efficient Conjugate Gradient Solvers in Large-Scale PDE Systems

DNN-MG: A Hybrid Neural Network/Finite Element Method with Applications to 3D Simulations of the Navier-Stokes Equations

DiffGCN: Graph Convolutional Networks via Differential Operators and Algebraic Multigrid Pooling

Learning Preconditioners for Conjugate Gradient PDE Solvers

Graph and convolutional neural network coupling with a high-performance large-eddy simulation solver

MgNet: A Unified Framework of Multigrid and Convolutional Neural Network

A Neural Multigrid Solver for Helmholtz Equations with High Wavenumber and Heterogeneous Media