Towards Universal Solvers: Using PGD Attack in Active Learning to Increase Generalizability of Neural Operators as Knowledge Distillation from Numerical PDE Solvers

TL;DR

This paper introduces an adversarial distillation approach using PGD-based active sampling to enhance the out-of-distribution generalization of neural operators for PDEs, achieving robustness without sacrificing efficiency.

Contribution

It presents a novel adversarial teacher-student framework with PGD active sampling to improve neural operators' OOD robustness for PDE solving.

Findings

Significant improvement in OOD robustness of neural operators.

Maintains low parameter count and fast inference.

Effective on Burgers and Navier-Stokes systems.

Abstract

Nonlinear PDE solvers require fine space-time discretizations and local linearizations, leading to high memory cost and slow runtimes. Neural operators such as FNOs and DeepONets offer fast single-shot inference by learning function-to-function mappings and truncating high-frequency components, but they suffer from poor out-of-distribution (OOD) generalization, often failing on inputs outside the training distribution. We propose an adversarial teacher-student distillation framework in which a differentiable numerical solver supervises a compact neural operator while a PGD-style active sampling loop searches for worst-case inputs under smoothness and energy constraints to expand the training set. Using differentiable spectral solvers enables gradient-based adversarial search and stabilizes sample mining. Experiments on Burgers and Navier-Stokes systems demonstrate that adversarial distillation substantially improves OOD robustness while preserving the low parameter cost and fast inference of neural operators.