Causal Deciphering and Inpainting in Spatio-Temporal Dynamics via Diffusion Model

TL;DR

CaPaint introduces a causal dec deciphering and diffusion-based inpainting framework to address data scarcity and interpretability in spatio-temporal forecasting. It identifies causal ST patches via self-supervised ViT reconstruction and applies front-door-adjusted diffusion inpainting to intervene on non-causal regions, reducing data-generation complexity from exponential to quasi-linear. Empirical results across five real-world ST benchmarks show consistent improvements in MAE, MSE, and SSIM across multiple backbones, with notable gains in data-scarce scenarios and long-term predictions. This approach offers a practical, interpretable, and scalable pathway to enhance ST prediction by combining causal reasoning with diffusion-based data augmentation and ST sequence modeling.

Abstract

Spatio-temporal (ST) prediction has garnered a De facto attention in earth sciences, such as meteorological prediction, human mobility perception. However, the scarcity of data coupled with the high expenses involved in sensor deployment results in notable data imbalances. Furthermore, models that are excessively customized and devoid of causal connections further undermine the generalizability and interpretability. To this end, we establish a causal framework for ST predictions, termed CaPaint, which targets to identify causal regions in data and endow model with causal reasoning ability in a two-stage process. Going beyond this process, we utilize the back-door adjustment to specifically address the sub-regions identified as non-causal in the upstream phase. Specifically, we employ a novel image inpainting technique. By using a fine-tuned unconditional Diffusion Probabilistic Model (DDPM) as the generative prior, we in-fill the masks defined as environmental parts, offering the possibility of reliable extrapolation for potential data distributions. CaPaint overcomes the high complexity dilemma of optimal ST causal discovery models by reducing the data generation complexity from exponential to quasi-linear levels. Extensive experiments conducted on five real-world ST benchmarks demonstrate that integrating the CaPaint concept allows models to achieve improvements ranging from 4.3% to 77.3%. Moreover, compared to traditional mainstream ST augmenters, CaPaint underscores the potential of diffusion models in ST enhancement, offering a novel paradigm for this field. Our project is available at https://anonymous.4open.science/r/12345-DFCC.

Figures (11)

• Figure 1: Illustration of the CaPaint overview and advantage across SOTA ST causal model on complexity.
• Figure 2: The details of CaPaint. (Upper.) The initial phase of discovering causal patches. (Bottom.) The update phase designed to eliminate spurious correlation shifts. Following the upstream training of the ViT, a diffusion model is trained in parallel. Using the identified causal patches as conditions, this generative model then performs inpainting for generating multiple sequences.
• Figure 3: Different SCM architectures of SOTA and CaPaint.
• Figure 4: Visualization of prediction results for TaxiBJ+ and SEVIR datasets. The left side shows the predicted results of the last 5 frames for TaxiBJ+. The middle presents the results of long-term predictions for SEVIR, displaying the last five frames from step 10 $\rightarrow$ step 20. The right side compares SSIM metrics with and without the incorporation of CaPaint.
• Figure 5: SSIM improvement across different datasets using the Mmvp model