SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Amortized specialization Static node features offer the opportunity to design and obtain node identification mechanisms across different time windows to tailor predictions to a specific node. However, in most settings, node features are either unavailable or insufficient to characterize the node dynamics. A possible solution consists of resorting to learnable node embeddings, i.e., a table of learnable parameters  $\Theta = V \in \mathbb{R}^{N \times d_v}$ . Rather than interpret these learned representations as positional encodings, our proposal is to consider them as a way of amortizing the learning of node-level specialized models. More specifically, instead of learning a local model for each time series, embeddings fed into modules of a global STGNN and learned end-to-end with the forecasting architecture allow for specializing predictions by simply relying on gradient descent to find a suitable encoding.
The most straightforward option for feeding embeddings into the processing is to update the template model by changing the encoder and decoder as
$$
\boldsymbol {h} _ {t} ^ {i, 0} = \operatorname {E N C O D E R} \left(\boldsymbol {x} _ {t - 1} ^ {i}, \boldsymbol {u} _ {t - 1} ^ {i}, \boldsymbol {v} ^ {i}\right), \quad (1 8) \quad \hat {\boldsymbol {x}} _ {t: t + H} ^ {i} = \operatorname {D E C O D E R} \left(\boldsymbol {h} _ {t} ^ {i, L}, \boldsymbol {v} ^ {i}\right). \tag {19}
$$
which can be seen as amortized versions of the encoder and decoder in Eq. 16-17. The encoding scheme of Eq. 18 also facilitates the propagation of relevant information by identifying nodes, an aspect that can be particularly significant as message-passing operators – in particular isotropic ones – can act as low-pass filters that smooth out node-level features [24, 25].
Tab. 1 (light gray block) reports empirical results that show the effectiveness of embeddings in amortizing the learning of local components (S2), with a negligible increase in the number of trainable parameters w.r.t. the base global model. In particular, feeding embeddings to the encoder, instead of conditioning the decoding step only, results in markedly better performance, hinting at the impact of providing node identification ahead of MP (additional empirical results provided in Sec. 7).
5.1 Structuring the embedding space
The latent space in which embeddings are learned can be structured and regularized to gather benefits in terms of interpretability, transferability, and generality (S5). In fact, accommodating new embeddings can be problematic, as they must fit in a region of the embedding space where the trained model can operate, and, at the same time, capture the local effects at the new nodes. In this setting, proper regularization can provide guidance and positive inductive biases to transfer the learned model to different node sets. As an example, if domain knowledge suggests that neighboring nodes have similar dynamics, Laplacian regularization [26, 27] can be added to the loss. In the following, we propose two general strategies based respectively on variational inference and on a clustering loss to impose soft constraints on the latent space. As shown in Sec. 7, the resulting structured space additionally allows us to gather insights into the features encoded in the embeddings.
Variational regularization As a probabilistic approach to structuring the latent space, we propose to consider learned embeddings as parameters of an approximate posterior distribution - given the training data - on the vector used to condition the predictions. In practice, we model each node embedding as a sample from a multivariate Gaussian  $\pmb{v}^i \sim q^i(\pmb{v}^i|\mathcal{D}) = \mathcal{N}\left(\pmb{\mu}_i, \mathrm{diag}(\pmb{\sigma}_i^2)\right)$  where  $(\pmb{\mu}_i, \pmb{\sigma}_i)$  are the learnable (local) parameters. Each node-level distribution is fitted on the training data by considering a standard Gaussian prior and exploiting the reparametrization trick [28] to minimize under the sampling of  $t$
$$
\delta_ {t} \doteq \mathbb {E} _ {\boldsymbol {V} \sim Q} \left[ \left\| \widehat {\boldsymbol {X}} _ {t: t + H} - \boldsymbol {X} _ {t: t + H} \right\| _ {2} ^ {2} \right] + \beta D _ {\mathrm {K L}} (Q | P),
$$
where  $P = \mathcal{N}(\mathbf{0},\mathbb{I})$  is the prior,  $D_{\mathrm{KL}}$  the Kulback-Leibler divergence, and  $\beta$  controls the regularization strength. This regularization scheme results in a smooth latent space where it is easier to interpolate between representations, thus providing a principled way for accommodating different node embeddings.
Clustering regularization A different (and potentially complementary) approach to structuring the latent space is to incentivize node embeddings to form clusters and, consequently, to self-organize into different groups. We do so by introducing a regularization loss inspired by deep  $K$ -means algorithms [29]. In particular, besides the embedding table  $\mathbf{V} \in \mathbb{R}^{N \times d_v}$ , we equip the embedding module with a matrix  $\mathbf{C} \in \mathbb{R}^{K \times d_v}$  of  $K \ll N$  learnable centroids and a cluster assignment matrix  $\mathbf{S} \in \mathbb{R}^{N \times K}$  encoding scores associated to each node-cluster pair. We consider scores as logits of a categorical (Boltzmann) distribution and learn them by minimizing the regularization term
$$
\mathcal {L} _ {r e g} \doteq \mathbb {E} _ {\boldsymbol {M}} \left[ \| \boldsymbol {V} - \boldsymbol {M} \boldsymbol {C} \| _ {2} \right], \quad p (\boldsymbol {M} _ {i j} = 1) = \frac {e ^ {\boldsymbol {S} _ {i j} / \tau}}{\sum e ^ {\boldsymbol {S} _ {i k} / \tau}},
$$
where  $\tau$  is a hyperparameter. We minimize  $\mathcal{L}_{reg}$  - which corresponds to the embedding-to-centroid distance - jointly with the forecasting loss by relying on the Gumbel softmax trick [30]. Similarly to the variational inference approach, the clustering regularization gives structure to embedding space and allows for inspecting patterns in the learned local components (see Sec. 7).
Transferability of graph-based predictors Global models based on GNNs can make predictions for never-seen-before node sets, and handle graphs of different sizes and variable topology. In practice, this means that the graph-based predictors can easily handle new sensors being added to the network over time and be used for zero-shot transfer. Clearly, including in the forecasting architecture node-specific local components compromises these properties. Luckily, if local components are replaced by node embedding, adapting the specialized components is relatively cheap since the number of parameters to fit w.r.t. the new context is usually contained, and – eventually – both the graph topology and the structure of the embedding latent space can be exploited (S4). Experiments in Sec. 7 provide an in-depth empirical analysis of transferability within our framework and show that the discussed regularizations can be useful in this regard.
GNNs have been remarkably successful in modeling structured dynamical systems [31-33], temporal networks [34-36] and sequences of graphs [37, 38]. For what concerns time series processing, recurrent GNNs [5, 6] were among the first STGNNs being developed, followed by fully convolutional models [7, 39] and attention-based solutions [40-42]. Among the methods that focus on modeling node-specific dynamics, Bai et al. [11] use a factorization of the weight matrices in a recurrent STGNN to adapt the extracted representation to each node. Conversely, Chen et al. [43] use a model inspired by Wang et al. [44] consisting of a global GNN paired with a local model conditioned on the neighborhood of each node. Node embeddings have been mainly used in structure-learning modules to amortize the cost of learning the full adjacency matrix [39, 45, 12] and in attention-based approaches as positional encodings [46, 47, 41]. Shao et al. [48] observed how adding spatiotemporal identification mechanisms to the forecasting architecture can outperform several state-of-the-art STGNNs. Conversely, Yin et al. [49] used a cluster-based regularization to fine-tune an AGCRN-like model on different datasets. However, none of the previous works systematically addressed directly the problem of globality and locality in STGNNs, nor provided a comprehensive framework accounting for learnable node embeddings within different settings and architectures. Finally, besides STGNNs, there are several examples of hybrid global and local time series forecasting models. Wang et al. [44] propose an architecture where  $K$  global models extract dynamic global factors that are then weighted and integrated with probabilistic local models. Sen et al. [50] instead use a matrix factorization scheme paired with a temporal convolutional network [51] to learn a multivariate model then used to condition a second local predictor.
This section reports salient results of an extensive empirical analysis of global and local models and combinations thereof in spatiotemporal forecasting benchmarks and different problem settings;
complete results of this systematic analysis can be found in Appendix B. Besides the reference architectures, we consider the following baselines and popular state-of-the-art architectures.
RNN: a global univariate RNN sharing the same parameters across the time series.
FC-RNN: a multivariate RNN taking as input the time series as if they were a multivariate one.
LocalRNNs: local univariate RNNs with different sets of parameters for each time series.
DCRNN [6]: a recurrent T&S model with the Diffusion Convolutional operator.
AGCRN [11]: the T&S global-local Adaptive Graph Convolutional Recurrent Network.
GraphWaveNet: the deep T&S spatiotemporal convolutional network by Wu et al. [39].
We also consider a global-local RNN, in which we specialize the model by using node embeddings in the encoder. Note that among the methods selected from the literature only DCRNN can be considered fully global (see Sec. 6). Performance is measured in terms of mean absolute error (MAE).
Synthetic data We start by assessing the performance of hybrid global-local spatiotemporal models in a controlled environment, considering a variation of GPVAR [52], a synthetic dataset based on a polynomial graph filter [53], that we modify to include local effects. In particular, data are generated from the spatiotemporal process
$$
\boldsymbol {H} _ {t} = \sum_ {l = 1} ^ {L} \sum_ {q = 1} ^ {Q} \Theta_ {q, l} \boldsymbol {A} ^ {l - 1} \boldsymbol {X} _ {t - q},
$$
$$
\boldsymbol {X} _ {t + 1} = \boldsymbol {a} \odot \tanh  \left(\boldsymbol {H} _ {t}\right) + \boldsymbol {b} \odot \tanh  \left(\boldsymbol {X} _ {t - 1}\right) + \eta_ {t}, \tag {20}
$$
where  $\Theta \in \mathbb{R}^{Q\times L}$ ,  $a\in \mathbb{R}^N$ ,  $b\in \mathbb{R}^N$  and  $\eta_t\sim \mathcal{N}(\mathbf{0},\sigma^2\mathbb{I})$ . We refer to this dataset as GPVAR-L: note that  $a$  and  $b$  are node-specific parameters that inject local effects into the spatiotemporal process. We indicate simply as GPVAR the process obtained by fixing
Table 2: One-step-ahead forecasting error (MAE) of the different models in GPVAR datasets (5 runs).
MODELS	GPVAR	GPVAR-L
FC-RNN	.4393±.0024	.5978±.0149
LocalRNNs	.4047±.0001	.4610±.0003
Global	RNN	.3999±.0000	.5440±.0003
TTS-IMP	.3232±.0002	.4059±.0032
TTS-AMP	.3193±.0000	.3587±.0049
Emb.	RNN	.3991±.0001	.4612±.0003
TTS-IMP	.3195±.0000	.3200±.0002
TTS-AMP	.3194±.0001	.3199±.0001
Optimal model	.3192	.3192
$\pmb{a} = \pmb{b} = \pmb{0.5}$ , i.e., by removing local effects. We use as the adjacency matrix the community graph used in prior works, increasing the number of nodes to 120 (see Appendix A.1).
Tab. 2 shows forecasting accuracy for reference architectures with a 6-steps window on data generated from the processes. In the setting with no local effects, all STGNNs achieve performance close to the theoretical optimum, outperforming global and local univariate models but also the multivariate FC-RNN that – without any inductive bias – struggles to properly fit the data. In GPVAR-L, global and univariate models fail to match the performance of STGNNs that include local components (S1); interestingly, the global model with anisotropic MP outperforms the isotropic alternative, suggesting that the more advanced MP schemes can lead to more effective state identification.
Table 3: TTS-IMP one-step-ahead MAE on GPVAR-L with varying window length  $W$  and capacity  ${d}_{h}$  (5 runs).
Global	Embeddings
W	dh=16	dh=32	dh=64	dh=16	dh=32	dh=64
2	.5371±.0014	.4679±.0016	.4124±.0021	.3198±.0001	.3199±.0001	.3203±.0001
6	.4059±.0032	.3578±.0031	.3365±.0006	.3200±.0002	.3201±.0001	.3209±.0002
12	.3672±.0035	.3362±.0012	.3280±.0003	.3200±.0001	.3200±.0000	.3211±.0003
24	.3485±.0032	.3286±.0005	.3250±.0001	.3200±.0002	.3200±.0000	.3211±.0003
Finally, Tab. 3 shows that to match the performance of models with local components, in fully global models both window size and capacity should be increased (S3). This result is in line with what we should expect by considering the theory of local and global models and suggests that similar trade-offs might also happen in practical applications where both computational and sample complexity are a concern.
Benchmarks We then compare the performance of reference architectures and baselines with and without node embeddings at the encoding and decoding steps. Note that, while reference architectures and DCRNN are purely global models, GraphWaveNet and AGCRN use node embeddings to obtain an adjacency matrix for MP. AGCRN, furthermore, uses embeddings to make the convolutional filters
Table 4: Forecasting error (MAE) on 4 benchmark datasets (5 runs). The best result between each model and its variant with embeddings is in bold. N/A indicates runs exceeding resource capacity.
MODELS	METR-LA	PEMS-BAY	CER-E	AQI	METR-LA	PEMS-BAY	CER-E	AQI
Reference arch.	Global models	Global-local models (with embeddings)
RNN	3.54±.00	1.77±.00	456.98±0.61	14.02±.04	3.15±.03	1.59±.00	421.50±1.78	13.73±.04
T&S-IMP	3.35±.01	1.70±.01	443.85±0.99	12.87±.02	3.10±.01	1.59±.00	417.71±1.28	12.48±.03
TTS-IMP	3.34±.01	1.72±.00	439.13±0.51	12.74±.02	3.08±.01	1.58±.00	412.44±2.80	12.33±.02
T&S-AMP	3.22±.02	1.65±.00	N/A	N/A	3.07±.02	1.59±.00	N/A	N/A
TTS-AMP	3.24±.01	1.66±.00	431.33±0.68	12.30±.02	3.06±.01	1.58±.01	412.95±1.28	12.15±.02
Baseline arch.	Original	Embeddings at Encoder and Decoder
DCRNN	3.22±.01	1.64±.00	428.36±1.23	12.96±.03	3.07±.02	1.60±.00	412.87±1.51	12.53±.02
GraphWaveNet	3.05±.03	1.56±.01	397.17±0.67	12.08±.11	2.99±.02	1.58±.00	401.15±1.49	11.81±.04
AGCRN	3.16±.01	1.61±.00	444.80±1.25	13.33±.02	3.14±.00	1.62±.00	436.84±2.06	13.28±.03
(a) Average cluster load by day of the week.
(b) Node embeddings.
Figure 1: Time series clusters in CER-E, obtained by regularizing the embedding space. (a) Average load for each clusters. (b) t-SNE plot of the corresponding node embeddings.
adaptive w.r.t. the node being processed. We evaluate all models on real-world datasets from three different domains: traffic networks, energy analytics, and air quality monitoring. Besides the already introduced traffic forecasting benchmarks (METR-LA and PEMS-BAY), we run experiments on smart metering data from the CER-E dataset [54] and air quality measurements from AQI [55] by using the same pre-processing steps and data splits of previous works [20, 47]. The full experimental setup is reported in appendix A. Tab. 4 reports forecasting mean absolute error (MAE) averaged over the forecasting horizon. Global-local reference models outperform the fully global variants in every considered scenario (S1). A similar observation can be made for the state-of-art architectures, where the impact of node embeddings (at encoding and decoding) is large for the fully global DCRNN and more contained in models already equipped with local components. Note that hyperparameters were not tuned to account for the change in architecture. Surprisingly, the simple TTS-IMP model equipped with node embeddings achieves results comparable to that of state-of-the-art STGNNs with a significantly lower number of parameters and a streamlined architecture. Interestingly, while both global and local RNNs models fail, the hybrid global-local RNN ranks among the best-performing models, outperforming graph-based models without node embeddings in most settings.
Structured embeddings To test the hypothesis that structure in embedding space provides insights on the local effects at play (S5), we consider the clustering regularization method (Sec. 5.1) and the reference TTS-IMP model trained on the CER-E dataset. We set the number of learned centroids to  $K = 5$  and train the cluster assignment mechanism end-to-end with the forecasting architecture. Then, we inspect the clustering assignment by looking at intra-cluster statistics. In particular, for each load profile, we compute the weekly average load curve, and, for each hour, we look at quantiles of the energy consumption within each cluster. Fig. 1a shows the results of the analysis by reporting the median load profile for each cluster; shaded areas correspond to quantiles with  $10\%$  increments. Results show that users in the different clusters have distinctly different consumption patterns. Fig. 1b
Table 5: Forecasting error (MAE) in the transfer learning setting (5 runs average). Results refer to a 1-week fine-tuning set size on all PEMS datasets.
TTS-IMP	PEMS03	PEMS04	PEMS07	PEMS08
Fine-tuning	Global	15.30 ± 0.03	21.59 ± 0.11	23.82 ± 0.03	15.90 ± 0.07
Embeddings	14.64 ± 0.05	20.27 ± 0.11	22.23 ± 0.08	15.45 ± 0.06
- Variational	14.56 ± 0.03	20.19 ± 0.05	22.43 ± 0.02	15.41 ± 0.06
- Clustering	14.60 ± 0.02	19.91 ± 0.11	22.16 ± 0.07	15.41 ± 0.06
Zero-shot	18.20 ± 0.09	23.88 ± 0.08	32.76 ± 0.69	20.41 ± 0.07
Table 6: Forecasting error (MAE) on PEMS04 in the transfer learning setting by varying fine-tuning set size (5 runs average).
Model	Training set size
TTS-IMP	2 weeks	1 week	3 days	1 day
Global	20.86 ± 0.03	21.59 ± 0.11	21.84 ± 0.06	22.26 ± 0.10
Embeddings	19.96 ± 0.08	20.27 ± 0.11	21.03 ± 0.14	21.99 ± 0.13
- Variational	19.94 ± 0.08	20.19 ± 0.05	20.71 ± 0.12	21.20 ± 0.15
- Clustering	19.69 ± 0.06	19.91 ± 0.11	20.48 ± 0.09	21.91 ± 0.21
shows a 2D t-SNE visualization of the learned node embeddings, providing a view of the latent space and the effects of the cluster-based regularization.
Transfer In this experiment, we consider the scenario in which an STGNN for traffic forecasting is trained by using data from multiple traffic networks and then used to make predictions for a disjoint set of sensors sampled from the same region. We use the PEMS03, PEMS04, PEMS07, and PEMS08 datasets [56], which contain measurements from 4 different districts in California. We train models on 3 of the datasets, fine-tune on 1 week of data from the target left-out dataset, validate on the following week, and test on the week thereafter. We compare variants of TTS-IMP with and without embeddings fed into encoder and decoder. Together with the unconstrained embeddings, we also consider the variational and clustering regularization approaches introduced in Sec. 5.1. At the fine-tuning stage, the global model updates all of its parameters, while in the hybrid global-local approaches only the embeddings are fitted to the new data. Tab. 5 reports results for the described scenario. The fully global approach is outperformed by the hybrid architectures in all target datasets (S4). Besides the significant improvement in performance, adjusting only node embeddings retains performance on the source datasets. Furthermore, results show the positive effects of regularizing the embedding space in the transfer setting (S5). This is further confirmed by results in Tab. 6, which report, for PEMS04, how forecasting error changes in relation to the length of the fine-tuning window. We refer to Appendix B for an in-depth analysis of several additional transfer learning scenarios.
We investigate the impact of locality and globality in graph-based spatiotemporal forecasting architectures. We propose a framework to explain empirical results associated with the use of trainable node embeddings and discuss different architectures and regularization techniques to account for local effects. The proposed methodologies are thoroughly empirically validated and, although not inductive, prove to be effective in a transfer learning context. We argue that our work provides necessary and key methodologies for the understanding and design of effective graph-based spatiotemporal forecasting architectures. Future works can build on the results presented here and study alternative, and even more transferable, methods to account for local effects.
This research was partly funded by the Swiss National Science Foundation under grant 204061: High-Order Relations and Dynamics in Graph Neural Networks.
The following sections provide additional details on the experimental setting and computational platform used to gather the results presented in the paper. Additionally, we also include supplementary empirical results and analysis.
Experimental setup and baselines have been developed with Python [57] by relying on the following open-source libraries:
PyTorch [58];
PyTorch Lightning [59];
PyTorch Geometric [60]
- Torch Spatiotemporal [61];
- numpy [62];
- scikit-learn [63].
Experiments were run on a workstation equipped with AMD EPYC 7513 processors and four NVIDIA RTX A5000 GPUs. The code needed to reproduce the reported results is available online<sup>2</sup>.
Table 7: Statistics of datasets used in the experiments.
DATASETS	Type	Time steps	Nodes	Edges	Rate
GPVAR	Undirected	30,000	120	199	N/A
GPVAR-L	Undirected	30,000	120	199	N/A
METR-LA	Directed	34,272	207	1515	5 minutes
PEMS-BAY	Directed	52,128	325	2369	5 minutes
CER-E	Directed	25,728	485	4365	30 minutes
AQI	Undirected	8,760	437	2699	1 hour
PEMS03	Directed	26,208	358	546	5 minutes
PEMS04	Directed	16,992	307	340	5 minutes
PEMS07	Directed	28,224	883	866	5 minutes
PEMS08	Directed	17,856	170	277	5 minutes
In this section, we provide additional information regarding each one of the considered datasets. Tab. 7 reports a summary of statistics.
GPVAR(-L) For the GPVAR datasets we follow the procedure described in Sec. 7 to generate data and then partition the resulting time series in  $70\% /10\% /20\%$  splits for training, validation and testing, respectively. For GPVAR-L the parameters of the spatiotemporal process are set as
$$
\Theta = \left[ \begin{array}{c c c} 2. 5 & - 2. 0 & - 0. 5 \\ 1. 0 & 3. 0 & 0. 0 \end{array} \right], \qquad \boldsymbol {a}, \boldsymbol {b} \sim \mathcal {U} (- 2, 2),
$$
$$
\eta \sim \mathcal {N} (\mathbf {0}, \operatorname {d i a g} \left(\sigma^ {2}\right)), \quad \sigma = 0. 4.
$$
Fig. 2 shows the topology of the graph used as a support to the process. In particular, we considered a network with 120 nodes with 20 communities and added self-loops to the graph adjacency matrix.
$^{2}$ https://github.com/Graph-Machine-Learning-Group/taming-local-effects-stgmns
Figure 2: GPVAR community graph. We used a graph with 20 communities resulting in a network with 120 nodes.
Spatiotemporal forecasting benchmarks We considered several datasets coming from relevant application domains and different problem settings corresponding to real-world application scenarios.
Traffic forecasting We consider two popular traffic forecasting datasets, namely METR-LA and PEMS-BAY [6], containing measurements from loop detectors in the Los Angeles County Highway and San Francisco Bay Area, respectively. For the experiment on transfer learning, we use the PEMS03, PEMS04, PEMS07, and PEMS08 datasets from Guo et al. [56] each collecting traffic flow readings, aggregated into 5-minutes intervals, from different areas in California provided by Caltrans Performance Measurement System (PeMS) [64].
Electric load forecasting We selected the CER-E dataset [54], a collection of energy consumption readings, aggregated into 30-minutes intervals, from 485 smart meters monitoring small and medium-sized enterprises.
Air quality monitoring The AQI [55] dataset collects hourly measurements of pollutant PM2.5 from 437 air quality monitoring stations in China. Note that all of these datasets have been previously used for spatiotemporal forecasting and imputation [47, 65].