【论文翻译】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$$