Physics-informed Temporal Alignment for Auto-regressive PDE Foundation Models

Congcong Zhu, Xiaoyan Xu, Jiayue Han, Jingrun Chen

TL;DR: A self-supervised learning framework for solving the error accumulation issue in PDE foundation models

Abstract: Auto-regressive partial differential equation (PDE) foundation models have shown great potential in handling time-dependent data. However, these models suffer from error accumulation caused by the shortcut problem deeply rooted in auto-regressive prediction. The challenge becomes particularly evident for out-of-distribution data, as the pretraining performance may approach random model initialization for downstream tasks with long-term dynamics. To deal with this problem, we propose physics-informed temporal alignment (PITA), a self-supervised learning framework inspired by inverse problem solving. Specifically, PITA aligns the physical dynamics discovered at different time steps on each given PDE trajectory by integrating physics-informed constraints into the self-supervision signal. The alignment is derived from observation data without relying on known physics priors, indicating strong generalization ability to out-of-distribution data. Extensive experiments show that PITA significantly enhances the accuracy and robustness of existing foundation models on diverse time-dependent PDE data. The code is available at \url{https://github.com/SCAILab-USTC/PITA}.

Lay Summary: Many deep learning models that predict the evolution of physical systems step by step suffer from accumulating errors over long-term forecasts. To address this, we propose physics-informed temporal alignment (PITA), a self-supervised learning framework inspired by inverse problem solving. Rather than relying on predefined physical equations, PITA automatically uncovers the underlying dynamics from observed trajectory data and uses these learned dynamics to align and correct predictions across successive time steps. By referencing the discovered physical patterns at each prediction stage, the model self-corrects and substantially mitigates error accumulation. Extensive experiments demonstrate that PITA significantly improves both the accuracy and robustness of existing auto-regressive PDE foundation models, even when applied to out-of-distribution scenarios. The code is available at \url{https://github.com/SCAILab-USTC/PITA}.

Link To Code: https://github.com/ SCAILab-USTC/PITA

Primary Area: Deep Learning->Foundation Models

Keywords: PDE foundation model, Physics-informed constrain, Inverse problem, Self-supervised learning

Submission Number: 6799