Safeguarded Neurally-Accelerated Preconditioning (SNAP) framework for Nonlinear Optimization

Abstract

We introduce the Safeguarded Neurally-Accelerated Preconditioning (SNAP) framework for hybridizing nonlinear optimization methods with pre-trained neural operators (NOs). SNAP targets parametric optimization problems arising from variational formulations of nonlinear partial differential equations (PDEs), where many parameter instances require repeated solves. Inspired by multilevel line-search methods, SNAP operates in cycles, combining classical optimization steps with a search direction generated by a pre-trained NO. Since NO-generated search directions carry no a priori descent guarantee, SNAP employs an a posteriori safeguard by accepting only corrections that satisfy prescribed descent and acceptance conditions. We prove global convergence of the SNAP algorithm and show that it preserves the convergence rate of the underlying optimization method independently of the employed NO architecture. Moreover, we also show that when the NO corrections are sufficiently accurate, SNAP converges faster, with the improvement directly related to the NO training error. To demonstrate SNAP's efficiency in practice, we run numerical experiments on a diverse set of nonlinear parametric optimization problems, including parametric minimal surface, Allen--Cahn, and nonlinear Helmholtz problems. The efficiency of SNAP is demonstrated by comparing two NO architectures, Fourier neural operator (FNO) and Transolver, and by using SNAP as a standalone solver and as a nonlinear preconditioner. Our numerical results show a substantial reduction in the number of required iterations, while demonstrating robustness and generalization across a wide range of problem parameters and mesh resolutions.