diffIRM: A Diffusion-Augmented Invariant Risk Minimization Framework for Spatiotemporal Prediction over Graphs

2024-12-31

AI Summary

Addressing out-of-distribution (OOD) generalization in spatiotemporal graph forecasting, this paper proposes an invariant learning framework integrating causal inference and graph diffusion augmentation. Methodologically, it jointly incorporates environmental diversity construction and the principle of invariance existence—novel in spatiotemporal graph prediction. A causal-mask-guided graph diffusion mechanism is introduced to enable interpretable data augmentation while jointly optimizing for invariant risk minimization (IRM). The architecture unifies a causal mask generator, a graph diffusion model, IRM regularization, and a spatiotemporal graph neural network. Evaluated on three real-world human mobility datasets—SafeGraph, PeMS04, and PeMS08—the method significantly outperforms existing OOD baselines, achieving superior robustness and cross-environment generalization performance.

Technology Category

• Planning, Routing, and Scheduling: Optimization of Spatio-temporal Systems
• Machine Learning: Causal Learning
• Computer Vision: Diffusion Models for Vision

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Spatiotemporal prediction over graphs (STPG) is challenging, because real-world data suffers from the Out-of-Distribution (OOD) generalization problem, where test data follow different distributions from training ones. To address this issue, Invariant Risk Minimization (IRM) has emerged as a promising approach for learning invariant representations across different environments. However, IRM and its variants are originally designed for Euclidean data like images, and may not generalize well to graph-structure data such as spatiotemporal graphs due to spatial correlations in graphs. To overcome the challenge posed by graph-structure data, the existing graph OOD methods adhere to the principles of invariance existence, or environment diversity. However, there is little research that combines both principles in the STPG problem. A combination of the two is crucial for efficiently distinguishing between invariant features and spurious ones. In this study, we fill in this research gap and propose a diffusion-augmented invariant risk minimization (diffIRM) framework that combines these two principles for the STPG problem. Our diffIRM contains two processes: i) data augmentation and ii) invariant learning. In the data augmentation process, a causal mask generator identifies causal features and a graph-based diffusion model acts as an environment augmentor to generate augmented spatiotemporal graph data. In the invariant learning process, an invariance penalty is designed using the augmented data, and then serves as a regularizer for training the spatiotemporal prediction model. The real-world experiment uses three human mobility datasets, i.e. SafeGraph, PeMS04, and PeMS08. Our proposed diffIRM outperforms baselines.

Problem

• Out-of-Distribution (OOD) Prediction
• Invariant Risk Minimization (IRM)
• Spatiotemporal Prediction

Innovation

• diffIRM
• invariant learning
• data augmentation