Graph Neural Networks (GNNs) represent a powerful learning paradigm that have achieved great success (Scarselli et al. 2008;Li et al. 2016;Kipf and Welling 2017;Hamilton, Ying, and Leskovec 2017;Velickovic et al. 2018;Xu et al. 2019;You, Ying, and Leskovec 2020). Among these models, messaging passing GNNs, such as GCN (Kipf and Welling 2017), GraphSAGE (Hamilton, Ying, and Leskovec 2017), and GAT (Velickovic et al. 2018), are dominantly used today due to their simplicity, efficiency and strong performance in realworld applications (Zitnik and Leskovec 2017;Ying et al. 2018;You et al. 2018aYou et al. , 2019bYou et al. , 2020a,b),b). The central idea behind message passing GNNs is to learn node embeddings Copyright c 2021, Association for the Advancement of Artificial Intelligence (www.aaai.org). All rights reserved.
via the repeated aggregation of information from local node neighborhoods using non-linear transformations (Battaglia et al. 2018).
Although GNNs represent a powerful learning paradigm, it has been shown that the expressive power of existing GNNs is upper-bounded by the 1-Weisfeiler-Lehman (1-WL) test (Xu et al. 2019). Concretely, a fundamental limitation of existing GNNs is that two nodes with different neighborhood structure can have the same computational graph, thus appearing indistinguishable. Here, a computational graph specifies the procedure to produce a node's embedding. Such failure cases are abundant (Figure 1): in node classification tasks, existing GNNs fail to distinguish nodes that reside in d-regular graphs of different sizes; in link prediction tasks, they cannot differentiate node candidates with the same neighborhood structures but different shortest path distance to the source node; and in graph classification tasks, they cannot differentiate d-regular graphs (Chen et al. 2019;Murphy et al. 2019). While task-specific feature augmentation can be used to mitigate these failure modes, the process of discovering meaningful features for different tasks is not generic and can, for example, hamper the inductive power of GNNs.
Several recent methods aim to overcome these limitations in existing GNNs. For graph classification tasks, a collection of works propose novel architectures more expressive than the 1-WL test (Chen et al. 2019;Maron et al. 2019a;Murphy et al. 2019). For link level tasks, P-GNNs are proposed to overcome the limitation of existing GNNs (You, Ying, and Leskovec 2019). While these methods have a rich theoretical grounding, they are often task specific (either graph or link level) and often suffer from increased complexity in computation or implementation. In contrast, message passing GNNs have a track record of high predictive performance across node, link, and graph level tasks, while being simple and efficient to implement. Therefore, extending message passing GNNs beyond the expressiveness of 1-WL test, to overcome current GNN limitations, is a problem of high importance. Present work. Here we propose Identity-aware Graph Neural Networks (ID-GNNs), a class of message passing GNNs with expressive power beyond the 1-WL test 1 . ID-GNN provides a universal extension and makes any existing message passing GNN more expressive. ID-GNN embeds each node 1 Project website with code: http://snap.stanford.edu/idgnn Across all examples, the task requires an embedding that allows for the differentiation of nodes labeled A vs. B in their respective graphs. However, across all tasks, existing GNNs, regardless of their depth, will always assign the same embedding to both nodes A and B, because for all tasks the computational graphs are identical (middle row). In contrast, the colored computational graphs provided by ID-GNN allow for clear differentiation between the nodes of label A and label B, as the colored computational graphs are no longer identical across the tasks.
by inductively taking into account its identity during message passing. The approach is different from labeling each node with a one-hot encoding, which is transductive (cannot generalize to unseen graphs). As shown in Figure 1, we use an inductive identity coloring technique to distinguish a node itself (the root node in the computational graph) from other nodes in its local neighborhood, within its respective computational graph. This added identity information allows ID-GNN to distinguish what would be identical computational graphs across node, edge and graph level tasks, and this way overcome the previously discussed limitations.
We propose two versions of ID-GNN. As a general approach, identity information is incorporated by applying rounds of heterogeneous message passing. Specifically, to embed a given node, ID-GNN first extracts the ego network centered at that node, then applies message passing, where the messages from the center node (colored nodes in Figure 1) and the rest of the nodes are computed using different sets of parameters. This approach naturally applies to applications involving node or edge features. We also consider a simplified version of ID-GNN, where we inject identity information via cycle counts originating from a given node as augmented node features. These cycle counts capture node identity information by counting the colored nodes within each layer of the ID-GNN computational graph, and can be efficiently computed by powers of a graph's adjacency matrix.
We compare ID-GNNs against GNNs across 8 datasets and 6 different tasks. First, we consider a collection of challenging graph property prediction tasks where existing GNNs fail, including predicting node clustering coefficient, predicting shortest path distance, and differentiating ran-dom d-regular graphs. Then, we further apply ID-GNNs to real-world datasets. Results show that transforming existing GNNs to their ID-GNN versions yields on average 40% accuracy improvement on challenging node, edge, and graph property prediction tasks; 3% accuracy improvement on node and graph classification benchmarks; and 15% ROC AUC improvement on real-world link prediction tasks. Additionally, we compare ID-GNNs against other expressive graph networks that are specifically designed for edge or graph-level tasks. ID-GNNs demonstrate improved or comparable performance over these models, further emphasizing the versatility of ID-GNNs.
Our key contribution includes: (1) We show that message passing GNNs can have expressive power beyond 1-WL test.
(2) We propose ID-GNNs as a general solution to the limitations in existing GNNs, with rich theoretical and experimental results. (3) We present synthetic and real world tasks to reveal the failure modes of existing GNNs and demonstrate the superior performance of ID-GNNs over both existing GNNs and other powerful graph networks.
Expressive neural networks beyond 1-WL test. Recently, many neural networks have been proposed with expressive power beyond the 1-WL test, including (Chen et al. 2019;Maron et al. 2019a;Murphy et al. 2019;You, Ying, and Leskovec 2019;Li et al. 2020). However, these papers introduce extra, often task/domain specific, components beyond standard message passing GNNs. For example, P-GNN's embeddings are tied with random anchor-sets and, thus, are not applicable to node/graph level tasks which require deter-ministic node embeddings (You, Ying, and Leskovec 2019). In this paper we emphasize the advantageous characteristics of message passing GNNs, and show that GNNs, after incorporating inductive identity information, can surpass the expressive power of the 1-WL test while maintaining benefits of efficiency, simplicity, and broad applicability. Graph Neural Networks with inductive coloring. Several models color nodes with augmented features to boost existing GNNs' performance (Xu et al. 2020;Veličković et al. 2020;Zhang and Chen 2018). However, existing coloring techniques are problem and domain-specific (i.e. link prediction, algorithm execution), and are not generally applicable to node and graph-level tasks. In contrast, ID-GNN is a general model that can be applied to any node, edge, and graph level task. It further adopts a heterogeneous message passing approach, which is fully compatible to cases where nodes or edges have rich features. GNNs with anisotropic message passing. We emphasize that ID-GNNs are fundamentally different from GNNs based on anisotropic message passing, where different attention weights are applied to different incoming edges (Bresson and Laurent 2017;Hamilton, Ying, and Leskovec 2017;Monti et al. 2017;Velickovic et al. 2018). Adding anisotropic message passing does not change the underlying computational graph because the same message passing function is symmetrically applied across all nodes. Therefore, these models still exhibit the limitations summarized in Figure 1.
A graph can be represented as G = (V, E), where V = {v 1 , ..., v n } is the node set and E ⊆ V × V is the edge set. Nodes can be paired with features X = {x v |∀v ∈ V}, and edges can have features F = {f uv |∀e uv ∈ E}. As discussed earlier, we focus on message passing GNNs throughout this paper. We follow the definition of GNNs in (Xu et al. 2019). The goal of a GNN is to learn meaningful node embeddings h v based on an iterative aggregation of local network neighborhoods. The k-th iteration of message passing, or the k-th layer of a GNN, can be written as:
where h
v is the node embedding after k iterations, h
The node embeddings h
v , ∀v ∈ V are then used for node, edge, and graph level prediction tasks.
We design ID-GNN so that it can make any message passing GNN more expressive. ID-GNN is built with two important
(1) inductive identity coloring where identity information is injected to each node, and (2) heterogeneous message passing where the identity information is utilized in message passing. Algorithm 1 provides an overview.
Inductive identity coloring. To embed a given node v ∈ G using a K-layer ID-GNN, we first extract the K-hop ego network
We then assign a unique coloring to the central node of the ego network
can be categorized into two types throughout the embedding process: nodes with coloring and nodes without coloring. This coloring technique is inductive because even if nodes are permuted, the center node of the ego network can still be differentiated from other neighboring nodes. Heterogeneous message passing. K rounds of message passing are then applied to all the extracted ego networks. To embed node u ∈ G (K) v , we extend Eq. 1 to enable heterogeneous message passing:
where only h
is used as the embedding representation for node v after applying K rounds of Eq. 3. Different from Eq. 1, two sets of MSG (k) functions are used, where MSG
is applied to nodes with identity coloring, and MSG
these functions. This way, the inductive identity coloring is encoded into the ID-GNN computational graph.
A benefit of this heterogeneous message passing approach is that it is applicable to any message passing GNN. For example, consider the following message passing scheme, which extends the definition of GNNs in Eq. 3 by including edge attributes f su during message passing:
Algorithmic complexity. Besides adding the identity coloring and applying two types of message passing instead of one, the computation of ID-GNN is almost identical to the widely used mini-batch version of GNNs (Hamilton, Ying, and Leskovec 2017;Ying et al. 2018). In our experiments, by matching the number of trainable parameters, the computation FLOPS used by ID-GNNs and mini-batch GNNs can be the same (shown in Table 4).
Extension to edge-level tasks. Here we discuss how to extend the ID-GNN framework to properly resolve existing GNN limitations in edge-level tasks (Figure 1, middle). Suppose we want to predict the edge-level label for a node pair u, v. For ID-GNN, the prediction is made from a conditional node embedding h u|v , which is computed by assigning node v, rather than u, identity coloring in node u's computation graph, as illustrated in Figure 1. In the case where node v does not lie within u's K-hop ego network, no identity coloring is used and ID-GNNs will still suffer from existing failure cases of GNNs. Therefore, we use deeper ID-GNNs for edge-level prediction tasks in practice.
1 (•), Eq. 3 becomes identical to Eq. 1 which trivially proves the first part. The d-regular graph example given in Figure 1 then proves the second part. ID-GNNs can count cycles. Proposition 1 provides an overview of the added expressive power of ID-GNNs. Here, we reveal one concrete aspect of this added expressive power, i.e., ID-GNN's capability to count cycles. We observe that the ability of counting cycles is intuitive to understand; moreover, it is crucial for useful tasks such as predicting node clustering coefficient, which we elaborate in the next section. Proposition 2. For any node v, there exists a K-layer ID-GNN instantiation that can learn an embedding h (K) v where the j-th dimension h (K) v [j] equals the number of length j cycles starting and ending at node v, for j = 1, ..., K.
We prove this by showing that ID-GNNs can count paths from any node u to the identity node v. Through induction, we show that a 1-layer ID-GNN embedding h
(1) u can count length 1 paths from u to v. Then, given a K-layer ID-GNN embedding h (K) u that counts paths of length 1, . . . , K between u and v, we show the K + 1-th layer of ID-GNN can accurately update h
Node-level: Predicting clustering coefficient. Here we show that existing message passing GNNs fail to inductively predict clustering coefficients purely from graph structure, while ID-GNNs can. Clustering coefficient is a widely used metric that characterizes the proportion of closed triangles in a node's 1-hop neighborhood (Watts and Strogatz 1998). The node classification failure case in Figure 1 demonstrates GNNs' inability to predict clustering coefficients, as GNNs fail to differentiate nodes v 1 and v 2 with clustering coefficient 1 and 0 respectively. By using one-hot node features, GNNs can overcome this failure mode (Hamilton, Ying, and Leskovec 2017). However, in this case GNNs are memorizing the clustering coefficients for each node, since one-hot encodings prevent generalization to unseen graphs.
Based on Proposition 2, ID-GNNs can learn node embeddings h
equals the number of length j cycles starting and ending at node v. Given these cycle counts, we can then calculate clustering coefficient c v of node v:
v , we can approximate it to an arbitrary precision with an MLP due to the universal approximation theorem (Hornik et al. 1989). Edge-level: Predicting reachability or shortest path distance. Vanilla GNNs make edge-level predictions from pairs of node embeddings (Hamilton, Ying, and Leskovec 2017). However, this type of approaches fail to predict reachability or shortest path distance (SPD) between node pairs. For example, two nodes can have the same GNN node embedding, independent of whether they are located in the same connected component. Although (Veličković et al. 2020) shows that proper node feature initialization allows for the prediction of reachability and SPD, ID-GNNs present a general solution to this limitation through the use of conditional node embeddings. As discussed in "Extension to edge-level tasks", we re-formulate edge-level prediction as conditional node-level prediction; consequently, a K-layer ID-GNN can predict if node u ∈ G is reachable from v ∈ G within K hops by using the conditional node embedding h
where h (0) u|v = 0, ∀u ∈ G, and the output h
Graph-level: Differentiating random d-regular graphs. As is illustrated in Figure 1, existing message passing GNNs cannot differentiate random d-regular graphs purely from graph structure, as the computation graphs for each node are identical, regardless of the number of layers. Here, we show that ID-GNNs can differentiate a significant proportion of random d-regular graphs. Specifically, we generate 100 nonisomorphic random d-regular graphs and consider 3 settings with different graph sizes (n) and node degree (d). We use up to length K cycle counts, which a K-layer ID-GNN can
Given that: (1) mini-batch implementations of GNNs have computational overhead when extracting ego networks, which is required by ID-GNNs with heterogeneous message passing, and (2) cycle count information explains an important aspect of the added expressive power of ID-GNNs over existing GNNs, we propose ID-GNN-Fast, where we inject identity information by using cycle counts as augmented node features. Similar cycle count information is also shown to be useful in the context of graph kernels (Zhang et al. 2018).
Following the definition in Proposition 3, we use the count of cycles with length 1, . . . , K starting and ending at the node v as augmented node feature x + v ∈ R K . These additional features x + v can be computed efficiently with sparse matrix multiplication via
where A is the adjacency matrix. We then update the input node attributes for all nodes by concatenating this augmented feature We bin over continuous clustering coefficients to make task
(1) and (3) 10-way classification tasks and threshold the shortest path distance to make task (2) a 5-way classification task. We also consider more common tasks with real-world labels, including (4) node classification, (5) link prediction, and (6) graph classification. For the ogbg-molhiv dataset we use provided splits, while for all the other tasks, we use a random 80/20% train/val split and average results over 3 random splits. Validation accuracy (multi-way classification) or ROC AUC (binary classification) in the final epoch is reported. Models. We present a standardized framework for fairly comparing ID-GNNs with existing GNNs. We use 4 widely adopted GNN models as base models: GAT (Velickovic et al. 2018), GCN (Kipf and Welling 2017), GIN (Xu et al. 2019), and GraphSAGE (Hamilton, Ying, and Leskovec 2017). We then transform each GNN model to its ID-GNN variants, ID-GNN-Full (based on heterogeneous message passing) and ID-GNN-Fast, holding all the other hyperparameters fixed.
To further ensure fairness, we adjust layer widths, so that all the models match the number of trainable parameters of a standard GCN model (i.e., match computational budget).
In summary, we run 12 models for each experimental setup, including 4 types of GNN architectures, each with 3 versions. We use 3-layer GNNs for node and graph level tasks, and 5layer GNNs for edge level tasks, where GCNs with 256-dim hidden units are used to set the computational budget for all 12 model variants. For ID-GNNs-Full, each layer has 2 sets of weights, thus each layer has fewer number of hidden units; for ID-GNNs-Fast, 10-dim augmented cycle counts features are used. We use ReLU activation and Batch Normalization for all the models. We use Adam optimizer with learning rate 0.01. Due to the different nature of these tasks, tasks (1)(3)( 6) excluding the ogbg-molhiv dataset, are trained for 1000 epochs, while the rest are trained for 100 epochs. For nodelevel tasks, GNN / ID-GNN node embeddings are directly used for prediction; for edge-level tasks, ID-GNNs-Full make predictions with conditional node embeddings, while GNNs and ID-GNNs-Fast make predictions by concatenating pairs of node embeddings and then passing the result through a 256-dim MLP; for graph-level tasks, predictions are based on a global sum pooling over node embeddings.
Overall, these comprehensive and consistent experimental settings reveal the general improvement of ID-GNNs compared with existing GNNs.
Node clustering coefficient prediction. In Table 2 we observe that across all models and datasets, both ID-GNN formulations perform at the level of or significantly outperform GNN counterparts, with an average absolute performance gain of 36.8% between the best ID-GNN and best GNN. In each dataset, both ID-GNN methods perform with near 100% accuracy for at least one GNN architecture. ID-GNN-Fast shows the most consistent improvements across models with greatest improvement in GraphSAGE. These results align with the previous discussion of using cycle counts alone to learn clustering coefficients. We defer discussion until later on ID-GNN-Full sometimes showing minimal improvement, to present a general understanding of this behavior. Shortest path distance prediction. In the pairwise shortest path prediction task, ID-GNNs-Full outperform GNNs by an average of 39.9%. Table 2 reveals that ID-GNN-Full performs with 100% or near 100% accuracy under all GNN architectures, across all datasets. This observation, along with the comparatively poor performance of ID-GNNs-Fast and GNNs, confirms the previously discussed conclusion that traditional edge-level predictions, through pairwise node embeddings, fail to accurately make edge-level predictions.
Average clustering coefficient prediction for random graphs. In Table 2, we observe that adding identity information results in a 55% and 42.3% increase in best performance over ScaleFree and SmallWorld graphs respectively. ID-GNN-Fast shows the most consistent improvement (56.9% avg. model gain), which aligns with previous intuitions about the utility of cycle count information in predicting clustering coefficients and differentiating random graphs.
Node classification. In node classification we see smaller but still significant improvements when using ID-GNNs. Table 3 shows an overall 1% and 1.6% improvement for Cora and CiteSeer respectively. In all cases except for GIN and GraphSAGE on Cora, adding identity information improves performance. In regards to the relatively small improvements, we hypothesize that the richness of node features (over 1000dim for both datasets) greatly dilutes the importance of graph structure in these tasks, and thus the added expressiveness from identity information is diminished. Link prediction. As shown in Table 3, we observe consistent improvement in ID-GNNs over GNNs, with 9.2% and 20.6% ROC AUC improvement on synthetic and real-world graphs respectively. Moreover, we observe that ID-GNN-Full nearly always performs the best, aligning with previous edge-level task results in Table 2 and intuitions on the importance of reformulating edge-level tasks as conditional node prediction tasks. We observe that performance improves less for random graphs, which we hypothesize is due to the randomness within these synthetic graphs causing the distinction between positive and negative edges to be much more vague. Graph classification. Across each dataset, we observe that the best ID-GNN consistently outperforms the best GNN of the same computational budget. However, model to model improvement is less clear. For the ENZYMES dataset, ID-GNN-Fast shows strong improvements under each GNN architecture, with gains as large as 10% in accuracy for the
We compare the runtime complexity (excluding mini-batch loading time) of ID-GNNs vs. existing GNNs, where we hold the computational budget constant across all models. Table 4 reveals that when considering the forward and backward pass, ID-GNN-Full runs 3.8x slower than its GNN equivalent but has an equivalent runtime complexity to the mini-batch implementation of GNN, while ID-GNN-Fast runs with essentially zero overhead over existing GNN implementations.
Overall, ID-GNN-Full and ID-GNN-Fast demonstrate significant improvements over their message passing GNN counterparts, of the same computational budget, on a variety of tasks. In all tasks, the best ID-GNNs outperforms the best GNNs; moreover, out of 160 model-task combinations, ID-GNNs fail to improve accuracy in fewer than 10 cases. For the rare cases where there is no improvement from ID-GNN-Full, we suspect that the model underfits since we control the complexity of models: given that ID-GNN-Full has two sets of weights (heterogeneous message passing), fewer weights are used for each message passing. For verification, if we double the computational budget, we observe that ID-GNN versions again outperform GNN counterparts.
We provide additional experimental comparisons against other expressive graph networks in both edge and graphlevel tasks. For edge-level task, we further compare with P-GNN (You, Ying, and Leskovec 2019) over the ENZYMES and PROTEINS datasets using the protocol introduced previously. For graph-level comparison, we include experimental results over 3 datasets: MUTAG with 182 graphs (Debnath et al. 1991), PTC with 344 graphs (Helma and Kramer 2003), and PROTEINS. We follow PPGN's (Maron et al. 2019a) 10-fold 90/10 data splits and compare against 5 other expressive graph networks. We report numbers in the corresponding papers, and report the best ID-GNNs out of the 4 variants. Link prediction. We compare against P-GNNs on 2 link prediction datasets. As shown in the Table 5, we observe significant improvements using ID-GNNs compared to both its GNN counterpart and P-GNNs. These results both demonstrate ID-GNNs' competitive performance as a general graph learning method against a task-specific model, while also
We have proposed ID-GNNs as a general and powerful extension to existing GNNs with rich theoretical and experimental results. Specifically, ID-GNNs have expressive power beyond the 1-WL test. When runtime efficiency is the primary concern, we also present a feature augmented version of ID-GNN that maintains theoretical guarantees and empirical success of heterogeneous message passing, while only requiring onetime feature pre-processing. We recommend that this cyclecount feature augmentation be the new go-to node feature initialization when additional node attributes are not available. Additionally, as direct extensions to message passing GNNs, ID-GNNs can be easily implemented and extended via existing code platform. Overall, ID-GNNs outperform corresponding message passing GNNs, while both maintaining the attractive proprieties of message passing GNNs and demonstrating competitive performance compared with other powerful/expressive graph networks. We hope ID-GNNs' added expressive power and proven practical applicability can enable exciting new applications and further development of message passing GNNs.
GNNs represent a promising family of models for analyzing and understanding relational data. A broad range of application domains, such as network fraud detection (Akoglu, Chandy, and Faloutsos 2013;Kumar, Cheng, and Leskovec 2017;Akoglu and Faloutsos 2013), molecular drug structure discovery (You et al. 2018a,b;Jin, Barzilay, and Jaakkola 2018), recommender systems (Ying et al. 2018;You et al. 2019a), and network analysis (Kumar, Cheng, and Leskovec 2017;Morris et al. 2019;Fan et al. 2019;Ying et al. 2019) stand to be greatly impacted by the use and development of GNNs. As a direct extension of existing message passing GNNs, ID-GNNs represent a simple but powerful transformation to GNNs that re-frames the discussion on GNN expressive power and thus their performance in impactful problem domains. In comparison to other models that have expressive power beyond 1-WL tests, ID-GNNs are easy to implement with existing graph learning packages; therefore, ID-GNNs can be easily used as extensions of existing GNN models for tackling important real-world tasks, as well as themselves extended and further explored in the research space.
The simplicity of ID-GNNs presents great promise for further exploration into the expressiveness of GNNs. In particular, we believe that our work motivates further research into heterogeneous message passing and coloring schemes, as well as generic, but powerful forms of feature augmentation. By further increasing the expressiveness of message passing GNNs, we hopefully enable new, important tasks to be solved across a wide range of disciplines or significant improvement on previously defined and widely adopted GNN models. Through ease of use and strong preliminary results, we believe that our work the doors for new explorations into the study of graphs and graph based tasks, with the potential for great improvement in existing GNN models.