Towards Generalizable PDE Dynamics Forecasting via Physics-Guided Invariant Learning

TL;DR

iMOOE is proposed, featuring an Invariance-aligned Mixture Of Operator Expert architecture and a frequency-enriched invariant learning objective, which points out that ingredient operators and their composition relationships remain invariant across different domains and PDE system evolution.

Abstract

Advanced deep learning-based approaches have been actively applied to forecast the spatiotemporal physical dynamics governed by partial differential equations (PDEs), which acts as a critical procedure in tackling many science and engineering problems. As real-world physical environments like PDE system parameters are always capricious, how to generalize across unseen out-of-distribution (OOD) forecasting scenarios using limited training data is of great importance. To bridge this barrier, existing methods focus on discovering domain-generalizable representations across various PDE dynamics trajectories. However, their zero-shot OOD generalization capability remains deficient, since extra test-time samples for domain-specific adaptation are still required. This is because the fundamental physical invariance in PDE dynamical systems are yet to be investigated or integrated. To this end, we first explicitly define a two-fold PDE invariance principle, which points out that ingredient operators and their composition relationships remain invariant across different domains and PDE system evolution. Next, to capture this two-fold PDE invariance, we propose a physics-guided invariant learning method termed iMOOE, featuring an Invariance-aligned Mixture Of Operator Expert architecture and a frequency-enriched invariant learning objective. Extensive experiments across simulated benchmarks and real-world applications validate iMOOE's superior in-distribution performance and zero-shot generalization capabilities on diverse OOD forecasting scenarios.

Figures (18)

• Figure 1: (a) The SCM diagram for the formation process PDE dynamics. It illustrates prescribed two-level PDE invariance and potential distribution shifts on exogenous inputs. (b) A case study by varying physical parameters of DR dynamics to compare the zero-shot OOD performance of four methods. Without the guidance of formalized PDE invariance, previous methods can not achieve better OOD results on unseen environments. (c) The ID-OOD correlations of two neural operators. Based on the slope of two linear positive ID-OOD lines, FNO equipped with PDE invariance learning can capture more transferrable knowledge from limited training domains and achieve better OOD robustness. Refer to Appendix \ref{['subsec:id_ood_corr']} for more details.
• Figure 2: (a) Overview of iMOOE method, which can capture the physics-guided PDE invariance by the mixture of operator experts architecture and frequency-enriched multi-context ($|\mathcal{E}_{tr}|=2$ here) training. (b) The structure of single operator expert, which can well fit in diverse neural operators.
• Figure 3: Zero-shot OOD results on SST dynamics.
• Figure 4: Test SST sample showcase.
• Figure 5: Zero-shot OOD results on SSE dynamics.