Graph Condensation for Graph Neural Networks

Given the prevalence of large-scale graphs in real-world applications , the storage and time for training neural models have raised increasing concerns . To alleviate the concerns , we propose and study the problem of graph condensation for graph neural networks ( GNNs ) . Specifically , we aim to condense the large , original graph into a small , synthetic and highly-informative graph , such that GNNs trained on the small graph and large graph have comparable performance . We approach the condensation problem by imitating the GNN training trajectory on the original graph through the optimization of a gradient matching loss and design a strategy to condense node futures and structural information simultaneously . Extensive experiments have demonstrated the effectiveness of the proposed framework in condensing different graph datasets into informative smaller graphs . In particular , we are able to approximate the original test accuracy by 95.3 % on Reddit , 99.8 % on Flickr and 99.0 % on Citeseer , while reducing their graph size by more than 99.9 % , and the condensed graphs can be used to train various GNN architectures . 1 INTRODUCTION . Many real-world data can be naturally represented as graphs such as social networks , chemical molecules , transportation networks , and recommender systems ( Battaglia et al. , 2018 ; Wu et al. , 2019b ; Zhou et al. , 2018 ) . As a generalization of deep neural networks for graph-structured data , graph neural networks ( GNNs ) have achieved great success in capturing the abundant information residing in graphs and tackle various graph-related applications ( Wu et al. , 2019b ; Zhou et al. , 2018 ) . However , the prevalence of large-scale graphs in real-world scenarios , often on the scale of millions of nodes and edges , poses significant computational challenges for training GNNs . More dramatically , the computational cost continues to increase when we need to retrain the models multiple times , e.g. , under incremental learning settings , hyperparameter and neural architecture search . To address this challenge , a natural idea is to properly simplify , or reduce the graph so that we can not only speed up graph algorithms ( including GNNs ) but also facilitate storage , visualization and retrieval for associated graph data analysis tasks . There are two main strategies to simplify graphs : graph sparsification ( Peleg & Schäffer , 1989 ; Spielman & Teng , 2011 ) and graph coarsening ( Loukas & Vandergheynst , 2018 ; Loukas , 2019 ) . Graph sparsification approximates a graph with a sparse graph by reducing the number of edges , while graph coarsening directly reduces the number of nodes by replacing the original node set with its subset . However , these methods have some shortcomings : ( 1 ) sparsification becomes much less promising in simplifying graphs when nodes are also associated with attributes as sparsification does not reduce the node attributes ; ( 2 ) the goal of sparsification and coarsening is to preserve some graph properties such as principle eigenvalues ( Loukas & Vandergheynst , 2018 ) that could be not optimal for the downstream performance of GNNs . In this work , we ask if it is possible to significantly reduce the graph size while providing sufficient information to well train GNN models . Motivated by dataset distillation ( Wang et al. , 2018 ) and dataset condensation ( Zhao et al. , 2021 ) which generate a small set of images to train deep neural networks on the downstream task , we aim to condense a given graph through learning a synthetic graph structure and node attributes . Correspondingly , we propose the task of graph condensation1 . It aims to minimize the performance gap between GNN models trained on a synthetic , simplified graph and the original training graph . In this work , we focus on attributed graphs and the node classification task . We show that we are able to reduce the number of graph nodes to as low as 0.1 % while training various GNN architectures to reach surprisingly good performance on the synthetic graph . For example , in Figure 1 , we condense the graph of the Reddit dataset with 153,932 training nodes into only 154 synthetic nodes together with their connections . In essence , we face two challenges for graph condensation : ( 1 ) how to formulate the objective for graph condensation tractable for learning ; and ( 2 ) how to parameterize the to-be-learned node features and graph structure . To address the above challenges , we adapt the gradient matching scheme in ( Zhao et al. , 2021 ) and match the gradients of GNN parameters w.r.t . the condensed graph and original graph . In this way , the GNN trained on condensed graph can mimic the training trajectory of that on real data . Further , we carefully design the strategy for parametrizations for the condensed graph . In particular , we introduce the strategy of parameterizing the condensed features as free parameters and model the synthetic graph structure as a function of features , which takes advantage of the implicit relationship between structure and node features , consumes less number of parameters and offers better performance . Our contributions can be summarized as follows : 1 . We make the first attempt to condense a large-real graph into a small-synthetic graph , such that the GNN models trained on the large graph and small graph have comparable performance . We introduce a proposed framework for graph condensation ( GCOND ) which parameterizes the condensed graph structure as a function of condensed node features , and leverages a gradient matching loss as the condensation objective . 2 . Through extensive experimentation , we show that GCOND is able to condense different graph datasets and achieve comparable performance to their larger counterparts . For instance , GCOND approximates the original test accuracy by 95.3 % on Reddit , 99.8 % on Flickr and 99.0 % on Citeseer , while reducing their graph size by more than 99.9 % . Our approach consistently outperforms coarsening , coreset and dataset condensation baselines . 3 . We make several important findings towards the condensation problem : ( a ) condensed graphs can generalize well to different GNN test models ; ( b ) condensing node features and structural information simultaneously can benefit the performance ; and ( c ) the condensed graphs are able to extract meaningful patterns hidden in the original data . 2 RELATED WORK . Dataset Distillation & Condensation . Dataset distillation ( DD ) ( Wang et al. , 2018 ; Bohdal et al. , 2020 ; Nguyen et al. , 2021 ) aims to distill knowledge of a large training dataset into a small synthetic dataset , such that a model trained on the synthetic set is able to obtain the comparable performance to that of a model trained on the original dataset . To improve the efficiency of DD , dataset condensation ( DC ) ( Zhao et al. , 2021 ; Zhao & Bilen , 2021 ) is proposed to learn the small synthetic dataset 1We aim to condense both graph structure and node attributes . A formal definition is given in Section 3. by matching the gradients of the network parameters w.r.t . large-real and small-synthetic training data . However , these methods are designed exclusively for image data and are not applicable to non-Euclidean graph-structured data where samples ( nodes ) are interdependent . In this work , we generalize the problem of dataset condensation to graph domain and we seek to jointly learn the synthetic node features as well as graph structure . Additionally , our work relates to coreset methods ( Welling , 2009 ; Sener & Savarese , 2017 ; Rebuffi et al. , 2017 ) , which seek to find informative samples from the original datasets . However , they rely on the presence of representative samples , and tend to give suboptimal performance . Graph Sparsification & Coarsening . Graph sparsification and coarsening are two means of reducing the size of a graph . Sparsification reduces the number of edges while approximating pairwise distances ( Peleg & Schäffer , 1989 ) , cuts ( Karger , 1999 ) or eigenvalues ( Spielman & Teng , 2011 ) while coarsening reduces the number of nodes with similar constraints ( Loukas & Vandergheynst , 2018 ; Loukas , 2019 ; Deng et al. , 2020 ) , typically by grouping original nodes into super-nodes , and defining their connections . Cai et al . ( 2021 ) proposes a GNN-based framework to learn these connections to improve coarsening quality . Huang et al . ( 2021 ) adopts coarsening as a preprocessing method to help scale up GNNs . Graph condensation also aims to reduce the number of nodes , but aims to learn synthetic nodes and connections in a supervised way , rather than unsupervised grouping as in these prior works . Graph pooling is also related to our work , but it targets at improving graph-level representation learning ( see Appendix D ) . Graph Neural Networks . Graph neural networks ( GNNs ) are a modern way to capture the intuition that inferences for individual samples ( nodes ) can be enhanced by utilizing graph-based information from neighboring nodes ( Kipf & Welling , 2016 ; Hamilton et al. , 2017 ; Klicpera et al. , 2018 ; Veličković et al. , 2018 ; Wu et al. , 2019b ; a ) . Due to their prevalence , various real-world applications have been tremendously facilitated including recommender systems ( Ying et al. , 2018a ; Fan et al. , 2019 ) , computer vision ( Li et al. , 2019 ) and drug discovery ( Duvenaud et al. , 2015 ) . Graph Structure Learning . Our work is also related to graph structure learning , which explores methods to learn graphs from data . One line of work ( Dong et al. , 2016 ; Egilmez et al. , 2017 ) learns graphs under certain structural constraints ( e.g . sparsity ) based on graph signal processing . Recent efforts aim to learn graphs by leveraging GNNs ( Franceschi et al. , 2019 ; Jin et al. , 2020 ; Chen et al. , 2020 ) . However , these methods are incapable of learning graphs with smaller size , and are thus not applicable for graph condensation . 3 METHODOLOGY . In this section , we present our proposed graph condensation framework , GCOND . Consider that we have a graph dataset T = { A , X , Y } , where A ∈ RN×N is the adjacency matrix , N is the number of nodes , X ∈ RN×d is the d-dimensional node feature matrix and Y ∈ { 0 , . . . , C − 1 } N denotes the node labels over C classes . Graph condensation aims to learn a small , synthetic graph dataset S = { A′ , X′ , Y′ } with A′ ∈ RN ′×N ′ , X′ ∈ RN ′×D , Y′ ∈ { 0 , . . . , C − 1 } N ′ and N ′ N , such that a GNN trained on S can achieve comparable performance to one trained on the much larger T . Thus , the objective can be formulated as the following bi-level problem , min S L ( GNNθS ( A , X ) , Y ) s.t θS = arg min θ L ( GNNθ ( A′ , X′ ) , Y′ ) , ( 1 ) where GNNθ denotes the GNN model parameterized with θ , θS denotes the parameters of the model trained on S , and L denotes the loss function used to measure the difference between model predictions and ground truth , i.e . cross-entropy loss . However , optimizing the above objective can lead to overfitting on a specific model initialization . To generate condensed data that generalizes to a distribution of random initializations Pθ0 , we rewrite the objective as follows : min S Eθ0∼Pθ0 [ L ( GNNθS ( A , X ) , Y ) ] s.t . θS = arg min θ L ( GNNθ ( θ0 ) ( A ′ , X′ ) , Y′ ) . ( 2 ) where θ ( θ0 ) indicates that θ is a function acting on θ0 . Note that the setting discussed above is for inductive learning where all the nodes are labeled and test nodes are unseen during training . We can easily generalize graph condensation to transductive setting by assuming Y is partially labeled .

Summary: The paper proposes and addresses the problem of graph condensation. In a nutshell, provided a large graph G, the scope of the paper is to propose a solution able to generate a smaller synthetic graph G’, which effectively allows to train Graph Neural Networks (GNNs) able to achieve similar performance as if they were trained on G. In order to do that, the authors propose to match the gradients produced by a specific GNN on mini-batches extracted from G and G’. This avoids the need of constructing a minimization problem based on an outer and inner training loop which would be computationally heavy to handle. The node-wise features X’ of the smaller graph G’ are considered as free parameters in the optimization process and the adjacency matrix A’ is instead inferred from X’ via a learnable MLP which is trained alongside X’ with an alternating optimization schema. The authors evaluate the quality of their approach on the Cora, Citeseer, OGBN-Arxiv and Reddit datasets, condensing and evaluating the new graphs with a variety of different architectures.