Out-of-distribution generalization of deep-learning surrogates for 2D PDE-generated dynamics in the small-data regime

TL;DR

This work tackles the challenge of out-of-distribution generalization for autoregressive deep-learning surrogates of 2D PDE-generated dynamics in a strict small-data regime. It introduces me-UNet, a periodic-padding convolutional U-Net tailored for incremental time stepping, and benchmarks it against ViT, AFNO, PDE-Transformer, and KAN-UNet across six PDE families on 64×64 periodic grids, using a unified training protocol and physics-aware metrics. The study demonstrates that me-UNet achieves accurate, stable long-horizon rollouts in-distribution and qualitatively robust generalization to unseen initial conditions with as few as ~20 training simulations, outperforming more complex models while requiring substantially less training time. These findings highlight the value of locality-focused inductive biases and periodic boundary alignment for data-efficient surrogate modeling in scientific computing, and they provide a controlled benchmark for evaluating surrogates under realistic small-data constraints, with avenues for extension to 3D, non-periodic domains, and physics-informed constraints.

Abstract

Partial differential equations (PDEs) are a central tool for modeling the dynamics of physical, engineering, and materials systems, but high-fidelity simulations are often computationally expensive. At the same time, many scientific applications can be viewed as the evolution of spatially distributed fields, making data-driven forecasting of such fields a core task in scientific machine learning. In this work we study autoregressive deep-learning surrogates for two-dimensional PDE dynamics on periodic domains, focusing on generalization to out-of-distribution initial conditions within a fixed PDE and parameter regime and on strict small-data settings with at most $\mathcal{O}(10^2)$ simulated trajectories per system. We introduce a multi-channel U-Net [...], evaluate it on five qualitatively different PDE families and compare it to ViT, AFNO, PDE-Transformer, and KAN-UNet under a common training setup. Across all datasets, me-UNet matches or outperforms these more complex architectures in terms of field-space error, spectral similarity, and physics-based metrics for in-distribution rollouts, while requiring substantially less training time. It also generalizes qualitatively to unseen initial conditions with as few as $\approx 20$ training simulations. A data-efficiency study and Grad-CAM analysis further suggest that, in small-data periodic 2D PDE settings, convolutional architectures with inductive biases aligned to locality and periodic boundary conditions remain strong contenders for accurate and moderately out-of-distribution-robust surrogate modeling.

Figures (31)

• Figure 1: Examples of simulation results from all mathematical models. Datasets DS-1 and 2 are obtained from simple convection and diffusion equations, while the complexity of the and the number of involved state variables increases towards the right: DS-3a, 3b, and 4 are obtained from , DS-5 from Navier-Stokes equations, while the variants of DS-6 are obtained from the Gray-Scott model.
• Figure 2: Illustration of the me-UNet architecture, showing the encoder--decoder structure, periodic padding, and multi-channel input/output setup for autoregressive time-stepping.
• Figure 3: Comparison of average values over $T=100$ time steps: (a) and (b) cosine-similarity scores of the power spectral density (PSD). Rows correspond to architectures and columns to datasets; colors are shown on a logarithmic scale for .
• Figure 4: Performance of the different neural network architectures on the dataset (DS-4): (a) example autoregressive rollouts; (b) per-time-step ; and (c) cosine similarity of the PSD curves between prediction and reference.
• Figure 5: Performance of the different neural network architectures on the Gray--Scott dataset DS-6a: (a) example autoregressive rollouts; (b) per-time-step ; and (c) cosine similarity of the PSD curves between prediction and reference.