Neural Preconditioning Operator for Efficient PDE Solves

TL;DR

The paper presents the Neural Preconditioning Operator (NPO), a neural network-based method that accelerates PDE solvers by improving convergence and reducing computation time across various problems and mesh types.

Contribution

Introducing NPO, a neural operator that generalizes preconditioning for PDE solvers, combining algebraic multigrid with transformer architecture for improved efficiency.

Findings

NPO reduces iteration counts and runtime significantly.

NPO outperforms traditional and neural methods across resolutions.

NPO maintains robust convergence on large, irregular grids.

Abstract

We introduce the Neural Preconditioning Operator (NPO), a novel approach designed to accelerate Krylov solvers in solving large, sparse linear systems derived from partial differential equations (PDEs). Unlike classical preconditioners that often require extensive tuning and struggle to generalize across different meshes or parameters, NPO employs neural operators trained via condition and residual losses. This framework seamlessly integrates with existing neural network models, serving effectively as a preconditioner to enhance the performance of Krylov subspace methods. Further, by melding algebraic multigrid principles with a transformer-based architecture, NPO significantly reduces iteration counts and runtime for solving Poisson, Diffusion, and Linear Elasticity problems on both uniform and irregular meshes. Our extensive numerical experiments demonstrate that NPO outperforms traditional methods and contemporary neural approaches across various resolutions, ensuring robust convergence even on grids as large as 4096, far exceeding its initial training limits. These findings underscore the potential of data-driven preconditioning to transform the computational efficiency of high-dimensional PDE applications.