Learning Neural PDE Solvers with Convergence Guarantees

1 INTRODUCTION . Partial differential equations ( PDEs ) are ubiquitous tools for modeling physical phenomena , such as heat , electrostatics , and quantum mechanics . Traditionally , PDEs are solved with hand-crafted approaches that iteratively update and improve a candidate solution until convergence . Decades of research and engineering went into designing update rules with fast convergence properties . The performance of existing solvers varies greatly across application domains , with no method uniformly dominating the others . Generic solvers are typically effective , but could be far from optimal for specific domains . In addition , high performing update rules could be too complex to design by hand . In recent years , we have seen that for many classical problems , complex updates learned from data or experience can out-perform hand-crafted ones . For example , for Markov chain Monte Carlo , learned proposal distributions lead to orders of magnitude speedups compared to handdesigned ones ( Song et al. , 2017 ; Levy et al. , 2017 ) . Other domains that benefited significantly include learned optimizers ( Andrychowicz et al. , 2016 ) and learned data structures ( Kraska et al. , 2018 ) . Our goal is to bring similar benefits to PDE solvers . Hand-designed solvers are relatively simple to analyze and are guaranteed to be correct in a large class of problems . The main challenge is how to provide the same guarantees with a potentially much more complex learned solver . To achieve this goal , we build our learned iterator on top of an existing standard iterative solver to inherit its desirable properties . The iterative solver updates the solution at each step , and we learn a parameterized function to modify this update . This function class is chosen so that for any choice of parameters , the fixed point of the original iterator is preserved . This guarantees correctness , and training can be performed to enhance convergence speed . Because of this design , we only train on a single problem instance ; our model correctly generalizes to a variety of different geometries and boundary conditions with no observable loss of performance . As a result , our approach provides : ( i ) theoretical guarantees of convergence to the correct stationary solution , ( ii ) faster convergence than existing solvers , and ( iii ) generalizes to geometries and boundary conditions very different from the ones seen at training time . This is in stark contrast with existing deep learning approaches for PDE solving ( Tang et al. , 2017 ; Farimani et al. , 2017 ) that are limited to specific geometries and boundary conditions , and offer no guarantee of correctness . Our approach applies to any PDE with existing linear iterative solvers . As an example application , we solve the 2D Poisson equations . Our method achieves a 2-3× speedup on number of multiplyadd operations when compared to standard iterative solvers , even on domains that are significantly different from our training set . Moreover , compared with state-of-the-art solvers implemented in FEniCS ( Logg et al. , 2012 ) , our method achieves faster performance in terms of wall clock CPU time . Our method is also simple as opposed to deeply optimized solvers such as our baseline in FEniCS ( minimal residual method + algebraic multigrid preconditioner ) . Finally , since we utilize standard convolutional networks which can be easily parallelized on GPU , our approach leads to an additional 30× speedup when run on GPU . 2 BACKGROUND . In this section , we give a brief introduction of linear PDEs and iterative solvers . We refer readers to LeVeque ( 2007 ) for a thorough review . 2.1 LINEAR PDES . Linear PDE solvers find functions that satisfy a ( possibly infinite ) set of linear differential equations . More formally , let F = { u : Rk → R } be the space of candidate functions , and A : F → F be a linear operator ; the goal is to find a function u ∈ F that satisfies a linear equation Au = f , where f is another function Rk → R given by our problem . Many PDEs fall into this framework . For example , heat diffusion satisfies ∇2u = f ( Poisson equation ) , where ∇2 = ∂ 2 ∂x21 + · · · + ∂ 2 ∂x2k is the linear Laplace operator ; u maps spatial coordinates ( e.g . in R3 ) into its temperature , and f maps spatial coordinates into the heat in/out flow . Solving this equation lets us know the stationary temperature given specified heat in/out flow . Usually the equation Au = f does not uniquely determine u . For example , u = constant for any constant is a solution to the equation ∇2u = 0 . To ensure a unique solution we provide additional equations , called “ boundary conditions ” . Several boundary conditions arise very naturally in physical problems . A very common one is the Dirichlet boundary condition , where we pick some subset G ⊂ Rk and fix the values of the function on G to some fixed value b , u ( x ) = b ( x ) , for all x ∈ G where the function b is usually clear from the underlying physical problem . As in previous literature , we refer to G as the geometry of the problem , and b as the boundary value . We refer to the pair ( G , b ) as the boundary condition . In this paper , we only consider linear PDEs and boundary conditions that have unique solutions . 2.2 FINITE DIFFERENCE METHOD . Most real-world PDEs do not admit an analytic solution and must be solved numerically . The first step is to discretize the solution space F from Rk → R into Dk → R , where D is a discrete subset of R. When the space is compact , it is discretized into an n×n×n · · · ( k many ) uniform Cartesian grid with mesh width h. Any function in F is approximated by its value on the nk grid points . We denote the discretized function as a vector u in Rnk . In this paper , we focus on 2D problems ( k = 2 ) , but the strategy applies to any dimension . We discretize all three terms in the equation Au = f and boundary condition ( G , b ) . The PDE solution u is discretized such that ui , j = u ( xi , yj ) corresponds to the value of u at grid point ( xi , yj ) . We can similarly discretize fand b . In linear PDEs , the linear operator A is a linear combination of partial derivative operators . For example , for the Poisson equation A = ∇2 = ∑ i ∂2 ∂x2i . Therefore we can first discretize each partial derivative , then linearly combine the discretized partial derivatives to obtain a discretized A. Finite difference is a method that approximates partial derivatives in a discretized space , and as mesh width h → 0 , the approximation approaches the true derivative . For example , ∂ 2 ∂x2 u can be discretized in 2D as ∂ 2 ∂x2 u ≈ 1 h2 ( ui−1 , j − 2ui , j + ui+1 , j ) , the Laplace operator in 2D can be correspondingly approximated as : ∇2u = ∂ 2u ∂x2 + ∂2u ∂y2 ≈ 1 h2 ( ui−1 , j + ui+1 , j + ui , j−1 + ui , j+1 − 4ui , j ) ( 1 ) After discretization , we can rewrite Au = fas a linear matrix equation Au = f ( 2 ) where u , f ∈ Rn2 , and A is a matrix in Rn2×n2 ( these are n2 dimensional because we focus on 2D problems ) . In many PDEs such as the Poisson and Helmholtz equation , A is sparse , banded , and symmetric . 2.3 BOUNDARY CONDITION . We also need to include the boundary condition u ( x ) = b ( x ) for all x ∈ G. If a discretized point ( xi , yj ) belongs to G , we need to fix the value of ui , j to bi , j . To achieve this , we first define e ∈ { 0 , 1 } n2 to be a vector of 0 ’ s and 1 ’ s , in which 0 indicates that the corresponding point belongs to G. Then , we define a “ reset ” matrix G = diag ( e ) , a diagonal matrix Rn2 → Rn2 such that ( Gu ) i , j = { ui , j ( xi , yj ) 6∈ G 0 ( xi , yj ) ∈ G ( 3 ) Intuitively G ” masks ” every point in G to 0 . Similarly , I − G can mask every point not in G to 0 . Note that the boundary values are fixed and do not need to satisfy Au = f . Thus , the solution u to the PDE under geometry G should satisfy : G ( Au ) = Gf ( I −G ) u = ( I −G ) b ( 4 ) The first equation ensures that the interior points ( points not in G ) satisfy Au = f , and the second ensures that the boundary condition is satisfied . To summarize , ( A , G , f , b , n ) is our PDE problem , and we first discretize the problem on an n× n grid to obtain ( A , G , f , b , n ) . Our objective is to obtain a solution u that satisfies Eq . ( 4 ) , i.e . Au = f for the interior points and boundary condition ui , j = bi , j , ∀ ( xi , yj ) ∈ G . 2.4 ITERATIVE SOLVERS . A linear iterative solver is defined as a function that inputs the current proposed solution u ∈ Rn2 and outputs an updated solution u′ . Formally it is a function Ψ : Rn2 → Rn2 that can be expressed as u′ = Ψ ( u ) = Tu+ c ( 5 ) where T is a constant update matrix and c is a constant vector . For each iterator Ψ there may be special vectors u∗ ∈ Rn2 that satisfy u∗ = Ψ ( u∗ ) . These vectors are called fixed points . The iterative solver Ψ should map any initial u0 ∈ Rn2 to a correct solution of the PDE problem . This is formalized in the following theorem . Definition 1 ( Valid Iterator ) . An iterator Ψ is valid w.r.t . a PDE problem ( A , G , f , b , n ) if it satisfies : a ) Convergence : There is a unique fixed point u∗ such that Ψ converges to u∗ from any initialization : ∀u0 ∈ Rn2 , limk→∞Ψk ( u0 ) = u∗ . b ) Fixed Point : The fixed point u∗ is the solution to the linear systemAu = f under boundary condition ( G , b ) . Convergence : Condition ( a ) in Definition 1 is satisfied if the matrix T is convergent , i.e . T k → 0 as k → ∞ . It has been proven that T is convergent if and only if the spectral radius ρ ( T ) < 1 ( Olver , 2008 ) : Theorem 1 . ( Olver , 2008 , Prop 7.25 ) For a linear iterator Ψ ( u ) = Tu + c , Ψ converges to a unique stable fixed point from any initialization if and only if the spectral radius ρ ( T ) < 1 . Proof . See Appendix A . It is important to note that Condition ( a ) only depends on T and not the constant c. Fixed Point : Condition ( b ) in Definition 1 contains two requirements : satisfy Au = f , and the boundary condition ( G , b ) . To satisfyAu = f a standard approach is to design Ψ by matrix splitting : split the matrix A into A = M − N ; rewrite Au = f as Mu = Nu + f ( LeVeque , 2007 ) . This naturally suggests the iterative update u′ = M−1Nu+M−1f ( 6 ) Because Eq . ( 6 ) is a rewrite of Au = f , stationary points u∗ of Eq . ( 6 ) satisfy Au∗ = f . Clearly , the choices of M and N are arbitrary but crucial . From Theorem 1 , we must choose M such that the update converges . In addition , M−1 must easy to compute ( e.g. , diagonal ) . Finally we also need to satisfy the boundary condition ( I − G ) u = ( I − G ) b in Eq.4 . After each update in Eq . ( 6 ) , the boundary condition could be violated . We use the “ reset ” operator defined in Eq . ( 3 ) to “ reset ” the values of ui , j to bi , j by Gu+ ( I −G ) b . The final update rule becomes u′ = G ( M−1Nu+M−1f ) + ( I −G ) b ( 7 ) Despite the added complexity , it is still a linear update rule in the form of u′ = Tu + c in Eq . ( 5 ) : we have T = GM−1N and c = GM−1f + ( 1−G ) b . As long as M is a full rank diagonal matrix , fixed points of this equation satisfies Eq . ( 4 ) . In other words , such a fixed point is a solution of the PDE problem ( A , G , f , b , n ) . Proposition 1 . If M is a full rank diagonal matrix , and u∗ ∈ Rn2×n2 satisfies Eq . ( 7 ) , then u∗ satisfies Eq . ( 4 ) .

Summary: This paper develops a method to accelerate the finite difference method in solving PDEs. Basically, the paper proposes a revised framework for fixed point iteration after discretization. The framework introduces a free linear operator --- the choice of the linear operator will influence the convergence rate. The paper uses a deep linear neural network to learn a good operator. Experimental results on Poisson equations show that the learned operator achieves significant speed-ups. The paper also gives theoretical analysis about the range of the valid linear operator (convex open set) and guarantees of the generalization for the learned operator.