MORPH: Shape-Agnostic Foundation Models for Heterogeneous PDE Systems

We present MORPH, an autoregressive foundation model for partial differential equation (PDE) systems that enables unified handling of heterogeneous spatiotemporal data across varying dimensionalities (1D-3D), resolutions, and physical field configurations. The architecture employs a convolutional vision transformer backbone integrating three key innovations: (i) component-wise convolution layers that jointly process scalar and vector field channels to capture local physical interactions, (ii) inter-field cross-attention mechanisms that model and selectively propagate information between coupled physical fields, and (iii) axial attention factorization that decomposes full spatiotemporal self-attention along individual spatial and temporal axes, substantially reducing computational complexity while preserving representational capacity. We pretrain multiple model variants on a diverse collection of PDE datasets spanning fluid dynamics, mutli-physics, and other physical systems, demonstrating robust transfer learning capabilities to downstream prediction tasks. This work addresses a critical challenge in scientific machine learning: the development of general-purpose models that can flexibly adapt to diverse PDE systems without specialized architectural modifications for each problem domain. Our results indicate that foundation models can effectively capture universal patterns in physical systems while maintaining computational efficiency for high-resolution spatiotemporal forecasting.