Towards a Foundation Model for Partial Differential Equations Across Physics Domains

TL;DR

PDE-FM introduces a foundation-model-style architecture that unifies spatial, spectral, and temporal reasoning across heterogeneous PDEs by integrating dual-tokenization, physics-aware conditioning, a Mamba state-space backbone, and a spectral decoder. Pretrained on a diverse set of The Well datasets, it delivers state-of-the-art or near-state-of-the-art performance across 12 2D/3D PDE regimes, with strong cross-physics generalization, especially in turbulent and radiative flows. The results suggest that large-scale pretraining across multiple physical domains can yield transferable representations for multi-physics surrogates, advancing toward universal, physics-informed surrogates for scientific discovery. Limitations in memory-dominated and elastic systems point to future work on energy-preserving losses, adaptive spectral decoding, and curriculum-based multi-domain training to broaden coverage.

Abstract

We present PDE-FM, a modular foundation model for physics-informed machine learning that unifies spatial, spectral, and temporal reasoning across heterogeneous partial differential equation (PDE) systems. PDE-FM combines spatial-spectral tokenization, physics-aware conditioning, and a Mamba-based state-space backbone with an operator-theoretic decoder, enabling scalable and data-efficient modeling of complex physical dynamics. In contrast to task-specific neural operators, PDE-FM is pretrained once on diverse PDE datasets and can be transferred to new physical regimes without architectural or data-specific modifications. Evaluated on twelve 2D and 3D datasets from The Well benchmark - spanning hydrodynamic, radiative, elastic, and astrophysical phenomena - PDE-FM achieves state-of-the-art accuracy in six domains, reducing mean VRMSE by 46% relative to prior operator-learning baselines. The model demonstrates robust cross-physics generalization, excelling in turbulent and radiative systems while maintaining strong performance in linear and steady-state regimes. These results suggest that large-scale pretraining across diverse physical processes can yield transferable representations of dynamics, marking a step toward unified, foundation-level surrogates for multi-physics simulation and scientific discovery.

Figures (4)

• Figure 1: General architecture of PDE-FM.
• Figure 2: Parity plot comparing VRMSE of PDE-FM versus the best SOTA baseline. Points below the diagonal (gray line) indicate improved performance. Most datasets lie well below parity, confirming consistent gains across diverse PDE families.
• Figure 3: VRMSE heatmap across models and datasets. Blue regions denote low errors. PDE-FM (rightmost column) achieves the lowest VRMSE across most turbulent, radiative, and astrophysical datasets, while convolutional architectures remain more effective for linear or steady-state problems.
• Figure 4: Mean VRMSE across all PDE datasets. PDE-FM achieves the lowest average error (0.165), outperforming all operator-learning baselines. The improvement margin relative to the next-best model (CNextU-net, 0.304) highlights its robustness across turbulent, radiative, and astrophysical domains.