本文档发表于2023年ACM CIKM会议，标题为“STAEformer: Spatio-Temporal Adaptive Embedding Makes Vanilla Transformer SOTA for Traffic Forecasting”。作者团队来自南方科技大学、东京大学、悉尼科技大学。这项工作针对智能交通系统中的关键挑战——交通流预测，提出了一种名为“时空自适应嵌入”的新组件，并将其与标准的Transformer架构结合，构建了STAEformer模型。该模型在六个真实世界交通数据集上取得了最先进的性能，证明了通过改进输入表示而非复杂化网络结构可以更有效地提升预测精度。
一、 学术背景与研究目的
交通预测是智能交通系统的核心组成部分，其目标是根据历史观测数据，预测道路网络中未来的交通状态序列。该领域的核心科学问题在于如何准确捕获复杂的时空依赖关系。近年来，深度学习模型，特别是时空图神经网络（STGNNs）和基于Transformer的模型，因其出色的性能而备受关注。研究者们投入了大量精力设计各种复杂、精巧的模型架构，例如新型图卷积、动态图结构学习和高效注意力机制等。
然而，研究背景指出，网络架构的进步正遭遇性能提升的边际效益递减现象。这促使研究焦点应从复杂的模型设计转向对数据本身的有效表征技术。输入嵌入是一种广泛使用、简单但强大的表征技术，但常被研究者忽视其潜力。以往的工作中，STGNNs主要使用特征嵌入；Transformer模型由于注意力机制不保留序列位置信息，需要额外添加时序位置编码和周期性（如日、周、月）嵌入；近期一些模型开始引入空间嵌入。其中，STID是少数深入探索这些嵌入的研究之一，它结合了空间嵌入和时间周期性嵌入，并取得了显著成效。
在上述背景下，本研究旨在探索一个基本但关键的交通时间序列预测表征学习技术——输入嵌入。具体目标是提出一种新的嵌入方法，以更有效的方式捕获交通数据中固有的时空关系和时序信息，从而简化模型架构的同时提升预测性能。研究提出，与其设计更复杂的网络，不如通过改进输入表征来更直接地解决问题。
二、 研究方法与流程
本研究的工作流程主要包括问题定义、模型架构设计、实验验证和深入分析几个部分。
1. 问题定义： 给定过去T个时间步的交通序列(X_{t-T+1:t})，交通预测旨在通过训练一个参数为(\theta)的模型(f(\cdot))来推断未来(T’)个时间步的交通数据。其中，每个时间步的数据(X_i \in \mathbb{R}^{N \times d})，(N)是空间节点的数量，(d)是输入特征的维度，在本研究中，(d=1)，代表交通流量。
2. 模型架构（STAEformer）： STAEformer的架构简洁，主要由嵌入层、Transformer层和回归层构成，其核心创新在于嵌入层。 * 嵌入层：该层负责生成隐藏的时空表征(Z \in \mathbb{R}^{T \times N \times d_h})。它由三种嵌入拼接而成： * 特征嵌入：使用一个全连接层将原始输入(X)投影到隐藏空间，得到(E_f \in \mathbb{R}^{T \times N \times d_f})，以保留数据的原始信息。 * 周期性嵌入：引入可学习的周几嵌入字典和一天内时间戳嵌入字典。根据输入序列对应的周几和时刻信息，从字典中提取相应的嵌入向量，拼接并广播后得到(E_p \in \mathbb{R}^{T \times N \times 2d_f})。这用于捕捉交通数据中的日、周等周期性模式。 * 时空自适应嵌入：这是本研究提出的核心创新组件，记为(E_a \in \mathbb{R}^{T \times N \times d_a})。它旨在以一种统一的方式捕获复杂的时空关系。与预先定义或动态学习的邻接矩阵不同，(E_a)是一个可学习的、在所有交通时间序列间共享的嵌入。其设计动机是：时序关系不仅由周期性决定，也受时间序列中时间先后顺序的影响（例如，邻近的时间帧更相似）；同时，来自不同传感器的时序模式也不同。因此，需要一个能够自适应捕捉这些内在时空结构和顺序信息的嵌入。 * Transformer层：将上述拼接得到的隐藏表征(Z)分别输入到时序Transformer层和空间Transformer层。这两个层均采用标准的Vanilla Transformer架构。时序层捕捉每个空间节点在不同时间步之间的依赖关系，其自注意力分数(A^{(te)} \in \mathbb{R}^{N \times T \times T})反映了不同空间节点上的时序模式。空间层则捕捉同一时间步下不同空间节点之间的关系。模型使用了层归一化、残差连接和多头注意力机制。 * 回归层：将经过时空Transformer层处理后的输出(Z’)通过一个全连接层，将维度从(T \times d_h)回归到(T’ \times 1)，生成最终的预测结果(Y)。
3. 实验设计： * 数据集：研究在六个广泛使用的交通预测基准数据集上验证模型：METR-LA、PEMS-BAY、PEMS03、PEMS04、PEMS07、PEMS08。这些数据集的时间间隔均为5分钟，详细信息如原文表2所示。 * 实现细节：模型使用PyTorch实现。嵌入维度(d_f)设为24，(d_a)设为80。时空Transformer的层数(L)均设为3，注意力头数为4。输入和预测长度均设为12（即1小时）。使用Adam优化器，学习率从0.001开始衰减，批次大小为16。采用早停策略。 * 评估指标：使用平均绝对误差、均方根误差和平均绝对百分比误差三个常用指标进行评估。 * 基线模型：与多种广泛使用的基线模型进行比较，包括传统模型（HI）、多种STGNNs（GWNet, DCRNN, AGCRN, STGCN, GTS, MTGNN）、专注于时间序列分解的STNorm，以及同样基于Transformer并针对交通预测任务的GMAN和PDFormer，还有强调时空身份嵌入的STID。 * 消融研究：设计了四个模型变体来评估各个组件的有效性：移除时空自适应嵌入、移除周期性嵌入、移除时序Transformer层、同时移除时空Transformer层。 * 案例分析： * 为了验证时空自适应嵌入在捕获时间顺序信息方面的能力，研究在PEMS04和PEMS08数据集上进行了输入序列沿时间轴随机打乱的实验，并与使用空间嵌入(E_s)的模型进行了对比。 * 对PEMS08数据集的时空自适应嵌入进行了可视化分析。在空间轴上，使用t-SNE降维观察不同节点嵌入的聚类情况；在时间轴上，计算输入12个时间帧之间嵌入的相关系数并绘制热力图。
三、 主要研究结果
1. 性能评估结果： 如表1和表3所示，STAEformer在六个数据集上的绝大多数评估指标上都取得了最佳性能。具体而言： * 在METR-LA和PEMS-BAY数据集上，STAEformer在几乎所有预测步长（15、30、60分钟）的三个指标上均优于所有基线模型。 * 在PEMS03、04、07、08数据集上，STAEformer同样在多数指标上领先，尤其在PEMS04、PEMS07和PEMS08上表现突出。 * 重要的是，STAEformer在没有使用任何图建模（如复杂的图卷积）的情况下，其性能大幅超越了STGNNs。同时，它也比其他基于Transformer的模型（如PDFormer）表现更好，表明其是一种更简单但更有效的解决方案。STNorm和STID也取得了有竞争力的结果，但本研究的模型通过更有效的嵌入设计，在简洁的架构上实现了更优的性能。
2. 消融研究结果： 如表4所示，消融实验清晰地展示了各组件的重要性： * 移除时空自适应嵌入(E_a)会导致所有数据集上的性能显著下降，尤其是在PEMS04和PEMS07上，MAE、RMSE和MAPE的恶化非常明显。这证明了(E_a)对于交通建模至关重要。 * 移除周期性嵌入(E_p)也会导致性能下降，但程度通常比移除(E_a)要轻，表明(E_p)对捕捉日、周模式有贡献，但并非最核心的组件。 * 移除时序Transformer层或同时移除时空Transformer层都会造成严重的性能退化。这说明了两个层面Transformer的必要性，也表明本研究提出的嵌入能够有效建模数据中的时空模式，使得这些层能够提取出有价值的特征。
3. 案例分析结果： * 与空间嵌入的对比实验：如图3所示，当对原始输入沿时间轴进行随机打乱时，使用STAEformer（包含(E_a)）的模型性能下降幅度远大于使用空间嵌入(E_s)的模型。这表明时空自适应嵌入(E_a)使模型对时间顺序信息更加敏感，而仅使用空间嵌入的模型对此相对不敏感。这一结果有力地证明了(E_a)能够更好地建模原始输入中的时间先后顺序信息以及其他复杂的交通模式。 * 时空自适应嵌入的可视化：如图4所示，可视化分析进一步验证了(E_a)的有效性。 * 在空间轴上，通过t-SNE可视化，不同节点的嵌入自然地形成了多个簇，这与交通数据固有的空间区域特性相符，表明(E_a)隐式地学习到了空间结构信息。 * 在时间轴上，相关性热力图显示，每个时间帧的嵌入与其邻近时间帧高度相关，并且相关性随着时间距离的增加而逐渐减弱。这正确地反映了时间序列中的时间连续性（时间顺序）信息，表明(E_a)成功地建模了时序上的邻近性模式。
四、 研究结论与价值
本研究的结论是，通过专注于输入嵌入这一基础但强大的表征学习技术，提出了一种新颖的时空自适应嵌入(E_a)。将该嵌入与标准的Vanilla Transformer结合，构建的STAEformer模型在六个交通预测基准数据集上取得了最先进的性能。
该研究的科学价值和应用意义在于： 1. 提供了新的研究方向：研究结果表明，相比于无止境地设计更复杂的网络架构，改进数据的输入表征可能是一条更高效、更有前景的路径。这为交通预测乃至更广泛的时空序列预测领域提供了新的思路。 2. 验证了简单模型的有效性：研究证明了通过精心设计的嵌入，一个结构相对简洁的Vanilla Transformer模型可以超越许多复杂的专用模型（如STGNNs），这挑战了“模型越复杂性能越好”的常见观念，强调了表征学习的重要性。 3. 揭示了嵌入的关键作用：研究提出的时空自适应嵌入被证明能够有效捕获交通时间序列中固有的时空关系和时序顺序信息。这不仅提升了模型性能，而且通过可视化分析得到了直观解释，增强了模型的可解释性。 4. 具有实际应用潜力：STAEformer模型架构简洁，性能优越，易于实现和部署，为实际智能交通系统中的高精度、高效率交通预测提供了一个强有力的工具。
五、 研究亮点