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: This paper presents a framework that applies GNN to subgraphs around each node independently to enhance node representations, followed by a global pooling layer to get graph-level embedding. The framework can be applied to any GNNs. The paper also demonstrates it's relative improvement over several real-world datasets, although the improvement seems marginal. From theoretical side, the paper proved formally that this framework can distinguish k-regular graphs where most traditional message passing based neural network cannot.