Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Learning Neural PDE Solvers with Convergence Guarantees

Jun-Ting Hsieh, Shengjia Zhao, Stephan Eismann, Lucia Mirabella, Stefano Ermon

TL;DR

The paper tackles speeding up linear PDE solvers while preserving convergence guarantees by learning a correction to a standard iterative update. It builds on a valid base solver $\Psi$ and introduces a learned linear operator $H$ so that the new update $\Phi_H(u)=\Psi(u)+G H w$ remains convergent and fixes the same solution, with $w=\Psi(u)-u$. Trained on a single geometry, the approach generalizes to different boundary conditions and grid sizes, achieving 2-3x speedups for 2D Poisson problems and competitive wall-clock performance versus FEniCS, including a large GPU speedup of ~30x. The authors also interpret $H$ as approximating $T(I-T)^{-1}$ to accelerate convergence and explore two linear-deep-network architectures (Conv and U-Net) to realize the correction efficiently.

Abstract

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.

Learning Neural PDE Solvers with Convergence Guarantees

TL;DR

The paper tackles speeding up linear PDE solvers while preserving convergence guarantees by learning a correction to a standard iterative update. It builds on a valid base solver and introduces a learned linear operator so that the new update remains convergent and fixes the same solution, with . Trained on a single geometry, the approach generalizes to different boundary conditions and grid sizes, achieving 2-3x speedups for 2D Poisson problems and competitive wall-clock performance versus FEniCS, including a large GPU speedup of ~30x. The authors also interpret as approximating to accelerate convergence and explore two linear-deep-network architectures (Conv and U-Net) to realize the correction efficiently.

Abstract

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.

Paper Structure

This paper contains 24 sections, 5 theorems, 22 equations, 2 figures, 1 table.

Table of Contents

  1. Introduction
  2. Background
  3. Linear PDEs
  4. Finite Difference Method
  5. Boundary Condition
  6. Iterative Solvers
  7. Jacobi Method
  8. Multigrid Method
  9. Learning Fast and Provably Correct Iterative PDE Solvers
  10. Formulation
  11. Training and Generalization
  12. Interpretation of H
  13. Linear Deep Networks
  14. Experiments
  15. Setting
  16. ...and 9 more sections

Key Result

Theorem 1

olver2008numerical For a linear iterator $\Psi(u) = Tu + c$, $\Psi$ converges to a unique stable fixed point from any initialization if and only if the spectral radius $\rho(T) < 1$.

Figures (2)

  • Figure 1: The ground truth solutions of examples in different settings. We only train our models on the square domain, and we test on all 4 settings.
  • Figure 2: CPU runtime comparisons of our model with the FEniCS model. Our method is comparable or faster than the best solver in FEniCS in all cases. When run on GPU, our solver provides an additional $30\times$ speedup.

Theorems & Definitions (14)

  • Definition 1: Valid Iterator
  • Theorem 1
  • proof
  • Proposition 1
  • Lemma 1
  • proof
  • Theorem 2
  • proof
  • Proposition 2
  • proof
  • ...and 4 more