Improving graph neural networks through feature importance learning

Abstract Graph neural networks (GNNs) are among the most widely used methods for node classification in graphs. A common strategy to improve their predictive performance is to enrich nodes with additional features. A weakness of this method is that the set of appropriate features can vary from graph to graph. We address this shortcoming by proposing a novel method. In a preprocessing step, a first GNN is trained on a set of graphs with varying structural properties, using a candidate set of node features fixed in advance. The resulting GNN model is then used to predict the most relevant features from the candidate set for unseen target graphs, which are later processed for node classification. For each target graph, a second GNN is trained on the graph, which is enriched with the node feature vectors calculated for the features selected by the first GNN. A key advantage of the proposed method is that the features are selected without computing the candidate features for the target graph. Our experimental results on synthetic and real-world graphs show that even a few features selected in this way is sufficient to significantly improve the predictive performance of GNNs that use either none or all of the candidate features. Moreover, the time needed to learn the second GNN for the target graph can be reduced by up to two orders of magnitude.