APEBench: A Benchmark for Autoregressive Neural Emulators of PDEs

Felix Koehler, Simon Niedermayr, rüdiger westermann, Nils Thuerey

Keywords: PDEs, Neural Operators, Neural Solvers, Numerical Methods, Autoregressive Emulators, Scientific Machine Learning

TL;DR: A JAX-based benchmark for time-dependent neural PDE emulation with a focus on long-term accuracy and support for more than 46 distinct dynamics across 1D, 2D, and 3D.

Abstract: We introduce the **A**utoregressive **P**DE **E**mulator Benchmark (APEBench), a comprehensive benchmark suite to evaluate autoregressive neural emulators for solving partial differential equations. APEBench is based on JAX and provides a seamlessly integrated differentiable simulation framework employing efficient pseudo-spectral methods, enabling 46 distinct PDEs across 1D, 2D, and 3D. Facilitating systematic analysis and comparison of learned emulators, we propose a novel taxonomy for unrolled training and introduce a unique identifier for PDE dynamics that directly relates to the stability criteria of classical numerical methods. APEBench enables the evaluation of diverse neural architectures, and unlike existing benchmarks, its tight integration of the solver enables support for differentiable physics training and neural-hybrid emulators. Moreover, APEBench emphasizes rollout metrics to understand temporal generalization, providing insights into the long-term behavior of emulating PDE dynamics. In several experiments, we highlight the similarities between neural emulators and numerical simulators. The code is available at [github.com/tum-pbs/apebench](https://github.com/tum-pbs/apebench) and APEBench can be installed via `pip install apebench`.

Submission Number: 767