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

October 30, 2025 ยท View on GitHub

License Python PyTorch arXiv GitHub stars

๐Ÿ“„ Paper | ๐Ÿ“Š Datasets


๐Ÿ“‹ Table of Contents


โœจ Overview

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 STRAP, which enhances model generalization by integrating retrieval-augmented learning into the STGNN continue learning pipeline. Extensive experiments across multiple real-world streaming graph datasets show that \methodname consistently outperforms state-of-the-art STGNN baselines on STOOD tasks, demonstrating its robustness, adaptability, and strong generalization capability without task-specific fine-tuning.

STRAP Framework

STRAP Framework Architecture


๐Ÿ—๏ธ Repository Structure

STRAP/
โ”‚
โ”œโ”€โ”€ ๐Ÿ“„ README.md                    # Project documentation
โ”œโ”€โ”€ ๐Ÿ“„ LICENSE                      # Apache 2.0 License
โ”œโ”€โ”€ ๐Ÿ“„ environment.yaml             # Conda environment configuration
โ”œโ”€โ”€ ๐Ÿš€ main.py                      # Main entry point for experiments
โ”œโ”€โ”€ ๐Ÿš€ stkec_main.py               # STKEC experiments entry point
โ”œโ”€โ”€ ๐Ÿ“œ run.sh                       # Batch experiment execution script
โ”‚
โ”œโ”€โ”€ ๐Ÿ“ conf/                        # โš™๏ธ Configuration files
โ”‚   โ”œโ”€โ”€ AIR/                       # Air quality dataset configs
โ”‚   โ”œโ”€โ”€ ENERGY-Wind/               # Wind energy dataset configs
โ”‚   โ””โ”€โ”€ PEMS/                      # Traffic dataset configs
โ”‚       โ”œโ”€โ”€ strap.json
โ”‚       โ”œโ”€โ”€ ewc.json
โ”‚       โ””โ”€โ”€ ...
โ”‚
โ”œโ”€โ”€ ๐Ÿ“ src/                         # ๐Ÿ’ป Source code
โ”‚   โ”œโ”€โ”€ dataer/                    # Data loading and preprocessing
โ”‚   โ”‚   โ”œโ”€โ”€ ...
โ”‚   โ”‚
โ”‚   โ”œโ”€โ”€ model/                     # Model implementations
โ”‚   โ”‚   โ”œโ”€โ”€ ...             # Model components
โ”‚   โ”‚
โ”‚   โ””โ”€โ”€ trainer/                   # Training and evaluation
โ”‚       โ”œโ”€โ”€ ...
โ”‚
โ”œโ”€โ”€ ๐Ÿ“ utils/                       # ๐Ÿ› ๏ธ Utility functions
โ”‚   โ”œโ”€โ”€ ...
โ”‚
โ”œโ”€โ”€ ๐Ÿ“ font/                        # Font files for visualization
โ”œโ”€โ”€ ๐Ÿ“ log/                         # ๐Ÿ“Š Training logs and checkpoints
โ””โ”€โ”€ ๐Ÿ“ data/                        # ๐Ÿ’พ Dataset storage (create this)

๐Ÿš€ Getting Started

๐Ÿ“‹ Prerequisites

Before you begin, ensure you have the following installed:

  • Conda or Miniconda (Download)
  • NVIDIA GPU with CUDA support (recommended)
  • Python 3.8+

๐Ÿ’ป Usage

# ENERGY-Wind, the same for other datasets.
bash run.sh

๐Ÿ™ Acknowledgements

We would like to express our gratitude to:

  • EAC: We thank the authors for their excellent work. Our implementation builds upon their codebase: EAC Repository

๐Ÿ“ Citation

If you find this work useful for your research, please consider citing our paper:

@article{zhang2025strap,
  title={STRAP: Spatio-Temporal Pattern Retrieval for Out-of-Distribution Generalization},
  author={Zhang, Haoyu and Zhang, Wentao and Miao, Hao and Jiang, Xinke and Fang, Yuchen and Zhang, Yifan},
  journal={arXiv preprint arXiv:2505.19547},
  year={2025}
}

๐ŸŒŸ Star History

Star History Chart