EMAGN：基于学习聚类的可扩展交通预测高效多注意力图网络

# EMAGN：通过学习聚类实现高效多注意力图网络的可扩展交通预测
来源：https://arxiv.org/html/2607.13241
Mingxing Xu1\*,Rakesh Chowdary Machineni2\*,Ke Liu3\*,Xi Cheng4, † Chengqi Lu5,Xin Hu2,Lyuhao Chen6,Xiangyu Li7,Junwei You8,Oliver Gao4 1上海交通大学2密歇根大学安娜堡分校3加州大学伯克利分校4康奈尔大学5德累斯顿工业大学6卡内基梅隆大学7德克萨斯大学奥斯汀分校8威斯康星大学麦迪逊分校

###### 摘要

交通预测因复杂且非线性的时空依赖关系而极具挑战性。自注意力机制已被广泛用于建模动态和长距离依赖，取得了最先进的性能，但由于二次计算和内存复杂度的限制，其可扩展性有限。为解决这一问题，我们提出了一种高效多注意力图网络（EMAGN），它受快速高维高斯滤波理论的启发，对空间注意力机制本身进行了线性化。两个学习到的聚类矩阵 \(C_k\) 和 \(C_v\) 自适应地将键向量和值向量分组为 \(M\) 个超级簇，将复杂度从 \(O(N^2 d)\) 降低到 \(O(N M d)\)，同时不损失注意力在动态依赖建模中的灵活性。在 PEMS-BAY 和 METR-LA 上的实验结果表明，EMAGN 的准确率与全注意力 GMAN 相差在 2.7–3.2% MAE 以内，同时训练时间减少 32%，推理时间减少 38%，GPU 显存减少 58%。关键的是，在 \(K=16\) 个注意力头的情况下，全注意力 GMAN 在标准 11 GB GPU 上完全耗尽内存，而 EMAGN 仍然可以运行，这展示了可行模型配置的显著扩展。EMAGN 在准确性和效率上也超越了同一骨干网络下的 Linformer 和 Performer，这得益于其交通网络感知的自适应聚类。

††脚注文本：*同等贡献 ††脚注文本：†通讯作者：[email protected]
*关键词*交通预测 ⋅ 自注意力 ⋅ 图神经网络 ⋅ 线性复杂度 ⋅ 时空建模。

## 1 引言

随着低成本交通传感器技术的部署，交通数据正在快速增长，我们已进入交通大数据时代。智能交通系统（ITS）[Mori等人，2015](https://arxiv.org/html/2607.13241#bib.bib4)利用这些数据实现高效的城市交通控制和规划。交通预测，即根据历史观测预测未来交通状况（例如速度、流量和密度），在 ITS 中扮演着关键角色。经典统计方法[Liu等人，2011](https://arxiv.org/html/2607.13241#bib.bib29)；Lippi等人，2013](https://arxiv.org/html/2607.13241#bib.bib55)依赖于平稳性假设，而这与交通流的非线性、非平稳特性相悖。深度学习方法，从 CNN [Ma等人，2017](https://arxiv.org/html/2607.13241#bib.bib82)和 RNN [Wu和Tan，2016](https://arxiv.org/html/2607.13241#bib.bib25)到图神经网络[Yu等人，2018](https://arxiv.org/html/2607.13241#bib.bib27)；Li等人，2018](https://arxiv.org/html/2607.13241#bib.bib5)；Wu等人，2019](https://arxiv.org/html/2607.13241#bib.bib38)，通过融入空间结构逐步提高了预测精度。自注意力机制[Vaswani等人，2017](https://arxiv.org/html/2607.13241#bib.bib85)已被广泛用于建模动态和长距离的时空依赖关系[Zheng等人，2020](https://arxiv.org/html/2607.13241#bib.bib36)；Park等人，2020](https://arxiv.org/html/2607.13241#bib.bib37)，取得了最先进的性能。然而，这些基于自注意力的模型因自注意力机制的二次计算和内存复杂度而面临可扩展性限制。如图1所示，必须计算所有元素之间的成对交互，导致 \(O(N^2)\) 的复杂度。随着多头注意力和更深层架构的增加，挑战愈发严峻，尽管两者对于大型网络的高性能至关重要。虽然最近的方法如 STAEformer [Liu等人，2023](https://arxiv.org/html/2607.13241#bib.bib84)和 MLCAFormer [He等人，2025](https://arxiv.org/html/2607.13241#bib.bib77)通过架构重新设计、自适应嵌入、层次注意力和传播延迟建模继续提升准确性，但二次空间注意力本身的内存瓶颈仍未得到解决。尽管存在这一众所周知的限制，现有的交通预测基准在其实验评估中并未报告 GPU 显存消耗，这使得无法评估这些模型在资源受限硬件上的实际可部署性。

参照图注图1：注意力机制的示意图。(a) 标准自注意力，复杂度为 \(O(N^2 d)\)。(b) 快速高斯滤波将附近位置聚合到聚类中心。(c) 我们的线性复杂度自注意力使用学习到的聚类矩阵 \(C_k\) 和 \(C_v\) 压缩键和值，实现 \(O(N M d)\) 的复杂度，其中 \(M \ll N\)。
为缓解这一瓶颈，我们提出了高效多注意力图网络（EMAGN），在保持可比预测性能的同时，实现了更高的效率和可扩展性。我们开发了一种线性复杂度自注意力机制来建模跨元素的空间依赖关系。具体来说，我们自适应地将 \(N\) 个键向量分组并聚合为 \(M\) 个超级键向量（相应地生成超级值向量），然后计算查询与超级键向量之间的依赖关系，将计算复杂度从 \(O(N^2)\) 降低到 \(O(N M)\)，其中 \(M\) 是一个预定义超参数。因此，复杂度随交通网络的规模线性扩展。我们的想法源于经典信号处理中的快速高维高斯滤波技术[Adams等人，2010](https://arxiv.org/html/2607.13241#bib.bib59)。其关键见解是，与附近位置相关联的值向量可以合并为一个超级值向量，从而显著降低计算成本，同时实现近似精确的滤波性能。EMAGN 基于 GMAN [Zheng等人，2020](https://arxiv.org/html/2607.13241#bib.bib36) 作为基础架构，因为 GMAN 显式计算所有传感器对之间的完整 \(N \times N\) 空间注意力，这使得二次瓶颈最为明显，效率提升也最可衡量。为解决可扩展性问题，GMAN 提出了一种分组空间注意力机制，随机将顶点划分为多个组，并仅计算组内和组间的注意力。然而，随机分组无法保证性能，并且需要更多组才能达到理想的精度。相比之下，我们在训练过程中自适应地学习分组，以利用交通网络内固有的空间聚类。经验上，我们发现少量学习到的组就足以获得有竞争力的性能。本文的主要贡献如下：

- •我们提出了一种源于快速高维高斯滤波的线性复杂度自注意力机制，通过学习到的自适应聚类矩阵 \(C_k\) 和 \(C_v\)，将空间注意力复杂度从 \(O(N^2 d)\) 降低到 \(O(N M d)\)，这些矩阵与预测目标端到端联合训练。
- •我们开发了 EMAGN，一种用于交通预测的高效多注意力图网络，其准确率与全注意力 GMAN 相差在 2.7–3.2% MAE 以内，同时在 METR-LA 上将训练时间减少 32%，推理时间减少 38%，GPU 显存减少 58%。
- •我们展示了 EMAGN 扩展了基于注意力的交通预测的可行配置空间：当 \(K=16\) 个注意力头时，全注意力 GMAN 超出标准 11 GB GPU 的内存，而 EMAGN 仍在 8,568 MB 下正常运行。
- •我们首次提供了通用线性注意力机制（Linformer [Wang等人，2020a](https://arxiv.org/html/2607.13241#bib.bib51) 和 Performer [Choromanski等人，2022](https://arxiv.org/html/2607.13241#bib.bib52)）与全注意力在标准交通基准上的直接实证比较，证明了 EMAGN 在准确性和效率上同时优于这两种替代方案。

## 2 相关工作

### 2.1 交通流预测

早期用于交通预测的深度学习方法采用 CNN 从网格化转换的交通网络中提取空间特征[Ma等人，2017](https://arxiv.org/html/2607.13241#bib.bib82)；Zhang等人，2019](https://arxiv.org/html/2607.13241#bib.bib81)，但这种转换丢弃了道路图固有的不规则拓扑结构。图神经网络（GNN）通过直接在图结构上操作克服了这一限制[Scarselli等人，2008](https://arxiv.org/html/2607.13241#bib.bib83)；Kipf和Welling，2017](https://arxiv.org/html/2607.13241#bib.bib19)。STGCN [Yu等人，2018](https://arxiv.org/html/2607.13241#bib.bib27)应用谱图卷积，而 DCRNN 和 GraphWaveNet [Li等人，2018](https://arxiv.org/html/2607.13241#bib.bib5)；Wu等人，2019](https://arxiv.org/html/2607.13241#bib.bib38)在有向图上采用扩散卷积来捕获方向性流动。然而，这些模型依赖于固定的邻接矩阵，无法适应时变空间依赖关系。自适应图方法如 AGCRN [Bai等人，2020](https://arxiv.org/html/2607.13241#bib.bib46)和 MTGNN [Wu等人，2020](https://arxiv.org/html/2607.13241#bib.bib47)通过从数据中学习图结构来解决这一问题，但它们仍然依赖于局部消息传递，可能难以捕获长距离空间交互。

自注意力机制为处理图结构数据提供了一种自然的解决方案，因为它建模所有节点之间的成对依赖关系，而与图距离无关[Pan等人，2019](https://arxiv.org/html/2607.13241#bib.bib73)；Zheng等人，2020](https://arxiv.org/html/2607.13241#bib.bib36)；Park等人，2020](https://arxiv.org/html/2607.13241#bib.bib37)；Wang等人，2020b](https://arxiv.org/html/2607.13241#bib.bib39)。GMAN [Zheng等人，2020](https://arxiv.org/html/2607.13241#bib.bib36)引入了一种带有分组空间注意力的图多注意力机制；PDFormer [Jiang等人，2023](https://arxiv.org/html/2607.13241#bib.bib48)融入了传播延迟感知的注意力；STAEformer [Liu等人，2023](https://arxiv.org/html/2607.13241#bib.bib84)表明，一个具有设计良好的自适应嵌入的 vanilla Transformer 能够取得很强结果；MLCAFormer [He等人，2025](https://arxiv.org/html/2607.13241#bib.bib77)通过多级因果注意力捕获层次化的时空依赖关系。尽管这些模型有效，但它们都继承了标准自注意力的 \(O(N^2)\) 成本，这限制了它们在大规模网络上的可扩展性。最近的工作已经开始解决这一瓶颈。BigST [Han等人，2024](https://arxiv.org/html/2607.13241#bib.bib49)通过线性化图卷积实现了线性复杂度，STGformer [Wang等人，2024](https://arxiv.org/html/2607.13241#bib.bib50)通过单层高阶注意力设计降低了成本。然而，BigST 依赖图卷积而牺牲了注意力的灵活性，而 STGformer 仍然依赖于预计算的图结构。在本文中，我们采取不同的方法：我们从快速高维高斯滤波推导出一种线性复杂度自注意力机制，该机制在保持注意力全部建模灵活性的同时，随网络规模线性扩展。

### 2.2 高效注意力机制

标准自注意力的二次复杂度[Vaswani等人，2017](https://arxiv.org/html/2607.13241#bib.bib85)推动了对高效替代方案的广泛研究，这些方案大致可分为三类。**稀疏注意力**方法将注意力模式限制在元素对的子集上。例如，Informer [Zhou等人，2021](https://arxiv.org/html/2607.13241#bib.bib76)通过 KL 散度度量选择主导查询，实现 \(O(L \log L)\) 复杂度，但其稀疏模式与数据无关，需要特定于任务的启发式规则。**低秩和核近似**方法利用注意力矩阵的低秩结构。Linformer [Wang等人，2020a](https://arxiv.org/html/2607.13241#bib.bib51)将键和值投影到较低维度空间，Performer [Choromanski等人，2022](https://arxiv.org/html/2607.13241#bib.bib52)使用正正交随机特征近似 softmax 核，FMMformer [Nguyen等人，2021](https://arxiv.org/html/2607.13241#bib.bib56)受快速多极子方法启发，将注意力分解为近场和远场分量。虽然这些方法在理论上实现了线性复杂度，但它们采用固定或随机的投影策略，对所有序列位置一视同仁，而不考虑底层数据结构。在交通预测场景中，传感器表现出强烈的地理聚类和空间结构化依赖关系，这种位置无关的处理会导致全注意力矩阵的次优近似。**聚类和引导点**方法在概念上与我们最接近。Set Transformer [Lee等人，2019](https://arxiv.org/html/2607.13241#bib.bib86)引入可学习的引导点来减少成对计算，而 Nyströmformer [Xiong等人，2021](https://arxiv.org/html/2607.13241#bib.bib54)通过 Nyström 近似从地标点子集重建注意力矩阵。然而，这些方法依赖于与底层数据结构无关的通用聚类或随机采样策略。

尽管高效注意力机制取得了进展，但它们直接应用于时空图数据以进行交通预测仍然有限。交通数据还表现出独特的特征（例如不规则图拓扑、空间和时间维度联合动态演化），通用方法并未明确解决这些问题。我们通过从快速高维高斯滤波理论推导出一种线性复杂度自注意力机制来弥补这一差距，其中聚类矩阵是端到端学习的，以适应交通网络的空间结构。

## 3 提出的模型

我们首先从快速高维高斯滤波理论推导出一种线性复杂度自注意力机制，然后开发用于交通预测的完整 EMAGN 架构。

### 3.1 线性复杂度自注意力

标准自注意力计算 \(\hat{X} = \text{Softmax}(\alpha Q K^T) V\)，复杂度为 \(O(N^2 d)\)[Vaswani等人，2017](https://arxiv.org/html/2607.13241#bib.bib85)，其中 \(Q, K, V\) 由输入线性投影得到。这种二次成本在结构上类似于朴素的高维高斯滤波：

\[\hat{v}_i = \sum_j e^{\frac{\|p_i - p_j\|^2}{2}} v_j,\] (1)

其中 \(p_i\) 和 \(v_i\) 是高维位置和值向量，这也需要全对评估。快速滤波算法[Adams等人，2010](https://arxiv.org/html/2607.13241#bib.bib59)通过将附近位置分组为 \(M\) 个聚类并预先聚合它们的值贡献，将成本降低到线性。