Data-Free Adversarial Knowledge Distillation for Graph Neural Networks

1 Introduction

An increasing number of machine learning tasks require dealing with a large amount of graph data, which capture rich and complex relationships among potentially billions of elements. Graph Neural Networks (GNNs) have become an effective way to address the graph learning problem by converting the graph data into a low dimensional space while keeping both the structural and property information to the maximum extent. Recently, the rapid evolution of GNNs has led to a growing number of new architectures as well as novel applications [Lu et al., 2021; Sun et al., 2018; Wu et al., 2021].

However, training a powerful GNN often requires heavy computation and storage. Hence it is hard to deploy them into resource-constrained devices, such as mobile phones. There has been a large literature [Bahri et al., 2021] aiming to compress and speed-up the cumbersome GNNs into lightweight ones. Among those methods, knowledge distillation [Hinton et al., 2015] is one of the most popular paradigms for learning a portable student model from the pre-trained complicated teacher by directly imitating its outputs.

∗Corresponding Author.

Knowledge distillation (KD) is proposed by Hinton et al. [Hinton et al., 2015] for supervising the training of a compact yet efficient student model by capturing and transferring the knowledge from a large complicated teacher model. KD has received significant attention from the community in the recent years [Yang et al., 2021; Gou et al., 2021]. Despite its successes, KD in its classical form has a critical limitation. It assumes that the real training data is still available in the distillation phase. However, in practice, the original training data is often unavailable due to privacy concerns. Besides, many large models are trained on millions or even billions of graphs [Lu et al., 2021]. While the pre-trained models might be made available for the community at large, making training data available also poses a lot of technical and policy challenges.

An effective way to avert the above-mentioned issue is using the synthetic graphs, i.e., data-free knowledge distillation [Lopes et al., 2017; Liu et al., 2021]. Just as "data free" implies, there is no training data. Instead, the data is reversely generated from the pre-trained models. Datafree distillation has received a lot of attention in the field of computer vision [Fang et al., 2019; Lopes et al., 2017; Fang et al., 2021], which is however rarely been explored in graph mining. Note that Deng et al. [Deng and Zhang, 2021] have made some pilot studies on this problem and proposed graph-free knowledge distillation (GFKD). Unfortunately, GFKD is not an end-to-end approach. It only takes the fixed teacher model into account, ignoring the information from the student model when generating graphs. Moreover, their method is based on the assumption that an appropriate graph usually has a high degree of confidence in the teacher model. In fact, the model maps the graphs from the data space to a very small output space, thus losing a large amount of information. Therefore, these generated graphs are not very useful for distilling the teacher model efficiently which leads to an unsatisfactory performance.

In this work, we propose a novel data-free adversarial knowledge distillation framework for GNNs (DFAD-GNN). DFAD-GNN uses a knowledge distillation method based on GAN [Goodfellow et al., 2014]. As illustrated in Figure 1, DFAD-GNN contains one generator and two discriminators: one fixed discriminator is the pre-trained teacher model, the other is the compact student model that we aim to learn. The generator generates graphs to help transfer teachers' knowledge to students. Unlike previous work [Deng and Zhang, 2021], our generator can fully utilize both the intrinsic statistics from the pre-trained teacher model and the customizable information from the student model, which can help generate high quality and diverse training data to improve the student model's generalizability. The contributions of our proposed framework can be summarized as follows: 1) We study a valuable yet intractable problem: how to distill a portable and efficient student model from a pre-trained teacher model when the original training data is not available; 2) We propose a novel data-free adversarial knowledge distillation framework for GNNs (DFAD-GNN) in order to train a compact student model using the generated graphs and a fixed teacher model. To the best of our knowledge, DFAD-GNN is the first end-to-end framework for data-free knowledge distillation on graph structured data; and 3) Extensive experiments demonstrate that our proposed framework significantly outperforms existing state-of-the-art data-free methods. DFAD-GNN can successfully distill a student model with 81.8%-94.7% accuracy of the teacher model on all six datasets.

2 Preliminary and RelatedWork

2.1 Graph Neural Networks

Graph Neural Networks (GNNs) have received considerable attention for a wide variety of tasks [Wu et al., 2020; Zhou et al., 2020]. Generally, GNN models can be unified by a neighborhood aggregation or message passing schema [Gilmer et al., 2017], where the representation of each node is learned by iteratively aggregating the embeddings ("message") of its neighbors.

As one of the most influential GNN models, Graph Convolutional Network (GCN) [Kipf and Welling, 2016] performs a linear approximation to graph convolutions. Graph Attention Network (GAT) [Veliˇckovi´c et al., 2017] introduces an attention mechanism that allows weighing nodes in the neighborhood differently during the aggregation step. GraphSAGE [Hamilton et al., 2017] is a comprehensive improvement on the original GCN which replaced full graph Laplacian with learnable aggregation functions. Graph Isomorphism Network (GIN) [Xu et al., 2018] uses a simple but expressive injective multiset function for neighbor aggregation. These GNNs above will be employed as our teacher models and student models in the experiments.

2.2 Knowledge Distillation

Knowledge distillation (KD) aims to transfer the knowledge of a (larger) teacher model to a (smaller) student model [Hinton et al., 2015; Wu et al., 2022]. It was originally introduced to reduce the size of models deployed on devices with limited computational resources. Since then, this line of research has attracted a lot of attention [Gou et al., 2021]. Recently, there are a few attempts that try to combine knowledge distillation with graph convolutional networks (GCNs). Yang et al. [Yang et al., 2021] proposed a knowledge distillation framework which can extract the knowledge of an arbitrary teacher model and inject it into a well-designed student model. Jing et al. [Jing et al., 2021] proposed to learn a lightweight student GNN that masters the complete set of expertise of multiple heterogeneous teachers. These works aim to improve the performance of student models on semi-supervised node classification task, rather than the graph classification task we considered in this work. Furthermore, although the abovementioned methods obtained promising results, they cannot be effectively launched without the original training dataset. In practice, the training dataset could be unavailable for some reasons, e.g. transmission limitations, privacy, etc. Therefore, it is necessary to consider a data-free approach for compressing neural networks.

Techniques addressing data-free knowledge distillation have relied on training a generative model to synthesize fake data. One recent work named graph-free KD (GFKD) [Deng and Zhang, 2021] has proposed a data-free knowledge distillation for graph neural network. GFKD learns graph topology structures for knowledge distillation by modeling them with a multinomial distribution. The training processes include two independent steps: (1) first, it uses a pre-trained teacher model to generate fake graphs; (2) afterwards, it uses these fake graphs to distill knowledge into the compact student model. However, those fake graphs are optimized by an invariant teacher model without considering the student model. Therefore, they are not very useful for distilling the student model efficiently. In order to generate high quality diverse training data to improve the student model's generalizability, we propose a data-free adversarial knowledge distillation framework for GNNs (DFAD-GNN). Our generator is trained end-to-end which not only uses the pre-trained teacher's intrinsic statistics but also obtains the discrepancy between teacher model and student model. We remark that the key differences between GFKD and our model lies in the training process and the generator.

2.3 Graph Generation

Data-free knowledge distillation involves the generation of training data. Motivated by the power of Generative Adversarial Networks (GANs) [Goodfellow et al., 2014], researchers have used them for generating graphs. Bojchevski et al. proposed NetGAN [Bojchevski et al., 2018], which uses the GAN framework to generate random walks on graphs. De Cao and Kipf proposed MolGAN [De Cao and Kipf, 2018], which generates molecular graphs using a simple multi-layer perceptron (MLP).

In this work, we build on a min-max game between two adversaries who try to optimize opposite loss functions. This approach is analogous to the optimization performed in GANs to train the generator and discriminator. The key difference is that GANs are generally trained to recover an underlying fixed data distribution. However, our generator chases a moving target: the distribution of data which is most indicative of the discrepancies of the current student model and its teacher model.

3 DFAD-GNN Framework

As shown in Figure 1, DFAD-GNN mainly consists of three components: one generator and two discriminators. One fixed discriminator is the pre-trained teacher model T , the other is the compact student model S that we aim to learn. More specifically, the generator G takes samples z from a prior distribution and generates fake graphs. Then the generated graphs are used to train a student model under the supervision of the teacher model.

3.1 Generator

The generator G is used to synthesize fake graphs that maximize the disagreement between the teacher T and the student S. G takes D-dimensional vectors z ∈ RD sampled from a standard normal distribution z ∼ N (0, I) and outputs graphs. For each z, G outputs an object: F ∈ RN×T that defines node features, where N is the node number and T is node feature dimension. Then we calculate adjacency matrix A as follows:

where σ(·) is the logistic sigmoid function. We transform the range of A from [0, 1] to {0, 1} with a threshold τ. If the element in A is larger than τ, it is set as 1, otherwise 0.

The loss function used for G is the same as that used for S, except that the goal is to maximize it. In other words, the student is trained to match the teacher's predictions and the generator is trained to generate difficult graphs for the student. The training process can be formulated as:

where D (·) indicates the discrepancy between the teacher T and the student S. If generator keeps generating simple and duplicate graphs, student model will fit these graphs, resulting in a very low model discrepancy between student model and teacher model. In this case, the generator is forced to generate difficult and different graphs to enlarge the discrepancy.

3.2 Adversarial Distillation

Overall, the adversarial training process consists of two stages: the distillation stage that minimizes the discrepancy D; and the generation stage that maximizes the discrepancy D, as shown in Figure 1. We detail each stage as follows.

Algorithm 1 DFAD-GNN

Input: A pre-trained teacher model,T (X; θt )

Output: A comparable student model,S (X; θs )

• 1: Randomly initialize a student model S (X; θs ) and a generator G (z; θg )
• 2: for Epochs do
• 3: // Distillation Stage
• 4: for k steps do
• 5: Generate graphs X from z with G (z; θg )
• 6: Calculate model discrepancy with LDIS
• 7: Update θs to minimize discrepancy with ∇θsD
• 8: end for
• 9: //Generation Stage
• 10: Generate graphs X from z with G (z; θg );
• 11: Calculate negative discrepancy with LGEN
• 12: Update θg to maximize discrepancy with ∇θg -D
• 13: end for

Distillation Stage

In this stage, we fix the generator G and only update the student S in the discriminator. We sample a batch of random noises z and construct fake graphs with generator G. Then each graph X is fed to both the teacher and the student model to produce the output qt and qs , where q is a vector indicating the scores of different categories.

In our approach, the choice of loss involves similar factors to those outlined in GANs: multiple works have discussed the problem of vanishing gradients as the discriminator becomes strong in case of GAN training [Arjovsky and Bottou, 2017]. Most prior work in model distillation optimized over the KullbackLeibler Divergence (KLD) and Mean Square Error (MSE) between the student and the teacher model. However, as the student model matches more closely the teacher model, these two loss functions tend to suffer from vanishing gradients. Specifically, back-propagating such vanishing gradients through the generator can harm its learning. For our approach, we minimize the Mean Absolute Error (MAE) between qt and qs , which provides stable gradients for the generator so that the vanishing gradients can be alleviated. In our experiment, we empirically find that this significantly improves student's performance over other possible losses. Now we can define the loss function for distillation stage as follows:

Generation Stage

The goal of the generation stage is to push the generation of new graphs. In this stage, we fix the two discriminators and only update the generator. We encourage the generator to produce more confusing training graphs. The loss function used for generator is the same as for student except that the goal is to maximize it:

With the generation loss, the error first back-propagates through the discriminator (the teacher and the student model), then through the generator to optimize it.

3.3 Optimization

The whole distillation process is summarized in Algorithm 1. DFAD-GNN trains the student and the generator by iterating over the distillation stage and the generation stage. Based on the learning progress of the student model, the generator crafts new graphs to further estimate the model discrepancy. The competition in this adversarial game drives the generator to discover more knowledge. In the distillation stage, we update the student model for k times so as to ensure its convergence. Note that compared with the conventional method GFKD [Deng and Zhang, 2021], the time complexity of DFAD-GNN mainly lies on the matrix multiplication of the generator, i.e., O(TN 2 ) where N is the number of nodes and T is the node feature dimension. Although our method has higher time complexity, the performances are much better. In addition, because there are not too many nodes in most real world graph-level applications (usually less than 100), we remark that our complexity is acceptable in practice.

4 Experiments

4.1 Datasets

We adopt six graph classification benchmark datasets including three bioinformatics graph datasets, i.e., MUTAG, PTC MR, and PROTEINS, and three social network graph datasets, i.e., IMDB-BINARY, COLLAB, and REDDITBINARY. The statistics of these datasets are summarized in Table 1. To remove the unwanted bias towards the training data, for all experiments on these datasets, we evaluate the model performance with a 10-fold cross validation setting, where the dataset split is based on the conventionally used training/test splits [Niepert et al., 2016; Zhang et al., 2018; Xu et al., 2018] with LIBSVM [Chang and Lin, 2011]. We report the average and standard deviation of validation accuracies across the 10 folds within the cross-validation.

4.2 Generator Architecture

We adopt a generator with fixed architecture for all experiments. The generator takes a 32-dimensional vector sampled from a standard normal distribution z ∼ N(0, I). We process it with a 3-layer MLP of [64,128,256] hidden units respectively, tanh is taken as the activation function. Finally, a linear layer is used to map the 256-dimensional vectors to N × T-dimensional vectors and reshape them as node features F ∈ RN×T . Throughout our experiments, we take the average number of nodes in the training data as N and test the effect of N in the ablation experiment.

4.3 Teacher/Student Architecture

To demonstrate the effectiveness of our proposed framework, we consider four GNN models as teacher and student models for a thorough comparison, including: GIN, GCN, GAT and GraphSAGE. Although the GNN model does not always require a deep network to achieve good results, however, from Appendix B, there is no fixed layer and hidden units that can make six datasets achieve the best performance on four different models. For fair comparison, we use 5 layers with 128 hidden units for teacher models. For the student model, we conduct experiments to gradually reduce the number of layers l ∈ {5, 3, 2, 1} and gradually reduce the number of hidden units h ∈ {128, 64, 32, 16}. We use a graph classifier layer which first builds a graph representation by averaging all node features extracted from the last GNN layer and then passing this graph representation to an MLP.

4.4 Implementation

For training, we use Adam optimizer with weight decay 5e-4 to update student models. The generator is trained with Adam without weight decay. Both student and generator are using a learning rate scheduler that multiplies the learning rate by a factor 0.3 at 10%, 30%, and 50% of the training epochs. The number of updates k of the student model in Algorithm 1 is set to 5. The threshold τ is empirically set to 0.5.

4.5 Baselines

We compare with the following baselines to demonstrate the effectiveness of our proposed framework.

Teacher: The given pre-trained model which serves as the teacher in the distillation process.

KD: The generator is removed, and the student model is trained on 100% original training data in our framework.

RANDOM: The generator's parameters are not updated and the student model is trained on the noisy graphs generated by the randomly initialized generator.

GFKD: GFKD is a data-free KD for GNNs by modeling the topology of graph with a multinomial distribution [Deng and Zhang, 2021].

4.6 Experimental Results

We have pre-trained all datasets on GCN, GIN, GAT and GraphSAGE with 5 layers and 128 hidden units (5-128 for short), and found that GIN performs best on all datasets (Detailed experimental results can be found in the Appendix B). Therefore, we adopt GIN as the teacher model in Table 2. We choose two representative architecture 1-128 and 5-32 for four kinds of student models (More experiments with other architectures can be found in the Appendix D).

From Table 2, it can be observed that KD's performance is very close to even outperforms the teacher model. That's because KD is a data-driven method which uses the same training data as the teacher model for knowledge distillation. This also implies that the loss function of our DFAD-GNN is very effective in distilling knowledge from the teacher model to the student model, as we apply the same loss function in both KD and our DFAD-GNN.

We also observe that RANDOM delivers the worst performance as the generator is not updated during the training process, thus the generator will not be able to generate difficult graphs as the student model progresses. Consequently, the student model fail to learn enough knowledge from the teacher, resulting in poor results.

In terms of the efficacy of our DFAD-GNN, Table 2 shows that our DFAD-GNN consistently outperforms the recent data-free method GFKD [Deng and Zhang, 2021]. We conjecture the potential reason that DFAD-GNN can significantly outperform GFKD is the teacher encodes the distribution characteristics of the original input graphs under its own feature space. Simply inverting graphs in GFKD tends to overfit to the partial distribution information stored in this teacher model. As a consequence, their generated fake graphs are lacking of generalizability and diversity. In contrast, our generated graphs are more conducive to transferring the knowledge of the teacher model to the student model.

In terms of the stability, it can be seen from Table 2 that the standard deviation of our DFAD-GNN is the smallest among all the data-free baselines and across all datasets, indicating that our model can obtain relatively stable prediction results.

Another interesting observation is that the performance of the compressed model is not necessarily worse than the more complex model. As can be seen from Table 2, that performance of a more compressed student model with 5-32 is not necessarily worse than the student model with 1-128. Therefore, we speculate that the performance of the student model may have no obvious relationship with the degree of model compression, which requires further investigation.

4.7 Model Analysis

Model Comparison. Here, we select PROTEINS and COLLAB with the most training data on molecular datasets and social datasets respectively for cross-training among the four kinds of models. It can be seen from Figure 2, when GIN is used as a teacher model, the overall performance of the student model is better, no matter which type of the student model is adopted. When GIN is used as a student model, the performance of the student model is significantly improved compared to other student models. We speculate there may be two reasons: (1) Under the same number of layers and hidden units, the GIN model has more parameters than GCN and GraphSAGE, so GIN owns more powerful learning capability. Note that although GAT has more parameters than GIN, it may be better at calculating node attention weights and then performing node classification tasks; (2) GIN is proposed to solve the problem of graph isomorphism. For our small molecule graphs and social network graphs, GIN is more powerful than other models in graph classification tasks.

Choice of Loss Function. The choice of loss is of paramount importance to a successful distillation. We explore five potential loss functions, including logit-MAE (calculate MAE with pre-softmax activations, L-MAE for short), softmax-MAE (calculate MAE with softmax outputs, S-MAE for short), MSE, KLD, and Cross Entropy (CE for short). These losses are commonly used in the knowledge distillation literature [Gou et al., 2021]. Figure 3 shows that using the MAE achieves significantly better test accuracies compared to other loss functions. Generally speaking, it is better to calculate MAE before softmax, because logits contain more information. As mentioned earlier, when the student model matches more closely to the teacher model, other loss functions tend to suffer from vanishing gradients [Fang et al., 2019]. Specifically, backpropagating such vanishing gradients through the generator can harm its learning. To prevent gradients from vanishing, we use the MAE computed with the teacher and student logits.

Percentage of Training Data. Although in the above reported results, we assume that the student cannot obtain any training data. However, in practice, student may have partial access to the training data. To reflect the practical scenario of knowledge distillation, we conduct extra experiments on PTC MR and IMDB-B, by varying the percentage of training data from 0.5% to 100%, while keeping other hyperparameters the same as in the previous experiments. As illustrated in Figure 4, PTC MR requires 17.1% of real data to achieve our results, while IMDB-B requires only 5.6%. This is because PTC MR has less data, thus the percentage of real data required is higher than IMDB-B.

Number of generated nodes. To explore the influence of node number N on the model performance, we conduct experiments with varying sizes of N on PTC MR and IMDB-B. It can be seen from Figure 5 that the model performs better when N takes the value near the average number of nodes in the training set. Far from the average number, the performance will decrease correspondingly due to a large deviation from real data.

Visualization of Generated Graphs. The generated graphs and real graphs of IMDB-B and PTC MR are shown in Figure 6. Although the generated graphs are not exactly the same as the real graphs, they can be used to generate a student model with relatively good performance.

5 Conclusion

This paper introduces a data-free adversarial knowledge distillation framework on graph neural network for model compression. Without any access to real data, we successfully reduce the discrepancy and obtain a student model with relatively good performance. Our extensive experiments on graph classification demonstrate that our framework can be effectively applied to different network architectures. In the future, we will extend this work to multi-teacher scenarios and continue to explore how to generate more complex graphs under various generator structures.

Acknowledgments

This work is supported in part by the National Natural Science Foundation of China (No. 62192784, U20B2045, 62172052, 61772082, 62002029, U1936104). It is also supported in part by The Fundamental Research Funds for the Central Universities 2021RC28.