Temperature-based pruning for input features in Graph Neural Networks

Abstract

In the present work, we employ the concept of neural network temperature to prune unimportant features in input to a Graph Neural Network (GNN) architecture. In benchmark datasets for node and graph property prediction, each node comes equipped with a vector of numerous features. It is paramount to understand which information is actually necessary and which can be discarded, both for efficiency and explainability. The temperature is linked to the gradient activity due to the loss function minimization and leads to pruning of weight structures associated with small gradients. This study is done on different GNN architectures, one for node classification and another one for link prediction, and several benchmark datasets are employed. We compare the results with similar experiments previously conducted on the filters of Convolutional Neural Networks. Although still at the proof-of-concept stage, our temperature-based pruning technique stands as a promising alternative to state-of-the-art magnitude-based pruning techniques.