Augmented Graph Neural Network with hierarchical global-based residual connections

Abstract

Graph Neural Networks (GNNs) are powerful architectures for learning on graphs. They are efficient for predicting nodes, links and graphs properties. Standard GNN variants follow a message passing schema to update nodes representations using information from higher-order neighborhoods iteratively. Consequently, deeper GNNs make it possible to define high-level nodes representations generated based on local as well as distant neighborhoods. However, deeper networks are prone to suffer from over-smoothing. To build deeper GNN architectures and avoid losing the dependency between lower (the layers closer to the input) and higher (the layers closer to the output) layers, networks can integrate residual connections to connect intermediate layers. We propose the Augmented Graph Neural Network (AGNN) model with hierarchical global-based residual connections. Using the proposed residual connections, the model generates high-level nodes representations without the need for a deeper architecture. We disclose that the nodes representations generated through our proposed AGNN model are able to define an expressive all-encompassing representation of the entire graph. As such, the graph predictions generated through the AGNN model surpass considerably state-of-the-art results. Moreover, we carry out extensive experiments to identify the best global pooling strategy and attention weights to define the adequate hierarchical and global-based residual connections for different graph property prediction tasks. Furthermore, we propose a reversible variant of the AGNN model to address the extensive memory consumption problem that typically arises from training networks on large and dense graph datasets. The proposed Reversible Augmented Graph Neural Network (R-AGNN) only stores the nodes representations acquired from the output layer as opposed to saving all representations from intermediate layers as it is conventionally done when optimizing the parameters of other GNNs. We further refine the definition of the backpropagation algorithm to fit the R-AGNN model. We evaluate the proposed models AGNN and R-AGNN on benchmark Molecular, Bioinformatics and Social Networks datasets for graph classification and achieve state-of-the-art results. For instance the AGNN model realizes improvements of +39% on IMDB-MULTI reaching 91.7% accuracy and +16% on COLLAB reaching 96.8% accuracy compared to other GNN variants.