Message Passing Neural PDE Solvers

1 INTRODUCTION . In the sciences , years of work have yielded extremely detailed mathematical models of physical phenomena . Many of these models are expressed naturally in differential equation form ( Olver , 2014 ) , most of the time as temporal partial differential equations ( PDE ) . Solving these differential equations is of huge importance for problems in all numerate disciplines such as weather forecasting ( Lynch , 2008 ) , astronomical simulations ( Courant et al. , 1967 ) , molecular modeling ( Lelièvre & Stoltz , 2016 ) , or jet engine design ( Athanasopoulos et al. , 2009 ) . Solving most equations of importance is analytically intractable and necessitates falling back on numerical approximation schemes . Obtaining accurate solutions of bounded error with minimal computational overhead requires the need for handcrafted solvers , always tailored to the equation at hand ( Hairer et al. , 1993 ) . The design of “ good ” PDE solvers is no mean feat . The perfect solver should satisfy an almost endless list of conditions . There are user requirements , such as being fast , using minimal computational overhead , being accurate , providing uncertainty estimates , generalizing across PDEs , and being easy to use . Then there are structural requirements of the problem , such as spatial resolution and timescale , domain sampling regularity , domain topology and geometry , boundary conditions , dimensionality , and solution space smoothness . And then there are implementational requirements , such as maintaining stability over long rollouts and preserving invariants . It is precisely because of this considerable list of requirements that the field of numerical methods is a splitter field ( Bartels , 2016 ) , tending to build handcrafted solvers for each sub-problem , rather than a lumper field , where a mentality of “ one method to rule them all ” reigns . This tendency is commonly justified with reference to no free lunch theorems . In this paper , we propose to numerically solve PDEs with an end-to-end , neural solver . Our contributions can be broken down into three main parts : • An end-to-end fully neural PDE solver , based on neural message passing , which offers flexibility to satisfy all structural requirements of a typical PDE problem . This design is motivated by the insight that some classical solvers ( finite differences , finite volumes , and WENO scheme ) can be posed as special cases of message passing . • Temporal bundling and the pushforward trick , which are methods to encourage zerostability in training autoregressive models . • Generalization across multiple PDEs within a given class . At test time , new PDE coefficients can be input to the solver . 2 BACKGROUND AND RELATED WORK . Here in Section 2.1 we briefly outline definitions and notation . We then outline some classical solving techniques in Section 2.2 . Lastly , in Section 2.3 , we list some recent neural solvers and split them into the two main neural solving paradigms for temporal PDEs . 2.1 PARTIAL DIFFERENTIAL EQUATIONS . We focus on PDEs in one time dimension t = [ 0 , T ] and possibly multiple spatial dimensions x = [ x1 , x2 , . . . , xD ] > ∈ X . These can be written down in the form ∂tu = F ( t , x , u , ∂xu , ∂xxu , ... ) ( t , x ) ∈ [ 0 , T ] × X ( 1 ) u ( 0 , x ) = u0 ( x ) , B [ u ] ( t , x ) = 0 x ∈ X , ( t , x ) ∈ [ 0 , T ] × ∂X ( 2 ) where u : [ 0 , T ] × X → Rn is the solution , with initial condition u0 ( x ) at time t = 0 and boundary conditions B [ u ] ( t , x ) = 0 when x is on the boundary ∂X of the domain X . The notation ∂xu , ∂xxu , . . . is shorthand for partial derivatives ∂u/∂x , ∂2u/∂x2 , and so forth . Most notably , ∂u/∂x represents a gradient vector . The boundary operator B we consider in this paper are Dirichlet boundary conditions , where BD [ u ] = u − bD for fixed function bD and Neumann boundary conditions BN [ u ] = n > ∂xu− bN , for scalar-valued u where n is an outward facing normal on ∂X . Conservation form Among all PDEs , we hone in on solving those that can be written down in conservation form , because there is already precedent in the field for having studied these ( Bar-Sinai et al. , 2019 ; Li et al. , 2020a ) . Conservation form PDEs are written as ∂tu +∇ · J ( u ) = 0 , ( 3 ) where∇·J is the divergence of J . The quantity J : Rn → Rn is the flux , which has the interpretation of a quantity that appears to flow . Consequently , u is a conserved quantity within a volume , only changing through the net flux J ( u ) through its boundaries . 2.2 CLASSICAL SOLVERS . Grids and cells Numerical solvers partition X into a finite grid X = { ci } Ni=1 of N small nonoverlapping volumes called cells ci ⊂ X . In this work , we focus on grids of rectangular cells . Each cell has a center at xi . uki is used to denote the discretized solution in cell ci and time tk . There are two main ways to compute uki : sampling u k i = u ( tk , xi ) and averaging u k i = ∫ ci u ( tk , x ) dx . In our notation , omitting an index implies that we use the entire slice , so uk = ( uk1 , u k 2 , ... , u k N ) . Method of lines A common technique to solve temporal PDEs is the method of lines ( Schiesser , 2012 ) . The domain X is discretized into a grid X and the solution as a vector uk . We then solve ∂tu t ∣∣ tk = f ( t , uk ) t ∈ [ 0 , T ] , ( 4 ) where f is the form of F acting on the vectorized ut instead of the function u ( t , x ) . The only derivative operator is now in time , making it an ordinary differential equation ( ODE ) . This ODE can be solved with off-the-shelf ODE solvers ( Butcher , 1987 ; Everhart , 1985 ) . f can be formed in different ways by approximating spatial derivatives on the grid . Below are three classical techniques . Finite difference method ( FDM ) In FDM , spatial derivative operators ( e.g. , ∂x ) are replaced with difference operators , called stencils . For instance , ∂xuk|xi might become ( uki+1− uki ) / ( xi+1− xi ) . Principled ways to derive stencils can be found in Appendix A. FDM is simple and efficient , but suffers poor stability unless the spatial and temporal discretizations are carefully controlled . Finite volume method ( FVM ) FVM works for equations in conservation form . It can be shown using the divergence theorem that the integral of u over cell i increases only by the net flux into the cell . In 1D , this leads to f ( t , uki ) = 1 ∆xi ( Jki−1/2 − J k i+1/2 ) , ( 5 ) where ∆xi is the cell width , and Jki−1/2 , J k i+1/2 the flux at the left and right cell boundary at time tk , respectively . The problem thus boils down to estimating the flux at cell boundaries xi±1/2 . The beauty of this technique is that the integral of u is exactly conserved . FVM is generally more stable and accurate than FDM , but can only be applied to conservation form equations . Pseudospectral method ( PSM ) PSM computes derivatives in Fourier space . In practical terms , the mth derivative is computed as IFFT { ( ιω ) mFFT ( u ) } for ι = √ −1 . These derivatives obtain exponential accuracy ( Tadmor , 1986 ) , for smooth solutions on periodic domains and regular grids . For non-periodic domains , analogues using other polynomial transforms exist , but for non-smooth solution this technique can not be applied . 2.3 NEURAL SOLVERS . We build on recent exciting developments in the field to learn PDE solvers . These neural PDE solvers , as we refer to them , are laying the foundations of what is becoming both a rapidly growing and impactful area of research . Neural PDE solvers for temporal PDEs fall into two broad categories , autoregressive methods and neural operator methods , ( see Figure 1a ) . Neural operator methods Neural operator methods treat the mapping from initial conditions to solutions at time t as an input–output mapping learnable via supervised learning . For a given PDE and given initial conditions u0 , a neural operator M : [ 0 , T ] × F → F , where F is a ( possibly infinite-dimensional ) function space , is trained to satisfy M ( t , u0 ) = u ( t ) . ( 6 ) Finite-dimensional operator methods ( Raissi , 2018 ; Sirignano & Spiliopoulos , 2018 ; Bhatnagar et al. , 2019 ; Guo et al. , 2016 ; Zhu & Zabaras , 2018 ; Khoo et al. , 2020 ) , where dim ( F ) < ∞ are grid-dependent , so can not generalize over geometry and sampling . Infinite-dimensional operator methods ( Li et al. , 2020c ; a ; Bhattacharya et al. , 2021 ; Patel et al. , 2021 ) by contrast resolve this issue . Each network is be trained on example solutions of the equation of interest and is therefore locked to that equation . These models are not designed to generalize to dynamics for out-of-distribution t. Autogressive methods An orthogonal approach , which we take , is autoregressive methods . These solve the PDE iteratively . For time-independent PDEs , the solution at time t+ ∆t is computed as u ( t+ ∆t ) = A ( ∆t , u ( t ) ) , ( 7 ) where A : R > 0 × RN → RN is the temporal update . In this work , since ∆t is fixed , we just write A ( u ( t ) ) . Three important works in this area are Bar-Sinai et al . ( 2019 ) , Greenfeld et al . ( 2019 ) , and Hsieh et al . ( 2019 ) . Each paper focuses on a different class of PDE solver : finite volumes , multigrid , and iterative finite elements , respectively . Crucially , they all use a hybrid approach ( Garcia Satorras et al. , 2019 ) , where the solver computational graph is preserved and heuristically-chosen parameters are predicted with a neural network . Hsieh et al . ( 2019 ) even have convergence guarantees for their method , something rare in deep learning . Hybrid methods are desirable for sharing structure with classical solvers . So far in the literature , however , it appears that autoregressive methods are more the exception than the norm , and for those methods published , it is reported that they are hard to train . In Section 3 we explore why this is and seek to remedy it .

Summary: The paper proposes a framework to solve generic PDEs written in conservation form, in an autoregressive fashion. The grid domain is considered as a graph, whose nodes contain information about solutions to the PDE until a given times steps, node positions, time steps, and hyperparameters information about the PDE. The node’s labels are first embedded in a latent space. Then message passing is applied and the output is decoded to predict the next n time steps solutions together. The network is trained with an additional loss which promotes zero stability of the solver, by constraining the network to be invariant to small perturbations in the input solutions of previous time steps. The approach is validated on a parameterized family of 1d and 2d PDEs, showing accuracy and generalization results that are above the state of the art of neural solvers and classical solver schemes. In addition, the paper show promising results dealing with irregular domains with different types of boundary conditions.