Revealing the Power of Masked Autoencoders in Traffic Forecasting

TL;DR

This paper introduces STMAE, a novel masked autoencoder framework that pretrains on partially observed traffic data to improve the accuracy and stability of existing spatial-temporal traffic forecasting models.

Contribution

The paper presents a plug-and-play masked autoencoder framework with dual masking strategies that enhances the performance of existing traffic prediction models.

Findings

STMAE significantly improves forecasting accuracy on traffic benchmarks.

Pretraining with STMAE enhances model stability and robustness.

The framework is compatible with various existing models.

Abstract

Traffic forecasting, crucial for urban planning, requires accurate predictions of spatial-temporal traffic patterns across urban areas. Existing research mainly focuses on designing complex models that capture spatial-temporal dependencies among variables explicitly. However, this field faces challenges related to data scarcity and model stability, which results in limited performance improvement. To address these issues, we propose Spatial-Temporal Masked AutoEncoders (STMAE), a plug-and-play framework designed to enhance existing spatial-temporal models on traffic prediction. STMAE consists of two learning stages. In the pretraining stage, an encoder processes partially visible traffic data produced by a dual-masking strategy, including biased random walk-based spatial masking and patch-based temporal masking. Subsequently, two decoders aim to reconstruct the masked counterparts from both spatial and temporal perspectives. The fine-tuning stage retains the pretrained encoder and integrates it with decoders from existing backbones to improve forecasting accuracy. Our results on traffic benchmarks show that STMAE can largely enhance the forecasting capabilities of various spatial-temporal models.