Decoupling Universal Laws and Environmental Heterogeneity: A Physics-Inspired Framework for Robust Spatio-Temporal Forecasting

Most spatio-temporal forecasting models assume in-distribution data and can degrade sharply under non-stationary environments. Existing methods for handling distribution shift largely rely on discrete graph inference, making it difficult to disentangle universal dynamics from environment-specific changes and to respect the continuous physical nature of spatio-temporal fields. To this end, we propose STPDE, a general framework that reformulates spatio-temporal dynamics as the evolution of inhomogeneous partial differential equations. STPDE explicitly decomposes dynamics into an Invariant Diffusion Operator that captures universal mechanisms and an Environment Basis Manifold that parameterizes local heterogeneous media. We show that the Green's function of the Laplacian can be effectively approximated by linear attention, enabling global diffusion at scale. Combined with stochastic environment perturbations, STPDE improves robustness under heterogeneous and shifting environments. Extensive experiments on in-distribution forecasting, out-of-distribution generalization, few-shot cross-city transfer, and continual learning demonstrate consistent improvements over state-of-the-art baselines with competitive computational efficiency.

Many real-world systems, such as city traffic, change across both space and time. A forecasting model may work well in one setting but become less reliable when conditions change, such as in a different city, a new season, an expanding road network, or during unexpected events. This paper introduces STPDE, a new forecasting approach inspired by ideas from physics. Rather than only learning patterns from past data, it separates general movement rules that remain useful across environments from local factors that may vary from place to place. This allows the model to keep stable knowledge while still adapting to changing conditions. For traffic prediction, this means learning broad traffic patterns while adjusting to local roads, city-specific behavior, and new situations. Experiments on several traffic forecasting tasks show that the proposed approach makes more accurate and reliable predictions, especially when the test environment differs from the training data, while still remaining computationally efficient.