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 linearizes the spatial attention mechanism itself, inspired by the theory of fast high-dimensional Gaussian filtering.
EMAGN utilizes two learned clustering matrices $C_k$ and $C_v$ to 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 modeling. 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%. Notably, 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. Furthermore, EMAGN surpasses Linformer and Performer in both accuracy and efficiency within the same backbone, owing to its traffic-network-aware adaptive clustering.
Blogger's Review: EMAGN significantly enhances the scalability of traffic forecasting by linearizing the spatial attention mechanism, showcasing its potential in handling large-scale data, especially under resource constraints. Its adaptive clustering strategy offers a good balance between model flexibility and efficiency, making it worthy of further exploration and application.