Generalizing PDE Emulation with Equation-Aware Neural Operators

2025-11-12

AI Summary

Traditional numerical methods for solving partial differential equations (PDEs) incur high computational costs, while existing deep learning surrogate models exhibit poor generalization—particularly when extrapolating to unseen PDE structures or coefficients. Method: We propose an equation-aware neural operator framework that explicitly encodes PDE-specific differential operators and coefficient fields into a conditional vector, enabling the model to intrinsically capture structural inductive biases of the governing equations. Contribution/Results: Our approach achieves, for the first time, stable out-of-distribution generalization across PDE types, parametric variations, and equation forms—including reliable long-term rollout stability. Evaluated on the APEBench 1D PDE benchmark suite against four state-of-the-art baselines, our model significantly outperforms equation-specific surrogates, maintaining high accuracy and robustness on both unseen parameter regimes and entirely novel PDE families.

Technology Category

• Machine Learning: Deep Neural Architectures and Foundation Models
• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation

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.

Problem

• Solving PDEs with traditional methods is computationally expensive
• Existing neural surrogates specialize only in single fixed-parameter PDEs
• Developing generalized emulators for unseen PDEs with varying parameters

Innovation

• Equation-aware neural operators generalize to unseen PDEs
• Conditioning neural models on PDE term encodings
• Achieves strong performance on held-out parameter sets