Nested Graph Neural Networks

1 Introduction . Graph is an important tool to model relational data in the real world . Representation learning over graphs has become a popular topic of machine learning in recent years . While network embedding methods , such as DeepWalk [ 1 ] , can learn node representations well , they fail to generalize to whole-graph representations , which are crucial for applications such as graph classification , molecule modeling , and drug discovery . On the contrary , although traditional graph kernels [ 2–7 ] can be used for graph classification , they define graph similarity often in a heuristic way , which is not parameterized and lacks some flexibility to deal with features . In this context , graph neural networks ( GNNs ) have regained people ’ s attention and become the state-of-the-art graph representation learning tool [ 8–17 ] . GNNs use message passing to propagate features between connected nodes . By iteratively aggregating neighboring node features to the center node , GNNs learn node representations encoding their local structure and feature information . These node representations can be further pooled into a graph representation , enabling graph-level tasks such as graph classification . In this paper , we will use message passing GNNs to denote this class ∗Corresponding author : Muhan Zhang ( muhan @ pku.edu.cn ) . †Pan Li contributes Sec . 3.3 that proves the Theorem 1 and some implementation ideas . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . of GNNs based on repeated neighbor aggregation [ 18 ] , in order to distinguish them from some high-order GNN variants [ 19–21 ] where the effective message passing happens between high-order node tuples instead of nodes . GNNs ’ message passing scheme mimics the 1-dimensional Weisfeiler-Lehman ( 1-WL ) algorithm [ 22 ] , which iteratively refines a node ’ s color according to its current color and the multiset of its neighbors ’ colors . This procedure essentially encodes a rooted subtree around each node into its final color , where the rooted subtree is constructed by recursively expanding the neighbors of the root node . One critical reason for GNN ’ s success in graph classification is because two graphs sharing many identical or similar rooted subtrees are more likely classified into the same class , which actually aligns with the inductive bias that two graphs are similar if they have many common substructures [ 23 ] . Despite this , rooted subtrees are still limited in terms of expressing all possible substructures that can appear in a graph . It is likely that two graphs , despite sharing a lot of identical rooted subtrees , are not similar at all because their other substructure patterns are not similar . Take the two graphs G1 and G2 in Figure 1 as an example . If we apply 1-WL or a message passing GNN to them , the two graphs will always have the same representation no matter how many iterations/layers we use . This is because all nodes in the two graphs have identical rooted subtrees across all tree heights . However , the two graphs are quite different from a holistic perspective . G1 is composed of two triangles , while G2 is a hexagon . The intrinsic reason for such a failure is that rooted subtrees have limited expressiveness for representing general graphs , especially those with cycles . To address this issue , we propose Nested Graph Neural Networks ( NGNNs ) . The core idea is , instead of encoding a rooted subtree , we want the final representation of a node to encode a rooted subgraph ( local h-hop subgraph ) around it . The subgraph is not restricted to be of any particular graph type such as tree , but serves as a general description of the local neighborhood around a node . Rooted subgraphs offer much better representation power than rooted subtrees , e.g. , we can easily discriminate the two graphs in Figure 1 by only comparing their height-1 rooted subgraphs . To represent a graph with rooted subgraphs , NGNN uses two levels of GNNs : base ( inner ) GNNs and an outer GNN . By extracting a local rooted subgraph around each node , NGNN first applies a base GNN to each node ’ s subgraph independently . Then , a subgraph pooling layer is applied to each subgraph to aggregate the intermediate node representations into a subgraph representation . This subgraph representation is used as the final representation of the root node . Rather than encoding a rooted subtree , this final node representation encodes the local subgraph around it , which contains more information than a subtree . Finally , all the final node representations are further fed into an outer GNN to learn a representation for the entire graph . Figure 2 shows one NGNN implementation using message passing GNNs as the base GNNs and a simple graph pooling layer as the outer GNN . One may wonder that the base GNN seems to still learn only rooted subtrees if it is message-passingbased . Then why is NGNN more powerful than GNN ? One key reason lies in the subgraph pooling layer . Take the height-1 rooted subgraphs ( marked with red boxes ) around v1 and v2 in Figure 1 as an example . Although v1 and v2 ’ s height-1 rooted subtrees are still the same , their neighbors ( labeled by 1 and 2 ) have different height-1 rooted subtrees . Thus , applying a one-layer message passing GNN plus a subgraph pooling as the base GNN is sufficient to discriminate G1 and G2 . The NGNN framework has multiple exclusive advantages . Firstly , it allows freely choosing the base GNN , and can enhance the base GNN ’ s representation power in a plug-and-play fashion . Theoretically , we proved that NGNN is more powerful than message passing GNNs and 1-WL by being able to discriminate almost all r-regular graphs ( where 1-WL always fails ) . Secondly , by extracting rooted subgraphs , NGNN allows augmenting the initial features of a node with subgraphspecific structural features such as distance encoding [ 24 ] to improve the quality of the learned node representations . Thirdly , unlike other more powerful graph neural networks , especially those based on higher-order WL tests [ 19–21 , 25 ] , NGNN still has linear time and space complexity w.r.t . graph size like standard message passing GNNs , thus maintaining good scalability . We demonstrate the effectiveness of the NGNN framework in various synthetic/real-world graph classification/regression datasets . On synthetic datasets , NGNN demonstrates higher-than-1-WL expressive power , matching very well with our theorem . On real-world datasets , NGNN consistently enhances a wide range of base GNNs ’ performance , achieving highly competitive results on all datasets . 2 Preliminaries . 2.1 Notation and problem definition . We consider the graph classification/regression problem . Given a graph G = ( V , E ) where V = { 1 , 2 , . . . n } is the node set and E ⊆ V × V is the edge set , we aim to learn a function mapping G to its class or target value y . The nodes and edges in G can have feature vectors associated with them , denoted by xi ( for node i ) and eij ( for edge ( i , j ) ) , respectively . 2.2 Weisfeiler-Lehman test . The Wesfeiler-Lehman ( 1-WL ) test [ 22 ] is a popular algorithm for graph isomorphism checking . The classical 1-WL works as follows . At first , all nodes receive a color 1 . Each node collects its neighbors ’ colors into a multiset . Then , 1-WL will update each node ’ s color so that two nodes get the same new color if and only if their current colors are the same and they have identical multisets of neighbor colors . Repeat this process until the number of colors does not increase between two iterations . Then , 1-WL will return that two graphs are non-isomorphic if their node colors are different at some iteration , or fail to determine whether they are non-isomorphic . See [ 7 , 26 ] for more details . 1-WL essentially encodes the rooted subtrees around each node at different heights into its color representations . Figure 1 middle shows the rooted subtrees around v1 and v2 . Two nodes will have the same color at iteration h if and only if their height-h rooted subtrees are the same . 3 Nested Graph Neural Network . In this section , we introduce our Nested Graph Neural Network ( NGNN ) framework and theoretically demonstrate its higher representation power than message passing GNNs . 3.1 Limitations of the message passing GNNs . Most existing GNNs follow the message passing framework [ 18 ] : given a graph G , each node ’ s hidden state ht+1v is updated based on its previous state h t v and the messages m t+1 v from its neighbors ht+1v = Ut ( h t v , m t+1 v ) , where m t+1 v = ∑ u∈N ( v|G ) Mt ( h t v , h t u , evu ) . ( 1 ) Here Mt , Ut are the message and update functions at time stamp t , evu is the feature of edge ( v , u ) , and N ( v|G ) is the set of v ’ s neighbors in graph G. The initial hidden states h0v are given by the raw node features xv . After T time stamps ( iterations ) , the final node representations hTv are summarized into a whole-graph representation with a readout ( pooling ) function R ( e.g. , mean or sum ) : hG = R ( { hTv |v ∈ G } ) . ( 2 ) Such a message passing ( or neighbor aggregation ) scheme iteratively aggregates neighbor information into a center node ’ s hidden state , making it encode a local rooted subtree around the node . The final node representations will contain both the local structure and feature information around nodes , enabling node-level tasks such as node classification . After a pooling layer , these node representations can be further summarized into a graph representation , enabling graph-level tasks . When there is no edge feature and the node features are from a countable space , it is shown that message passing GNNs are at most as powerful as the 1-WL test for discriminating non-isomorphic graphs [ 27 , 19 ] . For an h-layer message passing GNN , it will give two nodes the same final representation if they have identical height-h rooted subtrees ( i.e. , both the structures and the features on the corresponding nodes/edges are the same ) . If two graphs have a lot of identical ( or similar ) rooted subtrees , they will also have similar graph representations after pooling . This insight is crucial for the success of modern GNNs in graph classification , because it aligns with the inductive bias that two graphs are similar if they have many common substructures . Such insight has also been used in designing the WL subtree kernel [ 7 ] , a state-of-the-art graph classification method before GNNs . However , message passing GNNs have several limitations . Firstly , rooted subtree is only one specific substructure . It is not general enough to represent arbitrary subgraphs , especially those with cycles due to the natural restriction of tree structure . Secondly , using rooted subtree as the elementary substructure results in a discriminating power bounded by the 1-WL test . For example , all n-node r-regular graphs can not be discriminated by message passing GNNs . Thirdly , standard message passing GNNs do not allow using root-node-specific structural features ( such as the distance between a node and the root node ) to improve the quality of the learned root node ’ s representation . We might need to break through such limitations in order to design more powerful GNNs .

Summary: Since GNNs based on rooted subtrees are known with limited expressive power, this paper proposes to represent nodes with their rooted subgraphs. Its expressive power with respect to graph isomorphism testing is proven beyond 2-WL. Such a method is compatible to plug into popular GNNs to enhance their representation power. Enhanced GNNs show strong performances on graph classification and regression tasks.