Efficient Prompt Learning for Traffic Forecasting

TL;DR

This paper introduces SimpleST, a lightweight prompt tuning framework that enhances the generalization of pre-trained spatio-temporal GNNs for traffic forecasting, achieving better accuracy and efficiency on real-world datasets.

Contribution

It presents a novel prompt-based adaptation method for spatio-temporal GNNs that maintains fixed model parameters while improving out-of-distribution performance.

Findings

SimpleST outperforms baseline models in accuracy on five datasets.

The approach reduces adaptation overhead and computational complexity.

Experiments demonstrate improved generalization to distribution shifts.

Abstract

Accurate traffic prediction is essential for optimizing transportation systems, enhancing resource allocation, and improving overall urban administration. Spatio-temporal graph neural networks (GNNs) have achieved state-of-the-art performance and have been widely used in various spatio-temporal prediction scenarios. However, these prediction methods often exhibit low generalization ability, struggling with distribution shifts caused by spatio-temporal dynamics. To address this challenge, we propose an approach to enhance the generalization and adaptation of spatio-temporal GNNs through efficient prompting. Specifically, we introduce a lightweight and model-agnostic prompt tuning framework for spatio-temporal GNNs, named SimpleST. It facilitates adapting pre-trained spatio-temporal GNNs to novel distributions while keeping the model parameters fixed. This prompt mechanism reduces the overhead and complexity of adaptation, enabling efficient utilization of pre-trained models for out-of-distribution generalization. Extensive experiments conducted on five real-world urban spatio-temporal datasets demonstrate the superiority of our approach in terms of prediction accuracy and computational efficiency.