# Efficient Spatio-Temporal Graph Neural Networks for Traffic Forecasting

Yackov Lubarsky, Alexei Gaissinski, Pavel Kisilev

► To cite this version:

Yackov Lubarsky, Alexei Gaissinski, Pavel Kisilev. Efficient Spatio-Temporal Graph Neural Networks for Traffic Forecasting. 19th IFIP International Conference on Artificial Intelligence Applications and Innovations (AIAI), Jun 2023, León, Spain. pp.109-120, 10.1007/978-3-031-34107-6_9. hal-05498095

HAL Id: hal-05498095

https://inria.hal.science/hal-05498095v1

Submitted on 6 Feb 2026

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_1.png" /></figure>

Distributed under a Creative Commons CC BY 4.0 - Attribution - International License
<figure><img src="image_2.png" /></figure>

ifip

International Federation for Information Processing

This document is the original author manuscript of a paper submitted to an IFIP
conference proceedings or other IFIP publication by Springer Nature. As such, there
may be some differences in the official published version of the paper. Such
differences, if any, are usually due to reformatting during preparation for publication or
minor corrections made by the author(s) during final proofreading of the publication
manuscript.
# Efficient Spatio-Temporal Graph Neural Networks for Traffic Forecasting

Yackov Lubarsky(✉), Alexei Gaissinski, and Pavel Kisilev

Toga Networks, Huawei Tel Aviv Research Center, Hod Hasharon, Israel
{yackov.lubarsky, alexei.gaissinski, pavel.kisilev}@huawei.com

**Abstract.** Urban Traffic Forecasting has recently seen a lot of research activity as it entails a compelling combination of multivariate temporal data with geo-spatial dependencies between multiple data collection sensors. Current top approaches to this task tend to use costly spatio-temporal pipelines, where the model complexities typically have linear dependency on the time-series length and quadratic on the number of nodes. In this paper, we propose a number of steps to dramatically improve the runtime efficiency of the traffic forecasting solutions. First, we use a temporal pooling stack prior to spatial processing to effectively eliminate the time dimension before applying the spatial components. This removes the linear dependency of the model on the length of the time series. Second, we construct learnable graph pooling blocks inside the spatial stack which progressively reduce the size of the graph and facilitate better data flow between far away nodes. Experimental results on the standard METR-LA and PEMSBAY benchmarks show that the proposed approach yields significant inference and training speedups of up to x5 in the 1-h prediction task and x27 in the 24-h prediction task, while keeping or surpassing the state-of-the-art results. Our findings call into question the need for time-consuming spatio-temporal processing blocks, used in many of the latest solutions for the traffic forecasting task.

## 1 Introduction

Spatio-temporal signal modeling has recently seen a lot of activity with the rise in popularity of Graph Neural Networks (GNNs) and Attention-based neural architectures. These architectures, typically model the given real world data as a discrete set of entities, where each entity has both spatial attributes (e.g. geo-spatial location) and a temporal signal such as a series of sensor readings. Such a model can be applied to many real-world problems including weather forecasting [13], telecommunications network analysis [20] and many others [26]. In this work, we focus on the Urban Traffic Forecasting task, which is one of the most popular applications of this kind of data modeling.

In the canonical setting, data is collected using multiple sensors deployed at different locations in a particular area, each one measuring a time-varying signal such as the speed at a road segment. Given a series of history measurements of the sensors, the task is to forecast the future measurements up to some prediction horizon. The data can

<footnote>Y. Lubarsky and A. Gaissinski—Equal contribution.</footnote>
be structured as a graph where nodes represent the sensors, each carrying a temporal signal, and edges reflect the distance measure between them.

The top scoring solutions for this kind of setup employ deep neural networks that process the data via a series of computational blocks. Some of these blocks process the data along the temporal dimension, e.g. temporal convolutions or LSTM blocks [7], while others work along the spatial (i.e. graph) dimension, such as GCN [12] or self-attention layers [21]. These components are typically interleaved to build dependencies between spatial and temporal processing steps.

However, such an approach often comes at a high computational cost. As an example, the complexity of a single block of the popular Graph WaveNet model [22] is $O(TN^2D)$ where $T$ is the time series length, $N$ is the number of nodes and $D$ is the internal representation size. Such a model may be computationally too heavy to operate in a real-life production system where there is often a shortage in resources such as time, power, compute and memory. In addition, when the time-series are very long, training the models becomes increasingly demanding and in the extreme case, may be unfeasible due to out of memory condition.

In this work, we propose a method to dramatically improve the scalability and running times of existing traffic forecasting solutions without sacrificing the accuracy. It involves applying two architecture modifications to a given spatio-temporal model. First, a stack of gated convolutions is employed prior to spatial processing, reducing the size of the temporal dimension, while also functioning as a feature extractor for the subsequent layers. Second, sparse graph pooling blocks are added in the spatial processing layers to progressively reduce the size of the graph and facilitate better information propagation between far away nodes. To the best of our knowledge this is the first time a sparse and learnable pooling operation is used in the spatio-temporal modeling domain.

To validate our method, we use the Graph WaveNet architecture [22], a SOTA spatio-temporal architecture as shown in [11]. Using the proposed method, we develop efficient variants GWNET-TSQ and GWNET-SPOOL-TSQ. Our experiments show that these variants yield substantial x5-x27 speedups, while keeping or surpassing the state-of-the-art results. Furthermore, the speedup advantage is increased with the length of the input sequences.

The contributions of this paper can be summarized as follows:

- We propose a general method for significantly improving the runtime efficiency of spatio-temporal GNN models. It consists of preprocessing the input using a temporal pooling mechanism to reduce the temporal dimension into a fixed size representation and a novel application of learnable sparse graph pooling operators within the spatial blocks of the model, to reduce the runtime of the model and improve data flow between far away nodes.

- Using our method, we develop GWNET-TSQ and GWNET-SPOOL-TSQ architectures, a highly efficient solution for urban traffic forecasting. The proposed models are comparable or surpassing SOTA in the standard traffic forecasting benchmarks, while achieving a significant x5 speedup in training and inference.

- We evaluate the scalability of our proposed architectures using a long time-series forecasting task of 6-h and 24-h and observe dramatic speedups of up to x13 and x27 respectively, without any accuracy degradation compared to the top solutions.
<figure><img src="image_3.png" /><figcaption>**Fig. 1.** The proposed GWNET-SPOOL-TSQ architecture. The temporal stack encodes the temporal signal for each node into a fixed size representation, removing the dependency of spatial processing on the length of time series. The subsequent spatial processing block employs sparse learnable U-Net pooling blocks to progressively reduce the graph size, improving the runtime and the data flow between far away nodes.</figcaption></figure>

# 2 Related Work

## 2.1 Traffic Forecasting

The majority of recent Traffic Forecasting solutions, can be divided to two groups. 1) The spatio-temporal data is represented as a graph, where each node stores temporal data, and the nodes are connected by weighted edges, i.e. the distance between the sensors. 2) The a priori spatial correlation is neglected, and the data is treated as a multivariate prediction task.

Early GNN based works, such as [24], used a GRU [3] based architecture. In these works, the spatial graph features were extracted by a GNN and fed to a GRU unit to learn the temporal mechanics, which also produced the future predictions. These architectures are generally slow and fail on long-term forecasting horizons.

Modern Traffic Forecasting architectures use an interleaving processing of spatial GNNs and temporal convolutions, making the whole network fully convolutional and fast to train. In addition, all the predicted timesteps are produced at once by a single prediction block. A good representative of such architecture is [22]. This architecture uses a gated convolution for temporal feature extraction iterleaved with a slightly modified GCN [12] for message passing along the spatial dimension.

Some recent works neglect the underlying graph data and substitute it with an attention block, claiming that the correlation between nodes is dynamic and should be learned in an end-to-end manner. One of such works is GMAN [25]. In this work, the convolutional and GCN layers are substituted with attention blocks. There is a spatial attention block that operates over the different nodes, and temporal attention block working separately over the temporal dimension of each node. The blocks are then fused using weighted sum, to produce a single spatio-temporal representation of the data.

FC-GAGA [18] neglect the spatial information, and approached it as a multivariate prediction challenge. A learnable node embedding is used to measure the correlation between all node pairs, then each node is concatenated with a weighted representation of all other nodes, and this representation goes through fully connected (FC) layers to
produce the forecasting. The algorithm is fast, as it only uses FC layers, but it doesn't scale well as its complexity is $O(TN^2)$.

An extensive survey carried out by [11] covers many of the recent works done in the urban traffic prediction field. In addition, they provide a framework for standardized benchmarking of these solutions.

## 2.2 Graph Pooling

Graph pooling is an essential component of many GNN architectures. Similarly to pooling operator in convolutional neural networks (CNNs) [1, 19], its role is to reduce the number of nodes in the graph, while preserving the graph representation. Graph data is missing the grid-like structure of images, making the pooling operation an open problem with a lot of research groups proposing novel pooling algorithms.

Early GNN works used simple pooling techniques as a post-processing step for obtaining a graph representation [4, 8, 15]. The GNN would extract a feature representation of each node in the graph, and the final representation layer would globally pool these features into a single representation vector using mean or max pooling. Then this vector would be used by the classification layer for the final prediction.

One of the first works to utilize a learnable graph pooling module for hierarchical coarsening of the GNN architectures was DiffPool [26]. In this work, the pooling decision was learned using an additional GNN, that would generate a dense assignment matrix $S$, effectively making a soft clustering of the nodes.

A later work, Graph U-Nets [6], adopted the U-Net style architecture [19] to the GNN domain, introducing a learnable sparse pooling and unpooling architecture. In this work, all nodes are ranked with an importance score, the top-K important nodes are pooled to the next layer and the rest are dropped. In the unpooling stage, features from higher hierarchy nodes are combined with their previous layer representations and then a GNN propagates the data to the rest.

Since then, many pooling approaches have been introduced as shown in [9, 16, 17]. In particular, [9] introduced a taxonomy of different pooling operators.

# 3 The Proposed Efficient Spatio-Temporal Net

## 3.1 Preliminaries

The data is represented in the form of a graph $G = (V, E, X)$ where $V$ is the set of nodes and $E$ is the set of edges. We denote $N = |V|$, the number of nodes. An adjacency matrix derived from the graph is denoted by $A \in \mathbb{R}^{N \times N}$, such that $A_{i,j} = 1$ if the edge $(v_i, v_j) \in E$. At each time step $t$ the graph node features are represented by $X^{(t)} \in \mathbb{R}^{N \times D}$. Given a graph $G$, time $t$, history length $H$ and prediction horizon $P$, the forecasting task is to learn a function which maps $X^{(t-H:t)}$ to $X^{(t+1:t+P)}$.

## 3.2 Temporal Pooling

Recently proposed spatio-temporal architectures process the temporal dependency with interleaved CNN-GNN operations, or using recurrent networks, such as GRU or LSTM,
<figure><img src="image_4.png" /><figcaption>**Fig. 2.** a) Temporal pooling module. It consists of *B* consecutive blocks. Each block has a total of *L* convolutions, out of which *L* − 1 are gated convolutions followed by a single pooling convolution with stride 2 and a temporary channel expansion. The temporal dimension is halved after each block. The final linear layer is used to reduce the size to exactly one timestep. b) Gated Convolution block. The block has two identical layers, with different activation functions. One serving as a multiplicative gate to control the output of the other. The gate can be modeled as $h = tanh(\Theta_a * X + \beta_a) \odot \sigma(\Theta_b * X + \beta_b)$ where *X* is the input, $\Theta_a, \Theta_b, \beta_a, \beta_b$ are learnable weights and biases of the convolutions and $\sigma$ the sigmoid activation.</figcaption></figure>

over the graph representation of the data. Contrary to this paradigm, we use a Temporal Pooling module to extract the temporal data before applying the spatial processing. The Temporal Pooling module compresses the information stored in the temporal dimension of a given signal into a fixed vector representation. This design choice makes the spatial processing much more efficient as it needs to process only a single feature vector per node, decreasing the time complexity and allowing an efficient modeling of very long input sequences.

We implement the Temporal Pooling module using the following guiding principles: 1) all time steps should be accounted for, 2) low computational overhead and 3) powerful feature extraction. A detailed diagram of our temporal pooling stack is shown in Fig. 2. We employ a stack of gated convolutions as our Temporal Pooling module, since it fits well the above principals. It offers faster run-times compared to the recurrent architectures, while the gating mechanism is useful for overcoming the vanishing gradients problem in long sequences [5, 10, 22]. We use solely convolutional layers to keep the computational overhead linear in the input size. To reduce the size of the temporal dimension, a periodic stride of two is used to halve the input. One final linear layer is placed at the end of the stack to make the temporal dimension size exactly one.

## 3.3 Graph Pooling

The graph pooling block in our architecture serves two purposes: 1) as each consequent GNN layer receives less nodes, it also has less calculations, making the network more
<figure><img src="image_5.png" /><figcaption>**Fig. 3.** (a) Graph U-Net architecture. The blue block represents the original Spatio-Temporal layer, enhanced with the node pooling mechanism. Each yellow block fuses the unpooled data, which is sparse, with the dense data coming from the skip connection. Lastly, each green block does a message passing step between all received nodes and un-pools the data before fusing it again with the higher block. (b) Un-Pooling architecture. In the UnPool block, the green features from the sparse layer are padded with zero features to accommodate for the pooled nodes. Then the unpooled features are fused with the nodes representation from the dense layer. In our architecture, the fusion step is performed by concatenation. (Color figure online)</figcaption></figure>

efficient, 2) the pooling and unpooling operators improve the data propagation between distant nodes.

In this work, we propose to use the top-K pooling [6] operator. It allows end-to-end training, and alleviates the unpooling operator to produce dense predictions for all the nodes. The graph pooling architecture is shown in Fig. 3(a).

During the feature extraction phase, each consequent spatial processing block receives only the top-K chosen nodes from the previous layer. The importance score is derived by projecting node features on a learnable vector. In the proposed architecture, each consequent block has progressively less nodes which need to be accounted for by the un-pooling operation.

The unpooling stage is shown in Fig. 3(b). In the unpooling stage the low resolution features are upsampled with zeros, similarly to [1], and concatenated to the features extracted from the higher resolution block. After concatenation, the data between nodes is propagated using another GNN layer.

## 3.4 Architecture

**Graph WaveNet.** To demonstrate the effectiveness of the proposed approach, we use the Graph WaveNet model [22]. Based on the survey [11], Graph WaveNet consistently shows top results on the traffic forecasting task. The base architecture consists of a series of spatio-temporal blocks applied sequentially. The temporal component used is gated convolutions with dilation. The spatial component is a graph diffusion convolution layer applied on normalized forward an backward adjacency matrix as well as
on a self-adaptive adjacency matrix. The model also has residual and skip connections for better gradient propagation. Due to use of non-sparse adjacency matrix in the GCN component, the operating cost of each layer is $O(TN|E|D)$ or $O(TN^2D)$ if using the non-sparse self-adaptive matrix.

Using the components proposed in previous sections, we develop three variants to the base architecture.

- GWNET-TSQ: Temporal pooling stack is used after the first convolution layer to reduce the temporal dimension size to one. The gated convolutions are removed from the spatio-temporal blocks as they are no longer needed. This improves the runtime to $O(N^2D)$.

- GWNET-SPOOL: A sparse graph pooling mechanism is added in each layer, reducing the size of the graph by a constant factor. This results in geometric decrease in the number of nodes processed in each following GCN layer.

- GWNET-SPOOL-TSQ: Combines the two mentioned modifications - both temporal stack and pooling modules are added, and the temporal gated convolutions are removed from the spatio-temporal layers. This variant is shown in Fig. 1.

**GMAN.** GMAN [25] is another popular model for the spatio-temporal traffic data. Here, the model relies on the powerful multi-head attention mechanism [21] for feature extraction. The model consists of an encoder processing the history signal and a decoder producing the predictions, connected by transform attention. Both encoder and decoder are stacks of fused spatio-temporal blocks where one self-attention layer is applied among the nodes and another among the time steps of each node. This is a canonical example of applying self-attention in an interleaved manner to a signal that has more than one dimension of interest.

We modify the GMAN architecture with the proposed Temporal pooling stack and develop the GMAN-TSQ variant. The temporal stack is applied before the encoder, eliminating the self-attention processing done along the temporal dimension. We avoid modifying the decoder since it relies on temporal attention along the predicted signal. Although, it is possible to separate temporal and spatial processing there as well, it may seem as too excessive a change from the original architecture. At any rate, replacing even just the encoder with a temporal pooling stack results in dramatic reduction in operating cost as shown in the next section.

# 4 Experiments

## 4.1 Datasets

We verify our approach on METR-LA and PEMS-BAY, two of the most widely used public traffic network datasets released by [14]. METR-LA contains four months of traffic speed recordings collected from 2012/3/1 to 2012/6/30 by 207 sensors on the highways of Los Angeles County. PEMS-BAY contains six months of traffic speed recordings, collected from 2017/1/1 to 2017/5/31 by 325 sensors in the Bay area. As mentioned in [11] these two datasets have been used in quite a lot of studies and can
<table><caption>**Table 1.** Performance comparison of our and baseline models on METR-LA and PEMS-BAY datasets with one hour history and prediction horizon and a five minute resolution.</caption><thead><tr><th rowspan="2">Model</th><th colspan="3">15 min./3 steps</th><th colspan="3">30 min./6 steps</th><th colspan="3">60 min./12 steps</th><th rowspan="2">MACs [G]</th><th rowspan="2">MAC compr</th></tr><tr><th>MAE</th><th>RMSE</th><th>MAPE</th><th>MAE</th><th>RMSE</th><th>MAPE</th><th>MAE</th><th>RMSE</th><th>MAPE</th></tr></thead><tbody><tr><td colspan="12"><strong>METR-LA</strong></td></tr><tr><td>Graph WaveNet</td><td><strong>2.73</strong></td><td>5.25</td><td><strong>7.04</strong></td><td><strong>3.12</strong></td><td>6.30</td><td><strong>8.48</strong></td><td>3.60</td><td>7.45</td><td><strong>10.15</strong></td><td>0.67</td><td>—</td></tr><tr><td>GWNET-TSQ</td><td>2.74</td><td><strong>5.22</strong></td><td>7.12</td><td>3.13</td><td><strong>6.21</strong></td><td>8.61</td><td>3.60</td><td><strong>7.30</strong></td><td>10.32</td><td><strong>0.14</strong></td><td>0.21</td></tr><tr><td>GWNET-SPOOL</td><td>2.77</td><td>5.30</td><td>7.16</td><td>3.16</td><td>6.32</td><td>8.65</td><td>3.63</td><td>7.43</td><td>10.40</td><td>0.45</td><td>0.68</td></tr><tr><td>GWNET-SPOOL-TSQ</td><td>2.75</td><td><strong>5.22</strong></td><td>7.13</td><td>3.13</td><td>6.25</td><td>8.49</td><td>3.60</td><td>7.38</td><td>10.19</td><td>0.23</td><td>0.34</td></tr><tr><td>GMAN</td><td>3.52</td><td>7.17</td><td>10.48</td><td>3.59</td><td>7.34</td><td>10.70</td><td>4.12</td><td>8.39</td><td>12.24</td><td>0.61</td><td>—</td></tr><tr><td>GMAN-TSQ</td><td>3.17</td><td>6.13</td><td>8.63</td><td>3.30</td><td>6.49</td><td>9.12</td><td>3.98</td><td>8.04</td><td>11.31</td><td>0.35</td><td>0.57</td></tr><tr><td>GMAN-TSQ-4H</td><td>3.12</td><td>6.06</td><td>8.70</td><td>3.24</td><td>6.40</td><td>9.18</td><td>3.91</td><td>7.91</td><td>11.32</td><td><strong>0.14</strong></td><td>0.23</td></tr><tr><td>FC-GAGA</td><td>2.81</td><td>5.48</td><td>7.43</td><td>3.19</td><td>6.49</td><td>8.81</td><td>3.65</td><td>7.59</td><td>10.53</td><td>0.88</td><td>—</td></tr><tr><td>MTGNN</td><td>2.82</td><td>5.57</td><td>7.56</td><td>3.18</td><td>6.53</td><td>8.94</td><td><strong>3.58</strong></td><td>7.45</td><td>10.32</td><td>—</td><td>—</td></tr><tr><td>AGCRN</td><td>2.88</td><td>5.59</td><td>7.63</td><td>3.25</td><td>6.56</td><td>9.00</td><td>3.65</td><td>7.51</td><td>10.31</td><td>—</td><td>—</td></tr><tr><td colspan="12"><strong>PEMS-BAY</strong></td></tr><tr><td>Graph WaveNet</td><td><strong>1.31</strong></td><td>2.76</td><td><strong>2.76</strong></td><td>1.64</td><td>3.74</td><td><strong>3.71</strong></td><td>1.95</td><td>4.53</td><td>4.59</td><td>1.43</td><td>—</td></tr><tr><td>GWNET-TSQ</td><td><strong>1.31</strong></td><td><strong>2.74</strong></td><td>2.77</td><td><strong>1.63</strong></td><td><strong>3.67</strong></td><td><strong>3.71</strong></td><td><strong>1.94</strong></td><td><strong>4.44</strong></td><td><strong>4.57</strong></td><td>0.29</td><td>0.20</td></tr><tr><td>GWNET-SPOOL</td><td>1.34</td><td>2.78</td><td>2.84</td><td>1.68</td><td>3.77</td><td>3.84</td><td>1.99</td><td>4.54</td><td>4.74</td><td>1.09</td><td>0.76</td></tr><tr><td>GWNET-SPOOL-TSQ</td><td>1.34</td><td>2.78</td><td>2.88</td><td>1.67</td><td>3.75</td><td>3.87</td><td>1.99</td><td>4.54</td><td>4.81</td><td>0.44</td><td>0.31</td></tr><tr><td>GMAN</td><td>1.76</td><td>4.04</td><td>4.17</td><td>1.79</td><td>4.11</td><td>4.22</td><td>2.24</td><td>5.22</td><td>5.29</td><td>1.09</td><td>—</td></tr><tr><td>GMAN-TSQ</td><td>1.59</td><td>3.54</td><td>3.71</td><td>1.66</td><td>3.71</td><td>3.86</td><td>2.16</td><td>4.97</td><td>5.08</td><td>0.61</td><td>0.56</td></tr><tr><td>GMAN-TSQ-4H</td><td>1.54</td><td>3.38</td><td>3.47</td><td>1.66</td><td>3.70</td><td>3.76</td><td>2.20</td><td>5.11</td><td>5.10</td><td><strong>0.26</strong></td><td>0.23</td></tr><tr><td>FC-GAGA</td><td>1.37</td><td>2.89</td><td>2.90</td><td>1.70</td><td>3.84</td><td>3.86</td><td>1.99</td><td>4.56</td><td>4.71</td><td>2.09</td><td>—</td></tr><tr><td>MTGNN</td><td>1.33</td><td>2.85</td><td>2.84</td><td>1.66</td><td>3.80</td><td>3.77</td><td>1.95</td><td>4.49</td><td>4.59</td><td>—</td><td>—</td></tr><tr><td>AGCRN</td><td>1.35</td><td>2.86</td><td>2.94</td><td>1.67</td><td>3.82</td><td>3.84</td><td>1.96</td><td>4.57</td><td>4.69</td><td>—</td><td>—</td></tr></tbody></table>

be considered a standard benchmark for the traffic forecasting task. The sensor readings in both datasets are aggregated into 5-minute windows and the data is split in chronological order with 70% for training, 10% for validation and 20% for testing. As in other studies using these datasets, we use one hour (12 time steps) history length and one hour prediction in our forecasting task. To analyze model performances at longer time sequences, we introduce METR-LA-6H and METR-LA-24H configurations of the METR-LA data, having six hour (72 time steps) and 24 h (288 time steps) history and prediction horizons respectively.

## 4.2 Experimental Settings

In our experiments, we use Graph WaveNet [22], GMAN [25], FC-GAGA [18], MTGNN [23] and AGCRN [2] models as well as the variants presented in Sect. 3. For the one-hour (12 step) prediction task, we use the configuration found in the corresponding original publications. The temporal pooling stacks in GWNET-TSQ and GWNET-SPOOL-TSQ use *B* = 2 blocks in METR-LA and *B* = 4 in PEMS-BAY, with *L* = 2 layers and 32 channels. The spatial pooling in GWNET-SPOOL and GWNET-SPOOL-TSQ uses a pooling factor of 0.7.
For the longer time sequence experiments, in METR-LA-6H (72 step) we modify Graph WaveNet structure in order to cover the increased time dimension of 72 steps. The number of spatio-temporal blocks is increased to 5 and the dilation in gated convolutions is set to 4. In the temporal pooling stacks we use $B = 5$ blocks of $L = 2$ layers with 32 channels. We keep GMAN architecture unchanged. In METR-LA-24 (288 step) experiments, the Graph WaveNet structure uses 5 blocks of 3 spatio-temporal layers with dilation 4. The temporal pooling stack uses $B = 6$ blocks with $L = 2$ layers. All of the reported results are averaged over 3 runs with different seeds.

To compare the runtime efficiency of different models we compute and report the number of forward pass Multiply-Accumulate Operations (MACs) [27]. This serves as a common comparison metric, agnostic to the specific type of CPU or GPU used to run the model. For our proposed models, we also report the ratio between MACs of our model and MACs of the original.

**Implementation.** Our changes are implemented using the framework published by [11] for a standardized comparison of traffic forecasting methods. We introduce two main changes to the standardized setup of the framework. First, we use masked MAE loss instead of plain MAE (as is the framework default), in order to prevent the model learning from timesteps with missing data labels. Second, for the Graph WaveNet model, we reintroduce dropout of 0.3 as was done in the original publication [22]. Both of these changes are consistent with the original publications of the models we use for comparison. MAC estimates are computed using the THOP library [27], which we extend to support the benchmarked models.

## 4.3 Experimental Results

**METR-LA.** Table 1 summarizes the results. Notably, all the proposed Graph WaveNet variants perform on par with the baseline model, in particular for the far horizon of 12 time steps, often considered an indication for good modeling of long time-series dependencies. Furthermore, all the proposed variants use substantially less MACs to operate. GWNET-SPOOL and GWNET-SPOOL-TSQ reduce MACs used to 68% and 34% respectively compared to the original model. GWNET-TSQ uses only 21% of the original - a x4.76 speedup. GMAN performance is consistently worse than that of Graph WaveNet. Notably, GMAN-TSQ variants actually improve the prediction accuracy, while also using less computations. GMAN-TSQ variant reduces MACs to 57% of the original. We also consider reducing the number of attention heads in the GMAN-TSQ model to 4. Such a model GMAN-TSQ-4H further reduces the MACs to 23% of the original. In our experiments, FC-GAGA model did not reach better results than the baseline Graph WaveNet model while the MACs used indicate a slightly higher operational cost. It is worth noting that we were able to reproduce the results published in [18] when using their particular data setup. However, we couldn't reproduce the claimed results with a standard data pipeline.

**PEMS-BAY.** Table 1 summarizes the results. Here, the proposed variants show similar behavior to the METR-LA prediction task. GWNET-TSQ achieves the lowest errors
while running 5 times faster than Graph WaveNet. GWNET-SPOOL and GWNET-SPOOL-TSQ reduce MACs used to 76% and 31% respectively and get slightly higher errors. GMAN-TSQ and GMAN-TSQ-4H variants result in lower errors than the original GMAN architecture using only 56% and 23% of the original MACs.

<table><caption>**Table 2.** Performance comparison on METR-LA with a 6-h history and 6-h prediction horizons at a five minute resolution.</caption><thead><tr><th rowspan="2">Model</th><th colspan="3">2 h/24 steps</th><th colspan="3">4 h/48 steps</th><th colspan="3">6 h/72 steps</th><th rowspan="2">MACs</th><th rowspan="2">MAC<br>compr</th></tr><tr><th>MAE</th><th>RMSE</th><th>MAPE</th><th>MAE</th><th>RMSE</th><th>MAPE</th><th>MAE</th><th>RMSE</th><th>MAPE</th></tr></thead><tbody><tr><td>Graph WaveNet</td><td>4.17</td><td>8.73</td><td>12.75</td><td>4.31</td><td>9.00</td><td>13.64</td><td>4.49</td><td>9.29</td><td>14.49</td><td>4.90</td><td>-</td></tr><tr><td>FC-GAGA</td><td>4.33</td><td>8.56</td><td>13.35</td><td>4.49</td><td>8.89</td><td>14.52</td><td>4.72</td><td>9.17</td><td>15.34</td><td>4.88</td><td>-</td></tr><tr><td>GWNET-TSQ</td><td>4.11</td><td>8.40</td><td>12.51</td><td>4.29</td><td>8.80</td><td>13.56</td><td>4.50</td><td>9.17</td><td>14.74</td><td><strong>0.38</strong></td><td>0.08</td></tr><tr><td>GWNET-SPOOL</td><td>4.10</td><td><strong>8.30</strong></td><td>12.48</td><td>4.30</td><td><strong>8.67</strong></td><td><strong>13.35</strong></td><td>4.49</td><td><strong>9.00</strong></td><td>14.45</td><td>1.70</td><td>0.35</td></tr><tr><td>GWNET-SPOOL-TSQ</td><td><strong>4.09</strong></td><td>8.45</td><td><strong>12.38</strong></td><td><strong>4.26</strong></td><td>8.86</td><td>13.60</td><td><strong>4.40</strong></td><td>9.10</td><td><strong>14.09</strong></td><td>0.44</td><td>0.09</td></tr><tr><td>GMAN</td><td>5.24</td><td>10.71</td><td>18.55</td><td>5.19</td><td>10.60</td><td>18.29</td><td>6.47</td><td>12.55</td><td>22.17</td><td>4.05</td><td>-</td></tr><tr><td>GMAN-TSQ</td><td>5.17</td><td>10.52</td><td>18.34</td><td>5.16</td><td>10.42</td><td>18.17</td><td>6.45</td><td>12.37</td><td>22.15</td><td>2.21</td><td>0.55</td></tr><tr><td>GMAN-TSQ-4H</td><td>5.15</td><td>10.49</td><td>18.29</td><td>5.19</td><td>10.47</td><td>18.18</td><td>6.57</td><td>12.55</td><td>22.39</td><td>0.91</td><td>0.22</td></tr></tbody></table>

<table><caption>**Table 3.** Performance comparison on METR-LA with a 24-h history and 24-h prediction horizons at a five minute resolution.</caption><thead><tr><th rowspan="2">Model</th><th colspan="3">6 h/72 steps</th><th colspan="3">12 h/144 steps</th><th colspan="3">24 h/288 steps</th><th rowspan="2">MACs</th><th rowspan="2">MAC<br>compr</th></tr><tr><th>MAE</th><th>RMSE</th><th>MAPE</th><th>MAE</th><th>RMSE</th><th>MAPE</th><th>MAE</th><th>RMSE</th><th>MAPE</th></tr></thead><tbody><tr><td>Graph WaveNet</td><td>4.32</td><td>8.86</td><td>13.60</td><td>4.44</td><td>9.04</td><td>13.82</td><td>4.59</td><td>9.32</td><td>15.35</td><td>32.18</td><td>-</td></tr><tr><td>GWNET-TSQ</td><td>4.31</td><td>8.66</td><td>14.05</td><td>4.39</td><td>8.77</td><td>14.01</td><td>4.57</td><td>9.20</td><td>15.49</td><td><strong>1.18</strong></td><td>0.04</td></tr><tr><td>GWNET-SPOOL-TSQ</td><td><strong>4.21</strong></td><td><strong>8.36</strong></td><td><strong>13.26</strong></td><td><strong>4.30</strong></td><td><strong>8.55</strong></td><td><strong>13.31</strong></td><td><strong>4.44</strong></td><td><strong>8.78</strong></td><td><strong>14.18</strong></td><td>1.81</td><td>0.06</td></tr></tbody></table>

**METR-LA-6H.** Table 2 summarizes the results. Here the time dimension is 6 times larger than in 1-h prediction task. The baseline models Graph WaveNet, FC-GAGA and GMAN run at a roughly similar computational cost of 4-5 GMACs. In comparison, the proposed variants run at a dramatically lower operational cost. GWNET-SPOOL reduces MACs used to 35% of the original, while GWNET-SPOOL-TSQ and GWNET-TSQ reduce to 9% and 8% with speedups of x11.1 and x12.9. Error-wise, we note that the proposed variants either improve on the baseline or perform on par. In particular, for the longest horizon of 72 steps our variants give the lowest errors compared to the baseline. Notably, while both GWNET-TSQ and GWNET-SPOOL reduce the error compared to the original models, the best variant is GWNET-SPOOL-TSQ which combines the temporal pooling stack and the spatial pooling components. GMAN-TSQ gives lower errors than the GMAN baseline while using only 55% of the original MACs. The GMAN-TSQ-4H variant performs on par or better than the GMAN baseline and 2-h and 4-h horizons and has a slightly higher errors in the 6-h horizon, while using only 22% MACs compared to the original.
**METR-LA-24H.** Table 3 summarizes the results. Here, we again note the dramatic difference in MACs used in our proposed variants compared to the baseline, paired with similar or better accuracies. GWNET-TSQ runs at only 4% MACs compared to original - a x27.3 speedup - with slightly lower MAE scores and slightly higher MAPE (less than 0.5%). GWNET-SPOOL-TSQ runs at 6% MACs compared to baseline - a x17.8 speedup - and yields lower errors in all prediction horizons - a reduction of 0.15 in MAE and 1.17 MAPE in the longest 288-step score. This again suggests that the combination of Temporal Pooling stack with the spatial pooling produces the most accurate model for the longer time horizons.

# 5 Conclusion

In this paper, we proposed a new framework to dramatically improve the runtime efficiency of the traffic forecasting solutions. We deploy a new temporal pooling block prior to spatial processing, which effectively compresses the time dimension prior to applying the message passing modules. We also embed graph pooling blocks to progressively reduce the size of the graph and to facilitate better information propagation between distant nodes. Our approach effectively reduces the linear dependency of the spatio-temporal network on the time dimension. Therefore, the achieved speedup scales with the size of the input time-series, and allows effective modeling of long-range dependencies. Our experiments demonstrate that the proposed method yields substantial speedups in both training and inference times, ranging from x5 speedup in the standard 1-h forecasting task and up to x27 for longer time series forecast of 24-h, while keeping or surpassing the accuracy of state-of-the-art methods.

# References

1. Badrinarayanan, V., Kendall, A., Cipolla, R.: Segnet: a deep convolutional encoder-decoder architecture for image segmentation. IEEE Trans. Pattern Anal. Mach. Intell. 39(12), 2481-2495 (2017)

2. Bai, L., Yao, L., Li, C., Wang, X., Wang, C.: Adaptive graph convolutional recurrent network for traffic forecasting. Adv. Neural. Inf. Process. Syst. 33, 17804-17815 (2020)

3. Cho, K., Van Merriënboer, B., Bahdanau, D., Bengio, Y.: On the properties of neural machine translation: encoder-decoder approaches. arXiv preprint arXiv:1409.1259 (2014)

4. Dai, H., Dai, B., Song, L.: Discriminative embeddings of latent variable models for structured data. In: International Conference on Machine Learning, pp. 2702-2711. PMLR (2016)

5. Dauphin, Y.N., Fan, A., Auli, M., Grangier, D.: Language modeling with gated convolutional networks. In: International Conference on Machine Learning, pp. 933-941. PMLR (2017)

6. Gao, H., Ji, S.: Graph u-nets. In: International Conference on Machine Learning, pp. 2083-2092. PMLR (2019)

7. Gers, F.A., Schmidhuber, J., Cummins, F.: Learning to forget: continual prediction with LSTM. Neural Comput. 12(10), 2451-2471 (2000)

8. Gilmer, J., Schoenholz, S.S., Riley, P.F., Vinyals, O., Dahl, G.E.: Neural message passing for quantum chemistry. In: International Conference on Machine Learning, pp. 1263-1272. PMLR (2017)

9. Grattarola, D., Zambon, D., Bianchi, F.M., Alippi, C.: Understanding pooling in graph neural networks. IEEE Trans. Neural Netw. Learn. Syst. (2022)
10. Heinrich, K., Zschech, P., Janiesch, C., Bonin, M.: Process data properties matter: introducing gated convolutional neural networks (GCNN) and key-value-predict attention networks (KVP) for next event prediction with deep learning. Decis. Support Syst. 143, 113494 (2021)

11. Jiang, R., et al.: DL-Traff: survey and benchmark of deep learning models for urban traffic prediction. In: Proceedings of the 30th ACM International Conference on Information & Knowledge Management, pp. 4515–4525 (2021)

12. Kipf, T.N., Welling, M.: Semi-supervised classification with graph convolutional networks. arXiv preprint arXiv:1609.02907 (2016)

13. Lam, R., et al.: Graphcast: learning skillful medium-range global weather forecasting. arXiv preprint arXiv:2212.12794 (2022)

14. Li, Y., Yu, R., Shahabi, C., Liu, Y.: Diffusion convolutional recurrent neural network: data-driven traffic forecasting. In: International Conference on Learning Representations (ICLR 2018) (2018)

15. Li, Y., Tarlow, D., Brockschmidt, M., Zemel, R.: Gated graph sequence neural networks. arXiv preprint arXiv:1511.05493 (2015)

16. Liu, C., et al.: Graph pooling for graph neural networks: progress, challenges, and opportunities. arXiv preprint arXiv:2204.07321 (2022)

17. Mesquita, D., Souza, A., Kaski, S.: Rethinking pooling in graph neural networks. Adv. Neural. Inf. Process. Syst. 33, 2220–2231 (2020)

18. Oreshkin, B.N., Amini, A., Coyle, L., Coates, M.: FC-GAGA: fully connected gated graph architecture for spatio-temporal traffic forecasting. In: Proceedings of the AAAI Conference on Artificial Intelligence, vol. 35, pp. 9233–9241 (2021)

19. Ronneberger, O., Fischer, P., Brox, T.: U-Net: convolutional networks for biomedical image segmentation. In: Navab, N., Hornegger, J., Wells, W.M., Frangi, A.F. (eds.) MICCAI 2015. LNCS, vol. 9351, pp. 234–241. Springer, Cham (2015). https://doi.org/10.1007/978-3-319-24574-4_28

20. Rusek, K., Suárez-Varela, J., Almasan, P., Barlet-Ros, P., Cabellos-Aparicio, A.: Routenet: leveraging graph neural networks for network modeling and optimization in SDN. IEEE J. Sel. Areas Commun. 38(10), 2260–2270 (2020)

21. Vaswani, A., et al.: Attention is all you need. In: Advances in Neural Information Processing Systems, vol. 30 (2017)

22. Wu, Z., Pan, S., Long, G., Jiang, J., Zhang, C.: Graph wavenet for deep spatial-temporal graph modeling. In: The 28th International Joint Conference on Artificial Intelligence (IJCAI). International Joint Conferences on Artificial Intelligence Organization (2019)

23. Wu, Z., Pan, S., Long, G., Jiang, J., Chang, X., Zhang, C.: Connecting the dots: multivariate time series forecasting with graph neural networks. In: Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 753–763 (2020)

24. Zhao, L., et al.: T-GCN: a temporal graph convolutional network for traffic prediction. IEEE Trans. Intell. Transp. Syst. 21(9), 3848–3858 (2019)

25. Zheng, C., Fan, X., Wang, C., Qi, J.: GMAN: a graph multi-attention network for traffic prediction. In: Proceedings of the AAAI Conference on Artificial Intelligence, vol. 34, pp. 1234–1241 (2020)

26. Zhou, J., et al.: Graph neural networks: a review of methods and applications. AI Open 1, 57–81 (2020)

27. Zhu, L.: THOP: PyTorch-OpCounter (2019). https://github.com/Lyken17/pytorch-OpCounter