Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

RAST: A Retrieval Augmented Spatio-Temporal Framework for Traffic Prediction

Weilin Ruan, Xilin Dang, Ziyu Zhou, Sisuo Lyu, Yuxuan Liang

TL;DR

RAST tackles accurate traffic forecasting under limited contextual capacity and heterogeneous spatio-temporal patterns by introducing a retrieval-augmented spatio-temporal framework. It decouples spatial and temporal encodings, maintains dual memory banks, and retrieves context-relevant historical patterns via an information-theoretic ST-Retriever, then fuses retrieved signals with a universal backbone predictor through cross-attention. Extensive experiments on six real-world datasets show state-of-the-art performance and favorable efficiency, with ablations validating each component (notably the query generator) and hyperparameters that balance retrieval quality with computation. The framework offers a lightweight, plug-in enhancement for existing pre-trained STGNNs and holds promise for broader spatio-temporal domains beyond traffic forecasting.

Abstract

Traffic prediction is a cornerstone of modern intelligent transportation systems and a critical task in spatio-temporal forecasting. Although advanced Spatio-temporal Graph Neural Networks (STGNNs) and pre-trained models have achieved significant progress in traffic prediction, two key challenges remain: (i) limited contextual capacity when modeling complex spatio-temporal dependencies, and (ii) low predictability at fine-grained spatio-temporal points due to heterogeneous patterns. Inspired by Retrieval-Augmented Generation (RAG), we propose RAST, a universal framework that integrates retrieval-augmented mechanisms with spatio-temporal modeling to address these challenges. Our framework consists of three key designs: 1) Decoupled Encoder and Query Generator to capture decoupled spatial and temporal features and construct a fusion query via residual fusion; 2) Spatio-temporal Retrieval Store and Retrievers to maintain and retrieve vectorized fine-grained patterns; and 3) Universal Backbone Predictor that flexibly accommodates pre-trained STGNNs or simple MLP predictors. Extensive experiments on six real-world traffic networks, including large-scale datasets, demonstrate that RAST achieves superior performance while maintaining computational efficiency.

RAST: A Retrieval Augmented Spatio-Temporal Framework for Traffic Prediction

TL;DR

RAST tackles accurate traffic forecasting under limited contextual capacity and heterogeneous spatio-temporal patterns by introducing a retrieval-augmented spatio-temporal framework. It decouples spatial and temporal encodings, maintains dual memory banks, and retrieves context-relevant historical patterns via an information-theoretic ST-Retriever, then fuses retrieved signals with a universal backbone predictor through cross-attention. Extensive experiments on six real-world datasets show state-of-the-art performance and favorable efficiency, with ablations validating each component (notably the query generator) and hyperparameters that balance retrieval quality with computation. The framework offers a lightweight, plug-in enhancement for existing pre-trained STGNNs and holds promise for broader spatio-temporal domains beyond traffic forecasting.

Abstract

Traffic prediction is a cornerstone of modern intelligent transportation systems and a critical task in spatio-temporal forecasting. Although advanced Spatio-temporal Graph Neural Networks (STGNNs) and pre-trained models have achieved significant progress in traffic prediction, two key challenges remain: (i) limited contextual capacity when modeling complex spatio-temporal dependencies, and (ii) low predictability at fine-grained spatio-temporal points due to heterogeneous patterns. Inspired by Retrieval-Augmented Generation (RAG), we propose RAST, a universal framework that integrates retrieval-augmented mechanisms with spatio-temporal modeling to address these challenges. Our framework consists of three key designs: 1) Decoupled Encoder and Query Generator to capture decoupled spatial and temporal features and construct a fusion query via residual fusion; 2) Spatio-temporal Retrieval Store and Retrievers to maintain and retrieve vectorized fine-grained patterns; and 3) Universal Backbone Predictor that flexibly accommodates pre-trained STGNNs or simple MLP predictors. Extensive experiments on six real-world traffic networks, including large-scale datasets, demonstrate that RAST achieves superior performance while maintaining computational efficiency.

Paper Structure

This paper contains 40 sections, 18 equations, 4 figures, 7 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Spatio-temporal Forecasting (STF)
  4. Retrieval-Augmented Generation (RAG)
  5. Methodology
  6. Problem Formulation
  7. Data Encoding and Query Construction
  8. Dual-dimension Feature Disentanglement.
  9. Context-Aware Query Generation.
  10. Spatio-temporal Retrieval Store
  11. Pattern Indexing and Storage.
  12. Information-Theoretic ST-Retriever.
  13. Momentum-Based Memory Management.
  14. Cross-Attention Knowledge Fusion.
  15. Prediction and Optimization
  16. ...and 25 more sections

Figures (4)

  • Figure 1: Motivation of RAST. Inspired by retrieval-augmented generation for large language models, we design a specialized retrieval-augmented framework for STF tasks.
  • Figure 2: The spatial and temporal heterogeneity problem (a and b). Our methods utilize a retrieval-augmented mechanism with dual-dimension vector storage. By extracting low-predictable data points and retrieving history patterns, we improve model capability in those challenging cases.
  • Figure 3: Overview of the RAST framework. RAST integrates retrieval mechanisms with spatio-temporal modeling to enhance performance by maintaining and utilizing historical patterns. The proposed approach addresses the fundamental challenge of capturing long-term dependencies and complex spatio-temporal correlations in time series data.
  • Figure 4: Hyperparameter sensitivity analysis for RAST on the PEMS04 dataset. We chose 4 crucial hyperparameters and measured results by MAE and MAPE(%). Vertical dashed lines indicate optimal values for each parameter.