---
title: 'APEBench: Autoregressive PDE Emulator'
url: https://www.emergentmind.com/topics/apebench
type: topic
---

# APEBench: Autoregressive PDE Emulator

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 [2411.00180]. 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 [2605.15399] [2606.17460] [2511.09729].

## 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 [2411.00180]. 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 [2511.09729]. 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 [2411.00180].

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 [2411.00180]. 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 [2411.00180]. 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}),
\]
\[
\hat{u}^{n+1}
=
\hat{u}^{*}
+
\phi_2(\hat{L}\Delta t)\odot
\bigl[\hat{N}(\hat{u}^{*})-\hat{N}(\hat{u}^{n})\bigr],
\]
with
\[
\phi_1(z)=\frac{e^z-1}{z},
\qquad
\phi_2(z)=\frac{e^z-1-z}{z^2}.
\]

The reported complexity is
\[
O(N^D\,D\,\log N)
\]
per step via FFTs, and the implementation is described as trivially auto-differentiable in JAX [2411.00180]. 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 [2411.00180]. 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
\[
\partial_t u = -b\,\partial_x(u^2/2) + \xi\,\partial_x^3 u - \zeta\,\partial_x^4 u,
\]
the conservative Kuramoto–Sivashinsky equation
\[
\partial_t u = -b\,\partial_x(u^2/2) - \nu\,\partial_x^2 u - \zeta\,\partial_x^4 u,
\]
the Gray–Scott reaction–diffusion system
\[
\partial_t u = \nu_0 \nabla^2 u - uv^2 + f(1-u),
\qquad
\partial_t v = \nu_1 \nabla^2 v + uv^2 - (f+k)v,
\]
and two Navier–Stokes variants, including a streamfunction–vorticity decaying-turbulence form and a Kolmogorov-flow form [2411.00180].

To identify scenarios, APEBench defines normalized linear and nonlinear coefficients
\[
\alpha_s = a_s \Delta t / L^s,
\qquad
\gamma_s = \alpha_s N^s 2^{s-1} D,
\]
where \(\gamma_1 = c\Delta t/\Delta x\) is the classic CFL number for \(s=1\). The benchmark states that, more generally, \(\gamma_s \approx 1\) marks the stability limit of the most compact explicit finite-difference stencil. For nonlinear terms it defines
\[
\delta_{l_{\mathrm{pre}},p,l_{\mathrm{post}}}
=
b\Delta t
/
L^{\,l_{\mathrm{pre}}p+l_{\mathrm{post}}}
\cdot
N^{\,l_{\mathrm{pre}}p+l_{\mathrm{post}}}
D\,m,
\]
with \(m=\max|u|\). The tuple \(\{\gamma_s,\delta_j,N,D\}\) is presented as an unambiguous experiment identifier directly tied to numerical-stability intuition [2411.00180].

## 4. Training taxonomy, benchmark interfaces, and evaluation metrics

A central contribution of APEBench is its taxonomy for unrolled training. Denoting by \(P_h\) the reference simulator for one step and by \(f_\theta\) the neural emulator, and using trajectories \(u_h \sim \mathcal{D}_h\), the benchmark defines
\[
L(\theta)
=
\mathbb{E}_{u_h\sim\mathcal{D}_h}
\sum_{t=0}^{T-B}
\sum_{b=1}^{B}
l\!\left(
f_\theta^{\,t+b}(u_h),
P_h^{\,b}\!\left(f_\theta^{\,t}(u_h)\right)
\right),
\]
where \(T\) is the number of main-chain unrolled steps and \(B\) is the branch length of the reference rollout [2411.00180].

Three canonical regimes are distinguished. **One-step supervised** training sets \(T=B=1\), yielding
\[
L=\mathbb{E}[l(f_\theta(u_h),P_h(u_h))].
\]
**Supervised unrolling** uses \(T=B>1\), corresponding to full rollout versus precomputed reference. **Diverted-chain** training uses \(T>B\) and \(B=1\), so that at each main-chain step the one-step reference \(P_h(f_\theta^t)\) is compared to \(f_\theta^{t+1}\). 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 [2411.00180].

APEBench also exposes interfaces for neural–hybrid emulators. With a coarse solver \(\tilde{P}_h\) and a corrector network \(\hat{g}_\theta\), it defines **sequential correction**
\[
u^{n+1} = \hat{g}_\theta(\tilde{P}_h(u^n))
\]
and **parallel correction**
\[
u^{n+1} = \tilde{P}_h(u^n) + \hat{g}_\theta(u^n).
\]
Both layouts can be trained under any unrolling taxonomy [2411.00180].

For evaluation, APEBench emphasizes autoregressive rollout errors. Its primary metric is the mean normalized RMSE over \(M\) samples and \(N\) grid points,
\[
\mathrm{nRMSE}^{[t]}
=
\frac{1}{M}
\sum_{j=1}^{M}
\sqrt{
\frac{\sum_{i=1}^{N}(\hat{u}_{j,i}-u_{j,i})^2}
{\sum_{i=1}^{N}u_{j,i}^2}
},
\]
and temporal aggregation is given by the geometric mean
\[
\mathrm{Agg\text{-}nRMSE}
=
\exp\!\left[
\frac{1}{T}
\sum_{t=1}^{T}\log\bigl(\mathrm{nRMSE}^{[t]}\bigr)
\right].
\]
The benchmark also supplies Fourier-space RMSE/MAE, Sobolev \(H^1\) metrics, correlation, and band-limited metrics [2411.00180]. 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 [2411.00180]. 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 \(\gamma_1<1\), 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 [2411.00180]. In one-dimensional advection at \(\gamma_1\in\{0.5,2.5,10.5\}\), ConvNets succeed only when depth+1 is at least \(\gamma_1\), while FNO is described as agnostic to \(\gamma_1\); ResNets and dilated convolutions can handle moderate \(\gamma_1\), and unrolling during training markedly improves temporal generalization at high \(\gamma_1\).

For nonlinear one-dimensional Burgers, Kuramoto–Sivashinsky, and Korteweg–de Vries dynamics, all rollouts benefit from multi-step training, and diverted-chain training with \(T=5\), \(B=1\) yields one-step accuracy close to supervised one-step while matching or exceeding 5-step supervised rollouts in long-term RMSE, particularly for KdV [2411.00180]. In a two-dimensional advection experiment with \(\gamma_1=10.5\), 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** [2411.00180]. 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” [2605.15399] | Three PDEs on 2D periodic domains from APEBench: Navier–Stokes, Kuramoto–Sivashinsky, Gray–Scott | On “toy” periodic PDEs neural solvers need \(10^5\)–\(10^6\) calls to pay off; on harder settings that threshold falls to \(10^3\)–\(10^4\) |
| “Operator Boosting Produces Pareto-Efficient PDE Surrogates” [2606.17460] | Two-dimensional incompressible Navier–Stokes from APEBench | FNO boosting gives \(0.0805 \to 0.0188\) RelL2 with 94.5% parameter reduction; Pareto status is “Dominates” |
| “Generalizing PDE Emulation with Equation-Aware Neural Operators” [2511.09729] | 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 \([0,2]^2\), **Kuramoto–Sivashinsky** on \([0,50]^d\), and **Gray–Scott reaction–diffusion** on \([0,2]^2\) [2605.15399]. The reported setup includes \(200{,}000\) training and \(1000\) test trajectories at high fidelity for the periodic cases, followed by spectral down-sampling to \(64^2\) or \(64^d\). The paper’s central result is that, on the 2D periodic tasks, even the best supervised models typically require between \(10^5\) and \(10^6\) forward calls before they amortize training cost at moderate budgets \(B \approx 8\,\mathrm{k\,s}\), while increasing problem difficulty lowers the breakeven complexity \(N^\star\). For example, for FFNO on Kuramoto–Sivashinsky, \(N^\star_{\mathrm{avg}}=\infty\) in 1D, approximately \(5\times 10^5\) in 2D at \(B=1\,\mathrm{k\,s}\), and approximately \(6.5\times 10^3\) in 3D at \(B=4\,\mathrm{k\,s}\) [2605.15399]. 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
\[
a=\bigl(u(t),\Delta t\bigr)\longmapsto u(t+\Delta t),
\]
with
\[
\partial_t u + (u\cdot \nabla)u = -\nabla p + \nu \Delta u,
\qquad
\nabla\!\cdot u = 0
\]
on a rectangular domain [2606.17460]. Inputs and targets are normalized per component by the training-split mean and standard deviation, and test error is the denormalized relative \(L^2\) error
\[
\mathrm{RelL2}(\hat{u},u)
=
\frac{\|\hat{u}-u\|_2}{\|u\|_2+\varepsilon}.
\]
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 \([69.2\%,80.4\%]\)**, **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 [2606.17460]. 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 [2511.09729]. 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 \(x\in[0,1]\), discretized into **160 grid points**, with uniform time step \(\Delta t=1\); training trajectories have **50 time steps**, and test data use **30 random initial conditions per PDE**, rolled out for **200 steps** [2511.09729]. 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** [2411.00180]. 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** [2411.00180]. The equation-aware work reinforces this point by stating that the reference solvers are accurate, differentiable, and run “on the fly” during training [2511.09729].

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 [2411.00180]. 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 \(10^5\)–\(10^6\) forward calls before becoming cost-effective [2605.15399]. 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.

Source: https://www.emergentmind.com/topics/apebench