APEBench: Autoregressive PDE Emulator

Updated 7 July 2026

APEBench is an autoregressive PDE emulator benchmark that integrates differentiable simulations via JAX and pseudo-spectral methods.
It supports 46 distinct PDE dynamics across 1D, 2D, and 3D, enabling pure prediction, neural–hybrid correction, and reference-solver training.
The benchmark provides standardized rollout metrics and flexible training regimes to bridge surrogate modeling with classical numerical analysis.

Searching arXiv for papers on APEBench to ground the article in current literature. arxiv_search(query="APEBench benchmark autoregressive neural emulators PDEs", max_results=10) APEBench, the Autoregressive PDE Emulator Benchmark, is a JAX-based suite for systematically evaluating autoregressive neural emulators of time-dependent PDEs. It provides a seamlessly integrated differentiable simulation framework employing efficient pseudo-spectral methods, enabling 46 distinct PDEs across 1D, 2D, and 3D, and it is designed to support pure prediction, neural–hybrid correction, and differentiable reference-solver training (Koehler et al., 2024). In subsequent work, APEBench has also functioned as a canonical source of benchmark tasks for cost-aware evaluation of learned solvers, compact neural-operator design, and equation-aware generalization studies (Zhang et al., 14 May 2026, Shikhman, 16 Jun 2026, Zhu et al., 12 Nov 2025).

1. Purpose and benchmark scope

APEBench was introduced to fill the gap between purely data-driven benchmarks and differentiable-physics approaches. Its stated goal is to provide a common environment for systematically evaluating autoregressive neural emulators of time-dependent PDEs, while preserving tight integration with a high-performance reference solver (Koehler et al., 2024). This design makes the benchmark suitable not only for conventional supervised training, but also for workflows in which the simulator itself participates in training.

The benchmark is organized around the task of learning autoregressive solution operators. In the later equation-aware study, this task is summarized as follows: given the state $u(x,t)$ on a uniform spatial grid, a neural emulator must predict $u(x,t+\Delta t)$ repeatedly to roll out trajectories; the reference solvers are accurate, differentiable, and run “on the fly” during training (Zhu et al., 12 Nov 2025). A common misconception is therefore to treat APEBench as merely a fixed dataset. The benchmark description instead emphasizes its tight integration of the solver, support for differentiable physics training and neural-hybrid emulators, and an explicit focus on rollout metrics for temporal generalization (Koehler et al., 2024).

APEBench is also intended to make learned emulators comparable to classical numerical methods in a way that preserves numerical intuition. Its “unique normalized identifiers for dynamics” are defined so that they directly relate to stability criteria of classical numerical methods, rather than serving only as arbitrary dataset labels (Koehler et al., 2024). This suggests a benchmark philosophy in which surrogate modeling and numerical analysis are treated as closely connected rather than separate activities.

2. Numerical core and differentiable simulator

The numerical core of APEBench is a Fourier pseudo-spectral Exponential Time Differencing Runge–Kutta (ETDRK) solver that is fully differentiable (Koehler et al., 2024). The benchmark implements semi-linear PDEs of the form

$\partial_t u = \mathcal{L}u + \mathcal{N}(u),$

where $\mathcal{L}$ is a constant-coefficient linear operator and $\mathcal{N}(u)$ is a lower-order nonlinearity. Under periodic boundaries, $\mathcal{L}$ diagonalizes in Fourier space, allowing exact integration of the linear part.

The benchmark description gives the spectral derivative approximation as

$\partial_x^k u(x) \approx \mathcal{F}^{-1}\!\left[(ik)^k \mathcal{F}(u)\right].$

Using the integrating-factor formulation with time step $\Delta t$ , APEBench writes

$\hat{u}^{n+1} = \exp(\hat{L}\Delta t)\odot \hat{u}^{n} + \exp(\hat{L}\Delta t)\odot \int_0^{\Delta t} \exp(-\hat{L}\tau)\,\hat{N}(u(\tau))\,d\tau .$

An explicit ETDRK2 example is also provided: $\hat{u}^{*} = e^{\hat{L}\Delta t}\odot \hat{u}^{n} + \phi_1(\hat{L}\Delta t)\odot \hat{N}(\hat{u}^{n}),$

$u(x,t+\Delta t)$ 0

with

$u(x,t+\Delta t)$ 1

The reported complexity is

$u(x,t+\Delta t)$ 2

per step via FFTs, and the implementation is described as trivially auto-differentiable in JAX (Koehler et al., 2024). This numerical substrate is central to the benchmark’s role: it makes the reference dynamics both accurate enough for benchmarking and differentiable enough for simulator-in-the-loop training.

3. Catalog of supported dynamics and experiment identifiers

APEBench organizes its dynamics into a catalog of 46 scenarios derived from 20 equations, each instantiated in one or more spatial dimensions (Koehler et al., 2024). The benchmark groups them into linear and nonlinear dynamics.

Group	Representative dynamics	Coverage
Linear dynamics	Advection, Diffusion, Advection–Diffusion, Dispersion, Hyper-Diffusion, Unbalanced Advection, Diagonal Diffusion, Anisotropic Diffusion, Spatially Mixed Dispersion, Spatially Mixed Hyper-Diffusion	10 equations with multiple dimension instances
Nonlinear dynamics	Viscous Burgers, Burgers, Korteweg–de Vries, Kuramoto–Sivashinsky, Fisher–KPP, Gray–Scott, Swift–Hohenberg, Navier–Stokes (decaying turbulence), Navier–Stokes (Kolmogorov flow)	10 equations with multiple dimension instances

Among the listed nonlinear systems are the one-dimensional Korteweg–de Vries equation

$u(x,t+\Delta t)$ 3

the conservative Kuramoto–Sivashinsky equation

$u(x,t+\Delta t)$ 4

the Gray–Scott reaction–diffusion system

$u(x,t+\Delta t)$ 5

and two Navier–Stokes variants, including a streamfunction–vorticity decaying-turbulence form and a Kolmogorov-flow form (Koehler et al., 2024).

To identify scenarios, APEBench defines normalized linear and nonlinear coefficients

$u(x,t+\Delta t)$ 6

where $u(x,t+\Delta t)$ 7 is the classic CFL number for $u(x,t+\Delta t)$ 8. The benchmark states that, more generally, $u(x,t+\Delta t)$ 9 marks the stability limit of the most compact explicit finite-difference stencil. For nonlinear terms it defines

$\partial_t u = \mathcal{L}u + \mathcal{N}(u),$ 0

with $\partial_t u = \mathcal{L}u + \mathcal{N}(u),$ 1. The tuple $\partial_t u = \mathcal{L}u + \mathcal{N}(u),$ 2 is presented as an unambiguous experiment identifier directly tied to numerical-stability intuition (Koehler et al., 2024).

4. Training taxonomy, benchmark interfaces, and evaluation metrics

A central contribution of APEBench is its taxonomy for unrolled training. Denoting by $\partial_t u = \mathcal{L}u + \mathcal{N}(u),$ 3 the reference simulator for one step and by $\partial_t u = \mathcal{L}u + \mathcal{N}(u),$ 4 the neural emulator, and using trajectories $\partial_t u = \mathcal{L}u + \mathcal{N}(u),$ 5, the benchmark defines

$\partial_t u = \mathcal{L}u + \mathcal{N}(u),$ 6

where $\partial_t u = \mathcal{L}u + \mathcal{N}(u),$ 7 is the number of main-chain unrolled steps and $\partial_t u = \mathcal{L}u + \mathcal{N}(u),$ 8 is the branch length of the reference rollout (Koehler et al., 2024).

Three canonical regimes are distinguished. One-step supervised training sets $\partial_t u = \mathcal{L}u + \mathcal{N}(u),$ 9, yielding

$\mathcal{L}$ 0

Supervised unrolling uses $\mathcal{L}$ 1, corresponding to full rollout versus precomputed reference. Diverted-chain training uses $\mathcal{L}$ 2 and $\mathcal{L}$ 3, so that at each main-chain step the one-step reference $\mathcal{L}$ 4 is compared to $\mathcal{L}$ 5. The benchmark further notes that by cutting gradients through the main chain or reference branch one can recover variants such as the “pushforward trick” or implicit methods (Koehler et al., 2024).

APEBench also exposes interfaces for neural–hybrid emulators. With a coarse solver $\mathcal{L}$ 6 and a corrector network $\mathcal{L}$ 7, it defines sequential correction

$\mathcal{L}$ 8

and parallel correction

$\mathcal{L}$ 9

Both layouts can be trained under any unrolling taxonomy (Koehler et al., 2024).

For evaluation, APEBench emphasizes autoregressive rollout errors. Its primary metric is the mean normalized RMSE over $\mathcal{N}(u)$ 0 samples and $\mathcal{N}(u)$ 1 grid points,

$\mathcal{N}(u)$ 2

and temporal aggregation is given by the geometric mean

$\mathcal{N}(u)$ 3

The benchmark also supplies Fourier-space RMSE/MAE, Sobolev $\mathcal{N}(u)$ 4 metrics, correlation, and band-limited metrics (Koehler et al., 2024). This metric design reflects the benchmark’s stated emphasis on long-term rollout behavior rather than only one-step fit.

5. Architectures shipped with the benchmark and reported benchmark findings

APEBench ships dimension-agnostic implementations, via Equinox, of ConvNet, ResNet, UNet, Dilated ResNet, and Fourier Neural Operator (FNO) architectures (Koehler et al., 2024). The benchmark states that all support Dirichlet, Neumann, and periodic boundary conditions. This architecture layer is intended to make comparisons across PDE families and dimensions systematic rather than ad hoc.

The benchmark paper reports several characteristic findings. In a motivational advection experiment with $\mathcal{N}(u)$ 5, a learned 2-tap linear convolution outperforms the first-order upwind scheme for approximately the first 13 steps despite no consistency guarantee; training with 20-step unrolling pushes long-term error lower, with only an 11% short-term penalty, and the learned stencils gravitate towards first-order upwind as unroll length increases (Koehler et al., 2024). In one-dimensional advection at $\mathcal{N}(u)$ 6, ConvNets succeed only when depth+1 is at least $\mathcal{N}(u)$ 7, while FNO is described as agnostic to $\mathcal{N}(u)$ 8; ResNets and dilated convolutions can handle moderate $\mathcal{N}(u)$ 9, and unrolling during training markedly improves temporal generalization at high $\mathcal{L}$ 0.

For nonlinear one-dimensional Burgers, Kuramoto–Sivashinsky, and Korteweg–de Vries dynamics, all rollouts benefit from multi-step training, and diverted-chain training with $\mathcal{L}$ 1, $\mathcal{L}$ 2 yields one-step accuracy close to supervised one-step while matching or exceeding 5-step supervised rollouts in long-term RMSE, particularly for KdV (Koehler et al., 2024). In a two-dimensional advection experiment with $\mathcal{L}$ 3, a ResNet + coarse solver (50%) hybrid achieves the best gmean rollout error; convolutional hybrids profit strongly from unrolling, whereas FNO hybrids are relatively insensitive because of their global field.

Across the benchmark’s broad comparison over 46 PDEs from 1D to 3D, the paper reports that ResNets are the most robust across all dynamics and dimensions, ConvNets excel in low-order, low-difficulty linear or reaction-diffusion problems, FNO shines in spectral fluid flows such as Kolmogorov flow but struggles with high-frequency-rich reaction–diffusion patterns, and UNets offer consistent 3D performance (Koehler et al., 2024). It also reports that emulation accuracy often improves with resolution until model capacity saturates, and that all architectures benefit from more parameters and longer training while the ranking remains stable. These findings position APEBench as a benchmark for studying not just aggregate accuracy, but the interaction between PDE difficulty, numerical structure, receptive field, and training regime.

6. Use of APEBench in subsequent research

Later papers have used APEBench in three notably different ways: as a source of canonical periodic-domain PDEs for cost-aware evaluation, as a benchmark task for compact neural-operator construction, and as a parameterized one-dimensional suite for testing generalization across PDE families.

Study	APEBench subset	Reported result
“Breakeven complexity: A new perspective on neural partial differential equation solvers” (Zhang et al., 14 May 2026)	Three PDEs on 2D periodic domains from APEBench: Navier–Stokes, Kuramoto–Sivashinsky, Gray–Scott	On “toy” periodic PDEs neural solvers need $\mathcal{L}$ 4– $\mathcal{L}$ 5 calls to pay off; on harder settings that threshold falls to $\mathcal{L}$ 6– $\mathcal{L}$ 7
“Operator Boosting Produces Pareto-Efficient PDE Surrogates” (Shikhman, 16 Jun 2026)	Two-dimensional incompressible Navier–Stokes from APEBench	FNO boosting gives $\mathcal{L}$ 8 RelL2 with 94.5% parameter reduction; Pareto status is “Dominates”
“Generalizing PDE Emulation with Equation-Aware Neural Operators” (Zhu et al., 12 Nov 2025)	A 1D parameterized APEBench suite with KdV, cKS, Fisher’s Equation, Advection–Diffusion, and held-out Burgers’ Equation	All four generalized models produce stable 200-step rollouts on the unseen Burgers’ equation

In the breakeven-complexity study, the APEBench component comprises three canonical 2D periodic PDEs, all simulated via the GPU pseudo-spectral Exponax code: Navier–Stokes (vorticity form) on $\mathcal{L}$ 9, Kuramoto–Sivashinsky on $\partial_x^k u(x) \approx \mathcal{F}^{-1}\!\left[(ik)^k \mathcal{F}(u)\right].$ 0, and Gray–Scott reaction–diffusion on $\partial_x^k u(x) \approx \mathcal{F}^{-1}\!\left[(ik)^k \mathcal{F}(u)\right].$ 1 (Zhang et al., 14 May 2026). The reported setup includes $\partial_x^k u(x) \approx \mathcal{F}^{-1}\!\left[(ik)^k \mathcal{F}(u)\right].$ 2 training and $\partial_x^k u(x) \approx \mathcal{F}^{-1}\!\left[(ik)^k \mathcal{F}(u)\right].$ 3 test trajectories at high fidelity for the periodic cases, followed by spectral down-sampling to $\partial_x^k u(x) \approx \mathcal{F}^{-1}\!\left[(ik)^k \mathcal{F}(u)\right].$ 4 or $\partial_x^k u(x) \approx \mathcal{F}^{-1}\!\left[(ik)^k \mathcal{F}(u)\right].$ 5. The paper’s central result is that, on the 2D periodic tasks, even the best supervised models typically require between $\partial_x^k u(x) \approx \mathcal{F}^{-1}\!\left[(ik)^k \mathcal{F}(u)\right].$ 6 and $\partial_x^k u(x) \approx \mathcal{F}^{-1}\!\left[(ik)^k \mathcal{F}(u)\right].$ 7 forward calls before they amortize training cost at moderate budgets $\partial_x^k u(x) \approx \mathcal{F}^{-1}\!\left[(ik)^k \mathcal{F}(u)\right].$ 8, while increasing problem difficulty lowers the breakeven complexity $\partial_x^k u(x) \approx \mathcal{F}^{-1}\!\left[(ik)^k \mathcal{F}(u)\right].$ 9. For example, for FFNO on Kuramoto–Sivashinsky, $\Delta t$ 0 in 1D, approximately $\Delta t$ 1 in 2D at $\Delta t$ 2, and approximately $\Delta t$ 3 in 3D at $\Delta t$ 4 (Zhang et al., 14 May 2026). This usage treats APEBench as a controlled environment for evaluating full end-to-end economics rather than only rollout error.

In the operator-boosting study, the APEBench task is two-dimensional incompressible Navier–Stokes flow, formulated as a one-step time-advancement operator

$\Delta t$ 5

with

$\Delta t$ 6

on a rectangular domain (Shikhman, 16 Jun 2026). Inputs and targets are normalized per component by the training-split mean and standard deviation, and test error is the denormalized relative $\Delta t$ 7 error

$\Delta t$ 8

For this APEBench task, the paper reports that a 3-stage boosted tiny FNO stack improves from 0.0805 to 0.0188 RelL2, with 74.8% performance gain, 95% CI $\Delta t$ 9, 94.5% trainable-parameter reduction, and 10/10 wins over seeds; DeepONet boosting yields a modest gain; CNO boosting fails outright, worsening from 0.0822 to 0.2109 RelL2 (Shikhman, 16 Jun 2026). Here APEBench functions as a benchmark for within-family accuracy–parameter tradeoffs and Pareto judgments.

In the equation-aware generalization study, APEBench is described as a publicly available suite of parameterized, one-dimensional PDEs together with high-fidelity, differentiable reference solvers (Zhu et al., 12 Nov 2025). The reported experiments use four PDE families for training—Korteweg–de Vries, conserved Kuramoto–Sivashinsky, Fisher’s Equation, and Advection–Diffusion—and withhold Burgers’ Equation entirely for zero-shot testing. All simulations use a periodic domain $\hat{u}^{n+1} = \exp(\hat{L}\Delta t)\odot \hat{u}^{n} + \exp(\hat{L}\Delta t)\odot \int_0^{\Delta t} \exp(-\hat{L}\tau)\,\hat{N}(u(\tau))\,d\tau .$ 0, discretized into 160 grid points, with uniform time step $\hat{u}^{n+1} = \exp(\hat{L}\Delta t)\odot \hat{u}^{n} + \exp(\hat{L}\Delta t)\odot \int_0^{\Delta t} \exp(-\hat{L}\tau)\,\hat{N}(u(\tau))\,d\tau .$ 1; training trajectories have 50 time steps, and test data use 30 random initial conditions per PDE, rolled out for 200 steps (Zhu et al., 12 Nov 2025). The study evaluates both in-distribution and out-of-distribution parameter splits and reports that all four generalized models produce stable 200-step rollouts on the unseen Burgers’ equation, whereas the specialized Burgers baseline exhibits rapid error growth and eventual blow-up. In this setting, APEBench serves as an instrument for probing compositional generalization across PDEs.

7. Interpretation, misconceptions, and benchmark significance

Several recurring interpretations of APEBench are corrected by the literature surrounding it. First, APEBench is not only a one-dimensional benchmark: the original release explicitly supports 46 distinct PDE dynamics spanning 1D, 2D, and 3D (Koehler et al., 2024). Later papers may use restricted subsets—for example, a 1D parameterized family in equation-aware emulation or three 2D periodic systems in breakeven analysis—but those subsets do not exhaust the benchmark.

Second, APEBench is not only a static supervised-learning corpus. The original benchmark emphasizes a fully differentiable reference solver and support for differentiable reference-solver training and neural–hybrid emulators (Koehler et al., 2024). The equation-aware work reinforces this point by stating that the reference solvers are accurate, differentiable, and run “on the fly” during training (Zhu et al., 12 Nov 2025).

Third, APEBench is not only an accuracy leaderboard. Its native evaluation emphasizes rollout metrics, Fourier-space and Sobolev diagnostics, and experiment identifiers tied to numerical-stability limits (Koehler et al., 2024). The breakeven-complexity study extends that perspective by showing that end-to-end usefulness can depend on data-generation cost, training cost, inference cost, and the cost–fidelity tradeoff of classical solvers; on APEBench periodic tasks, models with acceptable rollout error may still require $\hat{u}^{n+1} = \exp(\hat{L}\Delta t)\odot \hat{u}^{n} + \exp(\hat{L}\Delta t)\odot \int_0^{\Delta t} \exp(-\hat{L}\tau)\,\hat{N}(u(\tau))\,d\tau .$ 2– $\hat{u}^{n+1} = \exp(\hat{L}\Delta t)\odot \hat{u}^{n} + \exp(\hat{L}\Delta t)\odot \int_0^{\Delta t} \exp(-\hat{L}\tau)\,\hat{N}(u(\tau))\,d\tau .$ 3 forward calls before becoming cost-effective (Zhang et al., 14 May 2026). This suggests that APEBench is especially valuable when used to connect surrogate-model quality, temporal stability, and computational economics rather than to isolate any one of those axes.

Within the arXiv literature represented here, APEBench therefore occupies a dual role. It is both a benchmark suite with a specific numerical and software design, and a reusable experimental substrate for broader questions about rollout stability, solver–architecture matching, parameter efficiency, zero-shot PDE transfer, and the practical breakeven point between learned and classical solvers.