Traffic forecasting with adaptive spatio-temporal masked pretraining

The complex spatio-temporal dependencies of urban traffic systems present a persistent challenge for accurate forecasting. Masked pretraining, particularly with autoencoders (MAEs), has emerged as a powerful paradigm for learning robust representations from long-term traffic data. However, the efficacy of these methods is often limited by simple random masking, which fails to account for the unique structural properties of traffic patterns. To address this, we propose the Adaptive Spatio-Temporal Masked Autoencoder (AST-MAE) for traffic forecasting with masked pretraining. Instead of masking inputs randomly, AST-MAE adaptively identifies and masks less informative tokens in both spatial and temporal dimensions. It employs specialized auxiliary networks to predict the contextual importance of each input, using graph neural networks to discern spatial correlations and temporal attention to capture sequential patterns. By selectively masking redundant or predictable data, our pretraining forces the encoder to learn a more sophisticated understanding of critical spatio-temporal dynamics. Models initialized with our pretrained representations achieve superior performance on downstream forecasting tasks. Extensive experiments on four real-world datasets validate that our adaptive pretraining with AST-MAE significantly outperforms state-of-the-art baselines, demonstrating its power to enhance model accuracy in traffic forecasting.