EMAGN: Efficient Multi-Attention Graph Network via Learned Clustering for Scalable Traffic Forecasting
July 14, 2026
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 \(\mathcal{O}(N^2d)\) to \(\mathcal{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.
1 Introduction↩︎
With the deployment of affordable traffic sensor technologies, traffic data are growing rapidly, bringing us to the era of transportation big data. Intelligent Transportation Systems (ITS) [1] leverage this data for efficient urban traffic control and planning. Traffic forecasting, namely predicting future traffic conditions (e.g., speeds, volumes, and density) from historical observations plays a critical role in ITS. Classical statistical methods [2], [3] rely on stationarity assumptions that are violated by the nonlinear, non-stationary nature of traffic flow. Deep learning approaches, from CNNs [4] and RNNs [5] to graph neural networks [6]–[8], have progressively improved forecasting accuracy by incorporating spatial structure. Self-attention mechanisms [9] have been widely adopted to model dynamic and long-range spatial-temporal dependencies [10], [11], achieving state-of-the-art performance. However, these self-attention-based models suffer from limited scalability due to the quadratic computational and memory complexity of the self-attention mechanism. As illustrated in Fig. 1, pairwise interactions across all elements must be computed, resulting in \(\mathcal{O}(N^2)\) complexity. The challenge grows with the addition of multi-head attention and deeper architectures, though both are essential for high performance on large-scale networks. While recent methods such as STAEformer [12] and MLCAFormer [13] continue to advance accuracy through architectural redesign, adaptive embeddings, hierarchical attention, propagation-delay modelling, the memory bottleneck of quadratic spatial attention itself remains unaddressed. Despite this well-known limitation, existing traffic forecasting benchmarks do not report GPU memory consumption in their experimental evaluations, making it impossible to assess the practical deployability of these models on resource-constrained hardware.
To alleviate this bottleneck, we propose an Efficient Multi-Attention Graph Network (EMAGN), achieving comparable forecasting performance while being much more efficient and scalable. A linear-complexity self-attention mechanism is developed to model spatial dependencies across elements. Specifically, we adaptively group and aggregate \(N\) key vectors into \(M\) super-key vectors (super-value vectors are generated accordingly), and then dependencies between query and super-key vectors are calculated, reducing the computational complexity from \(\mathcal{O}(N^2)\) to \(\mathcal{O}(NM)\), where \(M\) is a predefined hyperparameter. As a result, complexity scales linearly with the size of the traffic network. Our idea is derived from the fast high-dimensional Gaussian filtering technique in classical signal processing [14]. The key insight is that value vectors associated with nearby positions can be combined into a single super-value vector, significantly reducing computational cost while achieving near-exact filtering performance. EMAGN builds upon GMAN [10] as its base architecture, since GMAN explicitly computes full \(N \times N\) spatial attention across all sensor pairs, making the quadratic bottleneck most pronounced and the efficiency gain most measurable. To address the scalability problem, GMAN proposed a group-spatial attention mechanism that randomly partitions vertices into groups and computes only intra-group and inter-group attentions. However, random grouping does not guarantee performance, and more groups are needed to achieve desirable accuracy. In contrast, we learn the grouping adaptively during training to exploit the inherent spatial clustering within traffic networks. Empirically, we find that a small number of learned groups suffices for competitive performance. The main contributions of this paper are as follows:
We propose a linear-complexity self-attention mechanism derived from fast high-dimensional Gaussian filtering, reducing spatial attention complexity from \(\mathcal{O}(N^2d)\) to \(\mathcal{O}(NMd)\) through learned adaptive clustering matrices \(C_k\) and \(C_v\) that are trained end-to-end with the forecasting objective.
We develop EMAGN, an efficient multi-attention graph network for traffic forecasting that 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% on METR-LA.
We show that EMAGN expands the feasible configuration space of attention-based traffic forecasting: at \(K{=}16\) attention heads, full-attention GMAN exceeds the memory of a standard 11 GB GPU, while EMAGN continues to operate at 8,568 MB.
We provide the first direct empirical comparison of general linear attention mechanisms, Linformer [15] and Performer [16], against full attention on standard traffic benchmarks, demonstrating that EMAGN simultaneously achieves better accuracy and better efficiency than both alternatives.
2 Related Work↩︎
2.1 Traffic flow forecasting↩︎
Early deep learning approaches to traffic forecasting employ CNNs to extract spatial features from grid-converted traffic networks [4], [17], but this conversion discards the irregular topology inherent to road graphs. Graph neural networks (GNNs) overcome this limitation by operating directly on the network structure [18], [19]. STGCN [6] applies spectral graph convolutions, while DCRNN and GraphWaveNet [7], [8] adopt diffusion convolutions on directed graphs to capture directional flow. These models, however, rely on fixed adjacency matrices and cannot adapt to time-varying spatial dependencies. Adaptive-graph methods such as AGCRN [20] and MTGNN [21] address this by learning the graph structure from data, yet they still depend on local message-passing and may struggle to capture long-range spatial interactions.
Self-attention mechanisms offer a natural solution to these problems with graph structure data, as they model pairwise dependencies across all nodes regardless of graph distance [10], [11], [22], [23]. GMAN [10] introduces a graph multi-attention mechanism with group-based spatial attention; PDFormer [24] incorporates propagation delay-aware attention; STAEformer [12] shows that a vanilla Transformer with well-designed adaptive embeddings can achieve strong results; and MLCAFormer [13] captures hierarchical spatio-temporal dependencies via multi-level causal attention. Despite their effectiveness, all these models inherit the \(\mathcal{O}(N^2)\) cost of standard self-attention, which limits their scalability to large networks. Recent work has begun to tackle this bottleneck. BigST [25] achieves linear complexity through linearised graph convolutions, and STGformer [26] reduces cost via a single-layer high-order attention design. However, BigST sacrifices the flexibility of attention by relying on graph convolutions, while STGformer still depends on pre-computed graph structures. In this paper, we take a different approach: we derive a linear-complexity self-attention mechanism from fast high-dimensional Gaussian filtering that preserves the full modelling flexibility of attention while scaling linearly with network size.
2.2 Efficient attention mechanisms↩︎
The quadratic complexity of standard self-attention [9] has motivated extensive research on efficient alternatives, which can be broadly grouped into three families. Sparse attention methods restrict the attention pattern to a subset of element pairs. Informer [27], for example, selects dominant queries via a KL-divergence measure to achieve \(\mathcal{O}(L\log L)\) complexity, but its sparse patterns are data-agnostic and require task-specific heuristics. Low-rank and kernel approximation methods exploit the low-rank structure of the attention matrix. Linformer [15] projects keys and values into a lower-dimensional space, Performer [16] approximates the softmax kernel with positive orthogonal random features, and FMMformer [28] decomposes attention into near-field and far-field components inspired by the fast multipole method. While these methods achieve theoretical linear complexity, they apply fixed or random projection strategies that treat all sequence positions uniformly, without regard to the underlying data structure. In the traffic forecasting setting, where sensors exhibit strong geographic clustering and spatially structured dependencies, this position-agnostic treatment leads to suboptimal approximations of the full attention matrix. Clustering and inducing-point methods are conceptually closest to our work. Set Transformer [29] introduces learnable inducing points to reduce pairwise computations, while Nyströmformer [30] reconstructs the attention matrix from a subset of landmark points via the Nyström approximation. However, these methods rely on generic clustering or random sampling strategies that are agnostic to the underlying data structure.
Despite the advances of efficient attention mechanisms, their direct application to spatio-temporal graph data for traffic forecasting remains limited. Traffic data also shows unique characteristics (e.g., irregular graph topology, jointly and dynamically evolving spatial and temporal dimensions) that general-purpose methods do not explicitly address. We bridge this gap by deriving a linear-complexity self-attention mechanism from the theory of fast high-dimensional Gaussian filtering, where the clustering matrices are learned end-to-end to adapt to the spatial structure of traffic networks.
3 Proposed Model↩︎
We first derive a linear-complexity self-attention mechanism from the theory of fast high-dimensional Gaussian filtering, then develop the full EMAGN architecture for traffic forecasting.
3.1 Linear-Complexity Self-Attention↩︎
Standard self-attention computes \(\hat{X} = \text{Softmax}(\alpha QK^T)V\) with \(\mathcal{O}(N^2d)\) complexity [9], where \(Q, K, V\) are linearly projected from the input. This quadratic cost is structurally analogous to naive high-dimensional Gaussian filtering: \[\label{eq1} \hat{v}_i = \sum_j e^{\frac{\|p_i - p_j\|^2}{2}} v_j,\tag{1}\] where \(p_i\) and \(v_i\) are high-dimensional positions and value vectors, which also requires all-pairs evaluation. Fast filtering algorithms [14] reduce this to linear cost by grouping nearby positions into \(M\) clusters and pre-aggregating their value contributions, as displayed in Fig. 1. Specifically, given cluster centre \(p_c\), the contribution of all elements in a cluster to a query \(p_i\) is approximated via a truncated Taylor expansion as \[\label{eq3} \textstyle\sum_je^{{\|p_i-p_j\|^2}/{2}}{v_j}\;\approx\; e^{{\|p_i-p_c\|^2}/{2}}\,P_r(\|p_i-p_c\|^2),\tag{2}\] where \(P_r\) aggregates the \(r\)-th order Taylor terms and can be precomputed per cluster, so that each query attends only to \(M \ll N\) cluster summaries rather than all \(N\) elements.
Self-attention can be viewed as a generalisation of Gaussian filtering with directed dependencies and a scaled-product kernel rather than a Gaussian kernel. Following this analogy, we replace the hand-crafted clustering of classical filtering with two learnable clustering matrices \(C_k, C_v \in \mathbb{R}^{M \times N}\) that are optimised end-to-end, allowing the grouping to adapt to the spatial structure of traffic networks. The resulting linear-complexity self-attention is \[\hat{X} = \mathrm{Softmax}(\alpha Q (C_k K)^T) C_v V. \label{eq:linear95attention}\tag{3}\] The attention matrix becomes \(A = \alpha Q (C_k K)^T \in \mathbb{R}^{N \times M}\), reducing complexity from \(\mathcal{O}(N^2 d)\) to \(\mathcal{O}(NMd)\) (see Fig. 1). Since \(M\) is predefined and independent of \(N\), computation scales linearly with network size. This mechanism can be replicated across multiple heads for richer representational capacity.
3.2 EMAGN for Traffic Forecasting↩︎
3.2.0.1 Problem formulation.
We denote \(\mathcal{G}(\mathcal{V}, \mathcal{E}, W)\) as the input graph representing a traffic network, where \(\mathcal{V}\) represents the nodes set with \(N\) sensors, \(\mathcal{E}\) is the set of edges reflecting physical connectivity between sensors and \(W \in \mathbb{R}^{N \times N}\) is the adjacency matrix which is usually constructed with the Euclidean distances between sensors via Gaussian kernel. Formally, given \(P\) historical traffic speeds \([v^{\tau-P+1}, \cdots, v^{\tau}] \in \mathbb{R}^{P \times N \times 1}\) observed by the \(N\) sensors in the traffic network \(\mathcal{G}\), a traffic forecasting model \(\mathcal{F}\) is learned to predict the next \(T\) time step traffic speeds \([\hat{v}^{\tau+1}, \cdots, \hat{v}^{\tau+T}]\) as \[\hat{v}^{\tau+1}, \cdots, \hat{v}^{\tau+T} = \mathcal{F}(v^{\tau-P+1}, \cdots, v^{\tau}; \mathcal{G}). \label{eq:problem}\tag{4}\]
3.2.0.2 Overall architecture.
The architecture of the proposed model is illustrated in Fig. 2, consisting of an encoder and decoder. An alignment module that models the dynamical dependencies between historical and future time steps is also adopted. The STBlock consists of two modules: one for modeling spatial dependencies and another for modeling temporal dependencies. Spatial and temporal features are extracted independently within each STBlock and then fused using a gate. STBlocks are stacked to capture deep spatial and temporal features. To achieve spatial embedding, the node2vec algorithm [31] is used to learn the embeddings \(D_s\). For temporal embedding, we encode the day of the week and the time step in a day using one-hot encoding. These are then concatenated into a single vector, denoted as \(D_t\). The spatial and temporal embeddings are both projected into a fixed-dimensional vector \(D \in \mathbb{R}^{d_e}\) using two-layer fully connected networks. The learned spatial and temporal embeddings are concatenated with each input feature vector as positional embeddings.
3.2.0.3 Spatial dependencies modeling.
We propose adopting a linear-complexity self-attention mechanism to model spatial dependencies. For simplicity, we will use the example of time step \(t_i\) for notation. Especially, for graph representations \(X_{t_i} \in \mathbb{R}^{N \times (d_s + d_e)}\) that is either traffic speeds information or learned representation from the previous layer augmented with spatial-temporal embedding. Three linear projections \(W_q^s \in \mathbb{R}^{(d_s+d_e) \times d_q}\), \(W_k^s \in \mathbb{R}^{(d_s+d_e) \times d_k}\) and \(W_v^s \in \mathbb{R}^{(d_s+d_e) \times d_v}\) and two clustering matrices \(C_k \in \mathbb{R}^{M \times N}\) and \(C_v \in \mathbb{R}^{M \times N}\) are learned. Formally, the updated graph features could be written as \[\hat{X}_{t_i} = \mathrm{Softmax}\!\left(\frac{1}{\sqrt{d_k}} X_{t_i} W_q^s (C_k X_{t_i} W_k^s)^T\right) C_v X_{t_i} W_v^s. \label{eq:spatial}\tag{5}\] In Eq. (5 ), \(Q = X_{t_i} W_q^s\), \(K = C_k X_{t_i} W_k^s\) and \(V = C_v X_{t_i} W_v^s\) are the learned query, super-key and super-value vectors, respectively. The key and value vectors are firstly adaptively grouped into \(M\) clusters before modeling spatial dependencies. Compared to random grouping, we can achieve better performance with smaller groups. Multi-head attention can be further employed for various spatial patterns and the representations are further concatenated and fused with a two-layer fully connected network.
3.2.0.4 Temporal dependencies modeling.
In this study, we also employ a self-attention mechanism for modeling dynamic and long-range temporal dependencies. Specifically, we consider time sequence features \(H_{v_i} \in \mathbb{R}^{P \times (d_e + d_t)}\) of node \(v_i\). The updated time features can be formally written as \[\hat{H}_{v_i} = \mathrm{Softmax}\!\left(\frac{1}{\sqrt{d_k}} H_{v_i} W_q^t (H_{v_i} W_k^t)^T\right) H_{v_i} W_v^t \label{eq:temporal}\tag{6}\] where \(W_q^t \in \mathbb{R}^{(d_e+d_t) \times d_q}\), \(W_k^t \in \mathbb{R}^{(d_e+d_t) \times d_k}\) and \(W_v^t \in \mathbb{R}^{(d_e+d_t) \times d_v}\) are three linear projection matrices. As temporal sequences are fixed at \(T = P =12\) in our setting, the quadratic cost of temporal attention is negligible—\(\mathcal{O}(T^2 d)\) with \(T \ll N\), and the dominant complexity of EMAGN remains \(\mathcal{O}(NMd)\) in the spatial dimension. When the sequence grows much larger, linear complexity can be further considered. Similar to spatial dependencies modeling, multi-head attention can also be adopted and fused with fully connected networks.
3.2.0.5 Multi-step predictions.
To model the dynamic temporal dependencies between past and future time steps, we employ an additional attention mechanism. By utilizing spatial and temporal embeddings for input and output time steps as features, we can effectively capture the dynamic temporal dependencies for different nodes. Formally, let \(D_{v_i}^p \in \mathbb{R}^{T \times d_e}\) and \(D_{v_i}^f \in \mathbb{R}^{T \times d_e}\) denote the past and future embeddings of node \(v_i\), respectively, then the alignment matrix can be calculated with attention mechanism as \[A = \mathrm{Softmax}\bigl(f(D_{v_i}^p)(f(D_{v_i}^f))^T\bigr), \label{eq:alignment}\tag{7}\] where \(A\) is the alignment matrix and \(f\) are learnable linear projections. In our setting \(P = T\), so both past and future embeddings share the same dimensionality. The output of encoder \(X_{v_i}^{st}\) is then transformed with \(A\) as \(A X_{v_i}^{st}\) before fed into decoder.
4 Experiments↩︎
We evaluate the proposed model on two real-world benchmark datasets: PEMS-BAY and METR-LA. As METR-LA is the more challenging benchmark, with higher speed variance (std 19.50 vs. km/h) and stronger sensitivity to spatial structure, all ablation and scalability analyses are conducted on METR-LA unless otherwise stated.
4.1 Datasets and Data Preprocessing↩︎
PEMS-BAY consists of 6 months of speed data from 325 sensors in the Bay Area of California, starting from January 1st, 2017 through May 31st, 2017. METR-LA consists of traffic data from 207 sensor stations in the California state highway system during weekdays from May through June in 2012. The data statistics are presented in Table 1.
| Dataset | #nodes | #time steps | Mean speed | Std |
|---|---|---|---|---|
| PEMS-BAY | 325 | 52,116 | 62.74 km/h | 9.44 |
| METR-LA | 207 | 34,272 | 54.40 km/h | 19.50 |
Following the common data processing pipeline, traffic speeds are aggregated every five minutes and normalised with Z-score. Data are split into 70% training, 10% validation, and 20% testing. The road topology is encoded as a graph adjacency matrix using pairwise road distances via Gaussian kernel: \[W_{ij}= \begin{cases} \exp\!\left(\tfrac{-d_{ij}^2}{\sigma^2}\right), & \exp\!\left(\tfrac{-d_{ij}^2}{\sigma^2}\right)>\epsilon\\ 0, & \text{otherwise,} \end{cases}\] where \(d_{ij}\) is the distance between sensors \(v_i\) and \(v_j\), \(\sigma\) is the standard deviation and \(\epsilon\) controls the sparsity of the adjacency matrix.
4.2 Evaluation Metrics and Baselines↩︎
We report mean absolute error (MAE), root mean squared error (RMSE), and mean absolute percentage error (MAPE), respectively, at 15, 30, and 60 minute prediction horizons. Baselines include graph-based deep models: STGCN [6], DCRNN [7], GraphWaveNet [8], MTGNN [21], and STID [32]. Among Transformer based methods (\(\dagger\)), we compare against PDFormer [24], STAEformer [12], MLCAFormer [13], and GMAN [10], the full-attention model that EMAGN directly improves upon. For GMAN we adopt the published code with default settings; all remaining baseline results are reported from their original papers.
4.3 Experimental Settings and Model Configurations↩︎
All experiments are conducted on an NVIDIA RTX 2080Ti GPU. Models are trained using the Adam optimiser with an initial learning rate of \(10^{-3}\), decaying by a factor of 0.7 every 5 epochs, and weight decay \(10^{-3}\). Early stopping is applied after 10 consecutive epochs without validation improvement. We use 12 historical steps to predict the next 15-, 30-, and 60-minutes. Model depth is set proportional to dataset complexity, with PEMS-BAY using \(L=1\) STBlock and METR-LA using \(L=5\) STBlocks, both with \(K=8\) attention heads and 8 channels per head. Unless stated otherwise, EMAGN uses \(M=4\) clusters for METR-LA, determined by the ablation in Section 4.5.
| 15 min | 30 min | 60 min | |||||||
|---|---|---|---|---|---|---|---|---|---|
| 2-4(lr)5-7(lr)8-10 Model | MAE | RMSE | MAPE | MAE | RMSE | MAPE | MAE | RMSE | MAPE |
| PEMS-BAY | |||||||||
| STGCN | 1.36 | 2.96 | 2.90 | 1.81 | 4.27 | 4.17 | 2.49 | 5.69 | 5.79 |
| DCRNN | 1.38 | 2.95 | 2.90 | 1.74 | 3.97 | 3.90 | 2.07 | 4.74 | 4.90 |
| GWNet | 1.30 | 2.74 | 2.73 | 1.63 | 3.70 | 3.67 | 1.95 | 4.52 | 4.43 |
| MTGNN | 1.33 | 2.80 | 2.81 | 1.66 | 3.77 | 3.75 | 1.95 | 4.50 | 4.62 |
| STID | 1.31 | 2.79 | 3.78 | 1.64 | 3.73 | 3.73 | 1.91 | 4.42 | 4.55 |
| PDFormer\(\dagger\) | 1.32 | 2.83 | 2.78 | 1.64 | 3.79 | 3.71 | 1.91 | 4.43 | 4.51 |
| STAEformer\(\dagger\) | 1.31 | 2.78 | 2.76 | 1.62 | 3.68 | 3.62 | 1.88 | 4.34 | 4.41 |
| MLCAFormer\(\dagger\) | 1.28 | 2.73 | 2.66 | 1.59 | 3.63 | 3.50 | 1.86 | 4.30 | 4.30 |
| GMAN\(\dagger\) | 1.35 | 2.93 | 2.90 | 1.63 | 3.76 | 3.70 | 1.87 | 4.35 | 4.41 |
| EMAGN (ours) | 1.39 | 3.01 | 2.99 | 1.68 | 3.84 | 3.82 | 1.92 | 4.43 | 4.54 |
| STGCN | 2.88 | 5.74 | 7.62 | 3.47 | 7.24 | 9.57 | 4.59 | 9.40 | 12.70 |
| DCRNN | 2.77 | 5.38 | 7.30 | 3.15 | 6.45 | 8.80 | 3.60 | 7.59 | 10.50 |
| GWNet | 2.69 | 5.15 | 6.90 | 3.07 | 6.22 | 8.37 | 3.53 | 7.37 | 10.01 |
| MTGNN | 2.69 | 5.16 | 6.89 | 3.05 | 6.13 | 8.16 | 3.47 | 7.21 | 9.70 |
| STID | 2.82 | 5.53 | 7.75 | 3.19 | 6.57 | 9.39 | 3.55 | 7.55 | 10.95 |
| PDFormer\(\dagger\) | 2.83 | 5.45 | 7.77 | 3.20 | 6.46 | 9.19 | 3.62 | 7.47 | 10.91 |
| STAEformer\(\dagger\) | 2.65 | 5.11 | 6.85 | 2.97 | 6.00 | 8.13 | 3.34 | 7.02 | 9.70 |
| MLCAFormer\(\dagger\) | 2.62 | 5.05 | 6.72 | 2.93 | 5.96 | 7.95 | 3.30 | 6.97 | 9.47 |
| GMAN\(\dagger\) | 2.77 | 5.42 | 7.19 | 3.08 | 6.29 | 8.38 | 3.43 | 7.20 | 9.78 |
| EMAGN (ours) | 2.91 | 5.79 | 7.98 | 3.23 | 6.70 | 9.10 | 3.54 | 7.52 | 10.27 |
| 60 min | Efficiency | |||||
|---|---|---|---|---|---|---|
| 2-4(lr)5-7 Model | MAE | RMSE | MAPE | Train (s) | Infer (s) | Mem (MB) |
| GMAN | 3.43 | 7.20 | 9.78 | 721.5 | 36.9 | 10666 |
| EMAGN (ours) | 3.54 | 7.52 | 10.27 | 489.4 | 22.7 | 4472 |
| Linformer | 3.64 | 7.63 | 10.39 | 495.4 | 24.7 | 4562 |
| Performer | 3.62 | 7.62 | 10.78 | 496.6 | 24.7 | 4562 |
4.4 Experimental Results↩︎
Table 2 reports MAE, RMSE, and MAPE at the 15-, 30-, and 60-minute horizons on PEMS-BAY and METR-LA, respectively. We design the numerical experiments around three questions: (i) how does EMAGN compare in absolute accuracy ? (ii) what accuracy–efficiency trade-off does it achieve? and (iii) does its learned-clustering mechanism outperform other linear attention designs?
Accuracy comparison.↩︎
Quadratic-attention Transformers (i.e., MLCAFormer and STAEformer) achieve the highest overall accuracy on both benchmarks, benefiting from full \(\mathcal{O}(N^2)\) pairwise interactions, followed by graph-based models (GraphWaveNet, MTGNN) and the full-attention GMAN. EMAGN, operating at \(\mathcal{O}(NM)\) complexity, incurs a moderate accuracy gap relative to these \(\mathcal{O}(N^2)\) methods yet remains firmly within the range of established baselines. On METR-LA at the 60-minute horizon, EMAGN achieves an MAE of 3.54%, placing it within 3.2% of its full-attention parent GMAN (3.43%) and on par with GraphWaveNet (3.53%) and DCRNN (3.60%). On PEMS-BAY at 60 min, the gap narrows to 2.7% MAE (1.92% vs.GMAN’s 1.87%), while EMAGN substantially outperforms DCRNN (2.07%) and STGCN (2.49%). A noteworthy trend emerges across prediction horizons: the accuracy gap between EMAGN and GMAN consistently shrinks as the forecasting horizon increases—from 5.1% relative MAE at 15 min to 3.2% at 60 min on METR-LA, and from 3.0% to 2.7% on PEMS-BAY. This suggests that EMAGN’s learned clustering captures the dominant spatial dependencies most critical for long-range prediction, where modelling global structure matters more than fine-grained pairwise interactions.
Accuracy–efficiency trade-off.↩︎
Relative to GMAN, the full-attention architecture from which EMAGN is derived, our model reduces training time by 32% (\(721.5 \to 489.4\) s/epoch), inference time by 38% (\(36.9 \to 22.7\) s), and peak GPU memory by 58% (\(10{,}666 \to 4{,}472\) MB) on METR-LA (Table 4, row \(M{=}4\) vs.Full). As Section 4.6 demonstrates, these savings enable model configurations, such as \(K{=}16\) attention heads, that cause full-attention GMAN to exceed the memory of a standard 11 GB GPU, whereas EMAGN continues to operate at 8,568 MB. Quadratic-attention baselines that achieve higher accuracy (STAEformer, MLCAFormer) inherit the same \(\mathcal{O}(N^2)\) bottleneck and therefore face identical scalability limits. EMAGN thus occupies a distinct position on the accuracy–efficiency Pareto front: it is the only method that simultaneously achieves accuracy within the range of established graph-based baselines and guarantees linear-complexity attention.
Comparison with linear attention alternatives.↩︎
The efficiency gains above could, in principle, be obtained by any linear attention mechanism. To isolate the contribution of EMAGN’s learned clustering design, we replace its attention module with Linformer [15] and Performer [16], two widely adopted linear-attention methods, within the identical GMAN backbone. All three linear variants use the same network depth (\(L{=}5\)), head count (\(K{=}8\)), and projection dimension (\(m{=}4\)), ensuring a controlled comparison. Table 3 reports accuracy at 60 min and efficiency on METR-LA. EMAGN outperforms both alternatives across all reported metrics. In accuracy, it surpasses Linformer by 0.10 MAE (2.7% relative) and Performer by 0.08 MAE (2.2%); the gap is especially pronounced in MAPE, where EMAGN (10.27%) improves over Performer (10.78%) by 4.7% relative. In efficiency, EMAGN consumes 2% less memory (4,472 vs.562 MB) and runs 8% faster at inference (22.7 vs. s). The advantage stems from the fact that Linformer and Performer apply fixed or random projections that treat all positions uniformly, whereas EMAGN learns separate clustering matrices \(C_k\) and \(C_v\) end-to-end that can adapt to the spatial structure of the traffic network.
4.5 Ablation Study↩︎
Influence of the number of clusters.↩︎
A key hyperparameter in EMAGN is the cluster count \(M\), which controls the trade-off between approximation quality and computational efficiency. Table 4 lists accuracy and resource consumption as \(M\) varies from 2 to 128, benchmarked against full \(\mathcal{O}(N^2)\) attention.
| 15 min | 30 min | 60 min | Efficiency | |||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| 2-4(lr)5-7(lr)8-10(lr)11-13 \(M\) | MAE | RMSE | MAPE | MAE | RMSE | MAPE | MAE | RMSE | MAPE | Trn (s) | Inf (s) | Mem (MB) |
| Full | 2.77 | 5.42 | 7.19 | 3.08 | 6.29 | 8.38 | 3.43 | 7.20 | 9.78 | 721 | 36.9 | 10,666 |
| 128 | 2.96 | 5.77 | 8.23 | 3.29 | 6.69 | 9.54 | 3.61 | 7.46 | 10.80 | 619 | 30.2 | 8,543 |
| 64 | 2.93 | 5.83 | 8.14 | 3.26 | 6.76 | 9.47 | 3.60 | 7.60 | 10.81 | 593 | 28.7 | 7,728 |
| 32 | 2.95 | 5.88 | 8.35 | 3.29 | 6.83 | 9.68 | 3.61 | 7.64 | 10.89 | 567 | 27.2 | 7,728 |
| 16 | 2.91 | 5.80 | 8.08 | 3.24 | 6.75 | 9.47 | 3.58 | 7.61 | 10.85 | 541 | 25.7 | 6,100 |
| 8 | 2.90 | 5.76 | 7.91 | 3.22 | 6.65 | 9.15 | 3.58 | 7.52 | 10.52 | 515 | 24.2 | 5,286 |
| 4 | 2.91 | 5.79 | 7.98 | 3.23 | 6.70 | 9.10 | 3.54 | 7.52 | 10.27 | 489 | 22.7 | 4,472 |
| 2 | 2.92 | 5.75 | 8.03 | 3.22 | 6.56 | 9.21 | 3.59 | 7.44 | 10.72 | 478 | 21.9 | 4,250 |
Three findings are noteworthy. First, \(M=4\) achieves accuracy essentially on par with \(M=8\) (MAE 2.91 vs. at 15 min; 3.54 vs. at 60 min) while using 4,472 MB versus 5,286 MB, and both are far below full attention’s 10,666 MB. Second, larger cluster sizes (\(M=64\), \(M=128\)) do not consistently improve over smaller ones and in several cases perform worse. This behaviour is consistent with the Gaussian filtering analogy: once the learned clusters capture the dominant spatial groupings in the traffic network, additional clusters provide diminishing approximation benefit while increasing compute. Third, and most significantly, the fact that \(M=4\) matches or exceeds \(M=128\) on accuracy confirms that traffic sensor graphs have inherent low-rank spatial structure. A small number of learned clusters suffices because sensors naturally group into coherent geographic corridors, a structure that EMAGN discovers end-to-end, and that generic low-rank projections (Linformer, Performer) cannot exploit. We set \(M=4\) as the default for METR-LA.
4.6 Scalability Analysis↩︎
Beyond the efficiency gains at existing configurations, EMAGN’s linear complexity enables model configurations that are infeasible under full attention. We demonstrate this through three controlled scaling experiments on METR-LA.
Scalability with the Number of STBlocks \(L\)↩︎
Table 5 shows resource consumption as more STBlocks are stacked. Both models scale roughly linearly in \(L\), as expected. However, EMAGN’s costs are consistently and substantially lower at every depth. Notably, even at \(L=1\), EMAGN already consumes 58% less memory (895 MB vs.,133 MB) and trains 23% faster than GMAN, confirming the efficiency advantage is not a deep-model artefact but is present at every configuration. At \(L=5\), EMAGN requires 489.4s training time and 4,472 MB memory versus GMAN’s 721.5s and 10,666 MB, reductions of 32% and 58% respectively. Within a fixed memory budget, EMAGN can therefore deploy significantly deeper models than GMAN, enabling richer spatial-temporal feature extraction that is entirely inaccessible under full attention.
| GMAN (Full Attention) | EMAGN (Linear Attention) | |||||
|---|---|---|---|---|---|---|
| 2-4(lr)5-7 \(L\) | Train (s) | Infer (s) | Mem (MB) | Train (s) | Infer (s) | Mem (MB) |
| 1 | 172.4 | 8.7 | 2,133 | 133.0 | 6.2 | 895 |
| 2 | 311.1 | 15.8 | 4,447 | 224.0 | 11.1 | 1,850 |
| 3 | 446.8 | 22.8 | 6,495 | 315.4 | 16.0 | 2,760 |
| 4 | 579.4 | 29.7 | 8,543 | 406.8 | 20.9 | 3,680 |
| 5 | 721.5 | 36.9 | 10,666 | 489.4 | 22.7 | 4,472 |
Scalability with the Number of Attention Heads \(K\)↩︎
Figure 3 reports the effect of varying the number of attention heads \(K\). At \(K{=}16\), GMAN exceeds the 11 GB GPU memory (OOM), whereas EMAGN continues to operate at 954 s training time and 8,568 MB. At the shared configuration \(K{=}8\), EMAGN consumes 58% less memory and runs \(1.6{\times}\) faster at inference, confirming that the efficiency advantage grows with the number of attention heads.
Scalability with Traffic Network Size \(N\)↩︎
Figure 4 evaluates how costs scale with network size \(N\) by sub-sampling the METR-LA graph. At small \(N\) (\(N=10\)), both models exhibit similar resource requirements since the quadratic term \(N^2\) is negligible. As \(N\) grows, the gap widens significantly and follows the predicted pattern: GMAN’s costs grow steeply while EMAGN’s costs follow a near-linear trajectory, directly validating the theoretical \(\mathcal{O}(NM)\) complexity. At \(N=207\), EMAGN requires 32% less training time and 58% less memory. For deployment on larger real-world networks, city-scale systems with thousands of sensors, this gap would grow substantially further, making full-attention models practically infeasible while EMAGN remains deployable.
Together, the three scalability analyses consistently confirm that EMAGN achieves its theoretical \(\mathcal{O}(NM)\) advantage in practice across all three axes of scale. The efficiency gains are not merely cost reductions at existing configurations, they translate directly into new capabilities: deeper networks, richer multi-head configurations, and larger-scale deployments that full-attention models cannot support within standard hardware budgets.
5 Conclusion↩︎
In this paper, we proposed EMAGN, an efficient multi-attention graph network for traffic speed forecasting. Inspired by the high-dimensional Gaussian filtering, we developed a linear-complexity attention mechanism that replaces the \(\mathcal{O}(N^2)\) pairwise interactions of standard self-attention with \(\mathcal{O}(NM)\) operations through learned adaptive clustering. The EMAGN architecture stacks spatial-temporal blocks with parallel fusion of spatial and temporal attention, combined with a transform-attention alignment module for multi-step prediction.
Extensive experiments on PEMS-BAY and METR-LA demonstrate that EMAGN achieves accuracy closely competitive with the state-of-the-art full-attention model GMAN, with only a 2.7–3.2% MAE gap at the 60-minute horizon, while reducing training time by 32%, inference time by 38%, and memory consumption by 58%. Scalability experiments further show that EMAGN supports configurations (e.g., 16 attention heads) that exceed the memory capacity of full-attention models. Ablation studies confirm that as few as 4–8 learned clusters suffice to approximate full pairwise attention, validating the design principle derived from fast Gaussian filtering. A controlled comparison with Linformer and Performer further shows that learned clustering achieves better accuracy and efficiency than generic linear attention alternatives within the same backbone.
Future work includes extending the linear attention mechanism to temporal dependencies when dealing with longer historical sequences, exploring adaptive selection of cluster size \(M\) during inference, and evaluating EMAGN on larger city-scale traffic networks.