Partial differential equations (PDEs) are widely used across the physical and computational sciences. Decades of research and engineering went into designing fast iterative solution methods. Existing solvers are general purpose, but may be sub-optimal for specific classes of problems. In contrast to existing hand-crafted solutions, we propose an approach to learn a fast iterative solver tailored to a specific domain. We achieve this goal by learning to modify the updates of an existing solver using a deep neural network. Crucially, our approach is proven to preserve strong correctness and convergence guarantees. After training on a single geometry, our model generalizes to a wide variety of geometries and boundary conditions, and achieves 2-3 times speedup compared to state-of-the-art solvers. 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 BID20 BID12 . Other domains that benefited significantly include learned optimizers BID1 and learned data structures BID9 . 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 BID21 BID6 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 BID13 , 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. We presented a method to learn an iterative solver for PDEs that improves on an existing standard solver. The correct solution is theoretically guaranteed to be the fixed point of our iterator. We show that our model, trained on simple domains, can generalize to different grid sizes, geometries and boundary conditions. It converges correctly and achieves significant speedups compared to standard solvers, including highly optimized ones implemented in FEniCS.A PROOFS Theorem 1. For a linear iterator Ψ(u) = T u + c, Ψ converges to a unique stable fixed point from any initialization if and only if the spectral radius ρ(T ) < 1.Proof. Suppose ρ(T ) < 1, then (I−T ) −1 must exist because all eigenvalues of I−T must be strictly positive. Let u * = (I − T ) −1 c; this u * is a stationary point of the iterator Ψ, i.e. u * = T u * + c. For DISPLAYFORM0 Since ρ(T ) < 1, we know BID11 , which means the error e k → 0. Therefore, Ψ converges to u * from any u 0 . DISPLAYFORM1 Now suppose ρ(T ) ≥ 1. Let λ 1 be the largest absolute eigenvalue where ρ(T ) = |λ 1 | ≥ 1, and v 1 be its corresponding eigenvector. We select initialization u 0 = u * + v 1 , then e 0 = v 1 . Because |λ 1 | ≥ 1, we have |λ k 1 | ≥ 1, then T k e 0 = λ k 1 v 1 → k→∞ 0 However we know that under a different initializationû 0 = u * , we haveê 0 = 0, so T kê0 = 0. Therefore the iteration cannot converge to the same fixed point from different initializations u 0 and u 0 .Proposition 1 If M is a full rank diagonal matrix, and u * ∈ R n 2 ×n 2 satisfies Eq. (7) , then u * satisfies Eq. (4).Proof of Proposition 1. Let u * be a fixed point of Eq. (7) then The latter equation is equivalent to GM −1 (Au * − f ) = 0. If M is a full rank diagonal matrix, this implies G(Au * − f ) = 0, which is GAu * = Gf . Therefore, u * satisfies Eq.(4).Theorem 2. For fixed G, f, b, n, the spectral norm of Φ H (u; G, f, b, n) is a convex function of H, and the set of H such that the spectral norm of Φ H (u; G, f, b, n) < 1 is a convex open set.Proof. As before, denote Ψ(u) = T u + c. Observe that Φ H (u; G, f, b, n) = T u + c + GH(T u + c − u) = (T + GHT − GH)u + GHc + cThe spectral norm · 2 is convex with respect to its argument, and (T + GHT − GH) is linear in H. Thus, T + GHT − GH 2 is convex in H as well. Thus, under the condition that T + GHT − GH 2 < 1, the set of H must be convex because it is a sub-level set of the convex function T + GHT − GH 2 .To prove that it is open, observe that · 2 is a continuous function, so T + GHT − GH 2 is a continuous map from H to the spectral radius of Φ H . If we consider the set of H such that T + GHT − GH 2 < 1, this set is the preimage of (− , 1) for any > 0. As (− , 1) is open, its preimage must be open. Proposition 2. For fixed A, G, n and fixed H, if for some f 0 , b 0 , Φ H (u; G, f 0 , b 0 , n) is valid for the PDE problem (A, G, f 0 , b 0 , n), then for all f and b, the iterator Φ H (u; G, f, b, n) is valid for the PDE problem (A, G, f, b, n).Proof. From Theorem 1 and Lemma 1, our iterator is valid if and only if ρ(T + GHT − GH) < 1. The iterator T + GHT − GH only depends on A, G, and is independent of the constant c in Eq. (18). Thus, the validity of the iterator is independent with f and b. Thus, if the iterator is valid for some f 0 and b 0 , then it is valid for any choice of f and b.
We learn a fast neural solver for PDEs that has convergence guarantees. Develops a method to accelerate the finite difference method in solving PDEs and proposes a revised framework for fixed point iteration after discretization. The authors propose a linear method for speeding up PDE solvers.