FreeKD: Free-direction Knowledge Distillation for Graph Neural Networks

Graph neural networks have achieved promising results in processing graph data, whose goal is to learn node embeddings by aggregating nodes’ neighbor information. In recent years, lots of GNNs have been proposed (Kipf and Welling, 2016; Hamilton et al., 2017; Veličković et al., 2017). For instance, GCN (Kipf and Welling, 2016) designed a convolutional neural network architecture for graph data. GraphSAGE (Hamilton et al., 2017) proposed an efficient sample strategy to aggregate neighbor nodes. GAT (Veličković et al., 2017) applied a self-attention mechanism to GNN to assign different weights to different neighbors. SGC (Wu et al., 2019) simplified GCN by removing nonlinearities and weight matrices between consecutive convolutional layers. ROD (Zhang et al., 2021) proposed an ensemble learning based GNN model to fuse the knowledge in multiple hops. APPNP (Klicpera et al., 2018) analyzed the relationship between GCN and PageRank (Page et al., 1999), and proposed a propagation model combined with a personalized PageRank. Cluster-GCN (Chiang et al., 2019) built an efficient model based on graph clustering (Schaeffer, 2007). Being orthogonal to the above approaches developing different powerful GNN models, we concentrate on developing a new knowledge distillation framework on the basis of various GNNs.

Knowledge distillation (KD) has been widely studied in computer vision (Chen et al., 2017; Liu et al., 2018), natural language processing (Kim and Rush, 2016; liu, 2019), etc. Recently, a few KD methods have proposed for GNNs. LSP (Yang et al., 2020) transferred the topological structure knowledge from a pre-trained deeper teacher GNN to a shallower student GNN. CPF (Yang et al., 2021) designed a student architecture that is a combination of a parameterized label propagation and MLP layers. GFKD (Deng and Zhang, 2021) proposed a method to generate fake graphs and distilled knowledge from a teacher GNN model without any training data involved. The authors in (Yan et al., 2020) designed an efficient GNN model by utilizing the information from peer nodes to model the local structure explicitly and distilling the neighbor structure information from a deeper GNN implicitly. The work in (Chen et al., 2021a) studied a self-distillation framework, and proposed an adaptive discrepancy retaining regularizer to empower the transferability of knowledge. RDD (Zhang et al., 2020) was a semi-supervised knowledge distillation method for GNNs. It online learnt a complicated teacher GNN model by ensemble learning, and distilled knowledge from the generated teacher model into the student model. Different from them, we focus on studying a new free-direction knowledge distillation architecture, with the purpose of dynamically exchanging knowledge between two shallower GNNs.

Reinforcement learning aims at training agents to make optimal decisions by learning from interactions with the environment (Arulkumaran et al., 2017). Reinforcement learning mainly has two genres (Arulkumaran et al., 2017): value-based methods and policy-based methods. Value-based methods estimate the expected reward of actions (Mnih et al., 2015), while policy-based methods take actions according to the output probabilities of the agent (Williams, 1992). There is also a hybrid of this two genres, called the actor-critic architecture (Haarnoja et al., 2018). The actor-critic architecture utilizes the value-based method as a value function to estimate the expected reward, and employs the policy-based method as a policy search strategy to take actions. Until now, reinforcement learning has been taken as a popular tool to solve various tasks, such as recommendation systems (Zheng et al., 2018), anomaly detection (Oh and Iyengar, 2019), multi-label learning (Chen et al., 2018b), etc. In this paper, we explore reinforcement learning for graph data based knowledge distillation.

Let $\mathbf{G}=(\mathbf{V},\mathbf{E},\mathbf{X})$ denote a graph, where $\mathbf{V}$ is the set of nodes and $\mathbf{E}$ is the set of edge. $\ \mathbf{X}\in\mathbb{R}^{N\times d}$ is the feature matrix of nodes, where $N$ is the number of nodes and $d$ is the dimension of node features. Let $\mathbf{x}_{i}$ be the feature representation of node $i$ and $y_{i}$ be its class label. The neighborhood set of node $i$ is ${\mathcal{N}(i)}=\{\ j\in\mathbf{V}\ |\ (i,j)\in\mathbf{E}\ \}$. Currently, graph neural networks (GNNs) have become one of the most popular models for handling graph data. GNNs can learn the embedding $\mathbf{h}_{i}^{(l)}$ for node $i$ at the $l$-th layer by the following formula:

where $AGGREGATOR$ is an aggregation function, and it can be defined in many forms, e.g., mean aggregator (Hamilton et al., 2017). $\mathbf{W}^{(l)}$ is the learnt parameters in the $l$-th layer of the network. The initial feature of each node $i$ can be used as the input of the first layer, i.e., $\mathbf{h}_{i}^{(0)}=\mathbf{x}_{i}$.

Being orthogonal to those works developing various GNN models, our goal is to explore a new knowledge distillation framework for promoting the performance of GNNs, while addressing the issue involved because of producing a deeper teacher GNN model in the existing KD methods.

As shown in Figure 1, we observe typical GNN models often have different performances at different nodes during training. Based on this observation, we intend to dynamically exchange useful knowledge between two shallower GNNs, so as to benefit from each other. However, a challenging problem is attendant upon that: how to decide the directions of knowledge distillation for different nodes during training. To address this, we propose to manage the directions of knowledge distillation via reinforcement learning, where we regard the directions of knowledge transfer for different nodes as a sequential decision making problem (Pednault et al., 2002). Consequently, we propose a free-direction knowledge distillation framework via a hierarchical reinforcement learning, as shown in Figure 2. In our framework, the hierarchical reinforcement learning can be taken as a reinforced knowledge judge that consists of two levels of actions: 1) Level 1, called node-level action, is used to decide the distillation direction of each node for propagating the soft label; 2) Level 2, called structure-level action, is used to determine which of the local structures generated via node-level actions to be propagated.

Specifically, the reinforced knowledge judge (we call it agent for convenience) interacts with the environment constructed by two GNN models in each iteration, as in Figure 2. It receives the soft labels and cross entropy losses for a batch of nodes, and regards them as its node-level states. The agent then samples sequential node-level actions for nodes according to a learned policy network, where each action decides the direction of knowledge distillation for propagating node-level knowledge. Then, the agent receives the structure-level states and produces structure-level actions to decide which of the local structures generated on the basis of node-level actions to be propagated. After that, the two GNN models are trained based to the agent’s actions with a new loss function. Finally, the agent calculates the reward for each action to train the policy network, where the agent’s target is to maximize the expected reward. This process is repeatedly iterated until convergence.

We first give some notations for convenient presentation, before introducing how to distill both node-level and structure-level knowledge. Let $\Phi$ and $\Psi$ denote two GNN models, respectively. $\mathbf{h}_{i}^{\Phi}$ and $\mathbf{h}_{i}^{\Psi}$ denote the learnt representations of node $i$ obtained by $\Phi$ and $\Psi$, respectively. Let $\mathbf{p}_{i}^{\Phi}$ and $\mathbf{p}_{i}^{\Psi}$ be the predicted probabilities of the two GNN models for node $i$ respectively. We regard them as the soft labels. In addition, $L_{CE}^{\Phi}(i)$ and $L_{CE}^{\Psi}(i)$ denote the cross entropy losses of node $i$ in $\Phi$ and $\Psi$, respectively.

In this section, we introduce our reinforcement learning based strategy to dynamically distill node-level knowledge between two GNN models.

As we know, the structure information is important for graph learning. Thus, we attempt to dynamically transfer structure-level knowledge between two GNNs. It is worth noting that we don’t propagate all neighborhood information of one node as structure-level knowledge. Instead, we propagate a neighborhood subset of the node, which is comprised of agent-selected nodes. This is because we think agent-selected nodes contain more useful knowledge. We take Figure 3(b) as an example to illustrate it. $v_{1}$ is an agent-selected node in $\Phi$. When transferring its local structure information to $\Psi$, we only transfer the local structure composed of $\{v_{1},v_{4},v_{5}\}$. In other words, the local structure of node $v_{1}$ we consider to transfer is made up of agent-selected nodes. We call it agent-selected neighborhood set. Moreover, considering the knowledge of the local structure in graphs is not always reliable (Zhang et al., 2020; Chen et al., 2021b), we design a reinforcement learning based strategy to distinguish which of the local structures to be propagated. Next, we introduce it in detail.

In this section, we introduce the optimization procedure of our method. The detailed training procedure is in Appendix A.1.

In this subsection, we evaluate our method using three popular GNN models, GCN (Kipf and Welling, 2016), GAT (Veličković et al., 2017), and GraphSAGE (Hamilton et al., 2017). We arbitrarily select two networks from the above three models as our basic models $\Phi$ and $\Psi$, and perform our method FreeKD , enabling them to learn from each other. Note that we do not perform GCN on the PPI dataset, because of the inductive setting.

Table 1 and Table 2 report the experimental results. As shown in Table 1 and Table 2, our FreeKD can consistently promote the performance of the basic GNN models in a large margin on all the datasets. For instance, our method can achieve more than $4.5\%$ improvement by mutually learning from two GCN models on the Chameleon dataset, compared with the single GCN model. In summary, for the transductive learning tasks, our method improves the performance by 1.01% $\sim$ 1.97% on the Cora and Citeseer datasets and 1.08% $\sim$ 4.61% on the Chameleon and Texas datasets, compared with the corresponding GNN models. For the inductive learning task, our method improves the performance by 1.31% $\sim$ 3.11% on the PPI dataset dataset. In addition, we observe that two GNN models either sharing the same architecture or using different architectures can both benefit from each other by using our method, which shows the efficacy to various GNN models.

Since our method is related to knowledge distillation, we also compare with the existing knowledge distillation methods to further verify effectiveness of our method. In this experiment, we first compare with three traditional knowledge distillation methods, KD (Hinton et al., 2015), LSP (Yang et al., 2020), CPF (Yang et al., 2021) distilling knowledge from a deeper and stronger teacher GCNII model (Chen et al., 2020) into a shallower student GAT model. The structure details of GCNII and GAT could be found in the Appendix A.3. In addition, we also compare with an ensemble learning method, RDD (Zhang et al., 2020), where a complex teacher network is generated by ensemble learning for distilling knowledge. Finally, we take GNN-SD (Chen et al., 2021a) as another baseline, which distills knowledge from shallow layers into deep layers in one GNN. For our FreeKD, we take two GAT sharing the same structure as the basic models.

Table 3 lists the experimental results. Surprisingly, our FreeKD perform comparably or even better than the traditional knowledge distillation methods (KD, LSP, CPF) on all the datasets. This demonstrates the effectiveness of our method, as they distill knowledge from the stronger teacher GCNII while we only mutually distill knowledge between two shallower GAT. In addition, our FreeKD consistently outperforms GNN-SD and RDD, which further illustrates the effectiveness of our proposed FreeKD.

We perform ablation study to verify the effectiveness of the components in our method. We use GCN as the basic models $\Phi$ and $\Psi$ in our method, and conduct the experiments on two datasets of different domains, Chameleon and Cora. When setting $\rho=0$, this means that we only transfer the node-level knowledge. We denote it FreeKD-node for short. To evaluate our reinforcement learning based node judge module, we design three variants:

• •

  FreeKD-w.o.-judge: our FreeKD without using the agent. $\Phi$ and $\Psi$ distills knowledge for each node from each other.

• •

  FreeKD-loss: our FreeKD without using the reinforced knowledge judge. It determines the directions of knowledge distillation only relying on the cross entropy loss.

• •

  FreeKD-all-neighbors: our FreeKD selecting the directions of node-level knowledge distillation via node-level actions, but using all neighborhood nodes as the local structure.

• •

  FreeKD-all-structures: our FreeKD selecting the directions of node-level knowledge distillation, but without using structure-level actions for structure-level knowledge distillation.

Table 4 shows the results. FreeKD-node is better than GCN, showing that mutually transferring node-level knowledge via reinforcement learning is useful for boosting the performance of GNNs. FreeKD obtain better results than FreeKD-node. It illustrates distilling structure knowledge by our method is beneficial to GNNs. FreeKD achieves better performance than FreeKD-w.o.-judge, illustrating dynamically determining the knowledge distillation direction is important. In addition, FreeKD outperforms FreeKD-loss. This shows that directly using the cross entropy loss to decide the directions of knowledge distillation is sub-optimal. As stated before, this heuristic strategy only considers the performance of the node itself, but neglects the influence of the node on other nodes. Additionally, FreeKD has superiority over FreeKD-all-neighbors, demonstrating that transferring part of neighborhood information selected by our method is more effective than transferring all neighborhood information for GNNs. Finally, FreeKD obtains better performance than FreeKD-all-structures, which indicates our reinforcement learning based method can transfer more reliable structure-level knowledge. In summary, these results demonstrate our proposed knowledge distillation framework with a hierarchical reinforcement learning strategy is effective.

We further intuitively show the effectiveness of the reinforced knowledge judge to dynamically decide the directions of knowledge distillation. We set GCN as $\Phi$ and GraphSAGE as $\Psi$, and train our FreeKD on the Cora dataset. Then, we poison $\Phi$ by adding random Gaussian noise with a standard deviation $\sigma$ to its model parameters. Finally, we visualize the agent’s output, i.e., node-level policy probabilities $\pi_{\mathbf{\theta}}\left(\mathbf{s}_{i}^{[1]},0\right)$ and $\pi_{\mathbf{\theta}}\left(\mathbf{s}_{i}^{[1]},1\right)$ at node $i$ for $\Phi$ and $\Psi$, respectively. To better visualize, we show a subgraph composed of the first 30 nodes and their neighborhoods.

Figure 4 shows the results using different standard deviations $\sigma$. In Figure 4 (a), (b), and (c), the higher the probability output by the agent is, the redder the node is. And this means that the probability for the node in this network to serve as a distilled node to transfer knowledge to the corresponding node of the other network is higher. As shown in Figure 4 (a), when without adding noise, the degrees of the red color in $\Phi$ and $\Psi$ are comparable. As the noise is gradually increased in $\Phi$, the red color becomes more and more light in $\Phi$, but an opposite case happens in $\Psi$, as shown in Figure 4 (b) and (c). This is because the noise brings negative influence on the outputs of the network, leading to inaccurate soft labels and large losses. In such a case, our agent can output low probabilities for the network $\Phi$. Thus, our agent can effectively determine the direction of knowledge distillation for each node.

In our method, there are three main hyper-parameters, i.e., $\gamma$ in the reward function (12), $\mu$ and $\rho$ in the loss function (15) and (16). We study the sensitivity of our method to these hyper-parameters on the Cora dataset. First, we investigate the impact of $\gamma$ in the agent’s reward function on the performance of our method. As shown in Table 5, with the values of $\gamma$ increasing, the performance of our method will fall after rising. In the meantime, our method is not sensitive to $\gamma$ in a relatively large range. We also study the parameter sensitiveness of our method to $\mu$ and $\rho$. Figure 5(a) shows the results. Our method is still not sensitive to these two hyper-parameters in a relatively large range. Additionally, we analyze the convergence of our method. Figure 5(b) shows the reward convergence curve. Our method is convergent after 100 epochs.