Data-Free Adversarial Knowledge Distillation for 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čković 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.

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 above-mentioned 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.

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.

The generator $\mathcal{G}$ is used to synthesize fake graphs that maximize the disagreement between the teacher $\mathcal{T}$ and the student $\mathcal{S}$. $\mathcal{G}$ takes $D$-dimensional vectors ${z}\in\mathbb{R}^{D}$ sampled from a standard normal distribution ${z}\sim\mathcal{N}(0,I)$ and outputs graphs. For each $z$, $\mathcal{G}$ outputs an object: ${F}\in\mathbb{R}^{N\times 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 $\sigma(\cdot)$ is the logistic sigmoid function. We transform the range of $A$ from $[0,1]$ to $\{0,1\}$ with a threshold $\tau$. If the element in $A$ is larger than $\tau$, it is set as 1, otherwise 0.

The loss function used for $\mathcal{G}$ is the same as that used for $\mathcal{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 $\mathcal{D}\left(\cdot\right)$ indicates the discrepancy between the teacher $\mathcal{T}$ and the student $\mathcal{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.

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

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.

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 REDDIT-BINARY. 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.

We adopt a generator with fixed architecture for all experiments. The generator takes a 32-dimensional vector sampled from a standard normal distribution $\boldsymbol{z}\sim\mathcal{N}(\mathbf{0},\boldsymbol{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\times T$-dimensional vectors and reshape them as node features $\mathbf{F}\in\mathbb{R}^{N\times 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.

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\in\left\{5,3,2,1\right\}$ and gradually reduce the number of hidden units $h\in\left\{128,64,32,16\right\}$. 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.

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 $\tau$ is empirically set to 0.5.

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).

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.