Neural Preconditioning Operator for Efficient PDE Solves

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.