【论文翻译】IJCAI2024 | STD-MAE:用于时空预测的时空解耦掩码预训练框架

题目	Spatial-Temporal-Decoupled Masked Pre-training for Spatiotemporal Forecasting用于时空预测的时空解耦掩码预训练
论文链接	https://arxiv.org/abs/2312.00516
源码	https://github.com/Jimmy-7664/STD-MAE
发表会议	IJCAI 2024

摘要

时空预测技术在交通、能源和气象等领域具有重要意义。然而，由于时空异质性问题，准确预测时空序列仍然面临挑战。特别是，当前的端到端模型受到输入长度的限制，常常陷入时空海市蜃楼（即，类似的输入时间序列后跟随不相似的未来值，反之亦然）的困境。为了解决这些问题，我们提出了一种新的自监督预训练框架——时空解耦掩码预训练（STD-MAE），该框架采用两个解耦的掩码自编码器，在空间和时间维度上重构时空序列。通过这种重构学习到的丰富上下文表示可以无缝集成到任何架构的下游预测器中，从而增强它们的性能。

我们在六个广泛使用的基准数据集（PEMS03、PEMS04、PEMS07、PEMS08、METR-LA 和 PEMS-BAY）上进行了定量和定性评估，以验证STD-MAE的最先进性能。

1 引言

传感器网络收集的时空数据已经成为多个实际应用中的重要研究领域。时空数据受益于额外的空间上下文信息，例如传感器位置和道路网络，这些信息揭示了传感器之间的依赖关系。因此，与典型的多变量时间序列相比，时空数据的关键区别在于它表现出时空异质性。具体而言，虽然不同地点（如城市中心与郊区）和日期（如工作日与周末）的时间序列有所不同，但它们在相似的上下文中表现出一致且可预测的模式。因此，准确预测时空数据的关键在于有效捕捉这种异质性。

在时空数据集PEMS04中，图1a展示了传感器7的交通流量，揭示了工作日和周末在高峰时段的显著差异。工作日通常会有早高峰，而周末的流量则更加平均，表明了每周模式中的时间异质性。图1b则展示了同一时段内传感器5、86、155和177的交通流量。传感器5和177呈现出明显的峰值和谷值，而传感器86和155则全天保持相对稳定，展示了空间异质性。当数据量较小时，异质性清晰可见，但当数据规模较大时，空间和时间异质性会高度混合。

之前的研究者们提出了许多时空预测的尝试：将图卷积网络（GCN）嵌入到时序卷积网络（TCN）或递归神经网络（RNN）中，或者沿着时空轴应用Transformer。然而，这些模型在区分时空异质性方面存在困难。学习清晰的异质性仍然是时空预测的主要挑战。

此外，大多数现有模型都是端到端训练的。由于模型的高复杂度，它们的输入范围通常被限制为较短的值（通常为12步）。这种限制会使模型面临被称为时空海市蜃楼的问题：1）输入时间序列不相似，但未来值却相似；2）输入时间序列相似，但未来值不相似。我们通过图1c中传感器215和279的交通流量作为示例进行说明。在深夜，两个传感器的历史数据趋势分歧明显，但未来流量却相似；而在下午，它们的历史数据趋势非常相似，但未来数据差异显著。这背后的原因在于现有模型只能捕捉到碎片化的异质性，而非完整的异质性。因此，如何使这些模型在时空海市蜃楼问题上更加稳健是第二个挑战。

在本研究中，我们的目标是通过预训练来学习清晰且完整的时空异质性。特别地，掩码预训练在自然语言处理和计算机视觉中已显示出巨大的效果。核心思想是在预训练过程中掩盖输入序列的一部分，要求模型重建缺失的内容。通过这种方式，模型可以学习到丰富的上下文表示，进而增强各种下游任务的性能。

受到这些优点的启发，我们提出了一个新的时空解耦掩码预训练框架（STD-MAE）。它为学习清晰且完整的时空异质性提供了高效且有效的解决方案。这种学习到的异质性可以无缝集成到下游基线中，以穿透时空海市蜃楼。

总结来说，我们的主要贡献如下：

• 我们设计了一个时空数据预训练框架，该框架能够在不修改下游时空预测器原始结构的情况下大幅提高其性能。
• 我们提出了一种新颖的时空解耦掩码策略，通过在时空维度上捕捉长程上下文来有效学习时空异质性。
• 我们在六个基准数据集（PEMS03、PEMS04、PEMS07、PEMS08、METR-LA、PEMS-BAY）上验证了STD-MAE的性能。定量增强实验结果表明STD-MAE的性能优于基线模型，定性分析展示了其捕捉有意义长程时空模式的能力。

2 相关工作

2.1 时空预测

时空预测的目标是通过分析历史数据来预测未来的时空序列。早期的工作主要依赖于传统的时间序列模型。为了捕捉复杂的时间依赖关系，递归神经网络（RNN）和卷积神经网络（CNN）在时空数据建模中变得越来越流行，并取得了更好的预测效果。

然而，这些模型忽视了关键的空间关联，导致在网络化的道路系统上预测性能受限。为了联合捕捉时空特征，一些研究将图卷积网络（GCN）与时间模型结合。沿着这条研究路线，近年来提出了多种新的时空模型，展示了它们在捕捉时空关系中的优越性能。

注意力机制也对时空预测产生了深远的影响。一系列的Transformer模型被提出，表现出卓越的性能，突出了它们在捕捉时空关系中的有效性。然而，这些端到端模型仅关注短期输入，这限制了它们捕捉完整的时空依赖关系的能力。

2.2 掩码预训练

掩码预训练已成为一种在自然语言处理（NLP）和计算机视觉（CV）中用于自监督表示学习的高度有效的技术。核心思想是在预训练期间掩盖输入的一部分，要求模型根据可见的上下文来预测被掩盖的部分。在NLP中，BERT等模型使用掩码语言模型，依靠双向上下文来预测随机掩盖的词汇。随后的一些模型引入了更有效的掩码技术，并证明了通过更长时间的预训练可以获得显著的性能提升。在CV中，类似的掩码策略也被采用，诸如BEiT和掩码自编码器（MAE）等方法通过掩盖图像的随机块并根据未掩盖部分进行重建，获得了显著的改进。

最近，许多研究人员尝试在时间序列数据上使用预训练技术，以获得更好的隐藏表示。然而，这些方法要么是独立于通道的，要么忽略了空间维度的预训练。我们提出的STD-MAE引入了一种新的时空解耦掩码策略，在预训练过程中分别在空间和时间维度上进行掩码。通过这种方式，学习到的表示可以有效捕捉复杂的时空异质性。

3 问题定义

时空预测是一个专门的多变量时间序列预测问题。给定过去 $T$ 个时间步长的多变量时间序列 $X_{t-(T-1):t}$，我们的目标是预测未来 $\hat{T}$ 个时间步长的值：

$$[X_{t-(T-1)},\dots ,X_t]\to [X_{t+1},\dots ,X_{t+\hat{T}}]$$

其中，$X_i\in {\mathbb{R}}^{N\times C}$，$N$ 是空间节点的数量，$C$ 是信息通道的数量。在我们使用的数据集中，$C=1$。

4 方法

本节深入探讨了我们提出的时空解耦掩码预训练框架（STD-MAE）的技术细节，如图2所示。

4.1 时空掩码预训练

时空解耦掩码（Spatial-Temporal-Decoupled Masking）：在标准时空预测任务中，输入长度 $T$ 通常等于12（每个步长对应5分钟的间隔）。因此，端到端的模型经常会陷入图1c所描述的时空海市蜃楼问题。我们引入了一种长时输入的掩码预训练阶段。

由于时空数据相较于图像数据具有额外的时间维度，相较于语言数据具有额外的空间维度，直接应用原始的掩码预训练是不切实际的，因为这会导致时间和空间复杂度的平方级别增长。因此，我们提出了一种新方法，称为时空解耦掩码预训练（spatial-temporal-decoupled masking）。该方法在掩码预训练期间分别执行时间和空间维度上的掩码重建任务。这种解耦的掩码机制使模型能够学习到更清晰的异质性表示。

具体来说，给定输入的时空时间序列 $X\in {\mathbb{R}}^{T\times N\times C}$，我们提出了以下掩码策略：

1. 空间掩码（Spatial Masking, S-Mask）：随机掩盖 $N\times r$ 个传感器的数据，其中 $r$ 是掩码比例，取值范围为0到1。这样会得到一个空间掩码后的输入 ${\tilde{X}}^{(S)}\in {\mathbb{R}}^{T\times N(1-r)\times C}$。
2. 时间掩码（Temporal Masking, T-Mask）：随机掩盖 $T\times r$ 个时间步的数据，得到时间掩码后的输入 ${\tilde{X}}^{(T)}\in {\mathbb{R}}^{T(1-r)\times N\times C}$。

这两种掩码策略可以看作是从伯努利分布 $B(1-r)$ 中随机采样，其期望为 $1-r$，在对应的维度上进行如下公式的运算：

$${\tilde{X}}^{(S)}=\sum _{n=1}^N{B}_S(1-r)\cdot X[:,n,:]$$

$${\tilde{X}}^{(T)}=\sum _{t=1}^T{B}_T(1-r)\cdot X[t,:,:]$$

直观上，S-Mask 强制模型仅从其他可见传感器的数据中重建被掩盖的传感器数据，从而捕获长程的空间异质性。同样，通过利用可见序列的内在信息，T-Mask 可以学习到时间异质性并重建整个时间序列。

时空解耦掩码自编码器。在时空解耦掩码技术的基础上，我们进一步提出了时空解耦掩码自编码器（STD-MAE）。它由一个时间自编码器（T-MAE）和一个空间自编码器（S-MAE）组成，它们具有相似的架构。S-MAE在空间维度上应用自注意力机制，而T-MAE则在时间维度上应用自注意力机制。

具体来说，我们考虑长时间的输入，其长度为 $T_{long}$，通常跨越数天。然而，直接利用这样长的序列会导致计算和内存的挑战。为了解决这个问题，我们应用了patch嵌入技术。长序列被划分为不重叠的块，块的长度为 $T_p=T_{long}/L$，使用窗口 $L$ 进行划分。这会产生一个块状输入 $X_p\in {\mathbb{R}}^{T_p\times N\times L\times C}$。然后，我们将 $X_p$ 投影到一个全连接层，得到块嵌入 $E_p\in {\mathbb{R}}^{T_p\times N\times D}$，其中 $D$ 是嵌入维度。

此外，为了同时编码空间和时间的位置信息，我们实现了一个二维的位置编码。给定块嵌入 $E_p$，时空位置编码 $E_{pos}\in {\mathbb{R}}^{T_p\times N\times D}$ 计算如下：

$$\begin{array}{rl}{E}_{pos}[t,n,2i]& =\sin (t/10000^{4i/D})\\ E_{pos}[t,n,2i+1]& =\cos (t/10000^{4i/D})\\ E_{pos}[t,n,2j+D/2]& =\sin (n/10000^{4j/D})\\ E_{pos}[t,n,2j+1+D/2]& =\cos (n/10000^{4j/D})\end{array}$$

我们选择正弦位置编码而非学习的位置编码，因为它可以处理任意长度的输入。块嵌入 $E_p$ 和位置编码 $E_{pos}$ 相加以获得最终的输入嵌入 $E\in {\mathbb{R}}^{T_p\times N\times D}$。

接下来，使用S-Mask和T-Mask策略对输入嵌入 $E$ 进行掩码，分别获得可见的空间块嵌入 ${\tilde{E}}^{(S)}$ 和可见的时间块嵌入 ${\tilde{E}}^{(T)}$。S-MAE和T-MAE通过一系列Transformer层生成空间表示 $H^{(S)}\in {\mathbb{R}}^{T_p\times N(1-r)\times D}$ 和时间表示 $H^{(T)}\in {\mathbb{R}}^{T_p(1-r)\times N\times D}$。

轻量级解码器随后应用于S-MAE和T-MAE来重建被掩盖的输入。空间和时间解码器分别由一个padding层、标准Transformer层和回归层组成。我们使用共享的可学习掩码token $V\in {\mathbb{R}}^D$ 来表示缺失的块。给定空间表示 $H^{(S)}$ 和时间表示 $H^{(T)}$，空间和时间的padding会将 $V^{(S)}$ 和 $V^{(T)}$ 分别扩展到空间和时间掩码token。

最后，回归层用于在块级别重建时间序列。形式上，空间和时间的重建 $ \hat{Q}^{(S)} \in \mathbb{R}^{T_p \times N_r \times L}$ 和 ${\hat{Q}}^{(T)}\in {\mathbb{R}}^{T_m\times N\times L}$ 可以表示为：

$${\hat{Q}}^{(S)}=M^{(S)}\odot FC({\text{Attention}}^{(S)}(H^{(S)}))$$

$${\hat{Q}}^{(T)}=M^{(T)}\odot FC({\text{Attention}}^{(T)}(H^{(T)}))$$

其中，$M^{(S)}\in {\mathbb{R}}^{N_r}$ 是空间掩码索引，$M^{(T)}\in {\mathbb{R}}^{T_m}$ 是时间掩码索引。

我们计算掩码部分的损失函数，损失是通过计算重建结果与真实值的**平均绝对误差（MAE）**得到的。空间和时间的损失函数分别表示为：

$$L_S=\frac{1}{T_p{N}_{M}L}\sum _{t=1}^{T_p}\sum _{n=1}^{N_M}\sum _{l=1}^L\left|{\hat{Q}}^{(S)}[t,n,l]-Q^{(S)}[t,n,l]\right|$$

$$L_T=\frac{1}{T_{M}NL}\sum _{t=1}^{T_M}\sum _{n=1}^N\sum _{l=1}^L\left|{\hat{Q}}^{(T)}[t,n,l]-Q^{(T)}[t,n,l]\right|$$

综上所述，通过上述时空解耦掩码预训练，STD-MAE可以捕捉到完整且清晰的时空异质性。

4.2 下游时空预测

STD-MAE 可以无缝集成到现有的预测器结构中。该操作通过将 STD-MAE 生成的空间和时间表示添加到预测器的隐藏表示中实现。具体来说，我们首先将长时间输入 $T_{long}$ 的时间步长输入到预训练的空间和时间编码器中，以生成对应的空间表示 $H^{(S)}$ 和时间表示 $H^{(T)}$。然后，我们应用一个下游的时空预测器 $F_{\theta }$，其参数为 $\theta$，以获取广泛使用的短期输入 $X_{t-(T-1):t}$ 的隐藏表示 $H^{(F)}\in {\mathbb{R}}^{N\times D^{\prime }}$，计算公式如下：

$$H^{(F)}=F_{\theta }[X_{t-(T-1):t}]$$

其中，$D^{\prime }$ 是预测器的隐藏表示维度。为了与 $H^{(F)}$ 对齐，我们截取最后 $T^{\prime }$ 个块的空间表示 $H^{(S)}$ 和时间表示 $H^{(T)}$，并将这两个表示重新调整为 $H^{(S^{\prime })}\in {\mathbb{R}}^{N\times T^{\prime }\times D}$ 和 $H^{(T^{\prime })}\in {\mathbb{R}}^{N\times T^{\prime }\times D}$。接下来，我们通过两层 MLP 将这两个表示投影到 $D^{\prime }$ 维度。最后，通过将这些表示相加，得到增强表示 $H^{(Aug)}\in {\mathbb{R}}^{N\times D^{\prime }}$：

$$H^{(Aug)}=MLP(H^{(S^{\prime })})+MLP(H^{(T^{\prime })})+H^{(F)}$$

到目前为止，$H^{(Aug)}$ 包括了由预测器本身生成的表示以及来自 STD-MAE 的长时间空间和时间表示，这极大地增强了下游时空预测器的性能。

在我们的工作中，我们选择了 GWNet 作为我们的预测器，因为其表现优越。我们通过聚合 GWNet 中跨越多个时空层的跳跃连接状态以及 STD-MAE 生成的对应的空间和时间表示，来获取最终的表示。增强后的表示随后被输入到 GWNet 的回归层中进行预测。此外，我们还测试了其他经典的时空预测器，包括 DCRNN、MTGNN、STID 和 STAEformer。这些实验展示了 STD-MAE 的通用性。

5 实验

5.1 实验设置

数据集: 为了全面评估所提出的 STD-MAE 模型，我们在六个真实世界的时空基准数据集上进行了广泛的实验，如表1所示：PEMS03、PEMS04、PEMS07、PEMS08、METR-LA 和 PEMS-BAY。原始数据的时间分辨率为5分钟，每个时间步之间的间隔为5分钟。对于数据预处理，我们对原始输入进行了Z-score标准化。

基线模型。我们将 STD-MAE 与以下基线方法进行了对比：ARIMA、VAR、SVR、LSTM、TCN 和 Transformer 这些是时间序列模型。对于时空模型，我们选择了几个典型方法，包括 DCRNN、STGCN、ASTGCN、GWNet、STSGCN、STFGNN、STGODE、DSTAGNN、ST-WA、ASTGNN、EnhanceNet、AGCRN、Z-GCNETS、STEP、PDFormer 和 STAEformer。

设置。参考以往的工作，我们将 PEMS03、PEMS04、PEMS07 和 PEMS08 数据集按 6:2:2 的比例划分为训练、验证和测试集。对于 METR-LA 和 PEMS-BAY 数据集，训练、验证和测试的比例为 7:1:2。在预训练过程中，六个数据集的长时间输入 $T_{long}$ 分别设置为 864、864、864、2016、864 和 864 个时间步长。对于预测，我们将输入长度 $T$ 和输出 $\hat{T}$ 都设置为 12 个步长。嵌入维度 $D$ 设置为 96。编码器有 4 层 Transformer，而解码器有 1 层 Transformer。Transformer 层中的多头注意力头的数量为 4。我们使用的 patch 大小 $L$ 为 12，以对齐预测输入。$T^{\prime }$ 等于 1，这意味着我们截取并保留了 $H^{(S)}$ 和 $H^{(T)}$ 的最后一个块。掩码比例 $r$ 设置为 0.25。优化使用 Adam 优化器，初始学习率为 0.001，损失函数为平均绝对误差（MAE）。我们使用 MAE、均方根误差（RMSE）和平均绝对百分比误差（MAPE(%)）进行评估。所有实验主要在配备四个 NVIDIA GeForce RTX 3090 GPU 的 Linux 服务器上进行。为了进行公平和一致的比较，所有实验均在 BasicTS 平台上进行。

5.2 总体性能

模型的性能列在表 2、表 3 和表 4 中。为了进行公平的比较，基线模型的报告结果来自原始文献，并且在时空预测中被广泛引用和验证。在所有数据集中，STD-MAE 在所有评估指标上都显著优于基线模型。其他基线模型中，时空模型由于能够捕捉时空依赖关系，明显优于时间序列模型。总结来说，所提出的 STD-MAE 框架显著提升了时空预测的最先进水平，展示了其增强下游预测器能力的能力。

5.3 消融研究

掩码消融。我们设计了四个变体来验证我们时空掩码机制的有效性：

• S-MAE：仅在空间维度上掩码。
• T-MAE：仅在时间维度上掩码。
• STM-MAE：使用时空混合掩码。
• w/o Mask：不应用任何掩码预训练。

我们在 PEMS03 和 PEMS07 数据集上报告了实验结果。如图 3 所示，STD-MAE 通过时空解耦掩码显著优于消融版本。T-MAE 和 S-MAE 相比原始模型仍然有所提升，尽管它们只能部分捕捉异质性。对于 STM-MAE，我们在进行随机掩码操作前将空间和时间维度混合。然而，混合掩码任务过于简单，导致学习的表示语义较少。总的来说，结果突出了我们提出的时空解耦掩码预训练设计在时空预测中的价值。

预测器消融。为了评估 STD-MAE 的通用性，我们测试了五个不同的下游预测器，骨干网络包括 GCN+RNN、GCN+TCN、线性模型和 Transformer：

• STD-MAE-DCRNN：使用 DCRNN 作为预测器。
• STD-MAE-MTGNN：使用 MTGNN 作为预测器。
• STD-MAE-STID：使用 STID 作为预测器。
• STD-MAE-STAE：使用 STAEformer 作为预测器。
• STD-MAE：使用 GWNet 作为预测器。

这些实验是在 PEMS04 和 PEMS08 数据集上进行的。表5展示了在所有五个下游时空预测器上使用 STD-MAE 增强后的性能增益。这证明了 STD-MAE 生成的表示对所有类型的下游预测器都有显著的提升效果。

5.4 超参数研究

掩码比例。我们首先通过改变掩码比例 $r$ （$r$ 取值为 { 0.25 , 0.5 , 0.75 } \{0.25, 0.5, 0.75\} {0.25,0.5,0.75}）进行超参数研究，该值同样应用于 S-MAE 和 T-MAE。由图4所示，$r=0.25$ 时，在所有数据集上产生了最低的误差，表明该值是最优的。之前的研究表明，BERT 在掩码语言模型中使用相对较低的掩码比例（15%）。而图像重建中的掩码自编码器和视频建模则使用了更高的掩码比例（75%-90%）。我们发现，对于时空序列建模来说，最优掩码比例为25%，因为时空序列需要长时间输入以提供广泛的时间上下文。

预训练长度。我们还研究了 $T_{long}$ 的影响，即预训练时使用的输入长度。这里我们在四个数据集上测试了三种不同的预训练长度：一天、三天和一周。结果显示，在四个数据集中，有两个数据集预训练三天的效果最好。相比于以往的预训练方法，我们的方法展示了通过较短的预训练时间获得更好的性能的能力。

5.5 效率测试

由于 STD-MAE 引入了两个解耦的自编码器来编码空间和时间表示，因此效率可能会成为一个问题。然而，由于 STD-MAE 的解耦设计，它在效率方面仍然优于其他预训练模型，尤其是在具有大量传感器的数据集上。具体来说，与非解耦的预训练方法相比，我们的解耦时间复杂度为 $O(N^2+T_p^2)$。我们在表7中报告了四个数据集上每个样本的预训练和预测的总训练时间。

5.6 案例研究

预训练中的重建准确率。STD-MAE 展示了在仅依赖部分观察到的传感器记录情况下重建长时间序列的强大能力。图5a 和 5b 展示了这些效果，阴影区域表示被掩盖的区域。时间重建在周期性和趋势方面与真实值高度一致，表明 STD-MAE 成功地获取了时间模式的一般化知识。同样，STD-MAE 还可以根据空间相关传感器的上下文数据，恢复完全被掩盖的传感器数据，表明它也获得了有意义的空间关联知识。总的来说，STD-MAE 通过预训练学习到了丰富的时空表示。

时空海市蜃楼的鲁棒性。在图 6 中，我们展示了 GWNet 和 STD-MAE 对两个时空海市蜃楼的预测结果的对比分析。输入和预测窗口分别由紫色和绿色背景表示。一个重要的发现是，GWNet 在区分时空海市蜃楼方面存在局限性，如图 6a 和 6c 所示。相比之下，STD-MAE 在这些情况下表现出显著的准确性，如图 6b 和 6d 所示。STD-MAE 中的预训练组件显著增强了 GWNet 区分因异质性引起的时空海市蜃楼的能力。

6 结论

在本研究中，我们提出了 STD-MAE，一个用于时空预测的时空解耦掩码预训练框架。在预训练阶段，采用了一种新颖的时空解耦掩码方法，有效地建模时空数据的异质性。在预测阶段，利用 STD-MAE 生成的隐藏表示来提升下游时空预测器的性能。通过在六个基准数据集上进行的全面实验和深入分析，展示了 STD-MAE 的优越性。