From Numerical Simulators of PDEs to Neural Emulators and Back

Felix Koehler

cs.LGAug 25, 2026

From Numerical Simulators of PDEs to Neural Emulators and Back

Authors: Felix Koehler

Organizations: Technische Universität München TUM School of Computation, Information and Technology

Abstract

Simulation is central to modern engineering and science, but the cost of numerical solvers for partial differential equations (PDEs) remains a bottleneck whenever fast or many-query evaluations are required. Neural emulators trained on solver-generated data promise significant speedups, yet they are usually framed as opaque alternatives to the very methods that produce their training signal. This thesis argues the two paradigms are more alike than different: neural architectures mirror classical discretizations, their errors are amenable to the same spectral analysis, and insight flows profitably in both directions. We approach the relationship by disentangling the multiple roles a solver plays in the emulator learning pipeline. Mode-wise Fourier analysis then provides a common language in which solver errors, architectural inductive biases, and training objectives can all be read off simultaneously. Taken together, this allows synthesizing three contributions. (1) APEBench, a comprehensive benchmarking suite for autoregressive neural emulators of PDEs that uses fast differentiable pseudo-spectral solvers in JAX. (2) Progressively Refined Differentiable Physics, an investigation of the effect of unconverged solvers on surrogate training. (3) Neural Emulator Superiority, an analysis of the influence of numerical errors and architectural inductive biases.

Figures & tables

Panel of Figure 1 — Figure 1: Overview of the emulation pipeline. A PDE system $\mathbb{P}$ is discretized into a solver $\mathcal{P}_{h}$ by choosing a numerical method $\mathbb{M}$ , resolution $N$ , and time step $\Delta t$ . Initial conditions $\mathbf{u}_{h}^{[0]}$ sampled from a distribution $\mathcal{I}_{h}$ are rolled out via the trajectory operator $\mathcal{T}$ to produce data trajectories $\mathcal{D}_{h}$ , which are used to train and test the neural emulator $f_{\theta}$ .

Panel of Figure 2 — Figure 2: Overview of the three publications comprising this thesis. APEBench provides the benchmarking framework that enables systematic measurement of neural PDE emulators. In Emulator Superiority, we study the effect of numerical errors on surrogate quality with the surprising finding that they can sometimes exceed their training data fidelity. With PRDP, we demonstrate how to train emulators efficiently using progressively refined differentiable solvers. The bidirectional arc captures the thesis title: insights flow from numerical solvers to neural emulators and back.

Panel of Figure 3 — Figure 3: The method-of-lines pipeline for solving time-dependent PDEs. A continuous PDE is first discretized in space ( Section 5.1 ), yielding a semi-discrete system of ODEs. A time integration scheme ( Section 5.2 ) then defines the evolution operator $\mathcal{P}_{h}$ , which advances the discrete state $\mathbf{u}_{h}^{[n]}$ by one time step $\Delta t$ . An initial condition $\mathbf{u}_{h}^{[0]}\sim\mathcal{I}_{h}$ is sampled from a prescribed distribution, and repeated application of $\mathcal{P}_{h}$ produces a trajectory of states.

Panel of Figure 4 — Figure 4: Comparison of finite difference and Fourier-spectral spatial discretization. (a) Finite differences approximate derivatives via a local stencil; only neighbouring grid points contribute to each output, yielding polynomial-order accuracy. (b) Spectral methods transform to Fourier space ( $\mathcal{F}_{h}$ ), apply an exact pointwise multiplier per wavenumber, and transform back ( $\mathcal{F}_{h}^{-1}$ ); all grid points contribute globally, achieving exponential convergence for smooth, bandlimited functions on periodic domains.

Panel of Figure 5 — Figure 5: Connectivity patterns of neural architectures used as PDE emulators, shown as 1D schematics. Each panel highlights which input positions contribute to a single output position (dark blue) in one layer or block; shading intensity indicates contribution strength. Panels are ordered by locality of the receptive field: (a) ConvNet ( $\approx$ finite differences) uses a fixed local stencil; (b) SWIN attends within shifted windows; (c) Dilated ResNet grows the receptive field through dilation; (d) UNet ( $\approx$ multigrid) processes a hierarchy of resolutions with skip connections; (e) ViT applies global attention between patches; (f) FNO ( $\approx$ spectral methods) mixes all positions through the Fourier domain; (g) Transformer uses full self-attention over every position.

Panel of Figure 6 — Figure 6: Training and evaluation pipeline for autoregressive neural emulators. (a) During training, a numerical solver $\mathcal{P}_{h}$ produces reference trajectories from initial conditions sampled from $\mathcal{I}_{h}$ . The emulator $f_{\theta}$ is trained to match consecutive states; gradient flow (purple) may extend through multiple autoregressive steps. (b) At test time, the emulator is rolled out autoregressively for many steps and compared against a reference solver $\tilde{\mathcal{P}}_{h}$ . Tildes indicate that the test setup may differ from training, probing temporal, state-space, solver, or metric generalization.

Panel of Figure 7 — Figure 7: The $(H,C,L)$ temporal axis taxonomy. An autoregressive emulator $f_{\theta}$ maps $H$ input states to $L$ output states, potentially skipping $C{-}1$ intermediate time levels. (a) The standard Markovian baseline ( $H{=}C{=}L{=}1$ ). (b) Using $H{\geq}2$ past states as input enables higher-order temporal information, analogous to multi-step numerical methods. (c) A skip length $C{\geq}2$ coarsens the effective time step; skipped time levels (hollow circles) are not predicted. (d) Temporal bundling ( $L{\geq}2$ ) produces multiple future states simultaneously.

Panel of Figure 8 — Figure 8: The multiple roles of numerical solvers in the neural emulator pipeline. A continuous PDE system $\mathbb{P}$ admits many valid discretizations $\mathcal{P}_{h},\mathcal{P}_{H},\mathcal{P}_{h^{\prime}},\ldots$ obtained by choosing a numerical method $\mathbb{M}$ , resolution $N$ , and time step $\Delta t$ . Each solver may fill one or more of the five roles shown, and different experimental setups activate different subsets (see Table 1 ). Solid arrows denote data flow into the emulator, double-headed arrows denote comparison, and the dashed arrow denotes conceptual influence.

Setup	(1)	(2)	(3)	(4)	(5)
Default ( Section 11.1 )	✓	$=$ (1)	—	arch-dep	$=$ (1)
Neural-Hybrid	✓	$=$ (1)	✓	arch-dep	$=$ (1)
Superiority investigation	✓	✓	—	aligned	$=$ (1)
PRDP	✓	$=$ (1)	✓ $<\!\dots\!<$ ✓	arch-dep	$=$ (1)
Unfair benchmark	✓	(hidden)	—	arch-dep	$=$ (1)

Table 1: Which solver roles are active under different experimental setups. ✓ indicates the role is active and filled by a distinct solver; “

=

(k)” indicates the role collapses onto the solver filling role (k); “

<\!\dots\!<

” indicates progressive refinement of fidelity during training; “—” indicates the role is absent; “(hidden)” indicates the role is not exposed in the setup.

Panel of Figure 9 — Figure 9: Mode-wise multipliers for the one-dimensional diffusion equation with $\gamma_{2}=0.5$ (within the stability domain of the explicit scheme). The exact solution (black) decays exponentially with wavenumber. The explicit FTCS scheme (orange, solid) over-attenuates high frequencies, while the implicit BTCS scheme (orange, dashed) under-attenuates them. A learned emulator with FTCS-form inductive bias, trained on BTCS data at mode $\psi=0.15$ (blue), achieves a multiplier that lies closer to the exact solution than its training source across a range of higher modes.

Panel of Figure 10 — Figure 10: Conceptual sketch of the emulator superiority phenomenon over autoregressive rollout steps. The solver (orange) and emulator (blue) both accumulate errors over time, but the emulator’s more favorable error growth causes it to cross below the solver’s error curve. In the shaded region, the superiority ratio $\xi^{[t]}<1$ , meaning the emulator is closer to the high-fidelity reference than the solver it was trained on. This autoregressive superiority emerges even when the emulator is slightly less accurate at $t=1$ .

Panel of Figure 11 — Figure 11: Schematic of the PRDP progressive refinement strategy. Left axis (orange): Training begins with a coarse physics solver and monitors the validation loss. When the loss plateaus, solver fidelity is increased. The dashed gray curve shows training with fully converged physics ( $K_{\epsilon}$ ) throughout; both strategies converge to the same final accuracy. Right axis (blue): The PRDP staircase schedule of solver iterations $K$ . Two types of savings emerge: PR savings (blue hatched) from using fewer solver iterations during the early, coarse stages of training, and IC savings (green band spanning all stages) from the network achieving full accuracy at $K_{\max}$ well below the convergence threshold $K_{\epsilon}$ .

Figure 12: Qualitative overview of a subset of the PDE dynamics supported by APEBench, here in 2D displayed as volume-rendered spatiotemporal plots with time on the diagonal axis. Top row: linear dynamics (advection, diffusion, dispersion, hyper-diffusion). Middle row: nonlinear dynamics (Burgers, Kuramoto-Sivashinsky, Decaying Navier-Stokes, Kolmogorov Flow). Bottom row: reaction-diffusion dynamics (Fisher-KPP, Allen-Cahn, Gray-Scott, Swift-Hohenberg).

Panel of Figure 13 — Figure 13: The $(T,B)$ training taxonomy for autoregressive neural PDE emulators. $T$ denotes the number of autoregressive rollout steps, $B$ the branch length of the reference solver.

Panel of Figure 14 — Figure 14: Metric reduction order for the nRMSE. (a) Three reduction strategies for the prediction tensor $(T,C,X)$ . APEBench uses per-channel normalization (green), preserving sensitivity to weak channel. (b) Given a rollout metric $(T,)$ from (a), temporal aggregation via arithmetic vs. geometric mean. The geometric mean (APEBench choice, green) avoids domination by late-time divergent errors and is less sensitive to the choice of rollout horizon $T$ (sketch below). Batch averaging ( $\operatorname{mean}_{B}$ ) is applied last in all cases and omitted for clarity.

Panel of Figure 15 — Figure 15: The JAX software ecosystem developed as part of this thesis. Equinox (dashed, external) supplies the module system on which all packages except Chaotax build. The three pillars of APEBench are Exponax (spectral PDE solvers), PDEquinox (neural architectures), and Trainax (training methodologies). Picardax provides finite-difference solvers with the same equinox.Module interface and is hot-swappable with Exponax. This interchangeability underpins both the PRDP and the emulator superiority work. Chaotax operates on pure JAX callables (independent of Equinox) to quantify dynamical properties such as Lyapunov exponents. SuperAPE is a higher-level abstraction above APEBench that uses jax.lax.scan for the training loop, enabling fast 1D superiority studies across solver types.

PDE Surrogate Modeling

PDE surrogate modeling builds fast approximations of partial differential equation solutions for simulation, prediction, and design.

See all 319 papers

Breakeven complexity: A new perspective on neural partial differential equation solvers
May 14, 2026Yijing Zhang, Nicholas Roberts, Tanya Marwah +1Neural PDE Solvers Benchmark Design
University of Wisconsin–Madison · Google DeepMind
The Symbol of the Surrogate: Measuring Numerical Provenance in Neural PDE Solvers
Oct 7, 2026Ridham PatelPDE Surrogate Modeling
Computer Science & Engineering Indian Institute of Technology Gandhinagar Gandhinagar, India
BEACONS: Bounded-Error, Algebraically-Composable Neural Solvers for Partial Differential Equations
Feb 16, 2026Jonathan Gorard, Ammar Hakim, James JunoNeural PDE Solvers Partial Differential Equations
Princeton University Princeton, NJ 08544, USA · Princeton Plasma Physics Laboratory Princeton, NJ 08540, USA

May 14, 2026J/K move · Enter open · S save

Breakeven complexity: A new perspective on neural partial differential equation solvers

Yijing Zhang, Nicholas Roberts, Tanya Marwah, Mikhail Khodak

University of Wisconsin–Madison · Google DeepMind

Neural PDE Solvers Benchmark Design PDE Surrogate Modeling

Open paper AskPDF arXiv

Neural surrogate solvers of partial differential equations (PDEs) promise dramatic speedups over numerical methods, especially in scenarios requiring many solves. However, current accuracy-based evaluations do not fully consider two central issues: (1) neural solvers incur substantial up-front costs for data generation, training, and tuning; and (2) classical solvers can also generate low-fidelity solutions at a sufficiently low simulation cost. To explicitly account for these realities and fully incorporate end-to-end costs, we propose an evaluation framework centered on breakeven complexity, a metric that counts the forward solves before a learned solver is cost-effective relative to an error-equivalent traditional solver. To evaluate this measure, we apply scaling laws to determine how much training budget to allocate to data generation and discuss how to achieve smooth error-matching in diverse settings. We evaluate the breakeven complexity of multiple neural PDE solvers on three PDEs on 2D periodic domains from APEBench and a novel benchmark of flows past multiple obstacles generated by the GPU-native PyFR code. Among other findings, our results suggest that neural PDE solvers become more effective as problems get harder in terms of cost, dimension, rollout, physics regime (e.g. higher Reynolds number), etc.

Abstract

PDE Surrogate Modeling

Explore similar work

Breakeven complexity: A new perspective on neural partial differential equation solvers

The Symbol of the Surrogate: Measuring Numerical Provenance in Neural PDE Solvers

BEACONS: Bounded-Error, Algebraically-Composable Neural Solvers for Partial Differential Equations