Research map / Probabilistic and scientific modeling / Differential equations and PDEs
APEBench: A Benchmark for Autoregressive Neural Emulators of PDEs
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces APEBench, a JAX-based benchmark of 46 semi-linear PDEs in 1D-3D with a differentiable pseudo-spectral ETDRK solver, a unified unrolled-training taxonomy (including diverted-chain), and stability-based difficulty identifiers for evaluating autoregressive neural emulators.
Problem
Autoregressive neural emulators of time-dependent PDEs are evaluated on fixed pre-computed datasets with limited attention to rollout behavior, unrolled training strategies, and the interplay between emulators and numerical solvers.
Approach
Tightly integrate a differentiable Fourier pseudo-spectral ETDRK solver that procedurally generates data and serves as a differentiable reference during training, enabling diverted-chain unrolling and neural-hybrid correction, and define dimensionless difficulty coefficients (gamma) analogous to CFL-type stability criteria to identify dynamics.
Key results
- On 1D advection, purely linear convolutions surpass first-order upwind for the first 13 steps; 20-step unrolling gives 11% higher first-step error but better error at 30 steps.
- For 1D KdV, diverted-chain training of a ResNet improves long-term (100-step) accuracy over one-step training while keeping short-term accuracy close to one-step training.
- In 2D advection (gamma_1=10.5), ResNet with a coarse solver correcting 50% of difficulty performs best, and the FNO outperforms ResNet only in some settings while being largely insensitive to unrolling.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add 20-step unrolled training (ablation) | linear convolution emulator · 1D advection | rollout accuracy after 30 steps | Helped |
| add 20-step unrolled training (ablation) | linear convolution emulator · 1D advection | first-step error: 11% increased error | Hurt |
| increase training unrolling from 1 to 10 to 50 steps (ablation) | linear convolution emulator · 1D advection | distance to FOU stencil: 0.034 to 0.024 to 0.01 | Helped |
| add supervised unrolling (ablation) | ResNet emulator · 1D advection at highest difficulty | temporal generalization | Helped |
| add training rollout (5-step supervised unrolling or 5-step diverted chain) (ablation) | ResNet emulator · viscous Burgers, Kuramoto-Sivashinsky, and hyper-viscous KdV | 100-step rollout error | Helped |
| use 5-step diverted-chain training instead of one-step training (ablation) | ResNet emulator · hyper-viscous KdV | long-term accuracy | Helped |
| use 5-step diverted-chain training instead of supervised unrolling (ablation) | ResNet emulator · viscous Burgers and Kuramoto-Sivashinsky | long-term accuracy | Hurt |
| use 5-step diverted-chain training instead of supervised unrolling (ablation) | ResNet emulator · viscous Burgers and Kuramoto-Sivashinsky | short-term performance | Helped |
| add supervised unrolling (ablation) | ResNet and ResNet-hybrid · 2D advection, difficulty γ₁=10.5 | geometric mean rollout error over 100 time steps | Helped |
| add supervised unrolling (ablation) | FNO · 2D advection, difficulty γ₁=10.5 | geometric mean rollout error over 100 time steps | No clear effect |
| use a coarse solver for 50% of the difficulty in sequential correction (ablation) | ResNet-hybrid · 2D advection, difficulty γ₁=10.5 | geometric mean rollout error over 100 time steps | Helped |
| use ResNet instead of FNO | neural emulator · reaction-diffusion problems | rollout error | Helped |
Surprising findings
- Learned linear convolutions can outperform the first-order upwind numerical scheme for a limited rollout horizon, but diverge eventually because they are not consistent.
- ResNet and the deepest ConvNet fail at the highest advection difficulty despite having a sufficient receptive field, but unrolled training largely fixes this.
Limitations
- Only periodic boundary conditions and uniform Cartesian grids are supported.
- Solver is restricted to semi-linear PDEs with smooth, bandlimited solutions and no channel mixing in the linear part.
- Scope limited to the included solver and architectures; forced dynamics and specialized architectures are left to future work.
Experimental setup
Datasets: APEBench procedurally generated PDE trajectories (46 dynamics)
Benchmarks: APEBench (46 PDEs in 1D, 2D, 3D)
Metrics: normalized RMSE geometric mean rollout error over 100 steps median with 50% IQR across initializations
Baselines: FNO UNet Dilated ResNet ConvNet ResNet first-order upwind
Builds on: Fourier Neural Operator exponential time differencing Runge-Kutta PDEBench PDEArena Equinox
Closest papers by meaning
- Neural Emulator Superiority: When Machine Learning for PDEs Surpasses its Training Data (NeurIPS 2025 · similarity 0.89)
- INC: An Indirect Neural Corrector for Auto-Regressive Hybrid PDE Solvers (NeurIPS 2025 · similarity 0.88)
- Breakeven complexity: A new perspective on neural partial differential equation solvers (NeurIPS 2026 · similarity 0.88)
- SPDEBench: An Extensive Benchmark for Learning Stochastic PDEs (NeurIPS 2026 · similarity 0.87)
- Hierarchical Implicit Neural Emulators (NeurIPS 2025 · similarity 0.87)
- Text2PDE: Latent Diffusion Models for Accessible Physics Simulation (ICLR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.