通过自适应聚类让图注意力网络更高效,大幅降低内存与计算开销。
EMAGN: Efficient Multi-Attention Graph Network via Learned Clustering for Scalable Traffic Forecasting

- 用学习的聚类矩阵将键值向量分组,线性化空间注意力计算
- 在两个数据集上误差仅比全注意力模型高2.7%-3.2%,训练快32%
- 适合大规模交通预测场景,尤其适合显存受限的部署
交通预测因复杂的时空依赖关系而极具挑战。自注意力机制虽能建模动态长程依赖并达到顶尖性能,但受限于二次方计算与内存复杂度,难以扩展。为此,我们提出高效多注意力图网络(EMAGN),基于快速高维高斯滤波理论,将空间注意力机制线性化。通过两个可学习的聚类矩阵C_k和C_v,自适应地将键和值向量划分为M个超簇,使复杂度从O(N²d)降至O(NMd),同时保持注意力对动态依赖建模的灵活性。在PEMS-BAY和METR-LA数据集上的实验表明,EMAGN的平均绝对误差(MAE)仅比全注意力GMAN高出2.7%-3.2%,训练时间减少32%,推理时间降低38%,GPU内存占用下降58%。关键的是,在16头注意力配置下,标准11GB GPU上全注意力GMAN完全无法运行,而EMAGN仍可正常工作,显著拓展了可行模型配置范围。相比Linformer和Performer,EMAGN在相同骨干网络下兼具更高精度与效率,得益于其面向交通网络的自适应聚类设计。
原文摘要 · Abstract (English)
Traffic forecasting is highly challenging due to complex and nonlinear spatial and temporal dependencies. Self-attention mechanisms have been widely adopted to model dynamic and long-range dependencies, achieving state-of-the-art performance, but suffer from limited scalability due to quadratic computational and memory complexity. To address this, we propose an Efficient Multi-Attention Graph Network (EMAGN) that linearises the spatial attention mechanism itself, inspired by the theory of fast high-dimensional Gaussian filtering. Two learned clustering matrices C_k and C_v adaptively group key and value vectors into M super-clusters, reducing complexity from O(N^2 d) to O(NMd) without sacrificing the flexibility of attention for dynamic dependency modelling. Experimental results on PEMS-BAY and METR-LA show that EMAGN achieves accuracy within 2.7-3.2% MAE of full-attention GMAN while reducing training time by 32%, inference time by 38%, and GPU memory by 58%. Critically, at K=16 attention heads, full-attention GMAN runs out of memory on a standard 11 GB GPU entirely while EMAGN continues to operate, demonstrating a categorical expansion of feasible model configurations. EMAGN also surpasses Linformer and Performer in both accuracy and efficiency within the same backbone, owing to its traffic-network-aware adaptive clustering.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。