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

arXiv: https://arxiv.org/abs/2601.08404 · HTML: https://mlsanity.com/paper/2601.08404/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 45 references that are in the map.
- First posted on arXiv 13 Jan 2026.

## Builds on (most cited first, 21 of 21)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Perceptual Losses for Real-Time Style Transfer and Super-Resolution](https://mlsanity.com/paper/1603.08155/) (2016)
- [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://mlsanity.com/paper/1910.03193/) (2019)
- [FourCastNet: A Global Data-driven High-resolution Weather Model using Adaptive Fourier Neural Operators](https://mlsanity.com/paper/2202.11214/) (2022)
- [Machine learning accelerated computational fluid dynamics](https://mlsanity.com/paper/2102.01010/) (2021)
- [PDEBENCH: An Extensive Benchmark for Scientific Machine Learning](https://mlsanity.com/paper/2210.07182/) (2022)
- [Physics-Constrained Deep Learning for High-dimensional Surrogate Modeling and Uncertainty Quantification without Labeled Data](https://mlsanity.com/paper/1901.06314/) (2019)
- [Integrating Scientific Knowledge with Machine Learning for Engineering and Environmental Systems](https://mlsanity.com/paper/2003.04919/) (2020)
- [PhyCRNet: Physics-informed Convolutional-Recurrent Network for Solving Spatiotemporal PDEs](https://mlsanity.com/paper/2106.14103/) (2021)
- [Learning two-phase microstructure evolution using neural operators and autoencoder architectures](https://mlsanity.com/paper/2204.07230/) (2022)
- [Physics-informed neural network for ultrasound nondestructive quantification of surface breaking cracks](https://mlsanity.com/paper/2005.03596/) (2020)
- [Learning in Modal Space: Solving Time-Dependent Stochastic PDEs Using Physics-Informed Neural Networks](https://mlsanity.com/paper/1905.01205/) (2019)
- [Encoding physics to learn reaction-diffusion processes](https://mlsanity.com/paper/2106.04781/) (2021)
- [APEBench: A Benchmark for Autoregressive Neural Emulators of PDEs](https://mlsanity.com/paper/2411.00180/) (2024)
- [Towards a Foundation Model for Partial Differential Equations: Multi-Operator Learning and Extrapolation](https://mlsanity.com/paper/2404.12355/) (2024)
- [PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations](https://mlsanity.com/paper/2505.24717/) (2025)
- [Applying Convolutional Neural Networks to Data on Unstructured Meshes with Space-Filling Curves](https://mlsanity.com/paper/2011.14820/) (2020)
- [Efficient Surrogate Models for Materials Science Simulations: Machine Learning-based Prediction of Microstructure Properties](https://mlsanity.com/paper/2309.00305/) (2023)

## Built on by (most cited first, 3 of 3)

- [Diagnosing Failure Modes of Neural Operators Across Diverse PDE Families](https://mlsanity.com/paper/2601.11428/) (2026)
- [Semigroup Consistency as a Diagnostic for Learned Physics Simulators](https://mlsanity.com/paper/2605.26324/) (2026)
- [Into the danger zone: stable extrapolation in high-dimensional function and operator learning](https://mlsanity.com/paper/2609.36709/) (2026)

## Abstract (arXiv)

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 O(102) 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 ≈ 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.