Efficient Prompt Learning for Traffic Forecasting

2026-05-07

AI Summary

This work addresses the limited generalization of spatiotemporal graph neural networks (STGNNs) in traffic forecasting, which often stems from distribution shifts caused by complex spatiotemporal dynamics. To tackle this issue, the authors propose SimpleST—a lightweight, model-agnostic prompt-tuning framework that introduces prompt learning into STGNNs for the first time. By freezing the parameters of a pre-trained model and leveraging an efficient spatiotemporal prompting mechanism, SimpleST enables rapid adaptation to new data distributions without extensive retraining. Extensive experiments on five real-world urban traffic datasets demonstrate that SimpleST consistently achieves substantial improvements in both prediction accuracy and computational efficiency, highlighting its strong out-of-distribution generalization capabilities.

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.

Problem

• traffic forecasting
• spatio-temporal GNNs
• distribution shift
• generalization
• out-of-distribution

Innovation

• prompt tuning
• spatio-temporal GNNs
• out-of-distribution generalization
• model-agnostic adaptation
• traffic forecasting