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

TL;DR

RAST is a novel framework that enhances traffic prediction by integrating retrieval mechanisms with spatio-temporal models, addressing limitations in contextual understanding and fine-grained predictability.

Contribution

It introduces a retrieval-augmented approach with a decoupled encoder, retrieval store, and flexible predictor, improving accuracy and efficiency in traffic forecasting.

Findings

Outperforms existing models on six real-world datasets

Maintains computational efficiency with improved accuracy

Effectively captures complex spatio-temporal dependencies

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.