Neural Partial Differential Equations with Functional Convolution

1 INTRODUCTION ( 1+ ) Problem definition We aim to devise a learning paradigm to solve the inverse PDE identification problem . By observing a small data set in the PDE ’ s solution space with an unknown form of equations , we want to generate an effective neural representation that can precisely reconstruct the hidden structure of the target PDE system . This neural representation will further facilitate the prediction of the PDE solution with different boundary conditions . The right inset figure shows a typical example of our target problem : by observing a small part ( 4 samples in the figure ) of the solution space of a nonlinear PDE system F ( x ) = b , without knowing its analytical equations , our neural representation will depict the hidden differential operators underpinning F ( e.g. , to represent the unknown differential operator∇ · ( 1 + x2 ) ∇ by training the model on the solution of∇ · ( 1 + x2 ) ∇x = b . Challenges to solve The nonlinearity and the curse of dimensionality of the target PDE ’ s solution manifold are the two main challenges for the design of a high-performance neural discretization . An effective neural representation of a PDE system plays an essential role to solve these challenges . In retrospect , the design of neural PDE representations has been evolving from the raw , unstructured networks ( e.g. , by direct end-to-end data fitting ) to various structured ones with proper mathematical priors embedded . Examples include the residual-based loss function ( e.g. , physics-informed networks Raissi et al. , 2020 ; Lu et al. , 2019 ; Raissi et al. , 2019 ) , learnable convolution kernels ( e.g. , PDE-Nets Long et al. , 2018a ; b ; 2019 ) , and hybrid of numerical stencils and MLP layers ( e.g. , see Amos & Kolter , 2017 ; Pakravan et al. , 2020 ; Geng et al. , 2020 ; Stevens & Colonius , 2020 ) . Following this line of research , we aim to devise a lightweighted neural PDE representation that fuses the mathematical equation ’ s essential structure , the numerical solvers ’ computational efficiency , and the neural networks ’ expressive power . In particular , we want to aggressively reduce the scale of both model parameters and training data to some extremal extent , while extending the scope of the targeted PDE systems to a broad range , encompassing equations that are both linear and nonlinear , both steady-state and dynamic . i , j Flatten vector of , , . . . Output of , , Translational similarity of differential operators Our neural PDE representation design is motivated by the historical successes of the various sparse , iterative numerical solvers in solving nonlinear PDEs over the past decades . The key observation we have made is that the efficacy of a classical numerical PDE solver relies on the translational similarity of its discretized , local differential operators . Namely , the form of a differential operator can be written as a functional C ( x , p ) with respect to the the PDE unknown x and the local position p , which is showed in the right inset figure . For example , for a linear Poisson system ∇ · ∇x = b , C is a constant function ; for a linear Poisson system with embedded boundary conditions , C is a function of position p ; for a nonlinear PDE∇· ( 1 +x2 ) ∇x = b , C is a function of PDE unknown x ( or both x and p if it has embedded boundaries ) . For most numerical PDEs , these local functional operators can be parameterized and built on-the-fly within the solver iterations . Such operators ’ locality further inspired the design of a variety of computationally efficient PDE solvers , among which the most famous one is the matrix-free scheme that has been used widely in solving large-scale physical systems on GPU . These local procedures for stencil creation have demonstrated their extreme performance in accommodating PDE solvers . From a machine learning perspective , these “ translational similar ” differential operators resemble the concept of convolution operators that function as the cornerstone to embed the “ translational invariant ” priors into neural networks ( see LeCun et al. , 1995 ; 1998 ) . Method overview In this work , we leverage the PDE differential operators ’ “ translational similarity ” in a reverse manner by devising a local neural representation that can uncover and describe the global structure of the target PDE . At the heart of our approach lies in a differential procedure to simultaneously describe the spatial coupling and the temporal evolution of a local data point . Such procedure is implemented as a parameterized micro network , which is embedded in our iterative solving architecture , to learn the numerical process of converging from an initial guess to a final steady state for a PDE solution . We name these embedded micro networks “ functional convolutions , ” for two reasons . First , fitting the parameters of these local embedded networks amounts to the exploration of the optimal function that best describes the observed solution of the unknown nonlinear PDE within a functional space . Second , the local differential operators that span this functional space can be treated as numerically employing convolution kernels Hsieh et al . ( 2018 ) ; Lin et al . ( 2013 ) . Based on these functional convolutions , we are able to devise a learning framework by embedding the micro network architecture within an iterative procedure to 1 ) backwardly learn the underpinning , spatially varying structures of a nonlinear PDE system by iteratively applying the adjoint linear solvers and 2 ) forwardly predict the steady states of a PDE system by partially observing data samples of its equilibrium . We show that our model can simultaneously discover structures and predict solutions for different types of nonlinear PDEs . We particularly focus on solving elliptic boundary value problems that were less explored in the current literature . 2 MOTIVATING EXAMPLE : FORWARD NUMERICAL PDE . Naming convention We first show a motivating example to demonstrate the standard process of a forward numerical PDE solver . We take the simplest Poisson equation with Dirichlet boundary conditions as an example . The mathematical equation of a Poisson system can be written as ∇ · ∇x = b for x ∈ Ω , with x as the PDE unknowns , b as the right-hand side , and Ω as the problem ’ s domain . The boundary conditions are enforced in a Dirichlet way ( by assigning values directly ) as x = x̂ on the domain boundary , with x̂ as the specified boundary values . To create a discretized , numerical system to solve the equation , we use the symbol p to denote the position within the domain . The numerical solution of the PDE amounts to seeking an unknown function x ( p ) that can specify the value of x in an arbitrary position p within Ω . Linear PDE As shown in Figure 1 , we illustrate how to solve the Poisson system using a finitedifference method . We first subdivide the domain into n cell ( segment intervals in 1D and squares in 2D ) with the cell size of ∆p . Taking the 2D case for example , we can derive the discretized Poisson equation by approximating the Laplacian operator on each grid cell using the central finite difference method ( −xi−1 , j − xi+1 , j + 4xi , j − xi , j−1 − xi , j+1 ) /∆p2 = bi , j .The discretization of each cell forms one row in the linear system , and the combination of all the rows ( cells ) forms a sparse 1D Poisson equation 2D Poisson equation linear system Ax = b to solve . For a linear Poisson system , each row of A can be translated into a convolutional stencil instantiated with constant parameters , e.g. , ( -1,2,1 ) in 1D and ( -1 , -1,4 , -1 , -1 ) in 2D . This convolution perspective can be used to accelerate the numerical solver by maintaining a “ conceptual ” A matrix without storing the redundant matrix element values in memory ( we refer the readers to the matrix-free method Carey & Jiang ( 1986 ) for more details ) . This matrix-free nature indicates an extremely concise representation of a numerical Poisson system ( i.e. , the matrix A ) —using a 1× 3 or 3× 3 convolution kernel with fixed values . We use the symbol C to denote the convolutional representation of A . For a linear system , C is independent from the values of p and x. Nonlinear PDE and Picard interation The nonlinear case of the Poisson system is a bit complicated . We can still use the matrix-free representation to describe a nonlinear Poisson system , but the parameters of this convolutional stencil now depends on both the local position p and the local unknown p ( x ) . This dependency is nonlinear and therefore we can not find the solution by solving a single Ax = b . Here we present an iterative scheme—the Picard method—to solve a numerical nonlinear PDE system . ( see Picard , 1893 ; Bai , 2010 ) Let ’ s consider the nonlinear Poisson equation as : ∇ · α ( x ) ∇x = b for x ∈ Ω and Dirichlet boundary conditions on the boundary . The source of nonlinearity in this PDE includes the coefficient α ( x ) which is dependent on the solution of x , e.g. , α ( x ) = 1 + x2 . A simple and effective fixed-point procedure to discretize and solve the nonlinear Poisson equation , named Picard iteration , can be sketched as follows : while : |xn − xn−1| > xn+1 = A −1 ( xn ) b , ( 1 ) with the matrix A representing the current discretized nonlinear Laplacian operator approximated by the value of unknowns from the previous iteration . The key idea is to employ a linear approximation of a nonlinear problem and enforce such approximation iteratively to evolve to a steady state ( see Figure 2 ) . To uncover the underlying structure of A , which can evolve both spatially and temporally , we make a prior assumption that A can be described by a kernel function C ( x ( p ) , p ) . Such prior applies to most of the elliptic PDE systems where the spatial terms can be expressed as the combination of one or several differential operators . From a numerical perspective , C describes the local discretized interaction between an element and its neighbours . It amounts to a function that returns all the non-zero elements for each row i in A ( think of A in a matrix-free way ) . 3 METHOD : BACKWARD NEURAL PDE . Overview In this section , we present our neural PDE design motivated by the forward Picard solver with convolutional representation of the nonlinear system . Our framework consists of three key components : the neural representation of the convolution operator , the embedded Picard forward iterative solver , and the adjoint backward derivatives . The key idea is to differentiate the forward nonlinear Picard solver and build the neural representation for its sparse linearization step . This differentiation is implemented by our functional convolution scheme on the linearization level and the adjoint Picard for the nonlinear iterations . 3.1 FUNCTIONAL CONVOLUTION . In a conventional ML method , C can be approximated by the combination of a set of kernels Long et al . ( 2018a ) or by solving a deconvolution problem Xu et al . ( 2014 ) ; Jin et al . ( 2017 ) ; Zhang et al . ( 2017 ) . However , these strategies do not suit our scenario , where the instant linear effects of the system should be approximated by extracting the nonlinear effects of C. A natural choice to approximate this spatially-and-temporally varying kernel function is to devise a neural network , which takes the form of C ( x ( p ) , p , θ ) , with θ as the network parameters . Numerically , the global matrix A can be fully parameterized by C ( x ( p ) , p , θ ) by assuming the fact that C is a non-weight-shared convolution operator in the spatial domain . As illustrated in Figure 3 , such neural network can be further incorporated into a conventional nonlinear Picard solver to obtain the forward steady-state solution by solving the linear system A ( xn , θ ) xn+1 = b ( xn ) , where a black-box sparse linear solver can be used as in a conventional simulation program . The formal definition of functional convolution written in the kernel way is A ( x ( pi ) , θ ) = ∑ pj∈N ( pi ) [ C ( x ( N ( pi ) ) , N ( pi ) , θ ) ] e ( pj ) ( 2 ) where A ( x ( pi ) , θ ) is the ith row of matrix A , N ( pi ) is the neighboring positions of pi , x ( N ( pi ) ) is all the neighboring elements ( all channels ) around the position pi and e ( pj ) is the unit vector representing the jth column in A . To specify the 2D example , equation ( 2 ) has the following form A ( xm , n ) , θ ) = 1∑ i=−1 1∑ j=−1 [ C ( N ( xm , n ) , θ ) ] i , jem+i , n+j ( 3 ) where xm , n is the element that lies in row m , column n of the feature map . The input N ( xm , n ) = { xm+i , n+j for i , j = −1 , 0 , 1 } is the flatten vector of neighboring elements of pixel xm , n . After a simple neural network C with parameters θ , we obtain the output C ( N ( xm , n ) , θ ) , which is a vector with the same length of N ( xm , n ) .

Summary: Post-discussion update: The authors only partially adressed my concerns in their rebuttal. The paper suffers from lack of comparisons: only 2 baselines are compared, and only on few systems. Crucially the new Navier-Stokes experiment lacks comparisons. The authors also couldn't respond to my questions about research context or scope: it's difficult to assess what this work actually claims in relation to competing methods. For a machine learning paper this is not enough.