Spectrally Safe Neural Operator Warm-Starts for Large-Scale Newton Solvers

Neural operators serve as fast, data-driven surrogates for scientific modeling but typically rely on a monolithic, single-pass inference procedure that struggles to resolve high-frequency details, a limitation known as spectral bias. We introduce the Iterative Refinement Neural Operator (IRNO), which augments pre-trained operators with a learned refinement module iteratively applied via fixed-point iteration. IRNO decomposes the prediction into a coarse initialization followed by successive residual corrections, paralleling classical numerical solvers. Under local assumptions, we establish contraction of the induced operator, ensuring convergence to a unique fixed point. To explicitly target high-frequency errors, we propose a progressive spectral loss that adaptively increases penalty on high-frequency components over refinement steps during training. Across physical systems, IRNO consistently lowers error, with up to 56.05% improvement on turbulent flow. On Active Matter, spectral analysis reveals that, relative to base operator, the normalized error ratios decrease to 27.72-36.10% in low-, 5.07-6.68% in mid-, and 1.48-2.04% in high-frequencies, remaining stable beyond the trained iteration count. Code is available at https://github.com/xiaotianliu-dartmouth/Iterative_Refinement_Neural_Operator

Neural operators are widely used as fast surrogates for numerical PDE solvers, mapping input functions to solution functions. However, their generalizability and robustness are not yet clearly defined in the operator-learning setting, which differs from traditional adversarial robustness definitions. This paper studies the generalizability and robustness of a learned neural operator from a solver-integrated perspective, addressing the challenge that the output of a learned operator and a numerical solver tends to change in tandem under input perturbation. First, we formalize the definition of generalization and robustness through a model-solver error operator, identifying fixed-input model-solver loss as generalization metric, and norm-bounded adversarial attack loss increase and Jacobian-error function norm as robustness metric. Second, we identify the solver-integrated adversarial attack as appropriate for PDE operator learning and show why model-only or fixed-ground-truth attacks can be insufficient when the solver output also changes with the input. Third, we develop solver-integrated adversarial training methods for neural operators. Experiments on representative PDE benchmarks show that this solver-integrated adversarial training clearly improves both generalizability and robustness. Deeper solver integration yields more effective attacks, more informative samples, and more efficient training than less integrated alternatives. These results provide a general framework for robust operator training and automatic sample selection without heavy manual intervention. More broadly, the formulation applies to adversarial regression whenever a ground-truth oracle can evaluate, and ideally differentiate, the true input-output map; PDE operator learning is one such case.

Neural operators provide fast surrogates for partial differential equation (PDE) solvers, but their reliability can degrade for high-dimensional spatial inputs and inverse or repeated inference. State-only training constrains solution values but not the learned input--output response. We study sensitivity-constrained neural operators (SC-NOs), which augment standard training with sampled solver-derived Jacobian supervision. Selected sensitivities from differentiable solvers or discrete adjoints are matched during training, allowing response information to be amortized across minibatches without imposing the full Jacobian at every update. We evaluate SC-NO on advection--diffusion and RANS--Spalart--Allmaras benchmarks, input-dimensionality scaling tests, long-horizon autoregressive rollout, and a shallow-water Tohoku tsunami source-inversion case. Sensitivity supervision improves forward prediction and yields larger gains in gradient-based inverse reconstruction of distributed fields. Scaling experiments show an improved accuracy--cost tradeoff for high-dimensional gridded inputs, while ablations indicate that state values and Jacobian information provide complementary supervision. In the tsunami case, SC-FNO reconstructs gridded seafloor deformation from sparse early gauge observations and forecasts subsequent wave propagation in a near-real-time proof-of-concept workflow. These results support sampled sensitivity supervision as a practical way to improve neural PDE surrogates when forward accuracy, inverse stability, robustness, and computational cost must be considered together.