PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Neural PDE Solvers with Convergence Guarantees

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

arXiv:1906.01200v1math.NAstat.COstat.ML

TL;DR

Existing PDE solvers can be suboptimal for specific domains, motivating learned solvers that retain correctness and convergence guarantees. The paper learns modifications to an existing iterative solver, achieving faster convergence and generalization from one geometry to varied geometries and boundary conditions.

  • Problem

    Existing generic PDE solvers may be far from optimal for specific domains, while high-performing update rules can be difficult to design by hand.

  • Method

    The method learns a parameterized neural modification of an existing iterative solver while preserving its fixed point for every parameter choice.

  • Results

    The learned solver generalizes from a single training geometry to varied geometries and boundary conditions and achieves faster convergence than standard and highly optimized solvers.

  • Takeaways & Limitations

    Learning solver updates can provide provably correct iterative PDE solvers that retain broad generalization while improving convergence speed.

  • Takeaways & Limitations

    The method applies to PDEs with existing linear iterative solvers, limiting its current applicability to problems lacking such solvers.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Existing PDE solvers are effective but can be suboptimal for particular domains, motivating learned updates. The proposed solver modifies a standard iterative method while preserving correctness and convergence guarantees, generalizing across geometries and boundary conditions with substantial speedups.

  • Existing solver performance varies across application domains, so no method uniformly dominates and generic solvers may be suboptimal for specific domains.
  • The learned iterator modifies updates from an existing solver while preserving its fixed point for every parameter choice.This design provides correctness guarantees while allowing training to improve convergence speed.
  • Training on a single problem instance yields correct generalization across varied geometries and boundary conditions without observable performance loss.
  • 2-3× speedup in multiply-add operations is achieved over standard iterative solvers, with faster wall-clock CPU performance than state-of-the-art FEniCS solvers.GPU execution provides an additional 30× speedup, and the approach remains simpler than deeply optimized baselines.

2 BACKGROUND

Linear PDEs are discretized into finite-dimensional systems and solved with iterative updates whose convergence and fixed points can be analyzed. Boundary resetting preserves prescribed values, while Jacobi provides a valid baseline for Poisson problems.

  • 2.1 LINEAR PDES: Linear PDEs seek a function u satisfying Au = f, with boundary conditions added when the differential equation does not uniquely determine u.The paper considers linear PDEs and boundary conditions with unique solutions.
  • 2.2 FINITE DIFFERENCE METHOD: Finite-difference discretization represents a function by values on a uniform Cartesian grid, converting the PDE into a matrix equation.In two dimensions, the discretized solution and right-hand side have n^2 components, while the operator becomes a matrix.
  • 2.3 BOUNDARY CONDITION: Boundary conditions are enforced by masking boundary points and resetting their values to b after each iterative update.The interior points satisfy Au = f, while boundary points satisfy the prescribed values.
  • 2.4 ITERATIVE SOLVERS: A linear iterative solver maps the current solution u to an updated solution u′ through an affine rule u′ = Tu + c.Fixed points are vectors unchanged by the iterator and are required to solve the PDE under its boundary conditions.
  • 2.4 ITERATIVE SOLVERS: Convergence from any initialization is guaranteed when the update matrix has spectral radius ρ(T) < 1, while fixed-point validity links the limit to the PDE solution.For the boundary-aware update, a full-rank diagonal M ensures that fixed points satisfy the PDE system and boundary conditions.
  • 2.4.1 JACOBI METHOD: For Poisson equations, Jacobi uses M = I and produces a valid iterator for any geometry because ρ(T) < 1 and its fixed points satisfy the PDE.Jacobi updates can be implemented as convolutional neural-network layers followed by a boundary reset.

3 LEARNING FAST AND PROVABLY CORRECT ITERATIVE PDE SOLVERS

The paper learns a neural correction to an existing PDE iterator, preserving its fixed points while targeting faster convergence. The resulting method generalizes beyond training conditions and is supported by theoretical and empirical guarantees.

  • 3.1 FORMULATION: The method learns a parameterized correction to a valid linear solver, preserving the original solver’s fixed points by construction.The correction is applied to the existing update, and any fixed point of the original iterator remains a fixed point of the learned iterator.
  • 3.1 FORMULATION: The learned iterator includes the standard solver and can reproduce two standard iterations with comparable convolutional cost.Setting H = 0 recovers Ψ; setting H = T yields Ψ² using two convolutions, matching the learned iterator’s two-convolution computation.
  • 3.2 TRAINING AND GENERALIZATION: Training optimizes H to approach the ground-truth solution after a chosen number of learned iterations from random initialization.The training set contains problem instances, with ground-truth solutions obtained from the existing solver and k uniformly selected from [1, 20] in the experiments.
  • 3.2 TRAINING AND GENERALIZATION: The learned iterator has a convex open region of parameters with spectral norm below 1, which makes searching for rapidly convergent updates tractable.For symmetric iterators, spectral norm equals spectral radius, connecting this parameter region to the convergence criterion.
  • 3.2 TRAINING AND GENERALIZATION: Generalization across forcing terms and boundary conditions is guaranteed under the stated proposition, while generalization across geometries and grid sizes is only empirical.Although trained on one geometry and grid size, the iterator converged to the correct solution across varied geometries and grid sizes in the experiments.
  • 3.3 INTERPRETATION OF H: The learned correction targets T(I −T)−1, so better approximation of this operator yields faster convergence, although exact representation is generally impossible with convolutional H.A perfect approximation would converge in one iteration; the inverse is generally dense, whereas the implementation uses a linear deep convolutional network.

4 EXPERIMENTS

The experiments evaluate learned iterative solvers for the 2D Poisson equation across grid sizes, geometries, boundary conditions, and solver baselines. Models trained on square domains converge correctly while reducing computation and matching or exceeding FEniCS runtime.

  • 4.1 SETTING: The evaluation uses the 2D Poisson equation with Dirichlet boundary conditions and selects Jacobi as the standard solver baseline.
  • 4.1 SETTING: Models train on simple square domains and test on larger grids, L-shaped and cylindrical geometries, and nonzero forcing terms.Testing includes 256×256 grids for models trained on 64 × 64 grids.
  • 4.2 EVALUATION METRICS: Convergence is assessed by computation required to reduce error below a threshold, using convolutional layers on GPU and multiply-add operations on CPU.The multiply-add metric is biased in favor of Jacobi because CPU convolution implementations are less practical.
  • 4.3 CONV MODEL: ∼5× fewer layers and ∼2.5× fewer multiply-add operations make Conv3 faster than Jacobi while converging to the correct solution in all settings.
  • 4.4 U-NET MODEL: U-Net models outperform Multigrid models across settings, with U-Net3 converging faster than U-Net2 despite U-Net2 having lower relative computation cost against Multigrid2.
  • 4.5 COMPARISON WITH FENICS: The learned solver is comparable or faster than FEniCS in CPU wall-clock time and gains an additional 30× speedup on GPU.FEniCS uses minimal residual with an algebraic multigrid preconditioner as its fastest measured configuration.

5 RELATED WORK

Prior deep-learning PDE solvers directly generate solutions but lack correctness guarantees and generalization to arbitrary grids and boundary conditions. The paper distinguishes its approach from work on other PDEs that may lack existing linear iterative solvers.

  • Previous Poisson-solving deep networks directly generate solutions without correctness guarantees or generalization to arbitrary grid sizes and boundary conditions.
  • Many related studies address nonlinear PDEs, which may not have standard linear iterative solvers required by this method.The paper identifies extending the approach to different PDEs as future work.

6 CONCLUSION

The paper presents a learned iterative PDE solver that preserves correctness through a fixed-point guarantee and generalizes beyond its simple training domains. It reports correct convergence and speedups against standard and highly optimized solvers.

  • The method improves an existing standard PDE solver while theoretically guaranteeing that the correct solution is the iterator’s fixed point.
  • Training on simple domains generalizes across grid sizes, geometries, and boundary conditions, with speedups over standard solvers including optimized FEniCS implementations.

A PROOFS

The proofs characterize convergence through spectral properties of linear iterators and show how the learned update preserves validity across forcing terms and boundary values. They also establish convexity properties for selecting valid update operators.

  • Convergence of linear iterators: A linear iterator converges from any initialization to a unique stable fixed point exactly when its spectral radius is below 1.
  • Correctness: The learned iterator’s fixed point satisfies the PDE solution condition when the diagonal matrix used in the construction has full rank.
  • Validity of learned updates: The spectral norm of the learned iterator is convex in H, making the set of H values with norm below 1 a convex open set.
  • Generalization across problem data: For fixed A, G, n, and H, validity for one forcing term and boundary value implies validity for all f and b because the linear update operator is independent of c.

B PROOF OF CONVERGENCE OF JACOBI METHOD

The proof establishes that the Poisson-equation Jacobi update matrix has spectral radius below one for any geometry matrix G, guaranteeing convergence.

  • The Poisson-equation Jacobi update matrix is T = G(I −A).
  • The spectral-radius bound, symmetry of I −A, known convergence of I −A, and ∥G∥2 = 1 combine with sub-multiplicativity to bound T.These matrix properties provide the proof's key steps.
  • ρ(T) < 1 holds for any G.
  • Therefore, the standard Jacobi method is valid for the Poisson equation under any geometry.
Loading 1906.01200v1…