Inductive Lottery Ticket Learning for Graph Neural Networks

1 INTRODUCTION . Graph neural networks ( GNNs ) ( Kipf & Welling , 2016 ; Veličković et al. , 2017 ; Zhou et al. , 2018 ) have become a prevalent solution for machine learning tasks on graph-structured data . Such success is usually ascribed to the powerful representation learning of GNN , which incorporates the graph structure into the representations , such as aggregating neural messages from the neighboring nodes to update the ego node ’ s representation ( Veličković et al. , 2017 ; Kipf & Welling , 2016 ) . As the field grows , there is an increasing need of building deeper GNN architectures ( Li et al. , 2020a ; 2021 ) on larger-scale graphs ( Hu et al. , 2020 ) . While deepening GNNs shows potentials on large-scale graphs , it also brings expensive computations due to the increased scale of graph data and model parameters , limiting their deployment in resource-constrained applications . Taking fraud detection in a transaction network as an example , the scale of user nodes easily reaches millions or even larger , making a GNN-detector model prohibitive to stack deep layers and predict the malicious behaviors in real time . Hence , pruning over-parameterized GNNs is of great need , which aims to answer the question : Can we co-sparsify the input graphs and model parameters , while preserving or even improving the performance ? Recently , the pruning approach for GNNs , UGS , ( Chen et al. , 2021 ) , is proposed to find graph lottery tickets ( GLTs ) — smaller subsets of model parameters and input graphs . At its core is Lottery Ticket Hypothesis ( LTH ) ( Frankle & Carbin , 2018 ) speculating that any dense , randomly-initialized neural network contains a sparse subnetwork , which can be trained independently to achieve a matching performance as the dense network . Specifically , UGS employs trainable masks on each edge in the input graph and each weight in the model parameters , to specify their importance . When training the model with the masks , the strategy of iterative magnitude-based pruning ( IMP ) ( Frankle & Carbin , 2018 ) is used to discard the edges and weights with the lowest mask values at each iteration . Despite effectiveness , there exist the following limitations : ( 1 ) UGS focuses solely on providing the transductive graph masks by generating a painstakingly customized mask for a single edge 1 individually and independently . That is , the edge masks are limited to the given graph , making UGS infeasible to be applied in the inductive setting , since the edge masks hardly generalize to unseen edges or entirely new graphs . ( 2 ) Applying a mask for each edge alone only provides a local understanding of the edge , rather than the global view of the entire graph ( e.g. , in node classification ) or multiple graphs ( e.g. , in graph classification ) . Moreover , the way of creating trainable edge masks will double the parameters of GNNs , which violates the purpose of pruning somehow . As a result , these edge masks could be suboptimal to guide the pruning . ( 3 ) The unsatisfactory graph pruning will negatively influence the pruning of model weights . Worse still , low-quality weight pruning will amplify the misleading signal of edge masks in turn . They influence each other and form a vicious circle . We ascribe all these limitations of UGS to its transductive nature . Hence , conducting the combinatorial pruning in the inductive setting is crucial to high-quality winning tickets . In this work , we emphasize the inductive nature within the combinatorial pruning of input graphs and GNN parameters and present our framework , Inductive Co-Pruning of GNNs ( ICPG ) . It is an extremely simple but effective pruning framework that is applicable to any GNN in both inductive and transductive settings . Specifically , for the input graphs , we design a generative probabilistic model ( termed AutoMasker ) , which learns to generate edge masks from the observed graphs . It is parameterized with an additional GNN-based encoder , whose parameters are shared across the population of observed graphs . As a consequence , AutoMasker is naturally capable to specify the significance for each edge and extract core-subgraphs from a global view of the entire observations . For the model parameters , we simply exploit the magnitude of a model weight to assess whether it should be pruned , rather than training an additional mask ( Chen et al. , 2021 ) . Having established the edge masks and weight magnitudes , we can obtain high-quality GLTs by pruning the lowest-mask edges and lowest-magnitude weights . Experiments on ten graph-classification and two node-classification benchmarks consistently validate our framework ICPG by identifying high-quality GLTs . The visualizations show that ICPG always retains decisive subgraphs , such as edges located on digital pixels in MNIST graphs , which further illustrates the effectiveness and rationality . Moreover , we inspect the graph- and GNN-level transferability of GLTs , which promises for deploying our ICPG in the pre-training and fine-tuning paradigm . Our main contributions can be summarized as follows : • We proposed a simple but effective pruning framework , ICPG , to prune the GNN model and input graphs simultaneously . It can identify high-quality GLTs in diverse tasks of graph representation learning , under both inductive and transductive settings . • For graph classification tasks , ICPG can locate GLTs from small-scale ( TUDataset ) , mediumscale ( Superpixel graphs ) to challenging large-scale datasets ( OGB ) with graph sparsity from 22.62 % ∼43.12 % and GNN sparsity from 67.23 % ∼91.41 % with no degradation on performance . • For node classification tasks , ICPG can effectively develop on both transductive learning ( Cora dataset ) and inductive learning ( PPI dataset ) , which identifies the GLTs with graph sparsity from 22.62 % ∼26.49 % and GNN sparsity from 67.23 % ∼73.79 % without sacrificing performance . • The proposed AutoMasker promises for both GNN-level and graph-level transferability , which can achieve comparable or even better performance as compared with the original full graphs . In-depth analyses with visual inspections further demonstrate the effectiveness and rationality . 2 RELATED WORK . Graph Neural Networks ( GNNs ) ( Kipf & Welling , 2016 ; Veličković et al. , 2017 ; Xu et al. , 2019 ; Ying et al. , 2018 ) have emerged as a powerful tool for learning the representation of graph-structured data . The great success mainly comes from the structure-aware learning , which follows the iterative message-passing scheme ( Veličković et al. , 2017 ) . Specifically , we denote an undirected graph by G = ( A , X ) with the node set V and edge set E . A ∈ { 0 , 1 } |V|×|V| is the adjacency matrix , where A [ i , j ] = 1 denotes the edge between node vi and node vj , otherwise A [ i , j ] = 0 . X ∈ R|V|×d is the matrix of node features , where xi = X [ i , : ] is the d-dimensional feature of the node vi ∈ V . Given a K-layer GNN , its k-th layer generates the representation of node vi as : a ( k ) i = AGGREGATION ( k ) ( { h ( k−1 ) j |j ∈ N ( i ) } ) , h ( k ) i = COMBINE ( k ) ( h ( k−1 ) i , a ( k ) i ) , ( 1 ) where h ( k ) i and a ( k ) i are the representation of node vi and the message aggregated from its neighbor nodes set N ( i ) , respectively ; the AGGREGATION and COMBINE operators are the message 2 passing and update functions , respectively . After propagating through K layers , we get the final representations of nodes , which facilitate downstream node-level tasks , such as node classification and link prediction . As for graph-level tasks like graph classification and graph matching , we further hire the READOUT function to generate the representation of the whole graph G : ZG = READOUT ( { h ( k ) i |vi ∈ V , k ∈ { 1 , · · · , K } } ) . ( 2 ) Various GNNs , such as GIN ( Xu et al. , 2019 ) and GAT ( Veličković et al. , 2017 ) , implement different AGGREGATION , COMBINE and READOUT functions , so as to refine the desired information from graph structures . Graph Sparsification or Sampling ( Voudigari et al. , 2016 ; Leskovec & Faloutsos , 2006 ) aims to find small core-subgraphs from the original graph , which can remain effective for graph learning tasks . Numerous strategies ( Zeng et al. , 2019 ; Franceschi et al. , 2019 ; Ying et al. , 2019 ; Hamilton et al. , 2017 ; Chen et al. , 2018 ) are proposed to achieve efficient graph learning . For example , GraphSAGE ( Hamilton et al. , 2017 ) samples and aggregates feature from a node ’ s local neighborhood . FastGCN ( Chen et al. , 2018 ) adopts the global importance sampling , which is more efficient for training . DropEdge ( Rong et al. , 2019 ) randomly drops edges from the input graph , which can be seen as a data augmenter . Another research line selects the core-subgraph in an optimization way . SGCN ( Li et al. , 2020b ) adopts the ADMM optimization algorithm to sparsify the adjacency matrix . UGS ( Chen et al. , 2021 ) utilizes a trainable mask for each edge to remove the potential task-irrelevant edges . Distinct from them , our AutoMasker predicts the importance score of each edge from a global view , thus having better generalization ability in the inductive settings . Lottery Ticket Hypothesis ( LTH ) ( Frankle & Carbin , 2018 ) states that a sparse subnetwork exists in a dense randomly-initialized network that can be trained to achieve comparable performance to the full models . LTH is explored in many fields such as computer vision and natural language processing ( Chen et al. , 2020 ; Ma et al. , 2021 ; Yin et al. , 2020 ; Liu et al. , 2019 ; Wang et al. , 2020 ; Savarese et al. , 2020 ) . Recently , Chen et al . ( 2021 ) extends the LTH to the GNNs and proposes the Graph Lottery Ticket ( GLT ) , which includes subgraph and subnetwork pairs that can be trained independently to reach comparable performance to the dense pairs . However , due to the transductive nature of graph-specific masks , UGS ( Chen et al. , 2021 ) can not develop on inductive learning settings , such as graph classification tasks . To tackle the dilemma , we propose AutoMasker to globally learn the significance of each edge from training graphs and predict importance scores for new coming graphs , which is graph-independent and inductive . 3 METHODOLOGY . Here we first formulate the task of learning graph lottery tickets , and then present our inductive strategy of co-pruning the input graphs and model weights . 3.1 INDUCTIVE GRAPH LOTTERY TICKET . Without loss of generality , we consider the inductive task of graph classification as an example . Given a GNN classifier f ( · , Θg0 ) , it starts from the randomly-initialized parameters Θg0 before training and arrives at the well-optimized parameters Θg after training . Once trained , it takes any graph G = ( A , X ) as the input and yields a probability distribution over C classes ŷ = f ( G , Θg ) . Wherein , G is associated with the adjacency matrix A and the pre-existing node features X . The goal of learning graph lottery tickets is to make the input graph G and the model weights Θg0 sparse to reduce the computational costs , while preserving the classification performance . Formally , it aims to generate two masks mG and mΘ , which are applied on G and Θg0 correspondingly , so as to establish the sparser input graph G′ = ( mG A , X ) and initialized weights Θ′g0 = mΘ Θg0 . Hereafter , through retraining the subnetwork f ( · , Θ′g0 ) on the sparse versions { G ′ } of training graphs , we can get the new converged parameters Θ′g . If the well-optimized subnetwork can achieve comparable performance with full graphs and network , we term the pair of G′ and f ( · , Θ′g0 ) as graph lottery tickets ( GLTs ) . Although a very recent study , UGS ( Chen et al. , 2021 ) , has proposed an approach to learn the GLTs , it focuses solely on the transductive setting but leaves the inductive setting untouched . Specifically , 3 it assigns a trainable mask to each edge of the input graph and trains such graph-specific masks individually and independently . As a consequence , these edge-dependent masks are limited to the given graph , hardly generalizing to unseen edges or entirely new graphs . Distinct from UGS , we aim to uncover GLTs in inductive learning settings .

Summary: This paper is on generalizing UGS, (Chen et al., 2021) from transductive setting to inductive setting. The main idea is to learn the importance of each edge using the MLP over the learnt embedding from GNN, and pruning them through sorting the importance score. And weight pruning is done just by sorting the magnitude of weights in GCN.