NEURALPCG: LEARNING PRECONDITIONERS FOR SOLVING PARTIAL DIFFERENTIAL EQUATIONS WITH GRAPH NEURAL NETWORKS

Abstract

Fast and accurate partial differential equation (PDE) solvers empower scientific and engineering research. Classic numerical solvers provide unparalleled accuracy but often require extensive computation time. Machine learning solvers are significantly faster but lack convergence and accuracy guarantees. We present Neural-Network-Preconditioned Conjugate Gradient, or NeuralPCG, a novel linear second-order PDE solver that combines the benefits of classic iterative solvers and machine learning approaches. Our key observation is that both neural-network PDE solvers and classic preconditioners excel at obtaining fast but inexact solutions. NeuralPCG proposes to use neural network models to precondition PDE systems in classic iterative solvers. Compared with neural-network PDE solvers, NeuralPCG achieves converging and accurate solutions (e.g., 1e-12 precision) by construction. Compared with classic solvers, NeuralPCG is faster via data-driven preconditioners. We demonstrate the efficacy and generalizability of NeuralPCG by conducting extensive experiments on various 2D and 3D linear second-order PDEs. 1

1. INTRODUCTION

Partial differential equations (PDEs) are fundamental mathematical models with broad applications in science and engineering, for example, the Navier-Stokes equation in fluid dynamics, Poisson's equation in computational geometry, and the Black-Scholes equation in mathematical finance. Despite their powerful modeling ability and wide applications, it is notoriously difficult to find analytical solutions to a general PDE. Therefore, numerical solvers have long been the mainstay of solving PDEs. Classic PDE solvers provide accurate solutions to well-understood PDEs but typically at the cost of long computation time. Speeding up these classic solvers is non-trivial and often requires complex numerical techniques, e.g., multi-grid methods (Briggs et al., 2000) , domain decomposition (Smith, 1997) , and model reduction (Holmes et al., 2012) . Recently, several pioneering works (Li et al., 2018; Sanchez-Gonzalez et al., 2020) introduce machine learning techniques to solving PDEs, particularly in the field of physics simulation. While this line of methods typically outperforms classic solvers in speed by a large margin, it struggles with converging into a highly precise solution (e.g., 1e -12). A lack of theoretical analysis on convergence and accuracy inhibits neural-network PDE solvers' applications in mechanical engineering, structure analysis, and aerodynamics, where precise PDE solutions have a higher priority than fast yet inexact results. This work proposes NeuralPCG, a novel and hybrid method that combines the benefits of classic and machine-learning PDE solvers. Our key observation is that neural-network solvers are fast at estimating PDE solutions with a low to moderate accuracy. This property aligns with the preconditioning technique in numerical methods, which uses an easy-to-solve approximation of the original PDE to speed up numerical solvers. Based on this intuition, NeuralPCG proposes to learn a neural network that preconditions a classic iterative solver. This is in sharp contrast to existing neural-network PDE solvers that replace the numerical solver entirely with a neural network. The backbone of NeuralPCG remains to be a classic solver, empowering it to inherit convergence and accuracy guarantees. Moreover, the learned preconditioner even outperforms classic preconditioners because it adapts to data distributions from target PDE applications. To demonstrate the efficacy of NeuralPCG, we evaluate it on a set of 2D and 3D representative linear second-order PDEs. We compare its performance with two baselines of very different natures: (1) MeshGraphNet (MGN), a state-of-the-art neural-network PDE solver (Pfaff et al., 2020) ; (2) the preconditioned conjugate gradient (PCG) algorithm with classic preconditioners; Our experiments show that NeuralPCG has its unique advantage over both baselines: Compared with MGN and similar machine-learning solvers, NeuralPCG generates convergent and accurate solutions by construction and avoids error accumulation. On the other hand, NeuralPCG also outperforms classic PCG solvers in speed because our method learns a preconditioner tailored to the training data distribution, which classic preconditioners typically do not exploit. Finally, NeuralPCG generalizes to physics parameters or meshes with moderate differences (5σ) from the training data and is robust to outliers by design. In summary, our work makes the following contributions: 1. We propose NeuralPCG, a novel and efficient solver for linear second-order PDEs that combines the benefits of classic and machine-learning PDE solvers. 2. We present a framework for studying the performance of preconditioners in classic iterative solvers, including a benchmark problem set with carefully designed performance metrics. 3. We conduct extensive experiments to evaluate the efficacy and generalizability of Neu-ralPCG and demonstrate its advantages over existing methods.

2. RELATED WORK

Numerical PDE solvers Classic PDE solvers discretize the continuous PDEs into a numerical system. Popular discretization schemes include finite differences (Strikwerda, 2004) , finite elements (Hughes, 2012), and finite volumes (LeVeque et al., 2002) . Once the discretized numerical system is formulated, one often needs to solve a linear system. Solving these linear systems accurately and efficiently is the key to a robust PDE solver. Previous work has studied numerical linear algebra extensively (Trefethen & Bau III, 1997; Golub & Van Loan, 2013) . These solvers have formed the backbone of many successful, mature numerical PDE software, e.g., ANSYS (Madenci & Guven, 2015) , Abaqus (Helwany, 2007) , COMSOL (Pryor, 2009) , etc. These solvers are crucial for science and engineering, but extensive wall-clock computation time limits their wider deployment. This time bottleneck typically comes from (repeatedly) solving a linear system, and the problem becomes even more severe when the size of the problem scales up. To accelerate linear solvers, prior work has constructed efficient preconditioners for iterative linear solvers via matrix decomposition (Khare & Rajaratnam, 2012), sparsity control (Liu, 1990; Davis & Hager, 1999; Schäfer et al., 2021) , and multiscale techniques (Chen et al., 2021) . Machine learning (ML) methods for solving PDEs Researchers have attempted to incorporate machine learning techniques into solving PDEs. A particular line of work focuses on numerical simulation problems, e.g., solving the underlying PDEs for simulating fluid (Kim et al., 2019 ), cloth (Pfaff et al., 2020) , and elastoplastic solids (Sanchez-Gonzalez et al., 2020) . Typically, these approaches aim to train a surrogate neural network model that replaces the PDE solving process entirely (Brandstetter et al., 2022b; a) . Such methods enjoy a major speedup over classic numerical methods thanks to the fast inference time of a trained network model. However, while the network models are capable of producing visually plausible results, they yield no guarantee on the accuracy of the solution. As such, these methods often suffer from error accumulations, inhibiting their wider usage in applications where accuracy is a priority. Combining numerical methods with networks Given the pros and cons of classic and machine learning methods, researchers have been combining the classical numerical solver with neural networks (Belbute-Peres et al., 2020; Um et al., 2020; Li et al., 2020a; b) , with physics-informed neural networks being a particularly notable example (Raissi et al., 2019; Karniadakis et al., 2021) . Since neural networks serve as the backbones of these methods, they often have a hard time generating quantitatively accurate solutions (e.g., 1e -12 precision) and enforcing hard constraints (Márquez-Neila et al., 2017) , both of which are relatively simple with classic numerical solvers. Additionally,



Supplementary videos are available on the project webpage: https://sites.google.com/view/ neuralpcg

