NuwaDynamics+: A Causality-Aware Generative Framework for Spatio-Temporal Representation Learning

Spatio-temporal (ST) prediction is crucial in earth sciences, including meteorological forecasting and urban computing, to name just a few. Access to ample high-quality data, combined with deep models adept at inference, is essential for attaining significant outcomes. Yet, data scarcity and the substantial costs of sensor deployment result in notable data imbalances. Overly specialized models that lack causal linkages further undermine the generalizability of inference techniques. To address these challenges, we first introduce a causal framework for ST predictions, named $\mathtt{NuwaDynamics}$NuwaDynamicsNuwaDynamics, aimed at pinpointing causal regions in data and providing models with the capability for causal reasoning in a dual-phase process. Initially, we employ upstream self-supervision to identify causally significant patches, equipping the model with generalizable insights and performing targeted interventions on non-essential patches to approximate potential testing distributions. This stage is known as the discovery phase. Progressing from discovery, we apply the insights to downstream tasks tailored to specific ST goals, enhancing the model's recognition of a wider potential data distribution and augmenting its causal perceptual abilities (referred to as the Update phase). Additionally, we address environmental controllability and high computational complexity by implementing channel multiplication and conditional generation methods. This process, termed $\mathtt{NuwaDynamics+}$NuwaDynamics+NuwaDynamics+, can further be interpreted as the front-door adjustment technique in the causality domain. Through comprehensive experiments across ten real-world or simulated ST benchmarks, we demonstrate that integrating the $\mathtt{NuwaDynamics+}$NuwaDynamics+NuwaDynamics+ concept substantially improves various model performance. $\mathtt{NuwaDynamics+}$NuwaDynamics+NuwaDynamics+ concept also significantly enhances the versatility across various dynamic ST tasks, such as extreme weather forecasting and long-temporal-step super-resolution predictions.