Spatial-Temporal-Decoupled Masked Pre-training for Spatiotemporal Forecasting

1 Introduction

Spatiotemporal data collected by sensor networks has become a vital area of research with many real-world applications. It benefits from extra spatial context like sensor locations and road networks that reveal dependencies between sensors. Consequently, a key distinction from typical multivariate time series is that spatiotemporal data exhibits spatiotemporal heterogeneity. Specifically, while time series vary across different locations (urban centers against suburban areas) and day types (weekday versus weekend), they demonstrate consistent, predictable patterns within similar contexts. Hence, accurately predicting spatiotemporal data hinges on effectively capturing such heterogeneity. Figure 1a shows the traffic flow of sensor 7 in a spatiotemporal dataset PEMS04, revealing distinct weekday and weekend patterns especially during peak hours. Weekdays experience a morning rush hour peak, whereas weekend is more evenly distributed without significant peaks, illustrating temporal heterogeneity in weekly patterns. Additionally, Figure 1b displays the traffic flow of sensor 5, 86, 155, and 177 during the same period. Sensor 5 and 177 exhibit distinct peaks and troughs, while sensor 86 and 155 remain relatively stable throughout the day, demonstrating spatial heterogeneity. When the data scale is small, the heterogeneity is distinctly visible as shown above. However, when months of data from hundreds of sensors are given, the spatial and temporal heterogeneity becomes highly mixed. Previous researchers have done various attempts for spatiotemporal forecasting: embedding GCN into TCN [Yu et al., 2018; Wu et al., 2019] or RNN [Li et al., 2018; Bai et al., 2020], or applying transformer [Jiang et al., 2023a; Liu et al., 2023] along spatiotemporal axes. But these models often have difficulty in distinguishing the spatial and temporal heterogeneity in a clear way. So learning clear heterogeneity is still the primary challenge for spatiotemporal forecasting.

∗Corresponding author.

Moreover, most of the existing models are trained in an end-to-end manner. Due to their high model complexity, their input horizons are often restricted to a short value (usually 12 steps) [Yu et al., 2018; Wu et al., 2019; Song et al., 2020]. This limitation will make the models suffer from an issue denoted as spatiotemporal mirage: i) dissimilar input time series followed by similar future values; ii) similar input time series followed by dissimilar future values. We illustrate this phenomenon by taking traffic flow of sensor 215 and 279 as an example in Figure 1c. In the late evening, the two sensors show divergent historical data trends but similar future flow. Conversely, in the afternoon, their historical data trends align closely, yet their future data differ dramatically. Essentially, the reason behind this is that existing models can only capture the fragmented heterogeneity instead of the complete one. Therefore, how to make these models robust on such spatiotemporal mirage issue is the second challenge.

In this study, we focus our approach on learning clear and complete spatiotemporal heterogeneity through pre-training. In particular, masked pre-training has shown tremendous effectiveness in natural language processing [Kenton and Toutanova, 2019] and computer vision [Bao et al., 2021]. The core idea is to mask parts of the input sequence during pre-training, requiring the model to reconstruct the missing contents. In this way, the model learns context-rich representations which can augment various downstream tasks. Motivated by these benefits, we propose a novel spatialtemporal-decoupled masked pre-training framework called STD-MAE. It offers an effective and efficient solution for learning clear and complete spatiotemporal heterogeneity through pre-training. Such learned heterogeneity can be flawlessly integrated into downstream baselines to see through spatiotemporal mirages. In summary, our key contributions are as follows:

• We devise a pre-training framework on spatiotemporal data that can largely enhance downstream spatiotemporal predictors of arbitrary architectures without modifying their original structures.
• We propose a novel spatial-temporal-decoupled masking strategy to effectively learn spatial and temporal heterogeneity by capturing long-range context across spatial and temporal dimensions.
• We validate STD-MAE on four benchmarks (PEMS03, PEMS04, PEMS07, and PEMS08) with typical backbones. Quantitative enhancement on baselines highlights the exceptional performance of STD-MAE. Qualitative analyses demonstrate its power to capture meaningful long-range spatiotemporal patterns.

2 RelatedWork

2.1 Spatiotemporal Forecasting

Spatiotemporal forecasting [Jiang et al., 2021] aims to predict future spatiotemporal series by analyzing historical data.

Early work mainly relied on traditional time series models [Pan et al., 2012; Stock and Watson, 2001]. To capture the complex temporal dependencies, RNNs [Hochreiter and Schmidhuber, 1997; Chung et al., 2014] and CNNs [Oord et al., 2016] have gained popularity for better modeling spatiotemporal data and achieving improved predictions.

Nevertheless, these models overlook crucial spatial correlations, limiting predictive performance on networked road systems. To further capture spatiotemporal features jointly, some studies integrate Graph Convolutional Networks (GCNs) with temporal models [Yu et al., 2018; Li et al., 2018]. Following this line of research, several novel spatiotemporal models have been proposed in recent years [Wu et al., 2019; Xu et al., 2020; Bai et al., 2020; Cao et al., 2020; Wu et al., 2020; Guo et al., 2021a; Deng et al., 2022; Han et al., 2021; Jiang et al., 2023b; Deng et al., 2024]. Attention mechanisms [Vaswani et al., 2017] have also profoundly influenced spatiotemporal forecasting. A series of transformers [Xu et al., 2020; Jiang et al., 2023a; Liu et al., 2023] are proposed and exhibit superior performance, highlighting their effectiveness in capturing the spatiotemporal relations. However, these end-to-end models only focus on short-term input that limits them to capture complete spatiotemporal dependencies.

2.2 Masked Pre-training

Masked pre-training has emerged as a highly effective technique for self-supervised representation learning in both natural language processing (NLP) and computer vision (CV). The key idea is to train models to predict masked-out parts of the input based on visible context. In NLP, approaches like BERT [Kenton and Toutanova, 2019] use masked language modeling to predict randomly masked tokens with bidirectional context. Subsequent models [Liu et al., 2019; Lan et al., 2019] introduced more efficient masking techniques and demonstrated performance gains from longer pre-training. In CV, similar masking strategies have been adopted. Methods like BEiT [Bao et al., 2021] and Masked AutoEncoder (MAE) [He et al., 2022] mask out random patches of input images and do reconstruction based on unmasked patches. In both domains, masked pre-training produces substantial improvements on various downstream tasks.

Recently, many researchers have attempted to employ pretraining techniques on time series data to obtain superior hidden representations [Nie et al., 2022; Shao et al., 2022b; Li et al., 2023]. However, these methods are either channelindependent or neglect pre-training in the spatial dimension. Our proposed STD-MAE introduces a novel spatialtemporal-decoupled masking strategy during pre-training. By masking separately on spatial and temporal dimensions, the learned representations can effectively capture the intricate long-range heterogeneity in spatiotemporal data.

3 ProblemDefinition

Spatiotemporal forecasting is a specialized multivariate time series forecasting problem. Given multivariate time series Xt-(T-1):t in the past T time steps, our goal is to predict the future ̂ T time steps as: [Xt-(T-1), ...,Xt] →

[Xt+1, ...,Xt+̂ T ], where Xi ∈ RN×C for the i-th time step, N is the number of spatial nodes, and C is the number of the information channel. Here C=1 in our datasets.

4 Methodology

This section delves into the technical specifics of our proposed spatial-temporal-decoupled masked pre-training framework (STD-MAE) as delineated in Figure 2.

4.1 Spatial-Temporal Masked Pre-training

Spatial-Temporal-Decoupled Masking. In the standard spatiotemporal forecasting task, the input length T for Xt-(T-1):t is usually equal to 12 (each step corresponds to 5-minute interval) [Yu et al., 2018; Wu et al., 2019; Song et al., 2020], thus end-to-end models often fall into mirages outlined in Figure 1c. So we intend to introduce a masked pre-training phase involving long-range input. Since spatiotemporal data has an additional temporal dimension compared to image data and an extra spatial dimension over language data, a straightforward idea is to apply the original masked pre-training [Kenton and Toutanova, 2019; He et al., 2022] directly by mixing the temporal and spatial dimension as one. However, this is unfeasible due to the square-level time and space complexity. Therefore, we propose a novel method called spatial-temporal-decoupled masking during masked pre-training. This approach separately executes mask-reconstruction tasks along temporal and spatial dimension. Such decoupled masking mechanism allows the model to learn representation that can capture clearer heterogeneity. With this wider and clearer view, downstream predictors can see through spatiotemporal mirages. Consequently, our method presents an efficient and effective solution for pre-training on spatiotemporal data, enhancing model robustness against the challenges posed by complex spatiotemporal heterogeneity and mirage.

Specifically, given an input spatiotemporal time series X ∈ RT×N×C , we propose the following masking strategies: (1) Spatial Masking (S-Mask) randomly masks the time series of N × r sensors, where r is the masking ratio between 0 and 1. This results in a spatially masked input ˜ X (S) ∈ RT×N(1-r)×C . (2) Temporal Masking (T-Mask) randomly masks the time series of T × r time steps. This yields a temporally masked input ˜ X (T) ∈ RT(1-r)×N×C . Both masking strategies can be viewed as random sampling from a Bernoulli distribution B(1 - r) with expectation 1 - r in the corresponding dimensions:

Intuitively, S-Mask forces the model to reconstruct the data of masked sensors solely from the other visible sensors, thus capturing long-range spatial heterogeneity. Similarly, temporal heterogeneity can be learned by utilizing the intrinsic visible series to reconstruct the entire time series with T-Mask.

Spatial-Temporal-Decoupled Masked AutoEncoder. Building upon the spatial-temporal-decoupled masking technique, we further propose the spatial-temporal-decoupled masked autoencoder. It consists of a temporal autoencoder (T-MAE) and a spatial autoencoder (S-MAE), both having a similar architecture. S-MAE applies self-attention along spatial dimension, while T-MAE performs self-attention along temporal dimension. Specifically, we consider long input with length Tlong, typically spanning several days. However, directly utilizing such long sequences leads to computational and memory challenges. To address this, we apply a patch embedding technique [Nie et al., 2022]. The long input is divided into non-overlapping patches of length Tp = Tlong/L using a patch window L. This yields a patched input Xp ∈ RTp×N×LC . We then project Xp through a fully connected layer to obtain the patch embedding Ep ∈ RTp×N×D , where D is the embedding dimension. Moreover, to simultaneously encode spatial and temporal positional information, we implement a two-dimensional positional encoding [Wang and Liu, 2021]. Given the patch embedding Ep, the spatiotemporal positional encoding Epos ∈ RTp×N×D can be calculated as follows:

We choose sinusoidal positional encoding instead of learned positional encoding because it can handle inputs of arbitrary length. The patch embedding Ep and positional encoding Epos are summed to obtain the final input embedding E ∈ RTp×N×D .

This input embedding E is subsequently masked by SMask and T-Mask strategies to obtain the visible spatial patch embedding ˜ E(S) and the visible temporal patch embedding ˜ E(T) . S-MAE and T-MAE generate spatial representations H(S) ∈ RTp×N(1-r)×D and temporal representations H(T) ∈ RTp(1-r)×N×D through a series of transformer layers, respectively. We denote NM = N × r as the number of masked sensors and TM = Tp×r as count of masked patches. By focusing only on visible parts, such design could reduce time and memory complexity.

A lightweight decoder is then applied to S-MAE and TMAE to reconstruct the masked input. The spatial and temporal decoders each consists of a padding layer, a standard transformer layer, and a regression layer. Such asymmetrical design could dramatically reduce the pre-training time [He et al., 2022]. In the padding layer, we use a shared learnable mask token V ∈ RD to indicate missing patches. Given spatial representation H(S) and temporal representation H(T) , spatial padding expands V to spatial mask tokens V (S) ∈ RTp×NM×D while temporal padding expands V to temporal mask tokens V (T) ∈ RTM×N×D . The same spatiotemporal positional encoding as the encoders is added to V (S) and V (T) . Then we perform concatenation operations respectively as [H(S) ; V (S) ] and [H(T) ; V (T) ] to get the full set of patches H(S) ,H(T) ∈ RTp×N×L. Subsequently, they are passed to the transformer layer. Finally, a regression layer is used to reconstruct the time series at the patch level. Formally, the reconstruction ̂Q(S) ∈ RTp×NM×L and ̂Q(T) ∈ RTM×N×L for the spatially and temporally masked inputs can be derived by:

where M (S) ∈ RNM is the spatial masked index and M (T) ∈ RTM is the temporal masked index.

Following other masked pre-training architectures [He et al., 2022; Tong et al., 2022], we compute the loss on the masked part only. Our loss function is computed by calculating the mean absolute error (MAE) between the ground truth and the reconstruction result. Given reconstruction ̂Q(S) and ̂Q(T) , the corresponding ground truth are Q(S) ∈ RTp×NM×L and Q(T) ∈ RTM×N×L. The two loss functions in spatial and temporal can be calculated as:

In summary, through the above spatial-temporal-decoupled masked pre-training, STD-MAE can capture clear and complete spatial and temporal heterogeneity.

4.2 Downstream Spatiotemporal Forecasting

STD-MAE can be seamlessly integrated into existing predictor structures. This operation is done by adding the spatial and temporal representations generated by STD-MAE to the hidden representation of the predictor. Concretely, we first feed long-range input with Tlong time steps into pre-trained spatial and temporal encoders to generate the corresponding spatial representation H(S) and temporal representation H(T) . Then, we apply a downstream spatiotemporal predictor Fθ with parameter θ to obtain the hidden representation H(F) ∈ RN×D′ of the widely-used short input Xt-(T-1):t through the following:

where D′ is the hidden representation dimension of the predictor. To align with H(F) , we truncate the representation H(S) and H(T) of the last T ′ patches, and reshape these two representations to H′(S) ∈ RN×T ′D and H′(T) ∈ RN×T ′D . Next, we project these two representations into D′ dimension through a two-layer MLP. Finally, the augmented representation H(Aug) ∈ RN×D′ could be derived by adding these representations together:

By far, H(Aug) includes representations generated by the predictor itself as well as the long-range spatial and temporal representations from STD-MAE, which can largely enhance the performance of the downstream spatiotemporal predictor.

Specifically, in our work, we choose GWNet [Wu et al., 2019] as our predictor due to its superior performance. We obtain the final representation by aggregating the hidden states from the skip connections across the multiple spatiotemporal layers of GWNet, along with the corresponding spatial and temporal representations generated by the STD-MAE. The augmented representation is then fed into GWNet's regression layers for prediction. Furthermore, we also test other classical spatiotemporal predictors with a variety of structures, i.e., DCRNN [Li et al., 2018], MTGNN [Wu et al., 2020], STID [Shao et al., 2022a] and STAEformer [Liu et al., 2023]. These experiments demonstrate the generality of STD-MAE. Details could be found in our ablation study.

5 Experiment

5.1 Experimental Setup

Datasets. To thoroughly evaluate the proposed STD-MAE model, we conduct extensive experiments on four real-world spatiotemporal benchmark datasets as listed in Table 1: PEMS03, PEMS04, PEMS07, and PEMS08 [Song et al., 2020]. The raw data has a fine-grained time resolution of 5 minutes between consecutive time steps. For data preprocessing, we perform Z-score normalization on the raw inputs.

Baselines. We compare STD-MAE with the following baseline methods. ARIMA [Fang et al., 2021], VAR [Song et al., 2020], SVR [Song et al., 2020], LSTM [Song et al., 2020], TCN [Lan et al., 2022], and Transformer [Vaswani et al., 2017] are time series models. For spatiotemporal models, we select several typical methods including DCRNN [Li et al., 2018], STGCN [Yu et al., 2018], ASTGCN [Guo et al., 2019], GWNet [Wu et al., 2019], STSGCN [Song et al., 2020], STFGNN [Li and Zhu, 2021], STGODE [Fang et al., 2021], DSTAGNN [Lan et al., 2022], ST-WA [Cirstea et al., 2022], ASTGNN [Guo et al., 2021b], EnhanceNet [Cirstea et al., 2021], AGCRN [Bai et al., 2020], Z-GCNETs [Chen et al., 2021], STEP [Shao et al., 2022b], PDFormer [Jiang et al., 2023a] and STAEformer [Liu et al., 2023].

Settings. Following previous work [Song et al., 2020; Li and Zhu, 2021; Fang et al., 2021; Jiang et al., 2023a; Guo et al., 2021b], we divide the four datasets into training, validation, and test sets according to a 6:2:2 ratio. During pre-training, the long input Tlong of the four datasets are set to 864, 864, 864, and 2016 time steps, respectively. For prediction, we set the length of both input T and output ̂ T to 12 steps. The embedding dimension D is 96. The encoder has 4 transformer layers while the decoder has 1 transformer layer. The number of multi-attention heads in transformer layer is set to 4. We use a patch size L of 12 to align with the forecasting input. T ′ is equal to 1, which means we truncate and keep the last one patch of H(S) and H(T) . The masking ratio r is set to 0.25. Optimization is performed with Adam optimizer using an initial learning rate of 0.001 and mean absolute error (MAE) loss. For evaluation, we use MAE, root mean squared error (RMSE), and mean absolute percentage error (MAPE(%)). Performance is assessed by averaging over all 12 prediction steps. Experiments are mainly conducted on a Linux server with four NVIDIA GeForce RTX 3090 GPUs. To make fair and consistent comparison, they are all performed on BasicTS [Shao et al., 2023] platform.

5.2 Overall Performance

The performance of models is listed in Table 2. For a fair comparison, the reported results of the baseline models are taken from the original literature, which have been widely cited and validated in spatiotemporal forecasting. Across all datasets, our STD-MAE achieves superior performance over the baselines by a significant margin on all evaluation metrics. For other baselines, the spatiotemporal models clearly outperform the time series models due to their ability to capture spatiotemporal dependencies. In summary, the proposed STDMAE framework significantly advances the state-of-the-art in spatiotemporal forecasting, demonstrating its ability to augment downstream predictors.

5.3 Ablation Study

Masking Ablation. We design four variants to validate the effectiveness of our spatial-temporal masking mechanism:

• S-MAE: Only masking on the spatial dimension.
• T-MAE: Only masking on the temporal dimension.
• STM-MAE: Using spatial-temporal-mixed masking.
• w/o Mask: Without applying any masked pre-training.

We report the experimental results on the PEMS03 and PEMS07 datasets. As illustrated in Figure 3, STD-MAE with spatial-temporal-decoupled masking significantly outperforms the ablated versions. T-MAE and S-MAE still improve over the original model although they can only partly capture the heterogeneity. For STM-MAE, we mix the spatial and temporal dimensions before randomly masking operation. The task of mixed masking is trivial which would lead to learn representation with less rich semantics. Overall, the results highlight the value of our proposed spatial-temporaldecoupled masked pre-training design for spatiotemporal forecasting.

Predictor Ablation. To evaluate the generality of STDMAE, we test five downstream predictors with different backbones including GCN+RNN, GCN+TCN, Linear and Transformer:

• STD-MAE-DCRNN: Using DCRNN as the predictor.
• STD-MAE-MTGNN: Using MTGNN as the predictor.
• STD-MAE-STID: Using STID as the predictor.

• STD-MAE-STAE: Using STAEformer as the predictor.
• STD-MAE: Using GWNet as the predictor.

The experiments are conducted on the PEMS04 and PEMS08 datasets. Table 3 illustrates consistent and substantial performance gains across all five downstream spatiotemporal predictors when augmented with STD-MAE. This demonstrates the robustness of the representations generated by STD-MAE, which can benefit all kinds of downstream predictors regardless of architectures.

In conclusion, these ablation studies validate the effectiveness of STD-MAE's decoupled design and the generality for enhancing downstream baselines.

5.4 Hyper-parameter Study

Masking Ratio. We first conduct hyperparameter study by varying the masking ratio r in {0.25, 0.5, 0.75}, where the value is applied equally to S-MAE and T-MAE. As shown by Figure 4, a masking ratio of 0.25 yields the lowest error across all datasets, indicating optimal value. Prior work on masked language modeling with BERT [Kenton and Toutanova, 2019] utilizes a relatively low masking ratio of only 15% during pre-training. In contrast, masked autoencoders for image reconstruction [He et al., 2022] and video modeling [Tong et al., 2022] have found much higher optimal ratio of 75% and 90%, respectively. However, we find a somewhat lower optimal ratio of 25% for spatiotemporal time series modeling due to the fact that the long input time series required to provide extensive temporal context. This presents a challenge for reconstructing masked inputs, especially with spatial masking patterns that obstruct large contiguous blocks. Furthermore, our study demonstrates the importance of tuning masking specifically for spatiotemporal data. While an exact optimal is dataset-dependent, our results nonetheless show that relatively lower masking ratio is preferable for spatiotemporal time series.

Pre-training Length. We also study the effects of Tlong, which is the input length used in pre-training. Here we test different pre-training lengths of one day, three days, and a week on all four datasets. The results are shown in Table 4. In two out of the four datasets, a pre-training length of 3 days yields the best performance. Remarkably, compared to previous pre-training methodology, our approach demonstrates an enhanced capability to achieve superior performance with comparatively shorter pre-training lengths. These findings not only show the dynamic impact of pre-training length on the performance but also guide that the optimal pre-training length changes according to datasets.

5.5 Efficiency Test

Since STD-MAE introduce two decoupled autoencoders to encode spatial and temporal representation, efficiency might be a concern. However, due to the decoupled design of STDMAE, it still outperforms in efficiency compare to other pretrained model especially for datasets with a large amount of sensors. Specifically, compared to non-decoupled pretraining methods, our decoupled time complexity is O(N 2 + T 2 p ). We report the total training time for pre-training and forecasting per sample of each dataset as Table 5.

5.6 Case Study

Reconstruction Accuracy in Pre-training. STD-MAE exhibits a robust capacity for reconstructing long time series by relying solely on the partially observed sensor recordings, as shown in Figures 5a and 5b. The shaded area indicates the masked region. Temporal reconstruction closely matches the ground truth in terms of periodicity and trends. This suggests STD-MAE successfully acquires generalized knowledge of temporal patterns. Similarly, STD-MAE could restore entirely masked sensors based on contextual data from spatially related sensors, indicating it also gains meaningful spatial correlations. Overall, STD-MAE appears to learn rich spatiotemporal representations through pre-training.

Robustness on Spatiotemporal Mirage. In Figure 6, we present a comparative analysis of the prediction results from GWNet and STD-MAE for two spatiotemporal mirages in Figure 1c. The input and prediction windows are denoted by

(a) Reconstruction from T-MAE Pre-training

(b) Reconstruction from S-MAE Pre-training

Figure 6: Prediction under Spatiotemporal Mirage

purple and green backgrounds, respectively. A pivotal finding is that GWNet exhibits a limitation in distinguishing spatiotemporal mirages in Figures 6a and 6c. In contrast, STDMAE performs a significant accuracy in these situations as shown in Figures 6b and 6d. The pre-trained component in STD-MAE remarkably enhances GWNet's ability to distinguish spatiotemporal mirages arising from heterogeneity.

6 Conclusion

In this study, we propose STD-MAE, a novel spatialtemporal-decoupled masked pre-training framework for spatiotemporal forecasting. In the pre-training phase, a novel spatial-temporal-decoupled masking approach is utilized to effectively model the heterogeneity of spatiotemporal data. For forecasting phase, the hidden representations generated by STD-MAE are leveraged to boost the performance of downstream spatiotemporal predictors. Comprehensive experiments and in-depth analyses conducted on four benchmark datasets demonstrate the superiority of STD-MAE.

Acknowledgments

This work was supported by JST CREST Grant Number JPMJCR21M2 including AIP challenge program, JSPS KAKENHI Grant Number JP24K02996, and Guangdong Talent Program (No. 2021QN02X794).