Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Generalizing PDE Emulation with Equation-Aware Neural Operators

Qian-Ze Zhu, Paul Raccuglia, Michael P. Brenner

TL;DR

This work tackles generalizing PDE emulation across unseen equations and parameter regimes by introducing an equation-aware framework that conditions a neural operator on a compact 7-term equation encoding $c$. It trains four architectures—PI-FNO-UNET, LSC-FNO, PINO, and LC—to map $(u(t), c)$ to $u(t+dt)$, leveraging conditioning mechanisms like FiLM, spectral gating, and PDE residual regularization. The models achieve strong parameter generalization on 1D PDEs from APEBench and demonstrate zero-shot generalization to an unseen PDE (Burgers' equation), with the Learned Correction variant often performing best. An AI-assisted Tree Search accelerates architecture discovery, underscoring potential for automated creation of universal PDE emulators and faster scientific discovery, while noting challenges in extending to higher dimensions.

Abstract

Solving partial differential equations (PDEs) can be prohibitively expensive using traditional numerical methods. Deep learning-based surrogate models typically specialize in a single PDE with fixed parameters. We present a framework for equation-aware emulation that generalizes to unseen PDEs, conditioning a neural model on a vector encoding representing the terms in a PDE and their coefficients. We present a baseline of four distinct modeling technqiues, trained on a family of 1D PDEs from the APEBench suite. Our approach achieves strong performance on parameter sets held out from the training distribution, with strong stability for rollout beyond the training window, and generalization to an entirely unseen PDE. This work was developed as part of a broader effort exploring AI systems that automate the creation of expert-level empirical software for scorable scientific tasks. The data and codebase are available at https://github.com/google-research/generalized-pde-emulator.

Generalizing PDE Emulation with Equation-Aware Neural Operators

TL;DR

This work tackles generalizing PDE emulation across unseen equations and parameter regimes by introducing an equation-aware framework that conditions a neural operator on a compact 7-term equation encoding . It trains four architectures—PI-FNO-UNET, LSC-FNO, PINO, and LC—to map to , leveraging conditioning mechanisms like FiLM, spectral gating, and PDE residual regularization. The models achieve strong parameter generalization on 1D PDEs from APEBench and demonstrate zero-shot generalization to an unseen PDE (Burgers' equation), with the Learned Correction variant often performing best. An AI-assisted Tree Search accelerates architecture discovery, underscoring potential for automated creation of universal PDE emulators and faster scientific discovery, while noting challenges in extending to higher dimensions.

Abstract

Solving partial differential equations (PDEs) can be prohibitively expensive using traditional numerical methods. Deep learning-based surrogate models typically specialize in a single PDE with fixed parameters. We present a framework for equation-aware emulation that generalizes to unseen PDEs, conditioning a neural model on a vector encoding representing the terms in a PDE and their coefficients. We present a baseline of four distinct modeling technqiues, trained on a family of 1D PDEs from the APEBench suite. Our approach achieves strong performance on parameter sets held out from the training distribution, with strong stability for rollout beyond the training window, and generalization to an entirely unseen PDE. This work was developed as part of a broader effort exploring AI systems that automate the creation of expert-level empirical software for scorable scientific tasks. The data and codebase are available at https://github.com/google-research/generalized-pde-emulator.

Paper Structure

This paper contains 16 sections, 3 equations, 2 figures, 1 table.

Table of Contents

  1. Introduction
  2. Methodology
  3. Experiments and Results
  4. Generalization to Unseen Parameters
  5. Generalization to an Unseen PDE
  6. Automated Discovery via Tree Search
  7. Conclusion
  8. Experimental Setup
  9. Data generation
  10. Evaluation Metric
  11. Automated Discovery via Tree Search
  12. Model and Training Details
  13. Method 1: Physics-Informed U-Net with FNO layer (PI-FNO-UNet)
  14. Method 2: Latent Operator with Attention and Gating (LSC-FNO)
  15. Method 3: Physics-Informed Regularization (PINO)
  16. ...and 1 more sections

Figures (2)

  • Figure 1: Autoregressive rollout performance and OOD generalization. Mean normalized RMSE over 200 timesteps for KdV (a) and cKS (b), averaged over 30 random initializations. Red curves are the baselines from koehler2024apebench. The generalized models are evaluated on both the ID parameter (green curves) used for training and an OOD parameter via zero-shot inference (blue curves). Representative rollouts from PI-FNO-UNET, evaluated on the OOD parameter for KdV (c) and cKS (d). Zero-shot generalization error versus the dispersion coefficient for KdV (e) and hyper-diffusion coefficient for cKS (f). The error is the aggregated nRMSE (Geometric Mean) over the first 100 steps. The shaded area represents the parameter range used to train the generalized models; the vertical dotted line indicates the parameter used to train the baselines koehler2024apebench.
  • Figure 2: Zero-shot generalization on Burgers' equation. (a) nRMSE over 200 autoregressive time steps, averaged over 30 random initial conditions. Red curves are the baselines koehler2024apebench. Blue curves are our generalized models. (b) A representative autoregressive rollout from the LSC-FNO model.