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

Conference: AAAI2026 arXiv: 2508.16623
Code: RWLinno/RAST
Area: Autonomous Driving Keywords: traffic prediction, retrieval-augmented, spatio-temporal forecasting, memory bank, STGNN

TL;DR

This work introduces the RAG paradigm into spatio-temporal forecasting by maintaining a dual-dimensional memory bank to store historical spatio-temporal patterns and retrieving them at inference time for fusion. The resulting general-purpose retrieval-augmented spatio-temporal prediction framework, RAST, achieves state-of-the-art performance on six traffic datasets while requiring only 1/12 the GPU memory of competing methods.

Background & Motivation

State of the Field

Spatio-temporal graph neural networks (STGNNs) and pre-trained models continue to advance traffic prediction benchmarks. Methods such as DCRNN, STGCN, and DSTAGNN capture complex dependencies by stacking spatio-temporal attention layers. Large pre-trained models (iTransformer, TimeMixer) have also been explored for spatio-temporal forecasting, with mixed results.

Limitations of Prior Work

(1) Limited context capacity — bounded embedding lengths make it difficult to fully model complex spatio-temporal dependencies in large-scale traffic networks (thousands of nodes, tens of thousands of time steps); stacking more layers yields diminishing marginal returns. (2) Low fine-grained predictability — spatio-temporal data exhibits inherent heterogeneity (high variance across spatial nodes, irregular temporal periodicity), and existing methods offer limited improvement for low-predictability regions. (3) Complex STGNNs incur prohibitive memory and computational costs (e.g., D2STGNN requires 38.36 GB GPU memory).

Root Cause

Model parameter budgets are finite and cannot retain fine-grained information about all historical patterns. Increasing model complexity improves capacity at the cost of exploding computation and memory, while providing limited benefit for low-predictability patterns.

Solution

Goal: Compensate for the model's limited context capacity through external memory storage and retrieval, without increasing the number of model parameters. Key Insight: Inspired by RAG in NLP — where LLMs retrieve external knowledge bases to supplement parametric memory — this work draws an analogy to spatio-temporal forecasting by storing and retrieving historical fine-grained patterns. Core Idea: Maintain a dual-dimensional (temporal + spatial) memory bank, use FAISS for efficient similar-pattern retrieval, and fuse retrieved results via cross-attention to augment prediction.

Method

Overall Architecture

RAST consists of five core components forming a complete pipeline: Decoupled Encoder (input decomposition) → Query Generator (query construction) → Retrieval Store (memory bank management) → ST-Retriever (retrieval) → Backbone Predictor (forecasting). The input is a traffic graph \((\mathbf{X}, \mathcal{G})\) and the output is a future time-step forecast.

Key Designs

1. Decoupled Encoder + Query Generator:

   • Function: Decomposes the input into independent temporal and spatial features and generates a fused query.
   • Mechanism: Temporal features \(\mathbf{E}_{tp} = \sigma(\text{Conv2D}(\mathbf{X}))\) capture local temporal patterns; spatial features \(\mathbf{E}_{sp} = \sigma(\mathbf{W}_{sp}(\mathbf{X}, \mathcal{G}))\) exploit graph structure to capture spatial correlations. The concatenated features are passed through \(L\) residual FFN layers to produce a context-aware fused query \(\mathcal{Q}_{st}\).
   • Design Motivation: Decoupled processing allows temporal and spatial features to be learned independently; concatenation-based fusion ensures the query encodes information from both dimensions.

2. Spatio-temporal Retrieval Store + ST-Retriever:

   • Function: Efficiently stores and retrieves historical spatio-temporal patterns.
   • Mechanism: A dual-dimensional memory bank \(\mathcal{M} = \{\mathcal{M}_{sp}, \mathcal{M}_{tp}\}\) is maintained, storing chunked embedding vectors and associated metadata. FAISS indices enable L2-distance Top-\(k\) retrieval. An entropy-driven momentum score update is introduced: \(\omega'_i = \omega_i + \text{softmax}(s_i + \lambda \cdot \mathcal{H}(\mathbf{v}_i)) / \tau\), balancing pattern freshness against storage stability.
   • Design Motivation: FAISS indices support GPU acceleration and periodic rebuilding to keep retrieval latency tractable; momentum-based management prevents memory instability caused by frequent updates.

3. Cross-Attention Fusion + Backbone Predictor:

   • Function: Fuses retrieved patterns with the current query for prediction.
   • Mechanism: Cross-attention is applied separately over retrieved spatio-temporal embeddings: \(\mathcal{R}_f = \text{Attn}(\mathcal{Q}_{st}, \mathcal{R}_s, \mathcal{R}_t)\). The backbone is designed as a generic interface that can be plugged into existing pre-trained STGNNs or a simple MLP.
   • Design Motivation: Cross-attention allows the model to dynamically select useful information from retrieved results conditioned on the current context; the generic interface enables RAST to serve as an enhancement module for any STGNN.

Key Experimental Results

Main Results

Evaluations are conducted on six traffic datasets: PEMS03/04/07/08, SD (large-scale), and GBA (large-scale).

Dataset	Metric	RAST	Prev. SOTA	Gain
PEMS04	MAE	18.39	STDN 19.19	4.2%
PEMS07	MAE	19.52	DSTAGNN 21.42	8.9%
PEMS08	MAE	14.16	AGCRN 15.95	11.2%
SD	MAE	18.39	STGODE 19.55	5.9%
GBA	MAE	20.64	DCRNN 23.13	10.8%

Efficiency comparison on the SD dataset: RAST trains at 45.53 s/epoch, infers in 10.15 s, and uses 3.22 GB GPU memory, whereas D2STGNN requires 1,014.89 s/epoch and 38.36 GB memory.

Ablation Study

Configuration	PEMS08 MAE	Change
Full RAST	14.16	—
w/o Fusion Query	17.79	+25.6%, most critical component
w/o ST-Retriever	15.75	+11.2%
w/o Memory Update	14.89	+5.2%
MLP-only backbone	14.53	only +2.6% (retrieval matters more than backbone)

Key Findings

• Fusion Query contributes the most (removal causes a 25.6% MAE degradation), indicating that the integration of retrieved information is the core component.
• Using only an MLP as the backbone still outperforms most complex STGNNs, demonstrating that the retrieval mechanism rather than model complexity drives performance.
• Advantages are more pronounced on large-scale SD/GBA datasets, confirming the greater value of external memory at scale.
• GPU memory efficiency improves by 12× (3.22 GB vs. 38.36 GB) and training speed improves by 22×.

Highlights & Insights

• First RAG framework for spatio-temporal forecasting: Transfers the retrieval-augmented paradigm from NLP to spatio-temporal prediction with a conceptually clean and general design.
• External memory replaces parameter stacking: A lightweight MLP combined with retrieval surpasses complex STGNNs, revealing that "memory" is more important than "computation" in spatio-temporal forecasting.
• Universal interface design: RAST can serve as a plug-and-play enhancement module for existing pre-trained STGNNs.

Limitations & Future Work

• Performance on PEMS03 does not reach state of the art, suggesting that the retrieval advantage diminishes on small-scale or topologically specific networks.
• Validation is limited to traffic prediction; other spatio-temporal forecasting tasks (meteorology, energy, etc.) are not covered.
• The optimal update frequency and capacity of the Retrieval Store require manual tuning.
• Retrieval latency during online inference is not analyzed in detail.

Related Work & Insights

• vs. traditional STGNNs (DCRNN/DSTAGNN): RAST replaces layer stacking with external memory, achieving 12× efficiency gains alongside superior performance.
• vs. large pre-trained models (iTransformer/TimeMixer): These general-purpose time-series models underperform on traffic prediction; RAST surpasses them with a smaller, domain-specific model.
• vs. knowledge distillation (STKD): RAST acquires knowledge through retrieval rather than distillation, offering greater flexibility without requiring a teacher model.
• The RAG paradigm in vision and multimodal domains is worth monitoring; spatio-temporal retrieval stores, as a form of learned non-parametric memory, are generalizable to video understanding, trajectory prediction, and related tasks.

Rating

• Novelty: ⭐⭐⭐⭐ Novel entry point with RAG for spatio-temporal forecasting
• Experimental Thoroughness: ⭐⭐⭐⭐ 6 datasets + 21 baselines + complete ablation + efficiency analysis
• Writing Quality: ⭐⭐⭐⭐ Clear structure with explicit motivation
• Value: ⭐⭐⭐⭐ General-purpose framework with practical potential

Related Papers

• [CVPR 2026] RAG-TP: A General Framework for Vehicle Trajectory Prediction via Retrieval-Augmented Generation
• [AAAI 2026] Rethinking the Spatio-Temporal Alignment of End-to-End 3D Perception
• [CVPR 2026] Spatial Retrieval Augmented Autonomous Driving
• [AAAI 2026] CaTFormer: Causal Temporal Transformer with Dynamic Contextual Fusion for Driving Intention Prediction
• [AAAI 2026] Meta Dynamic Graph for Traffic Flow Prediction