Unlocking Out-of-Distribution Generalization in Dynamics through Physics-Guided Augmentation

TL;DR

SPARK tackles out-of-distribution generalization and data scarcity in dynamical systems by introducing physics-guided augmentation. It encodes physical priors into a discrete state dictionary via a reconstruction autoencoder and vector quantization, enabling principled latent-space interpolation for augmentation. Prediction uses a Fourier-enhanced Graph ODE with attention-based history encoding to capture long-term dynamics. The approach is supported by information-theoretic generalization bounds showing reduced dependence on training data when physical priors are included, and it demonstrates state-of-the-art performance across Prometheus, ERA5, Navier–Stokes, Spherical-SWE, and sea-ice tasks, including transfer to data-scarce domains. Overall, SPARK offers a robust, scalable framework for physics-informed dynamical modeling with strong OOD and transfer capabilities.

Abstract

In dynamical system modeling, traditional numerical methods are limited by high computational costs, while modern data-driven approaches struggle with data scarcity and distribution shifts. To address these fundamental limitations, we first propose SPARK, a physics-guided quantitative augmentation plugin. Specifically, SPARK utilizes a reconstruction autoencoder to integrate physical parameters into a physics-rich discrete state dictionary. This state dictionary then acts as a structured dictionary of physical states, enabling the creation of new, physically-plausible training samples via principled interpolation in the latent space. Further, for downstream prediction, these augmented representations are seamlessly integrated with a Fourier-enhanced Graph ODE, a combination designed to robustly model the enriched data distribution while capturing long-term temporal dependencies. Extensive experiments on diverse benchmarks demonstrate that SPARK significantly outperforms state-of-the-art baselines, particularly in challenging out-of-distribution scenarios and data-scarce regimes, proving the efficacy of our physics-guided augmentation paradigm.

Figures (4)

• Figure 1: Model overview -- The proposed SPARK consists of four steps: ❶ Physical prior incorporation, where physical parameters are encoded into the input observations; ❷ Dynamical discretization modeling through reconstruction to create a discrete physics-rich state dictionary; ❸ State dictionary guided augmenting on sampled training data; and ❹ Fourier-enhanced graph ODE for dynamical system prediction based on historical observations.
• Figure 2: Comparison of Prediction Performance. The figure shows the target values and the predictions from different models (Ours, FNO, ViT) at multiple time steps. It is evident that SPARK's predictions are closest to the target values, especially in the locally complex regions (highlighted in red or white boxes), demonstrating higher detail-capturing ability and accuracy.
• Figure 3: The figures in the left show the predictions of different models ( SPARK, FNO, and U-Net) for sea ice data at 10-th timestep. The figures in the right demonstrate the curves of loss and performance metrics during training.
• Figure 4: Energy Spectrum Comparison Results on Navier Stokes, Spherical-Shallow Water, and 3D Reaction–Diff datasets.

Theorems & Definitions (2)

• theorem 1: Enhancement of Model Generalization via Physical Priors from an Information-Theoretic Perspective
• theorem 2: Upper Bound on Generalization Error in Bayesian Learning with Physical Priors