FlashST: A Simple and Universal Prompt-Tuning Framework for Traffic Prediction

TL;DR

FlashST tackles the challenge of distribution shift in traffic forecasting by introducing a universal spatio-temporal prompt-tuning framework. It combines a spatio-temporal in-context learning pipeline with a distribution-mapping module to align pre-training and downstream data, enabling fast, parameter-efficient adaptation of pre-trained models to new urban datasets. Through a two-stage process—pre-training the prompt network with a multi-component ST embedding and then tuning only the prompts for downstream tasks—FlashST achieves strong generalization and significant efficiency gains across multiple datasets and backbone models. The approach is model-agnostic, demonstrates notable reductions in training time, and provides insights via ablations and case studies, suggesting practical impact for scalable, cross-domain traffic prediction. Future work may explore integrating large language models for knowledge guidance within the FlashST framework.

Abstract

The objective of traffic prediction is to accurately forecast and analyze the dynamics of transportation patterns, considering both space and time. However, the presence of distribution shift poses a significant challenge in this field, as existing models struggle to generalize well when faced with test data that significantly differs from the training distribution. To tackle this issue, this paper introduces a simple and universal spatio-temporal prompt-tuning framework-FlashST, which adapts pre-trained models to the specific characteristics of diverse downstream datasets, improving generalization in diverse traffic prediction scenarios. Specifically, the FlashST framework employs a lightweight spatio-temporal prompt network for in-context learning, capturing spatio-temporal invariant knowledge and facilitating effective adaptation to diverse scenarios. Additionally, we incorporate a distribution mapping mechanism to align the data distributions of pre-training and downstream data, facilitating effective knowledge transfer in spatio-temporal forecasting. Empirical evaluations demonstrate the effectiveness of our FlashST across different spatio-temporal prediction tasks using diverse urban datasets. Code is available at https://github.com/HKUDS/FlashST.

Figures (6)

• Figure 1: Motivations behind FlashST. The left figure illustrates the diverse data distributions across various ST datasets, while the right figure demonstrates that the end-to-end model's parameters are overfit to training set A and fail to generalize to test set B.
• Figure 2: Our proposed FlashST framework adopts an architecture that integrates spatio-temporal in-context learning and a unified distribution mapping mechanism, offering an efficient and effective approach for spatio-temporal prompt-tuning across diverse scenarios.
• Figure 3: The convergence efficiency of FlashST.
• Figure 4: Ablation study of FlashST.
• Figure 5: Hyperparameter study of $\tau$ and $\lambda$.