EMAGN: Efficient Multi-Attention Graph Network via Learned Clustering for Scalable Traffic Forecasting

2026-07-14

AI Summary

This work addresses the scalability limitations of self-attention mechanisms in traffic forecasting, which suffer from O(N²d) computational complexity and hinder application to large-scale road networks. The authors propose EMAGN, a novel model that linearizes spatial attention to O(NMd) by introducing learnable clustering matrices Cₖ and Cᵥ, substantially reducing both computational and memory overhead. EMAGN integrates linear multi-head graph attention based on high-dimensional Gaussian filtering, learnable hyper-cluster aggregation, and spatiotemporal graph neural networks. Evaluated on PEMS-BAY and METR-LA datasets, EMAGN achieves only 2.7–3.2% higher MAE than the full-attention GMAN while reducing training time by 32%, accelerating inference by 38%, and cutting GPU memory usage by 58%. Notably, it efficiently supports a 16-head attention configuration on GPUs with merely 11GB of memory.

Abstract

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.

Problem

• traffic forecasting
• self-attention
• scalability
• computational complexity
• memory efficiency

Innovation

• efficient attention
• learned clustering
• graph neural network
• traffic forecasting
• scalable modeling