TOWARDS EXPRESSIVE GRAPH REPRESENTATIONS FOR GRAPH NEURAL NETWORKS

Abstract

Graph Neural Network (GNN) aggregates the neighborhood information into the node embedding and shows its powerful capability for graph representation learning in various application areas. However, most existing GNN variants aggregate the neighborhood information in a fixed non-injective fashion, which may map different graphs or nodes to the same embedding, detrimental to the model expressiveness. In this paper, we present a theoretical framework to improve the expressive power of GNN by taking both injectivity and continuity into account. Based on the framework, we develop injective and continuous expressive Graph Neural Network (iceGNN) that learns the graph and node representations in an injective and continuous fashion, so that it can map similar nodes or graphs to similar embeddings, and non-equivalent nodes or non-isomorphic graphs to different embeddings. We validate the proposed iceGNN model for graph classification and node classification on multiple benchmark datasets. The experimental results demonstrate that our model achieves state-of-the-art performances on most of the benchmarks.

1. INTRODUCTION

Graph representation learning that maps graphs or their components to vector representations has attracted growing attentions for graph analysis. Recently, graph neural networks (GNN) that can learn a distributed representation for a graph or a node in a graph are widely applied to a variety of areas, such as social network analysis (Hamilton et al., 2017; Ying et al., 2018a) , molecular structure inference (Duvenaud et al., 2015; Gilmer et al., 2017 ), text mining (Yao et al., 2019; Peng et al., 2018) , clinical decision-making (Mao et al., 2022b; Li et al., 2018) and image processing (Mao et al., 2022a; Garcia & Bruna, 2018) . GNN recursively updates the representation of a node in a graph by aggregating the feature vectors of its neighbors and itself (Hamilton et al., 2017; Morris et al., 2019; Xu et al., 2019) . The graph-level representation can then be obtained through aggregating the final representations of all the nodes in the graph. The learned representations can be fed into a prediction model for different learning tasks, such as node classification and graph classification. In GNN, the aggregation rule plays a vital role in learning expressive representations for the nodes and the entire graph. There are many GNN variants with different aggregation rules proposed to achieve good performances for different tasks and different problems (Kipf & Welling, 2017; Hamilton et al., 2017; Zhang et al., 2018; Xinyi & Chen, 2019; Wang et al., 2020) . However, most of the existing GNN aggregation rules are designed based on a fixed non-injective pooling function, (e.g., max pooling and mean pooling) or on non-continuous node types (e.g., graph isomorphism test). The non-injective aggregation may map different (non-isomorphic) graphs or (non-equivalent) nodes to the same embedding; and the non-continuous aggregation may map similar graphs or nodes to quite different embeddings, both detrimental to the expressive power of GNN. For example, for the graph with attributed nodes in Figure 1 (a), mean pooling or sum aggregation on the neighborhoods generates the same neighborhood representation for all the nodes (Figure 1(d) ), thus cannot capture any meaningful structure information. Xu et al. (2019) showed that a powerful GNN can at most achieve the discriminative power of Weisfeiler-Lehman graph isomorphism test (WL test) which can discriminate a broad class of graphs (Weisfeiler & Lehman, 1968) , and proposed the powerful graph isomorphism network (GIN). However, the theoretical framework of GIN is under the assumption that the input feature space is countable, which makes GIN less expressive when applied to graphs with continuous attributes. We argue that the expressive power of a graph mapping should imply two aspects, injectivity and continuity: the injectivity ensures different graphs are mapped to different representations and the continuity ensures that similar graphs are mapped to similar representations. Most previous works only took either one into account for GNN design; few considered both injectivity and continuity. Here, we present a theoretical framework that can guide us to design highly expressive GNNs with both injectivity and continuity for general graphs with continuous attributes. We also present a necessary condition related to the representation dimension for a fully injective and continuous graph mapping. The general idea of our framework is illustrated in Figure 1 . Our main contributions are summarized as follows. (1) We present a theoretical framework to guide the design of expressive GNNs by ensuring the injectivity and continuity in the neighborhood aggregation process. (2) We present a limitation about the representation dimension for a fully injective and continuous graph mapping. (3) Based on the framework, we implement two injective and continuous expressive GNN (iceGNN) models with a fixed and learnable aggregation function, respectively. (4) We validate our models on multiple benchmark datasets including simple graphs and attributed graphs for graph classification and node classification, the experimental results demonstrate that our models can achieve state-of-the-art performances on most of the benchmarks. Our code is available in the Supplementary Material. Common notations used throughout the paper are found in Appendix A.1 Table 5 .

2. RELATED WORK

Many GNN variants with different aggregation rules are proposed in the literature to achieve good performances in different tasks. GIN proposed by Xu et al. ( 2019) is expected to be highly expressive for simple graphs where node attributes are one-hot encoders on which sum aggregation is injective. However, GIN cannot be directly extended to attributed graphs with the same expressive power, because the sum aggregation is no longer injective in uncountable cases. GCN is another GNN variant with simple element-wise mean pooling in a node's neighborhood, including the node itself (Kipf & Welling, 2017). Hamilton et al. ( 2017) tested 3 aggregators in GraphSAGE, including mean aggregator, LSTM aggregator and max pooling aggregator, they found no significant performance difference exists between the LSTM aggregators and pool aggregators, but GraphSAGE-LSTM is significantly slower than GraphSAGE-pool. Mean aggregation and max pooling are permutation invariant on sets, but the operation is not injective, which may result in the same embedding for different inputs. LSTM aggregation could have large expressive capacity, but it is not permutation invariant, this may cause equivalent nodes or isomorphic graphs to have different embeddings. Corso et al. ( 2020) combined multiple aggregators with degree-scalers and proposed PNA to improve the expressive power of GNN, but they did not provide a theoretical guidance on how to improve the expressive power of GNN. In this article, we present a theoretical framework to guide the design of expressive GNNs by ensuring the injectivity and continuity in the neighborhood aggregation process. PNA can exactly fall into our framework by a simple comparison analysis.



Figure 1: An overview of our framework on an exemplar attributed graph in one iteration. (a) Original graph with attributed nodes; (b) Graph nodes are represented by the corresponding attribute and neighborhood set through WL test; (c) The node vector representations after an injective set function on neighborhood sets, here the set function is f (X) = x∈X (1, x, x 2 , x 3 , x 4 ); (d) A non-injective alternative set function in GNNs, after aggregation, the node information remain unchanged, node B and D still have the same representation despite their different neighborhoods.

