GLASS: GNN with Labeling Tricks for Subgraph Representation Learning

1 INTRODUCTION . Graph is a natural tool for modeling objects with complex internal relationships , which is widely used in fields such as natural language processing ( Yao et al. , 2019 ) , biology ( Fout et al. , 2017 ) , and social network ( Chen et al. , 2018 ) . Among the various graph representation learning methods , GNN has achieved state-of-the-art performance on almost all sorts of tasks . Existing GNNs are mainly designed for node ( Dabhi & Parmar , 2020 ; Chen et al. , 2020 ) , edge ( Singh et al. , 2021 ; Galkin et al. , 2021 ) and whole graph ( Ying et al. , 2021b ; Yang et al. , 2020 ) property prediction tasks . An ordinary GNN produces embeddings of a node by aggregating the features from the ( multi-hop ) neighbors of the node , which is equivalent to encoding a breadth-first-search ( BFS ) tree rooted in the node ( Xu et al. , 2019 ) . Such embeddings can be used to predict node properties directly . As for edge and graph tasks , pooling the embeddings of nodes related to the structure is a prevailing method . Though node , edge , and graph tasks are the three most common graph representation learning tasks , properties of subgraphs are also worth predicting . Take company structure network as an example . The nodes are employees , and the edges between them represent cooperation relations . We want to predict the performance of a department , in other words , a subgraph in the network . Obviously , on the one hand , we need to consider the internal organization ( such as the collaboration within the department and the competence of individual employees ) . If the structure is disorganized or the employees are incompetent , we can expect the department to perform poorly . On the other hand , the external information of the department also deserves attention . A department is less likely to be productive if the company as a whole is facing bankruptcy . In contrast , close cooperation with ∗Corresponding author : Muhan Zhang other remarkable departments can be a sign of good performance . As illustrated by this example , a subgraph task is to predict the property of subgraphs in the whole graph . It needs to consider the topology both inside and outside the subgraph . It may also need to combine multi-level information of nodes , edges , higher-order substructures , and even the whole graph . Thus , a very general model is needed , and a natural idea is to extend ordinary GNNs to subgraph tasks . Figure 1 left shows a typical subgraph to predict . The target subgraph S is embedded in the whole graph and may have multiple connected components , and our task is to produce a subgraph representation , which can be used to predict properties of S. However , in experiments , we find that SubGNN ( Alsentzer et al. , 2020 ) , the current state-of-the-art method for subgraph tasks , only slightly outperforms plain GNNs which directly pool node embeddings within the subgraph as the subgraph representation . SubGNN replaces the message passing between nodes with a subgraph-level message passing framework and designs three channels . Each channel is further divided into an internal and a border module to aggregate subgraph features separately . Despite its performance , SubGNN is both space and time-consuming due to its lengthy precomputation . Furthermore , the units of message passing are subgraph patches sampled from the whole graph randomly , and there is no guarantee of the optimality of the samples , leading to high variance in performance and dubious robustness . Last but not least , the framework of SubGNN is overly complicated—many of its designs seem ad hoc or suboptimal , resulting in poor compatibility with recent advances in GNN research . Nevertheless , by comparing SubGNN with plain GNNs , we find that differentiating internal and external topology is crucial for subgraph tasks . Inspired by this insight , we introduce a max-zero-one labeling trick ( Zhang et al. , 2021 ) , which explicitly marks whether a node is within a subgraph or not , to augment GNNs and show that plain GNNs with this labeling trick are superior to SubGNN . Here we give a brief introduction to labeling trick . First proposed by Zhang et al . ( 2021 ) , labeling trick is a theoretical framework for using graph neural networks to produce multi-node representations , which show that producing expressive representations for high-order structures needs to capture the interaction among the different nodes within the structure . This theory shows that directly aggregating node representations to represent high-order structures is not expressive enough , and labeling trick can aid this problem . In implementation , a labeling trick assigns a label to each node and combines the node features and the labels as the new input node features to GNNs . Labeling trick have achieved great successes on graph representation learning in previous works . For example , the state-of-the-art link prediction method SEAL ( Zhang & Chen , 2018 ) gains better performance with carefully designed labeling trick . IDGNN ( You et al. , 2021 ) differentiates one center node from others , and Distance Encoding ( Li et al. , 2020 ) uses the distance to the target nodes to label other nodes , both of which gain improved performance on node , edge , and graph tasks . In this work , we for the first time introduce labeling trick to subgraph problems , and design an expressive and scalable labeling trick called max-zero-one . Max-zero-one is the first labeling trick that enables jointly predicting a batch of structures within the same graph . Present work We propose GLASS ( GNN with LAbeling trickS for Subgraph ) , a novel and simple graph neural network for subgraph tasks . To the best of our knowledge , GLASS is the first subgraph representation learning method using the ordinary message passing framework and a labeling trick . GLASS is more scalable , more expressive , and easier to implement and extend than the existing state-of-the-art method . Theoretically , we prove that GLASS is more expressive than plain GNNs , and can capture a range of important subgraph properties defined in SubGNN like density , cut ratio , border , positions , etc . Experiments on eight datasets show that GLASS achieves new state-of-theart performance . On synthetic datasets , GLASS with a single message passing layer beats SubGNN with three carefully designed channels by up to 48.6 % , illustrating the expressive power of nodelevel message passing augmented by labeling trick for subgraph tasks . On real-world datasets , GLASS also outperforms the strongest baseline SubGNN by up to 14.3 % . Moreover , training a GLASS model on average only takes 26 % time needed to train a SubGNN . With such a powerful model , we also prove the effectiveness of labeling trick on subgraph tasks . 2 RELATED WORK . Subgraph Representation Learning . Though some works have utilized subgraphs to perform other graph representation tasks ( Sun et al. , 2021 ; Wang et al. , 2021 ; Huang & Zitnik , 2020 ) or studied some specific tasks involving subgraphs ( Bordes et al. , 2014 ; Meng et al. , 2018 ; Ying et al. , 2020 ) , few works have studied the general subgraph representation learning problem . Alsentzer et al . ( 2020 ) introduced the problem formally and proposed SubGNN ( Alsentzer et al. , 2020 ) , the current state-of-the-art method , which samples patches from the whole graph and aggregates their features to produce subgraph representations . Before that , Sub2Vec ( Adhikari et al. , 2018 ) , designed for graph classification and community detection , samples random walks in subgraphs and feeds them to the language model Paragraph2vec ( Le & Mikolov , 2014 ) to generate embeddings for subgraphs . Labeling trick . SEAL ( Zhang & Chen , 2018 ) first introduces labeling trick to graph representation learning and applies them to link prediction . IDGNN ( You et al. , 2021 ) uses different message passing parameters for a target node and the other nodes , which is essentially a labeling trick assigning different labels to the target node and others . Distance Encoding ( Li et al. , 2020 ) uses distances to target nodes as node labels . Zhang et al . ( 2021 ) give a theoretical analysis of labeling trick and prove that they can produce the most expressive representations for substructures with a GNN expressive enough . These previous methods ignore subgraph tasks and have poor scalability due to the relabeling for every target substructure to predict . Besides these deterministic labeling methods , rGIN ( Sato et al. , 2021 ) assigns a random vector to each nodes as its label in each forward process . Similarly , GNN-RNI ( Abboud et al. , 2021 ) randomly initializes node embeddings and can approximate any functions mapping graphs to real numbers . Despite the theoretical power , random labels suffer from slow convergence and subpar performance . We also discuss other structural encoding methods in Appendix A.1 . Our GLASS is an application of labeling trick on subgraph tasks . Its success verifies the theory of using GNNs and labeling trick to produce multi-node representations . 3 PRELIMINARIES . Let G = ( V , E , X ) denote a graph with a finite node set V = { 1 , 2 , ... , n } , an edge set E ⊆ V × V and node feature matrix X , whose ith rowXi is the feature of node i. N ( v ) refers to the set of nodes adjacent to node v. S = ( VS , ES , XS ) is a subgraph of G if VS ⊆ V and ES ⊆ ( VS ×VS ) ∩E and XS is the stack of the rows of X corresponding to nodes in VS . In this paper , we focus on induced subgraphs , whose edge set ES = ( VS × VS ) ∩ E. Let S ⊆ G donate that S is a subgraph in G. Problem ( Subgraph Representation and Property Prediction ) . Given the whole graph G , its subgraphs S = { S1 , S2 , ... , Sn } and their target properties T = { tS1 , tS2 , ... , tSn } , the goal is to learn a representation vector hSi that can be used to predict tSi of Si . A Plain GNN for Subgraph Tasks . Message passing neural network ( MPNN ) ( Gilmer et al. , 2017 ) is a common framework of GNNs . A message passing layer aggregates embeddings from neighbors to update the representation of a node . The kth message passing layer can be formulated as follows . a ( k ) v = AGGREGATE ( k ) ( { h ( k−1 ) u |u ∈ N ( v ) } ) , ( 1 ) h ( k ) v = COMBINE ( k ) ( h ( k−1 ) v , a ( k ) v ) , ( 2 ) where h ( k ) v is the embedding of node v at the kth layer and h ( 0 ) v = Xv . The embeddings at the last layer can be used to predict node properties . As for edge or graph tasks , pooling the multiset of embeddings of nodes belonging to the edge or within the graph is a widely used method . Naturally , we can extend it to subgraphs , and the representation of a subgraph S can be learned by pooling the embeddings of nodes within the subgraph as follows , which we call a plain GNN . hS = READOUT ( { hu|u ∈ VS } ) . ( 3 )

Summary: In this paper, the authors propose GLASS, a GNN designed for subgraph tasks using labelling tricks. A simple zero-one trick can be used to differentiate nodes in and out of a target subgraphs. However, this would cause efficiency problem as different target subgraphs would require different labels. To fit the classic batch training, a batch of target subgraphs are used to label the nodes together. To resolve conflicts among subgraphs, a zero-max-one trick is proposed. GLASS is easier / intuitive to implement and more scalable. The zero-one trick version is also claimed to have more expressive power than plain GNNs and can capture the six properties in the state-of-the-art method SubGNN. Experiments have been conducted on both real-world and synthetic datasets with generally promising results.