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

TL;DR

This paper demonstrates that a multi-channel U-Net deep learning model can effectively generalize to out-of-distribution initial conditions in 2D PDE dynamics, especially in small-data regimes, outperforming more complex models in accuracy and efficiency.

Contribution

The study introduces a multi-channel U-Net architecture tailored for 2D PDE surrogate modeling, showing its superior performance and data efficiency compared to advanced models in small-data, out-of-distribution scenarios.

Findings

me-UNet matches or outperforms complex architectures in accuracy.

It generalizes to unseen initial conditions with as few as 20 training simulations.

Convolutional architectures with locality and periodic boundary biases are effective in small-data PDE settings.

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.