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 addresses the stability of training neural PDE solvers. It uses the "pushforward trick": the adversarial-style loss is used by perturbing the input at step $k$ by a certain noise (note, that this idea looks similar to denoising autoencoders). Another trick is to predict K steps simulaneously in time. The architecture itself is graph-based: the nodes are distributed through the domain, and it is an encoder-processsor-decoder architecture. This helps to be closer to classical numerical solvers. The experiments show that the pushforward trick helps to increase "survival time" (which is defined as a time when a certain error exceed the threshold) clearly helps.