NEURAL PARTIAL DIFFERENTIAL EQUATIONS WITH FUNCTIONAL CONVOLUTION

Abstract

We present a lightweighted neural PDE representation to discover the hidden structure and predict the solution of different nonlinear PDEs. Our key idea is to leverage the prior of "translational similarity" of numerical PDE differential operators to drastically reduce the scale of learning model and training data. We implemented three central network components, including a neural functional convolution operator, a Picard forward iterative procedure, and an adjoint backward gradient calculator. Our novel paradigm fully leverages the multifaceted priors that stem from the sparse and smooth nature of the physical PDE solution manifold and the various mature numerical techniques such as adjoint solver, linearization, and iterative procedure to accelerate the computation. We demonstrate the efficacy of our method by robustly discovering the model and accurately predicting the solutions of various types of PDEs with small-scale networks and training sets. We highlight that all the PDE examples we showed were trained with up to 8 data samples and within 325 network parameters.

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 + x 2 )∇ by training the model on the solution of ∇ • (1 + x 2 )∇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. 

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 + x 2 )∇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. 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 (-x i-1,j -x i+1,j + 4x i,j -x i,j-1 -x i,j+1 )/∆p 2 = b i,j .The discretization of each cell forms one row in the linear system, and the combination of all the rows (cells) forms a sparse



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 Ω.

