Generalizing PDE Emulation with Equation-Aware Neural Operators

TL;DR

This paper introduces a neural operator framework that generalizes PDE emulation to unseen equations by conditioning on PDE term encodings, enabling accurate and stable predictions beyond training data.

Contribution

It proposes a novel equation-aware neural operator that generalizes to unseen PDEs, outperforming traditional models limited to fixed equations.

Findings

Strong performance on unseen PDE parameter sets

Stable rollouts beyond training window

Generalization to completely new PDEs

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.