STRAP: Spatio-Temporal Pattern Retrieval for Out-of-Distribution Generalization

Yifan Zhang (Princeton University (yifzhang.com)) · Haoyu Zhang (Harbin Institute of Technology (Shenzhen)) · Hao Miao (Hong Kong Polytechnic University) · WentaoZhang (ShanghaiTech University) · Xinke Jiang (Peking University) · Yuchen Fang (University of Electronic Science and Technology of China)

catastrophic forgettingcontinue learning pipelinedynamic graph-structured datahistorical informationknowledge-balancing objectivepattern libraryplug-and-play prompting mechanismreal-world streaming graph datasetsretrieval-augmented learningsemantic informationspatio-temporal graph neural networksspatio-temporal out-of-distributionspatio-temporal patternsspatio-temporal representationsstrong generalization capabilitystructural information

Spatio-Temporal Graph Neural Networks (STGNNs) have emerged as a powerful tool for modeling dynamic graph-structured data across diverse domains. However, they often fail to generalize in Spatio-Temporal Out-of-Distribution (STOOD) scenarios, where both temporal dynamics and spatial structures evolve beyond the training distribution. To address this problem, we propose an innovative Spatio-Temporal Retrieval-Augmented Pattern Learning framework, STRAP, which enhances model generalization by integrating retrieval-augmented learning into the STGNN continue learning pipeline. The core of STRAP is a compact and expressive pattern library that stores representative spatio-temporal patterns enriched with historical, structural, and semantic information, which is obtained and optimized during the training phase. During inference, STRAP retrieves relevant patterns from this library based on similarity to the current input and injects them into the model via a plug-and-play prompting mechanism. This not only strengthens spatio-temporal representations but also mitigates catastrophic forgetting. Moreover, STRAP introduces a knowledge-balancing objective to harmonize new information with retrieved knowledge. Extensive experiments across multiple real-world streaming graph datasets show that STRAP consistently outperforms state-of-the-art STGNN baselines on STOOD tasks, demonstrating its robustness, adaptability, and strong generalization capability without task-specific fine-tuning.