Condensing Graphs via One-Step Gradient Matching
Abstract
As training deep learning models on large dataset takes a lot of time and resources, it is desired to construct a small synthetic dataset with which we can train deep learning models sufficiently.
There are recent works that have explored solutions on condensing image datasets through complex bi-level optimization.
For instance, dataset condensation (DC) matches network gradients w.r.t.
large-real data and small-synthetic data, where the network weights are optimized for multiple steps at each outer iteration.
However, existing approaches have their inherent limitations: (1) they are not directly applicable to graphs where the data is discrete; and (2) the condensation process is computationally expensive due to the involved nested optimization.
To bridge the gap, we investigate efficient dataset condensation tailored for graph datasets where we model the discrete graph structure as a probabilistic model.
We further propose a one-step gradient matching scheme, which performs gradient matching for only one single step without training the network weights.
Our theoretical analysis shows this strategy can generate synthetic graphs that lead to lower classification loss on real graphs.
Extensive experiments on various graph datasets demonstrate the effectiveness and efficiency of the proposed method.
In particular, we are able to reduce the dataset size by 90 % while approximating up to 98 % of the original performance and our method is significantly faster than multi-step gradient matching (e.g. 15 × in CIFAR10 for synthesizing 500 graphs).
Our code is available at https://github.com/amazon-science/doscond .
Introduction
Graph-structured data plays a key role in various real-world applications.
For example, by exploiting graph structural information, we can predict the chemical property of a given molecular graph [1], detect fraud activities in a financial transaction graph [2], or recommend new friends to users in a social network [3].
Due to its prevalence, graph neural networks (GNNs) [4, 5, 6, 7] have been developed to effectively extract meaningful patterns from graph data and thus tremendously facilitate computational tasks on graphs.
Despite their effectiveness, GNNs are notoriously datahungry like traditional deep neural networks: they usually require massive datasets to learn powerful representations.
Thus, training GNNs is often computationally expensive.
Such cost even becomes prohibitive when we need to repeatedly train GNNs, e.g., in neural architecture search [8] and continual learning [9].
One potential solution to alleviate the aforementioned issue is dataset condensation or dataset distillation .
It targets at constructing a small-synthetic training set that can provide sufficient information to train neural networks [10, 11, 12, 13, 14, 15, 16].
In particular, one of the representative methods, DC [11], formulates the condensation goal as matching the gradients of the network parameters between small-synthetic and large-real training data.
It has been demonstrated that such a solution can greatly reduce the training set size of image datasets without significantly sacrificing model performance.
For example, using 100 images generated by DC can achieve 97 .
4 % test accuracy on MNIST compared with 99 . 6 % on the original dataset ( 60 , 000 images).
These condensed samples can significantly save space for storing datasets and speed up retraining neural networks in many critical applications, e.g., continual learning and neural architecture search.
In spite of the recent advances in dataset distillation/condensation for images, limited attention has been paid on domains involving graph structures.
To bridge this gap, we investigate the problem of condensing graphs such that GNNs trained on condensed graphs can achieve comparable performance to those trained on the original dataset.
However, directly applying existing solutions for dataset condensation [10, 11, 12, 13] to graph domain faces some challenges.
First, existing solutions have been designed for images where the data is continuous and they cannot output binary values to form the discrete graph structure.
Thus, we need to develop a strategy that can handle the discrete nature of graphs.
Second, they usually involve a complex bi-level problem that is computationally expensive to optimize: they require multiple iterations (inner iterations) of updating neural network parameters before updating the synthetic data for multiple iterations (outer iterations).
It can be catastrophically inefficient for learning pairwise relations for nodes, of which the complexity is quadratic to the number of nodes.
To address the aforementioned challenges, we propose an efficient condensation method for graphs, where we follow DC [11] to match the gradients of GNNs between synthetic graphs and real graphs.
In order to produce discrete values, we model the graph structure as a probabilistic graph model and optimize the discrete structures in a differentiable manner.
Based on this formulation, we further propose a one-step gradient matching strategy which only performs gradient matching for one single step.
Consequently, the advantages of the proposed strategy are twofold.
First, it significantly speeds up the condensation process while providing reasonable guidance for synthesizing condensed graphs.
Second, it removes the burden of tuning hyper-parameters such as the number of outer/inner iterations of the bi-level optimization as required by DC.
Furthermore, we demonstrate the effectiveness of the proposed one-step gradient matching strategy both theoretically and empirically.
Our contributions can be summarized as follows:
We study a novel problem of learning discrete synthetic graphs for condensing graph datasets, where the discrete structure is captured via a graph probabilistic model that can be learned in a differentiable manner.
We propose a one-step gradient matching scheme that significantly accelerates the vanilla gradient matching process.
Theoretical analysis is provided to understand the rationality of the proposed one-step gradient matching.
We show that learning with one-step matching produces synthetic graphs that lead to a small classification loss on real graphs. 3.
Extensive experiments have demonstrated the effectiveness and efficiency of the proposed method.
Particularly, we are able to reduce the dataset size by 90 % while approximating up to 98 % of the original performance and our method is significantly faster than multi-step gradient matching (e
g. 15 × in CIFAR10 for synthesizing 500 graphs).
The Proposed Framework
Before detailing the framework, we first introduce the main notations used in this paper.
We majorly focus on the graph classification task where the goal is to predict the labels of given graphs.
Specifically, we denote a graph dataset as T =
{ G 1 , . .
. , G N } with ground-truth label set Y .
Each graph in T is associated with a discrete adjacency matrix and a node feature matrix.
Let A ( i ) , X ( i ) represent the adjacency matrix and the feature matrix of i -th real graph, respectively.
Similarly, we use S =
{ G ′ 1 , . . .
, G ′ N ′ } and Y ′ to indicate the synthetic graphs and their labels, respectively.
Note that the number of synthetic graphs N ′ is essentially much smaller than that of real graphs N .
We use d and n to denote the number of feature dimensions and the number of nodes in each synthetic graph, respectively 1 .
Let C denote the number of classes and ℓ denote the cross entropy loss.
The goal of our work is to learn a set of synthetic graphs S such that a GNN trained on S can achieve comparable performance to the one trained on the much larger dataset T .
2.1 Gradient Matching as the Objective
Since we aim at learning synthetic graphs that are highly informative, one solution is to allow GNNs trained on synthetic graphs to imitate the training trajectory on the original large dataset.
Dataset condensation [11, 12] introduces a gradient matching scheme to achieve this goal.
Concretely, it tries to reduce the difference of model gradients w.r.t.
large-real data and small-synthetic data for model parameters at every training epoch.
Hence, the model parameters trained on synthetic data will be close to these trained on real data at every training epoch.
Let θ t denote the network parameters at the t -th epoch and f θ t indicate the neural network parameterized by θ t .
The condensation objective is expressed as:
[Equation]
where D ( · , · ) is a distance function, T is the number of steps of the whole training trajectory and opt θ ( · ) is the optimization operator for updating parameter θ .
Note that Eq.
(1) is a bi-level problem where we need to learn the synthetic graphs S at the outer optimization and update model parameters θ t at the inner optimization.
To learn synthetic graphs that generalize to a distribution of model parameters P θ 0 , we sample θ 0 ∼ P θ 0 and rewrite Eq.
(1) as:
[Equation]
Discussion.
The aforementioned strategy has demonstrated promising performance on condensing image datasets [11, 12].
However, it is not clear how to model the discrete graph structure.
Moreover, the inherent bi-level optimization inevitably hinders its scalability.
To tackle these shortcomings, we propose DosCond that models the structure as a probabilistic graph model and is optimized through one-step gradient matching.
In the following subsections, we introduce the details of DosCond .
2.2 Learning Discrete Graph Structure
For graph classification, each graph in the dataset is composed of an adjacency matrix and a feature matrix.
For simplicity, we use X ′ ∈ R N ′ × n × d to denote the node features in all synthetic graphs
S and A ′ ∈ { 0 , 1 } N ′ × n × n to indicate the graph structure information in S .
Note that f θ t can be instantiated as any graph neural network and it takes both graph structure and node features as input.
Then we rewrite the objective in Eq.
(2) as follows:
[Equation]
where we aim to learn both graph structure A ′ and node features X ′ .
However, Eq.
(3) is challenging to optimize as it requires a function that outputs binary values.
To address this issue, we propose to model the graph structure as a probabilistic graph model with Bernoulli distribution.
Note that in the following, we reshape A ′ from N ′ × n × n to N ′ × n 2 for the purpose of demonstration only.
Specifically, for each entry A ′ ij ∈ { 0 , 1 } in the adjacency matrix A ′ , it follows a Bernoulli distribution: P Ω ij ( A ′ ij ) =
A ′ ij σ ( Ω ij )
+ (1 − A ′ ij ) σ ( − Ω ij ) , where σ ( · ) is the sigmoid function; Ω ij ∈ R is the success probability of the Bernoulli distribution and also the parameter to be learned.
Since A ′ ij is independent of all other entries, the distribution of A ′ can be modeled as: P Ω ( A ′ ) =
(cid:81) N ′ i =1 (cid:81) n 2 j =1 P Ω ij (cid:0)
A ′ ij (cid:1) .
Then, the objective in Eq.
(2) needs to be modified to
[Equation]
With the new parameterization, we obtain a function that outputs discrete values
but it is not differentiable due to the involved sampling process.
Thus, we employ the reparameterization method [17], binary concrete distribution, to refactor the discrete random variable into a differentiable function of its parameters and a random variable with fixed distribution.
Specifically, we first sample α ∼ Uniform(0 , 1) , and edge weight A ′ ij ∈ [0 , 1] is calculated by:
[Equation]
where τ ∈ (0 , ∞ ) is the temperature parameter that controls the continuous relaxation.
As τ → 0 , the random variable A ′ ij smoothly approaches the Bernoulli distribution.
In other words, we have lim
[Equation]
A
[Equation]
= σ [Equation] Ω ij ) . While small τ is necessary for obtaining discrete samples, large τ is useful in getting large gradients as suggested by [17].
In practice, we employ an annealing schedule [18] to gradually decrease the value of τ in training.
With the reparameterization trick, the objective function becomes differentiable w.r.t.
Ω ij with well-defined gradients.
Then we rewrite our objective as:
[Equation]
2.3 One-Step Gradient Matching
The vanilla gradient matching scheme in Eq.
(2) presents a bi-level optimization problem.
To solve this problem, we need to update the synthetic graphs S at the outer loop and then optimize the network parameters θ t at the inner loop.
The nested loops heavily impede the scalability of the condensation method, which motivates us to design a new strategy for efficient condensation.
In this work, we propose a one-step gradient matching scheme where we only match the network gradients for the model initializations θ 0 while discarding the training trajectory of θ t .
Essentially, this strategy approximates the overall gradient matching loss for θ t with the initial matching loss at the first epoch, which we term as one-step matching loss .
The intuition is: the one-step matching loss informs us about the direction to update the synthetic data, in which, we have empirically observed a strong decrease in the cross-entropy loss (on real samples) obtained from the model trained on synthetic data.
Hence, we can drop the summation symbol (cid:80) T − 1 t =0 in Eq.
(6) and simplify Eq.
(6) as follows:
[Equation]
where we sample θ 0 ∼ P θ 0 and α ∼ Uniform(0 , 1) .
Compared with Eq.
(6), one-step gradient matching avoids the expensive nested-loop optimization and directly updates the synthetic graph
S .
It greatly simplifies the condensation process.
In practice, as shown in Section 3.3, we find this strategy yields comparable performance to its bi-level counterpart while enabling much more efficient condensation.
Next, we provide theoretical analysis to understand the rationality of the proposed one-step gradient matching scheme.
Theoretical Understanding.
We denote the cross entropy loss on the real graphs as ℓ T ( θ ) =
(cid:80) i ℓ i ( A ( i ) ,
X ( i ) , θ ) and that on synthetic graphs as ℓ S ( θ ) =
ℓ S ( A ′ ( i ) , X
( i ) , θ ) .
Let θ ∗ denote the optimal parameter and θ t be the parameter trained on S at the t -th epoch by optimizing ℓ S ( θ ) .
For notation simplicity, we assume that A and A ′ are already normalized.
The matrix norm ∥ · ∥ is the
Frobenius norm.
We focus on the GNN of Simple Graph Convolutions (SGC) [19] to study our problem since SGC has a simpler architecture but shares a similar filtering pattern as GCN.
Theorem 1 When we use a K -layer SGC as the GNN used in condensation, i.e., f θ ( A ( i ) , X ( i ) )
=
Pool ( A K ( i ) X ( i ) W 1 ) W 2 with θ =
[ W 1 ; W 2 ] and assume that all network parameters satisfy ∥ θ ∥ 2 ≤
M 2 ( M > 0) , we have
[Equation]
where γ i = 1 if we use sum pooling in f θ ; γ i = 1 n i if we use mean pooling, with n i as the number of nodes in the i -th synthetic graph.
We provide the proof of Theorem 1 in Appendix C.1.
Theorem 1 suggests that the smallest gap between the resulted loss (by training on synthetic graphs) and the optimal loss has an upper bound.
This upper bound depends on two terms: (1) the difference of gradients w.r.t.
real data and synthetic data and (2) the norm of input matrices.
Thus, the theorem justifies that reducing the gradient difference w.r.t real and synthetic graphs can help learn desirable synthetic data that preserves sufficient information to train GNNs well.
Based on Theorem 1, we have the following proposition.
Proposition 1 Assume the largest gradient gap happens at 0 -th epoch, i.e., ∥∇ θ ℓ T ( θ 0 ) − ∇ θ ℓ S
θ 0 )
∥
= max t ∥∇ θ ℓ T ( θ t ) − ∇ θ ℓ S
( θ t ) ∥ with t = 0 , 1 , . .
, T − 1 , we have min
[Equation]
We omit the proof for the proposition since it is straightforward.
The above proposition suggests that the smallest gap between the ℓ T ( θ t ) and ℓ T ( θ ∗ ) is bounded by the one-step matching loss and the norm ∥ 1 ⊤ A ′ K ( i )
X ′ ( i ) ∥ 2 .
As we will show in Section B.1, when using mean pooling, the second term tend to have a smaller scale than the first one and can be neglected; the second term matters more when we use sum pooling.
Hence, we solely optimize the one-step gradient matching loss for
GNNs with mean pooling and additionally include the second term (the norm of input matrices) as a regularization for GNNs with sum pooling.
As such, when we consider the optimal loss ℓ T ( θ ∗ ) as a constant, reducing the one-step matching loss indeed learns synthetic graphs that lead to a small classification loss on real graphs.
This demonstrates the rationality of one-step gradient matching theoretically.
Remark 1.
Note that the spectral analysis from [19] demonstrated that both GCN and SGC share similar graph filtering behaviors.
Thus practically, we extend the one-step gradient matching loss from K -layer SGC to K -layer GCN and observe that it works well under the non-linear scenario.
Remark 2.
While we focus on the graph classification task, it is straightforward to extend our framework to node classification.
We obtain similar conclusions for node classification as shown in
Theorem 2 in Appendix C.2 and achieve impressive empirical performance in Appendix B.2.
2.4 Final Objective and Training Algorithm
In this subsection, we describe the final objective function and the detailed training algorithm.
We note that the objective in Eq.
(6) involves two nested expectations, we adopt Monte Carlo to approximately optimize the objective function.
Together with one-step gradient matching, we have min
[Equation]
where K 1 is the number of sampled model initializations and K 2 is the number of sampled graphs.
We find that K 2 = 1 is able to yield good performance in our experiments.
Regularization.
In addition to the one-step gradient matching loss, we note that the proposed
DosCond can be easily integrated with various priors as regularization terms.
In this work, we focus on exerting sparsity regularization on the adjacency matrix, since a denser adjacency matrix will lead to higher cost for training graph neural networks.
Specifically, we penalize the difference of the sparsity between σ (Ω) and a given sparsity ϵ :
[Equation]
We initialize σ ( Ω ) and X ′ as randomly sampled training graphs and set ϵ to the average sparsity of initialized σ ( Ω ) so as to maintain a low sparsity.
On top of that, as we discussed earlier in Section 2.3, we include the following regularization for GNNs with sum pooling:
[Equation]
Training Algorithm.
We provide the details of our proposed framework in Algorithm 1 in Appendix A.3.
Specifically, we sample K 1 model initializations θ 0 to perform one-step gradient matching.
Following the convention in DC [11], we match gradients and update synthetic graphs for each class separately in order to make matching easier.
For class c , we first retrieve the synthetic graphs of that class, denoted as ( A ′ c , X ′ c , Y ′ c ) ∼ S , and sample a batch of real graphs ( A c , X c , Y c ) .
We then forward them to the graph neural network and calculate the one-step gradient matching loss together with the regularization term.
Afterwards, Ω and X ′ are updated via gradient descent.
It is worth noting that the training process for each class can be run in parallel since the graph updates for one class is independent of another class.
Comparison with DC.
Recall that the gradient matching scheme in DC involves a complex bi-level optimization.
If we denote the number of inner-iterations as τ i and that of outer-iterations as τ o , its computational complexity can be τ i ×
τ o of our method.
Thus DC is significantly slower than
DosCond .
In addition to speeding up condensation, DosCond removes the burden of tuning some hyper-parameters, i.e., the number of iterations for outer/inner optimization and learning rate for updating f θ , which can save us enormous training time when learning larger synthetic sets.
Comparison with Coreset Methods.
Coreset methods [20, 21] select representative data samples based on some heuristics calculated on the pre-trained embedding.
Thus, it requires training the model first.
Given the cheap cost on calculating and ranking heuristics, the major computational bottleneck for coreset method is on pre-training the neural network for a certain number of iterations.
Likewise, our proposed DosCond has comparable complexity because it also needs to forward and backward the neural network for multiple iterations.
Thus, their efficiency difference majorly depends on how many epochs we run for learning synthetic graphs in DosCond and for pre-training the model embedding in coreset methods.
In practice, we find that DosCond even requires less training cost than the coreset methods as shown in Section 3.2.
Experiment 3.1 Experimental settings
Datasets.
To evaluate the performance of our method, we use multiple molecular datasets from
Open Graph Benchmark (OGB) [22] and TU Datasets (DD, MUTAG and NCI1) [23] for graph-level property classification, and one superpixel dataset CIFAR10 [24].
We also introduce a real-world e-commerce dataset.
In particular, we randomly sample 1,109 sub-graphs from a large, anonymized internal knowledge graph.
Each sub-graph is created from the ego network of a random selected product on the e-commerce website.
We form a binary classification problem aiming at predicting the product category of the central product node in each sub-graph.
We use the public splits for OGB datasets and CIFAR10.
For TU Datasets and the e-commerce dataset, we randomly split the graphs into 80%/10%/10% for training/validation/test.
Detailed dataset statistics are shown in Appendix A.2.
We compare our proposed methods with four baselines that produce discrete structures: three coreset methods ( Random , Herding [20] and K-Center [25, 21]), and a dataset condensation method DCG [11]: randomly graphs the training
Evaluation Protocol.
To evaluate the effectiveness of the proposed method, we test the classification performance of GNNs trained with condensed graphs on the aforementioned graph datasets.
Concretely, it involves three stages: (1) learning synthetic graphs, (2) training a GCN on the synthetic graphs and (3) test the performance of GCN.
We first generate the condensed graphs following the procedure in Algorithm 1.
Then we train a GCN classifier with the condensed graphs.
Finally we evaluate its classification performance on the real graphs from test set.
For baseline methods, we first get the selected/condensed graphs and then follow the same procedure.
We repeat the generation process of condensed graphs 5 times with different random seeds and train GCN on these graphs with 10 different random seeds.
We report the mean and standard deviation of these results.
3.2 Performance with Condensed Graphs
Classification Performance Comparison.
To validate the effectiveness of the proposed framework, we measure the classification performance of GCN trained on condensed graphs.
Specifically, we vary the number of learned synthetic graphs per class in the range of { 1 , 10 , 50 } ( { 1 , 10 , 20 } for MUTAG and E-commerce) and train a GCN on these graphs.
Then we evaluate the classification performance of the trained GCN on the original test graphs.
Following the convention in OGB [22], we report the
ROC-AUC metric for ogbg-molbace, ogbg-molbbbp and ogbg-molhiv; for other datasets we report the classification accuracy (%).
The results are summarized in Table 1.
Note that the Ratio column presents the ratio of synthetic graphs to original graphs and we name it as condensation ratio ;
the
[Table]
Whole Dataset column shows the GCN performance achieved by training on the original dataset.
From the table
[Table]
we make four observations:
(a) The proposed DosCond consistently achieves better performance than the baseline methods under different condensation ratios and different datasets.
Notably, when generating only 2 graphs on ogbg-molbace dataset (0.2%), we achieve an ROC-AUC of 0.657 while the performance on full training set is 0.714, which means we approximate 92% of the original performance with only 0.2% data.
Likewise, we are able to approximate 96.5% of the original performance on ogbg-molhiv with 0.3% data.
By contrast, baselines underperform our method by a large margin.
Similar observations can be made on other datasets, which demonstrates the effectiveness of learned synthetic graphs in preserving the information of the original dataset.
(b) Increasing the number of synthetic graphs can improve the classification performance.
For example, we can approximate the original performance by 89%/93%/98% with 0.2%/2.1%/10.6% data on DD.
More synthetic samples indicate more learnable parameters that can preserve the information residing in the original dataset and present more diverse patterns that can help train GNNs better.
This observation is in line with our experimental results in Section 3.3.
(c) The performance on CIFAR10 is less promising due to the limit number of synthetic graphs.
We posit that the dataset has more complex topology and feature information and thus requires more parameters to preserve sufficient information.
However, we note that our method still outperforms the baseline methods especially when producing only 1 sample per class, which suggests that our method is much more data-efficient.
Moreover, we are able to promote the performance on CIFAR10 by learning a larger synthetic set as shown in Section 3.3.
(d) Learning both synthetic graph structure and node features is necessary for preserving the information in original graph datasets.
By checking the performance DCG, which only learns node features based on randomly selected graph structure, we see that DCG underperforms DosCond by a large margin in most cases.
This indicates that learning node features solely is sub-optimal.
Efficiency Comparison.
Since one of our goals is to enable scalable dataset condensation, we now evaluate the efficiency of DosCond .
We compare DosCond with the coreset method Herding, as it is less time-consuming than DCG and generally achieves better performance than other baselines.
We adopt the same setting as in Table 1: 1000 iterations for DosCond , i.e., K 1 = 1000 , andepochs (100 epochs for ogbg-molhiv) for pre-training the graph convolutional network as required by
Herding.
We also note that pre-training the neural network need to go over the whole dataset at every
[Figure]
epoch while DosCond only processes a batch of graphs.
In Table 2, we report the running time on an
NVIDIA V100 GPU for CIFAR10, ogbg-molhiv and DD.
We make three observations:
(a) DosCond can be faster than Herding.
In fact, DosCond requires less training time in all the cases except in DD with 50 graphs per class.
Herding needs to fully train the model on the whole dataset to obtain good-quality embedding, which can be quite time-consuming.
On the contrary, DosCond only requires matching gradients for K 1 initializations and does not need to fully train the model on the large real dataset.
(b) The running time of DosCond increases with the increase of the number of synthetic graphs N ′ .
It is because DosCond processes the condensed graphs at each iteration, of which the time complexity is O ( N ′ L ( n 2 d + nd 2 )) for an L -layer GCN.
Thus, the additional complexity depends on N ′ .
By contrast, the increase of N ′ has little impact on Herding since the process of selecting samples based on pre-defined heuristic is very fast.
(c)
The average nodes in synthetic graph n also impacts the training cost of DosCond .
For instance, the training cost on ogbg-molhiv ( n =26) is much lower than that on DD ( n =285), and the gap of cost between the two methods on ogbg-molhiv and DD is very different.
As mentioned earlier, it is because the complexity of the forward process in GCN is O ( N ′ L ( n 2 d + nd 2 )) for N ′ condensed graphs with node size of n .
3.3 Further Investigation
Increasing the Number of Synthetic Graphs.
We study whether the classification performance can be further boosted when using larger synthetic size.
Concretely, we vary the size of the learned graphs from 1 to 300 and report the results of absolute and relative accuracy w.r.t.
whole dataset training accuracy for CIFAR10 in Figure 1a.
It is clear to see that both Random and DosCond achieve better performance when we increase the number of samples used for training.
Moreover, our method outperforms the random baseline under different condensed dataset sizes.
Note that the performance gap between the two methods diminishes with the increase of the number of samples.
This is because the random baseline will finally approach the whole dataset training if we continue to enlarge the size of the condensed set, in which the performance can be viewed as the upper bound of DosCond .
Ablation Study.
We perform ablation study on the proposed one-step gradient matching and regularization terms.
We create an ablation of our method, namely DosCond-Bi , which adopts the vanilla gradient matching scheme that involves a bi-level optimization.
Without loss of generality, we compare the training time and classification accuracy of DosCond and DosCond-Bi in the setting of learning 50 graphs/class synthetic graphs on CIFAR10 dataset.
The results are summarized in
Figure 1b and we can see that DosCond needs approximately 5 minutes to reach the performance of
DosCond-Bi trained for 75 minutes, which indicates that DosCond only requires 6.7% training cost.
It further demonstrates the efficiency of the proposed one-step gradient matching strategy.
Next we study the effect of sparsity regularization on DosCond .
Specifically, we vary the sparsity coefficient β in the range of { 0 , 0 . 001 , 0 . 01 , 0 .
1 , 1 , 10 } and report the classification accuracy and graph sparsity on DD and NCI datasets in Figure 1c and 1d.
As shown in the figure, when β gets larger, we exert a stronger regularization on the learned graphs and the graphs become more sparse.
Furthermore, the increased sparsity does not affect the classification performance.
This is a desired
[Figure]
property since sparse graphs can save much space for storage and reduce training cost for GNNs.
We also remove the regularization of Eq.
(11) for ogbg-molhiv, we obtain the performance of 0.724/ 0.727/0.731 for 1/10/50 graphs per class, slightly worse than the one with this regularization.
Visualization.
We further investigate whether GCN can learn discriminative representations from the synthetic graphs learned by DosCond .
Specifically, we use t-SNE [28] to visualize the learned graph representation from GCN trained on different condensed graphs.
We train a GCN on graphs produced by different methods and use it to extract the latent representation for real graphs from test set.
Without loss of generality, we provide the t-SNE plots on DD dataset with 50 graphs per class in Figure 2.
It is observed that the graph representations learned with randomly selected graphs are mixed for different classes.
Similarly, DCG graphs also resulted in poorly trained GCN that outputs indistinguishable graph representations.
By contrast, the representations are well separated for different classes when learned with DosCond graphs (Figure 2c) and they are as discriminative as those learned on the whole dataset (Figure 2d).
This demonstrates that the graphs learned by DosCond preserve sufficient information of the original dataset so as to recover the original performance.
Related Work
Graph Neural Networks.
As the generalization of deep neural network to graph data, graph neural networks (GNNs) [4, 29, 5, 7, 19, 30, 31, 32, 33] have revolutionized the field of graph representation learning through effectively exploiting graph structural information.
GNNs have achieved remarkable performances in basic graph-related tasks such as graph classification [34, 35], link prediction [3] and node classification [4].
Recent years have also witnessed their great success achieved in many realworld applications such as recommender systems [3, 36], computer vision [37], drug discovery [38] and single-cell analysis [39].
GNNs take both adjacency matrix and node feature matrix as input and output node-level representations or graph-level representations.
Essentially, they follow a message-passing scheme [40] where each node first aggregates the information from its neighborhood and then transforms the aggregated information to update its representation.
Dataset Distillation & Dataset Condensation.
It is widely received that training neural networks on large datasets can be prohibitively costly.
To alleviate this issue, dataset distillation (DD) [10] aims to distill knowledge of a large training dataset into a small number of synthetic samples.
DD formulates the distillation process as a learning-to-learning problem and solves it through bi-level optimization.
To improve the efficiency of DD, dataset condensation (DC) [11, 12] is proposed to learn the small synthetic dataset by matching the gradients of the network parameters w.r.t.
largereal and small-synthetic training data.
It has been demonstrated that these condensed samples can facilitate critical applications such as continual learning [11, 12, 41, 42, 43], neural architecture search [13, 14, 44] and privacy-preserving scenarios [45].
Recently, following the gradient matching scheme in DC, GCond [46] proposes to condense a large-scale graph to a small graph for node classification.
Different from GCond , we aim to solve the challenge of learning discrete structure and we majorly target at graph classification.
Our method avoids the costly bi-level optimization and is much more efficient than the previous work.
A detailed comparison is included in Appendix B.2.
Conclusion
Training graph neural networks on a large-scale graph dataset consumes high computational cost.
One solution to alleviate this issue is to condense the large graph dataset into a small synthetic dataset.
In this work, we propose a novel framework DosCond that adopts a one-step gradient matching strategy to efficiently condenses real graphs into a small number of informative graphs with discrete structures.
We further justify the proposed method from both theoretical and empirical perspectives.
Notably, our experiments show that we are able to reduce the dataset size by 90% while approximating up to 98% of the original performance.
In the future, we plan to investigate interpretable condensation methods and diverse applications of the condensed graphs.
A Experimental Setup A.1 Parameter Settings.
When learning the synthetic graphs, we adopt 3-layer GCN with 128 hidden units as the model for gradient matching.
The learning rates for structure and feature parameters are set to 1.0 (0.01 for ogbg-molbace and CIFAR10) and 0.01, respectively.
We set K 1 to 1000 and β to 0.1.
Additionally, we use mean pooling to obtain graph representation for all datasets except ogbg-molhiv.
We use sum pooling for ogbg-molhiv as it achieves better classification performance on the real dataset.
During the test stage, we use GCN with the same architecture and we train the model for 500 epochs ( epochs for ogbg-molhiv) with an initial learning rate of 0.001.
A.2 Dataset Statistics
Dataset statistics for node classification and graph classification are shown in Table 3 and 4, respectively.
[Table] [Table]
A.3 Algorithm
We provide the details of our proposed framework in Algorithm 1.
Specifically, we sample K 1 model initializations θ 0 to perform one-step gradient matching.
Following the convention in DC [11], we match gradients and update synthetic graphs for each class separately in order to make matching easier.
For class c , we first retrieve the synthetic graphs of that class, denoted as ( A ′ c , X ′ c , Y ′ c ) ∼ S , and sample a batch of real graphs ( A c , X c , Y c ) .
We then forward them to the graph neural network and calculate the one-step gradient matching loss together with the regularization term.
Afterwards, Ω and X ′ are updated via gradient descent.
It is worth noting that the training process for each class can be run in parallel since the graph updates for one class is independent of another class.
B More Experiments B.1 Scale of the two terms in Eq. (9) .
As mentioned earlier in Section 2.3, the scale of the first term is essentially larger than the second term in Eq.
We now perform empirical study to verify this statement.
Since both terms contain the factor M , we simply drop it and focus on studying ℓ 1 = √ 2 ∥∇ θ ℓ T ( θ 0 ) − ∇ θ ℓ S
θ 0 ) ∥ and
1: Input: Training data T =
( A , X , Y ) 2: Required: Pre-defined condensed labels Y ′ , graph neural network f θ , temperature τ , desired sparsity ϵ , regularization coefficient β , learning rates η 1 , η 2 , number of epochs K 1 .
3: Initialize Ω , X ′ 4: for k = 0 , . . .
, K 1 − 1 do 5:
Sample θ 0 ∼ P θ6:
Sample α ∼ Uniform(0 , 1) 7:
Compute A ′ =
σ ((log α − log(1 − α ) + Ω ) /τ ) 8: for c = 0 , . . .
, C − 1 do 9:
Sample ( A c , X c , Y c ) ∼ T and ( A ′ c , X ′ c , Y ′ c ) ∼ S 10:
Compute ℓ T = ℓ ( f θ
( A c , X c ) , Y c ) 11:
Compute ℓ S = ℓ ( f θ
( A ′ c , X ′ c ) , Y ′ c ) 12:
Compute ℓ reg = max( (cid:80) i,j σ ( Ω ij ) − ϵ, 0) 13:
Update Ω ← Ω − η 1 ∇ Ω ( D ( ∇ θ 0 ℓ T , ∇ θ 0 ℓ S ) + βℓ reg ) 14:
Update X ′ ← X ′ − η 2 ∇ X ′
( D ( ∇ θ 0 ℓ T , ∇ θ 0 ℓ S )
+ βℓ reg ) 15: end for 16: end for 17: Return: ( Ω , X ′ , Y ′ )
ℓ 2 = 3 2 √ T · C − 1 CN ′
(cid:113)(cid:80)
i γ i ∥
1 ⊤ A ′ K ( i ) X ′ ( i ) ∥ 2 .
Specifically, we set T to 500 and N ′ to 50, and plot the changes of these two terms during the training process of DosCond .
The results on DD (with mean pooling) and ogbg-molhiv (with sum pooling) are shown in Figure 3.
We can observe that the scale of ℓ 1 is much larger than ℓ 2 at the first few epochs when using mean pooling as shown in Figure 3a.
By contrast, ℓ 2 is not negligible when using sum pooling as shown in Figure 3b
and it is desired to include it as a regularization term in this case.
These observations provide support for ours discussion of theoretical analysis in Section 2.3.
[Figure]
B.2 Node Classification
Next, we investigate whether the proposed method works well in node classification so as to support our analysis in Theorem 2 in Appendix C.2.
Specifically, following GCond [46], a condensation method for node classification, we use 5 node classification datasets: Cora, Citeseer, Pubmed [4], ogbn-arxiv [22] and Flickr [47].
The dataset statistics are shown in 4.
We follow the settings in
GCond to generate one condensed graph for each dataset, train a GCN on the condensed graph, and evaluate its classification performance on the original test nodes.
To adopt DosCond into node classification, we replace the bi-level gradient matching scheme in GCond with our proposed one-step gradient matching.
The results of classification accuracy and running time per epoch are summarized in Table 5.
From the table, we make the following observations:
(a) The proposed DosCond achieves similar performance as GCond and the performance is also comparable to the original dataset.
For example, we are able to approximate the original training
[Table]
performance by 99% with only 2.6% data on Cora.
It demonstrates the effectiveness of DosCond in the node classification case and justifies Theorem 2 from an empirical perspective.
(b) The training cost of DosCond is essentially lower than GCond as DosCond avoids the expensive bi-level optimization.
By examining their running time, we can see that DosCond is up totimes faster than GCond .
We further note that GCond produces weighted graphs which require storing the edge weights in float formats, while DosCond outputs discrete graph structure which can be stored as binary values.
Hence, the graphs learned by DosCond are more memory-efficient.
C Proofs C.1 Proof of Theorem
Let A ( i ) , X ( i ) denote the adjacency matrix and the feature matrix of i -th real graph, respectively.
We denote the cross entropy loss on the real samples as ℓ T ( θ ) =
(cid:80) i ℓ i ( A ( i ) ,
X ( i ) , θ ) and denote that on synthetic samples as ℓ S ( θ ) =
ℓ S ( A ′ ( i ) , X
( i ) , θ ) .
Let θ ∗ denote the optimal parameter and let θ t be the parameter trained on condensed data at t -th epoch by optimizing ℓ S ( θ ) .
For simplicity of notations, we assume A and A ′ are already normalized.
Part of the proof is inspired from the work [48].
Theorem 1 When we use a linearized K -layer SGC as the GNN used in condensation, i.e., f θ ( A ( i ) , X ( i ) )
=
Pool ( A K ( i ) X
( i ) W 1 )
W 2 with θ =
[ W 1 ; W 2 ] and assume that all network parameters satisfy ∥ θ ∥ 2 ≤ M 2 ( M > 0) , we have
[Equation]
where γ i = 1 if we use sum pooling in f θ ; γ i = 1 n i if we use mean pooling, with n i being the number of nodes in i -th synthetic graph.
We start by proving that ℓ T ( θ ) is convex and ℓ S ( θ ) is lipschitz continuous when we use f θ ( A ( i ) , X ( i ) )
=
Pool ( A K ( i ) X
( i ) W 1 ) W 2 as the mapping function.
Before proving these two properties, we first rewrite f θ ( A ( i ) , X ( i ) ) as:
[Equation]
pooling
[Equation]
[Equation]
where n is the number of nodes in A ( i ) and 1 is an n i × 1 matrix filled with constant one.
From the above equation we can see that f θ with different pooling methods only differ in a multiplication factor 1 n i .
Thus, in the following we focus on f θ with sum pooling to derive the major proof.
I.
For f θ with sum pooling: Substitute W for W 1 W 2 and we have f θ ( A ( i ) , X ( i ) )
= 1 ⊤ A K ( i ) X
( i ) W for the case with sum pooling.
Next we show that ℓ T ( θ ) is convex and ℓ S ( θ ) is lipschitz continuous when we use f θ ( A ( i ) , X ( i ) )
= 1 ⊤ A K ( i ) X
( i ) W with θ = W .
(a) Convexity of ℓ T ( θ ) .
From chapter 4 of the book [49], we know that softmax classification f ( W ) =
XW with cross entropy loss is convex w.r.t.
the parameters W .
In our case, the mapping function f
θ ( A ( i ) , X ( i ) )
= 1 ⊤ A K ( i ) X
( i ) W applies an affine function on XW .
Given that applying affine function does not change the convexity, we know that ℓ T ( θ ) is convex.
(b) Lipschitz continuity of ℓ S ( θ ) .
In [50], it shows that the lipschitz constant of softmax regression with cross entropy loss is C − 1 Cm ∥ X ∥ , where X is the input feature matrix, C is the number of classes and m is the number of samples.
Since ℓ S ( θ ) is cross entropy loss and f θ is linear, we know that the f θ is lipschitz continuous and it satisfies:
[Equation]
With (a) and (b), we are able to proceed our proof.
First, from the convexity of ℓ T ( θ ) we have
[Equation]
We can rewrite ∇ θ ℓ T ( θ t ) T ( θ t − θ ∗ ) as follows:
[Equation]
Given that we use gradient descent to update network parameters, we have ∇ θ ℓ S ( θ t ) =
1 η ( θ t − θ t +1 ) where η is the learning rate.
Then we have,
[Equation]
Combining Eq.
(15) and Eq.
(17) we have,
[Equation]
We sum up the two sides of the above inequality for different values of t ∈ [0 , T − 1] :
[Equation]
Since 12 η ∥ θ T − θ ∗ ∥ 2 ≥
[Equation]
we have
[Equation] [Equation]
As we assume that ∥ θ ∥ 2 ≤ M 2 , we have ∥ θ − θ ∗ ∥ 2 ≤ 2 ∥ θ ∥ 2 = 2 M 2 .
Then Eq.
(20) can be rewritten as,
[Equation]
Recall that ℓ S ( θ ) is lipschitz continuous as shown in Eq.
(14), and combine
[Equation]
Then we choose η =
[Equation]
and we can get:
[Equation]
II. For f θ with mean pooling
Following similar derivation as in the case of sum pooling, we have
[Equation] [Equation]
where n i is the number of nodes in i -th synthetic graph.
C.2 Theorem for Node Classification Case
We adopt similar notations for representing the data in node classification but note that there is only one graph for node classification task.
Let A ∈ { 0 , 1 } N × N , A ′ ∈ { 0 , 1 } N ′ × N ′ denote the adjacency matrix for real graph and synthetic graph, respectively.
Let X ∈ R N × d , X ′ ∈ R N ′ ×
d denote the feature matrix for real graph and synthetic graph, respectively.
We denote the cross entropy loss on the real samples as ℓ T ( θ ) and denote that on synthetic samples as ℓ S ( θ ) .
Theorem 2 When we use a K -layer SGC as the model used in condensation, i.e., f θ ( A , X , θ ) =
A K XW with θ = W and assume that all network parameters satisfy ∥ θ ∥ 2 ≤ M 2 ( M > 0) , we have
[Equation]
We start by proving that ℓ T ( θ ) is convex and ℓ S ( θ ) is lipschitz continuous when f θ ( A , X , θ )
=
A K XW .
(a) Convexity of ℓ T ( θ ) :
Similar to the graph classification case, the Hessian matrix of ℓ T ( θ ) in node classification is positive semidefinite and thus ℓ T ( θ ) is convex.
(b) Lipschitz continuity of ℓ S ( θ ) :
As shown in [50], the lipschitz constant of softmax regression with cross entropy loss is C − 1 Cm ∥ X ∥ with C being the number of classes and m being the number of samples.
Thus, we know that the lipschitz constant of ℓ S ( θ ) is C − 1 CN ′ ∥
A ′ K X ′ ∥ , which indicates ∇ θ ℓ S ( θ ) ≤ C − 1 CN ′
From the convexity of ℓ T ( θ ) , we still have the following inequality (see Eq.
Then recall that ℓ S ( θ ) is lipschitz continuous and ∇ θ ℓ S ( θ ) ≤
[Equation]
1 CN ′ ∥
A
[Equation]
K X ′ ∥ , and combine min
[Equation] [Equation]
Then we choose η =
[Equation]
∥ and we can get:
[Equation]