论文标题： Efficient Large-Scale Traffic Forecasting with Transformers: A Spatial Data Management Perspectiveg
作者： 方雨晨，梁宇轩，惠博，邵泽志，邓力玮，柳旭，蒋新科，郑凯
机构： 电子科技大学，香港科技大学（广州），奥本大学（Auburn），中科院计算所，新加坡国立大学，北京大学
关键词： 时空预测，Patch，空间数据管理，高效
论文链接： https://arxiv.org/abs/2412.09972
代码链接： https://github.com/LMissher/PatchSTG
TL; DR: 此文从空间数据管理的角度出发研究不均匀空间分布节点的patch 化方法，以将Transformer 从空间的patch-level 而非point-level 应用于大规模交通预测。
挑战（Challenge）
作者首先指出当前大规模时空图预测任务仍存低效动态建模的挑战。动态空间建模技术已成为交通预测的主流研究方向，其旨在揭示不同时间段随时间演化的空间相关性并动态传播空间信息。然而大多数动态空间建模方法的二次复杂度已成为城市级交通预测算法部署瓶颈。如下表所示，三种广泛被使用的动态空间建模范式（点积、线性和低秩）以注意力机制形式被展示，也即首先使用查询和键动态计算空间相关性，然后根据计算出的空间相关性将空间信息与原始值进行传播。其中基于点积的方法（例如 D2STGNN 和 STAEformer）中的空间相关性需要在每一对点上计算注意力，因此在大规模交通预测中导致不可接受的二次复杂度。尽管基于线性的 BigST 和基于低秩的 Airformer 被提出以缓解大规模时空预测中动态空间建模的计算需求，但它们仍然存在一些局限性。其中线性方法缺乏可解释性，因为无法明确显示空间注意力矩阵。而低秩方法的缺陷是信息失真，也即关键信息无法保证能保留在被压缩后的低秩表示中，从而导致性能下降。
动机（Motivation）
虽然目前缺乏高效动态建模时空图数据的方法，如下图 (a)、(b) 所示，此文发现视觉以及时间序列领域如 ViT 与 PatchTST 可以通过将稠密原始数据进行 patch 化，从而减少参与注意力机制计算的 token 数量进而提升模型效率。交通和视觉及时序数据的区别在于，像素与时间点是均匀且连续的，而交通传感器在道路上的分布不均匀且在计算机中的存储不连续。也即道路上相邻的传感器在计算机中的存储下标不一定连续，使其不能像图片或时序数据一样，可以通过 reshape 操作轻松的将相同数量且连续的元素分割到 patch 块中。因此，时空图 Patching 的主要目标是设计一种可以先将计算机存储的空间坐标重新排列使其具有连续性也即存储相邻元素在实际空间中也相邻，然后能在新坐标上执行元素均衡且不重叠的空间点分割。
虽然目前聚类或图分割算法可以得到局部连续的空间新坐标，但聚类方法难以得到元素均衡分割，而图分割算法难以做到补全后元素不重叠。受到空间数据管理算法的启发，它们类似视觉图片 patching，显式的通过坐标轴分割使得相邻坐标连续并可以做到元素均衡且不重叠。如下图所示，可以先通过空间数据管理算法如KDTree 将均衡数量的点划分到小区域。然后将相似点填充到未达到最大点数的区域以以防重叠分割，再将小区域合并为 patch。最后，就可对同一 patch 中点以及其他 patch 中具有相同索引的点执行深度和广度注意力，以捕获局部和全局空间知识。
贡献（Contribution）
此文从空间数据管理角度出发研究非均匀分布节点 patch 化方法，提出了一个高效且准确的用于大规模交通预测的 Transformer 模型 PatchSTG。主要的贡献有三：
PatchSTG 是第一个从空间数据管理角度出发，从而将 KDTree 发展为不均匀空间点 Patching 技术的方法，其可将不均匀分布的节点均匀且不重叠的划分为 patch 块，以降低后续 Transformer 复杂度，并因其显式坐标轴划分而具有可解释性。 PatchSTG 对每个 patch 块中的点以及在不同 patch 块中具有相同索引的点交替使用深度和广度注意力机制，以高效且保真地动态挖掘局部和全局空间信息。 在大规模交通数据集上的综合实验结果表明，与动态空间建模基线相比，PatchSTG 不仅取得了最先进的性能，并且在减少 倍内存占用的情况下，提高了 倍训练速度。
方法（Method）
Overview：
下图展示了 PatchSTG 模型，它包含四个主要组件：
Spatio-Temporal Embedding：用于将交通数据预处理为高维嵌入； Irregular Spatial Patching：用于将相同数量的点分割到 patch； Dual Attention Encoder：用于提取空间信息； Projection Decoder： 用于预测未来值。
Spatio-Temporal Embedding：
此文首先对每个输入交通时间序列采用全连接层，将数值交通流转换为高维嵌入：
其中 和 是可学习参数。 是包含流量演化的投影嵌入。此外，此文考虑了日周期模式、周周期模式和空间异构模式。对于时间方面，一周中的第几天和一天中的第几个时间可以存储在数据驱动的字典 中，其中 和 表示一周中的天数和一天中的时间片数。因此，此文使用所有点的最后一个时间片作为索引，从字典中提取相应的星期几嵌入 和一天中的时间片嵌入。与时间方面类似，此文利用可学习的嵌入 作为独立表征来区分数据集中的空间节点。最后，可将上述嵌入拼接起来以得出时空嵌入：
Irregular Spatial Patching：
Leaf KDTree: 考虑到分割的平衡性、不重叠性和效率要求，此文采用简单而有效的 KDTree 来寻找分割不均匀分布交通数据的解决方案。如上图所示，此文提出 Leaf KDTree 以交通数据的经纬度位置视对节点进行交替均匀分割，直到划分到树中叶子节点。不难从上图中看出，属于同一子树的叶节点基于它们在现实世界中更近的距离保持更强的空间相关性，这为后续补全提供了可解释的回溯。在构建KDTree 之后，可对树进行广度优先搜索，以得出新的交通点坐标，从而确保属于同一子树的叶子节点在新坐标中相邻。整个过程基于纬度、经度 和叶子节点的容量（KDTree 中的叶子节点最多包含 个点为预定常数，上图中），可以表示如下：
其中。 和 表示叶 KDTree 构造和广度优先搜索操作。
Padding: 尽管 Leaf KDTree 可以提供均衡分割，但点的数量 不一定能被容量 整除，这会导致存在叶子节点不满。不一致的叶子节点所包含数量阻碍了基于 patch 的方法应用。虽然填充零或不相关的点可以缓解不可整除的问题，但会降低预测性能。为了避免重复使用叶子节点中点，此文从其他叶子节点中选择与未满叶子节点中点最相似的点填充以达最大容量，相似填充可以确保非重叠的 patch 分割：
其中 表示通过余弦相似度 查询每个未满叶节点的最相似点。 和 表示新坐标中应填充的位置以及查询点对应的原始坐标。因此，填充过程可以表述如下：
填充嵌入 有 个点且。
Patching: 尽管填充后可保证均衡分割，但直接使用容量较大的叶节点作为 patch 会导致填充不均衡的问题，即未满的叶节点可能会被与同一点的相似点所填充。又因为同一子树下的叶节点在新坐标中是强相关和相邻的，此文先将点划分到容量较小的叶节点中，然后从叶节点回溯到子树的根节点，直到满足子树的数量等于预先定义的 patch 数量。上述做法可以缓解填充不平衡的问题，因为未满的叶节点大多由与不同点的相似点填充。具体地，让 表示 patch 中的点数， 表示 patch 数，其中 是决定子树中有多少个叶节点的超参数，则。值得注意的是， 是 的幂，因为只有属于同一子树的叶节点才具有强的空间局部性，并且 Leaf KDTree是一棵二叉树。
在均衡且不重叠的空间 patch 化后，填充的嵌入 被转换为新的表示，作为后续模型输入。
Dual Attention Encoder：
如上图所示，此文使用双注意力编码器以动态捕获空间依赖关系。对于 patch 后输入，第一维中的 个点可以看作子树的根节点，第二维中的 个点可以看作子树中的点。因此，PatchSTG 首先对第二维中点也即 patch 里的点使用注意力机制来动态提取局部空间信息，这是因为子树中的点具有更强的局部相关性，又因为子树中的点向下展开也称为深度注意力。此外，由于全局依赖关系在流量预测中与局部信息同样重要，此文对第一维中的点也即每个 patch 坐标对应的点采用注意力机制来无损学习全局知识，这是因为 patch 中每个点都可从其他根节点中具有相同坐标的点接收信息，并且这些点在先前的深度注意之后已经具有局部融合信息，又因为同一层根节点平铺展开也称为广度注意力。此外，双重注意力机制可在编码器中交替堆叠 层以挖掘深层空间信息。
Projection Decoder：
此文通过双重注意力编码器的空间信息交互输出 来预测未来的流量。其首先通过将表示 reshape 为，并去除 padding 的点以保持与输入形状和坐标一致：
其中 表示删除填充位置上的点， 是具有原始坐标的表示。最后，采用全连接层将历史表示投影到未来：
其中 为可学习参数， 为预测流量。在训练阶段，此文使用 和 计算 L1 损失作为目标函数，以指导模型学习。
实验（Experiments）
Datasets：
此文使用常用的大规模交通数据集LargeST以及它的三个子集进行交通预测，分别是CA、GLA、GBA 和 SD，他们的详细信息如下表：
Main Results：
主要结果如下表所示，可以看出 PatchSTG 在几乎所有 task 都取得了 sota 的结果，并且在一小时的平均结果上取得 左右提升。
Ablation Study：
如下表所示，此文分别从消息汇聚、填充方法及分割方法进行消融实验。其中 w/o FGGC、w/o Depth 和 w/o Breadth 分别表示去除多样化全局融合、去除深度注意力及去除宽度注意力，可以看出局部、全局及不失真的全局信息汇聚都是有必要的。w/ PadDis 与 w/o PadSim 结果显示不填充相似点而导致重叠将会降低性能。w/ METIS、w/ KMeans 及 w/o LKDT 结果表明在不均匀空间节点 patch 中，近邻、均衡且不重叠的分割至关重要。
Efficiency Comparisons :
如下表所示，PatchSTG 在各个数据集上都拥有最高效的 GPU 内存利用率，并在训练速度方面也排名第一。尤其在 CA 与 GLA 这两个超大数据集上，PatchSTG 的高效性得到突出显示。
Hyper-parameter Study:
如下图所示，此文从 [8, 16, 32, 64, 128, 256, 512] 的搜索空间中搜索最佳 patch 块的数量，如图 8 所示。当 patch 数量分别为 16、16、64 和 512 个时，PatchSTG 可以在 SD、GBA、GLA 和 CA 数据集上取得最佳性能，这表明 patch 的数量与数据集的大小呈正相关。
Visualization:
下图可视化了 Leaf KDTree 在 GLA 数据集上的均衡分割。可以观察到，现实世界的相邻点显然被划分到同一个叶节点以保持空间局部性，与基于低秩的动态空间建模方法相比，这可以提供视觉上可解释的空间划分。为了展示 PatchSTG 可学习不失真全局空间知识，此文可视化了在 GLA 数据集上每个 patch 里第 23 和 56 号学习到的空间关联关系。如下图 6 所示，广度注意力学习到的全局相关性对于 patch 中的不同节点是不同的，这符合交通变化的多样与异质性。因此，与基于低秩的方法相比，PatchSTG 是具有保真度的，因为低秩方法对 patch 中所有点只能反映同一种空间相关性。
点击文末阅读原文跳转本文arXiv链接。
如果觉得有帮助还请分享，在看，点赞
