PruneGCRN: Minimizing and explaining spatio-temporal problems through node pruning

TL;DR

PruneGCRN addresses explainability in spatio-temporal forecasting by learning a pruned graph during training to identify the most informative nodes. It combines Node Adaptive Parameter Learning (NAPL) and Pruned Graph Learning (PGL) with a GRU-based temporal core to jointly optimize prediction accuracy and sparsity via a learnable mask, with parameter sharing via $\Theta = E_N W_N$. The model produces both predictions and an explainability mask, and validates on PeMS traffic datasets showing that high-sparsity pruning can preserve or improve accuracy while reducing compute and memory, thus enabling scalable, interpretable forecasting. These results demonstrate that problem-level explainability via node pruning can reveal key sensors and interdependencies, offering practical benefits for real-world spatio-temporal analytics.

Abstract

This work addresses the challenge of using a deep learning model to prune graphs and the ability of this method to integrate explainability into spatio-temporal problems through a new approach. Instead of applying explainability to the model's behavior, we seek to gain a better understanding of the problem itself. To this end, we propose a novel model that integrates an optimized pruning mechanism capable of removing nodes from the graph during the training process, rather than doing so as a separate procedure. This integration allows the architecture to learn how to minimize prediction error while selecting the most relevant nodes. Thus, during training, the model searches for the most relevant subset of nodes, obtaining the most important elements of the problem, facilitating its analysis. To evaluate the proposed approach, we used several widely used traffic datasets, comparing the accuracy obtained by pruning with the model and with other methods. The experiments demonstrate that our method is capable of retaining a greater amount of information as the graph reduces in size compared to the other methods used. These results highlight the potential of pruning as a tool for developing models capable of simplifying spatio-temporal problems, thereby obtaining their most important elements.

Figures (9)

• Figure 1: Diagram of the generation of $\tilde{M}$. The parameters are shown in green and the layers in orange.
• Figure 2: Distribution of the number of times nodes are used, for different pruning percentages applied. The nodes removed are selected randomly.
• Figure 3: Distribution of the number of times nodes are used, for different pruning percentages applied. The nodes removed are selected by the model during training.
• Figure 4: Map showing all the nodes in PeMS-Bay dataset.
• Figure 5: Comparison between the mask generated using random nodes and the mask learned during training.