Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Yifei Sun

TL;DR

The paper tackles the limited out-of-distribution generalization of neural operators for nonlinear PDEs by distilling a differentiable spectral PDE solver into a compact neural operator, augmented with a minimax, PGD-style adversarial search in function space to reveal worst-case inputs. By enabling gradients to propagate through the differentiable solver, the adversarial sampling is informed by true solver sensitivity, improving sample efficiency and robustness. The approach is validated on Burgers and Navier–Stokes problems, showing substantial gains in OOD robustness while preserving the speed and small parameter footprint of neural operators. This framework offers a practical path toward universal PDE solvers by expanding the training distribution in function space through solver-guided adversarial distillation.

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.

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

TL;DR

The paper tackles the limited out-of-distribution generalization of neural operators for nonlinear PDEs by distilling a differentiable spectral PDE solver into a compact neural operator, augmented with a minimax, PGD-style adversarial search in function space to reveal worst-case inputs. By enabling gradients to propagate through the differentiable solver, the adversarial sampling is informed by true solver sensitivity, improving sample efficiency and robustness. The approach is validated on Burgers and Navier–Stokes problems, showing substantial gains in OOD robustness while preserving the speed and small parameter footprint of neural operators. This framework offers a practical path toward universal PDE solvers by expanding the training distribution in function space through solver-guided adversarial distillation.

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.
Paper Structure (44 sections, 68 equations, 59 figures, 3 tables, 4 algorithms)

This paper contains 44 sections, 68 equations, 59 figures, 3 tables, 4 algorithms.

Table of Contents

  1. Classical and Differentiable Numerical Solvers for Nonlinear PDEs.
  2. Neural Operators: Learning Operators in Function Space.
  3. Active learning and adversarial sampling for PINNs.
  4. PGD attack on FNO.
  5. Knowledge Distillation and Representation Alignment.
  6. Adversarial Training for Better Generalizability and Robustness
  7. Neural Operator Model $G$
  8. Classical Numerical Solver $g$
  9. Back Propagation and Derivative
  10. Choice of Loss Function
  11. Batch Derivative
  12. PGD in the $L_\infty$ ball
  13. PGD in the $L_2$ ball
  14. Min–Max Optimization Framework
  15. 1D Burgers Equation via Pseudospectral Method
  16. ...and 29 more sections

Figures (59)

  • Figure 1: Unnormalized and Normalized FNO 1D Burgers OOD testing RMSE and MAE
  • Figure 2: Unnormalized FNO 2D NS OOD testing RMSE
  • Figure 3: 1D Burgers PGD attack illustration (with solver)
  • Figure 4: 2D NS PGD attack illustration (with solver)
  • Figure 5: Loss change of round-by-round method
  • ...and 54 more figures