APEBench: A Benchmark for Autoregressive Neural Emulators of PDEs

TL;DR

A novel taxonomy for unrolled training is proposed for unrolled training and a unique identifier for PDE dynamics that directly relates to the stability criteria of classical numerical methods is introduced.

Abstract

We introduce the Autoregressive PDE Emulator 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.

Figures (18)

• Figure 1: APEBench provides an efficient pseudo-spectral solver to simulate $46$ PDE dynamics across one to three spatial dimensions. Shown are examples visualized with APEBench's custom volume renderer.
• Figure 2: Test rollout performance of linear convolution emulators when learned with different training rollout lengths relative to the performance of a FOU method. All learned emulators surpass the numerical method in an initial operating regime. More unrolling improves long-term accuracy for a small sacrifice in short-term performance.
• Figure 3: (a) Performance of various neural emulator architectures on a 1D advection problem with increasing difficulty ($\gamma_1 = \text{CFL}$). If the demands on the receptive field (given by $\gamma_1$) are not fulfilled, emulators diverge immediately. (b) Unrolling improves accuracy at the highest difficulty.
• Figure 4: Comparison of training methodologies for a ResNet emulator on three nonlinear 1D dynamics. Emulators benefit from rollout training, the strongest visible for the KdV case. Diverted-Chain offers the advantage of long-term accuracy without sacrificing short-term performance.
• Figure 5: ResNet and FNO either as full prediction emulators or neural-hybrid emulators for 2D advection ($\gamma_1=10.5$) with a coarse solver doing 10% or 50% of the difficulty. The geometric mean of the rollout error over 100 time steps is shown. Training with unrolling benefits the ResNet yet only shows marginal improvement for the FNO. The ResNet can work in symbiosis with a coarse simulator.