Suit the Node Pair to the Case: A Multi-Scale Node Pair Grouping Strategy for Graph-MLP Distillation

1 Introduction

Graph, as an abstract data structure, is currently widely applied in various real-world applications, social network analysis, recommendation system, neuroscience, etc [Qiu et al., 2018; Wu et al., 2022; Ye et al., 2024]. Recently, the appearance of Graph Neural Networks (GNNs) has provided efficient solutions to graph-related tasks. Based on message passing mechanism where node can capture features from its neighborhood, GNNs are advantageous in capturing graph representations, and semi-supervised learning has witnessed a success in node-level classification task [Kipf and Welling, 2016]. However, when considering practical deployment, there exist certain gaps. One gap lies in its neighborhoodfetching mechanism, where the explicit utilization of structure information leads to latency during reference time. In contrast, Multi-Layer-Perceptrons (MLPs) achieve higher inference speed, but with poorer performance. One way to bridge its gap is through Knowledge Distillation, a model compression method aiming to transfer the knowledge from one complex teacher model to lightweight student model [Hinton et al., 2015; Joshi et al., 2022].

∗Corresponding author

Graph Knowledge Distillation is useful in bridging GNN and MLP [Liu et al., 2023; Tian et al., 2023; Gou et al., 2021]. Previous studies have explored different distillation frameworks from GNNs to MLPs. GLNN [Zhang et al., 2021] followed the traditional distillation paradigm between teacher and student models, leading to better performance and fast inference speed. [Huo et al., 2023] distills graph information from two aspects: feature and structure, which applies data augmentation during training. BGNN [Guo et al., 2023] improves the performance of MLPs by boosting multiple teacher GNNs. KRD [Wu et al., 2023b] is a framework which quantifies the knowledge and can make reliable distillation. NOSMOG [Tian et al., 2022] can learn better structure information and is robust to feature noise. VQGraph [Yang et al., 2024] proposed a VQ-VAE-based framework that can facilitate better structure-aware knowledge distillation. SimMLP [Wang et al., 2024] revisited distillation from mutual information and introduced a "pretrain-finetune" framework without teacher's supervision.

However, these researches have neglected following aspects: Most distillation methods just follow a direct "node-tonode" fashion, which directly aligns logits between teacher nodes and corresponding student nodes. While for graph data, information within GNN may be more complicated, especially considering the complex topological structure of graph, where the relationship between one node and its neighbors can be various, and simple distillation is not capable of capturing these comprehensive knowledge. According to [Wu et al., 2023a], GNN may carry two types of information from spectral perspective, and GNN's message passing mechanism promotes the similarity of representation between nodes and its neighbors. Thus, traditional "node-tonode" distillation may fail to transfer finer information at high frequency which is drown by low frequency. However, its FF-G2M method is a global solution for all nodes, ignoring difference in local structure for different node pairs in one graph. Specially, for one given node, its multiple relationships with different neighborhoods reflect features at different scales (shown in Fig. 1), and different distillation methods are required to capture corresponding different information. Moreover, compared to message passing mechanism for layers within GNN, the fully connected layers inside MLP just ignore the learning of graph structure. Current distillation methods just focus on alignment between different logits, ignoring the forward process within inside MLP layers.

To address the weakness mentioned above, we design a novel GNN-MLP distillation method. The motivation is that the relationships between one node and its neighborhoods can be multiple and different relationships among node pairs require different methods to extract corresponding information. The model we proposed, namely Multi-Scale Node-pair grouping Graph-Distillation-MLP (MSN-GDM), groups the node pairs at different scales and uses corresponding optim functions to capture multi-scale information. The main contribution of this paper is summarized as follows:

• We propose a multi-scale node pair grouping strategy for preprocess. Nodes are assigned to different groups according to the similarity metric between node pairs. The metrics calculation considers both node features and their structures.
• We design corresponding multi-scale distillationoptimization loss based on multi-scale node pair set groups in the previous step. The loss consists of intra-layer loss and inter-layer loss.
• Extensive experiments on eight real-world datasets demonstrate the effectiveness of our distillation method.

2 RelatedWork

2.1 Graph Neural Network

Recent studies have shown the superior performance of Graph neural networks (GNNs). Based on the Spectral Graph Theory, the original GNN captured the features of the graph from a spectral perspective [Bojchevski and G¨unnemann, 2017]. [Kipf and Welling, 2016] proposed a semi-supervised framework (GCN) which follows the paradigm of spatial domain by aggregating the neighborhood nodes and updating target nodes. Through this message passing mechanism, spatial GNN outperforms spectral GNN in its fast training time while making good use of graph structure. Recently, increasing GNN-based methods are dedicated to proposing novel architectures for better graph representation learning in spatial domain. GAT adds attention module for better aggregation [Veliˇckovi´c et al., 2017], GraphSAGE designs an inductive framework that can generate embeddings for unseen nodes [Hamilton et al., 2017].

2.2 Knowledge Distillation

Knowledge Distillation (KD) is a method to compress models at a large scale, through which the knowledge is transferred from a large-scale teacher model to a lightweight student model, thereby imporoving the performance of student model while keeping its low complexity of time and space. Current graph knowledge distillation methods can be divided into three types [Gou et al., 2021; Tian et al., 2023]: GNN-GNN distillation that intends to decrease the layers and dimensions of teacher GNNs [Yang et al., 2020; Yan et al., 2020]; GNN-MLP distillation that aims to transfer knowledge from GNNs to MLPs [Wu et al., 2022; Zhang et al., 2021; Wu et al., 2023b; Yang et al., 2024; Wang et al., 2024] and Self-Distillation which refines the transfer of knowledge inside one single model [Chen et al., 2021; Zhang et al., 2023; Deng and Zhang, 2021; Wu et al., 2024].

3 Preliminaries

Graph: We denote a Graph as G = (V, E), which stands for its two components: nodes and edges. Each node vi has a ddimension feature xi and the corresponding feature matrix of G is X ∈ RN×d . For node-level task, each node has a class label yi, and for graph-level task, the label yG is for the whole graph.

Graph Neural Network: For most graph neural networks (GNNs), they all follow message passing mechanism which aggregates its neighborhood nodes before updating the value of its node vi [Kipf and Welling, 2016; Veliˇckovi´c et al., 2017; Hamilton et al., 2017].

Where hli is the representation for node vi in l-th layer, AGG(.) stands for the aggregation of its neighborhood nodes Ni, and UPDATE(.) means the update operation of hi.

Knowledge Distillation: Knowledge Distillation aims to transfer the knowledge from a cumbersome teacher model to a lightweight student model [Hinton et al., 2015]. Knowledge Distillation for graphs mainly takes the form of imposing KL-divergence between the soft labels of teacher and student models. The distillation loss function can be defined as follows:

Where zT and hS means the soft labels of teacher and student respectively, τ is the temperature which controls the process of distillation.

4 Methods

This section focuses on details of the proposed MSN-GDM. Before formal introduction, we first make an analysis of measurement of different knowledge in GNNs. The framework of MSN-GDM is shown in Fig. 2, which is composed of two steps: In step 1, we preprocess the graph by grouping different neighbor node pairs of each node vi based on our node similarity metric. Then in step 2, a multi-scale optimization function is designed according to mean value of corresponding node pair set groups. The weighted sum of each layer will be the final output for MLP to capture comprehensive representations. More details will be presented in the following subsections.

4.1 Measuring Different Knowledge in GNN

According to [Wu et al., 2022], a trained GNN contains both low and high frequency features, while simple distillation method only captures low-frequency information, ignoring high-frequency information, which can be viewed as the difference between two nodes spatially. Two ways are proposed to measure different knowledge in GNN. The distance between connected nodes can be utilized as low-frequency measurement and the pairwise distance differences between teacher GNN and student MLP shows the capability for MLP to capture fine-grained information.

In general, the knowledge can be divided into two parts. Here we define these two types of knowledge in Eq. 3.

We use cosine similarity to measure low-frequency knowledge in GNN which represents similarities between node pairs (Mean-Sim-Metric, MSM), and use MSE(.) between teacher and student logits to measure those fine-grained information , which shows the student model's ability to keep its distance between two nodes (Mean-Differ-Metric, MDM). The following plots are shown in Fig. 1.

According to the Mean-Sim-Metric figure in Fig. 1, the MSM value of vanilla MLP is the lowest, showing its weakness in capturing graph information, while the value of GCN and its distilling student MLP increases from the beginning and converges to a certain value during the training process [Chen et al., 2020a; Keriven, 2022; Chen et al., 2021]. The reason lies in GNN's message passing mechanism, which finally minimizes the distance between one node and its neighbors. Current "node-to-node" fashion just ignores the difference in local features of graph, where the neighbors of one node may have different features. In other word, a multiscale distillation is necessary which is tailored for capturing features at different scales.

4.2 Multi-Scale Node Pair Grouping Strategy

According to our motivation, different node pairs contain different feature patterns, and we hope to tell if the node pair shares similar patterns or different ones before designing corresponding operations. Thus, we propose a multiscale node pair grouping strategy to assign node pairs into different groups. The grouping strategy consists of two parts, namely node pair similarity metric calculation and multiscale node pair grouping, and the whole process is shown in Fig. 3.

Node pair similarity metric calculation. Considering that the two attributes of the graph: feature and structure, together make the representation, here we define feature similarity (φf ) and structure similarity (φs) accordingly. For feature similarity, we utilize cosine similarity between given node pair (vi, vj). For structure similarity, intuitively, the representation of a node can also be influenced by its neighbors [Xing et al., 2024]. Nodes with more "similar" neighbors tend to maintain their own features while nodes with more "different" neighbors are more likely to be influenced. Thus, we consider the neighbor nodes for given vi (denote as Ni). We normalize the mean value of feature similarity for vi:

Here, the term degi refers to the degree of node vi. Based on above definitions, φf and φs for node pair (vi, vj) are as follows:

Where the ||.|| operator is the L1 Norm of the variable, and the smaller value means larger similarity between vi and vj .

The value of φf and φs together decide the type of knowledge between (vi, vj). We define knowledge similarity φk as the combination of the two values. Moreover, encouraged by [Wu et al., 2023a], the abstract value of φk should reflect the extent of knowledge (vi, vj) carries. So the final value of φk is as follows:

Here α is a hyper-parameter ranging [0, 1], by which the range of φk is normalized to [-1, 1], where smaller value means larger similarity between two nodes.

Multi-scale node pair grouping. The value of φk varies from -1 to 1, and for node vi and its neighbors {j|j ∈ Ni}, we assign vj to different groups according to φk(i, j).

For scale factor J, we firstly divide the interval of [-1,1] equally into J parts, where the range of j-th interval is between [-1 + 2(j-1) J ,-1 + 2j J ] → [ℓj, rj]. Then we assign node to the correct interval according to its φk. As is shown in Eq. 7, for each node vi, we maintain a J-scale node set (denote as MSi), where the j-th element contains nodes whose value is among [ℓj, rj). Finally, we merge the node set of each node as the multi-scale node pair set group for the given graph G:

After merging, two connected nodes may appear in the same group. Essentially, the grouping strategy can be viewed as operations on edges. For implementation, we do both calculation and grouping based on edges (E) of the graph. For graphs at larger scale where the expense to directly preprocess whole graph can be extremely great, we randomly sample edges from original E for scalability. Then we employ above preprocess operation on the edge-scaling graph.

Where ρ is a hyper-parameter controlling number of edges to sample on large-scale graphs. In this paper we set it as 0.3.

4.3 Multi-Scale Distillation-Optimization Loss

According to Eq. 1, the GNN forward process can be regarded as type of "aggregate then update " operation. Encouraged by this, Our optim loss is composed of two parts: Intra-Layer Loss within one single layer and Inter-Layer Loss between two layers, whose process in shown in Fig. 4.

Intra-Layer Optim Loss. The idea of intra-layer optim loss (Lintra) is from graph self distillation method, where a model distills knowledge from itself without participation of teacher model [Hu et al., 2021; Luo et al., 2021]. Instead, model distills knowledge from structure information by certain constraints [Wu et al., 2024].

Further, inspired by message passing mechanism, for MLP model with L layers, the intra layer optim function can be formulated as:

Where zl i means the representation of node vi at l-th layer. By adding this item, the node may be closer to nodes around, implicitly utilizing graph structure information in the same layer.

Inter-Layer Optim Loss. FF-G2M designed two types of loss function to capture low-frequency and high-frequency features respectively. Our inter-layer optim loss derives from its loss function, yet we step further by applying it into multiscale node groups. Briefly, we select which type of knowledge to be passed to deeper layer. Corresponding to our analysis before, we consider two situations: aggregate knowledge distillation (LAKD) for similar node pairs and different knowledge distillation (LDKD) for different node pairs.

For graph G and the corresponding node pair set group MSG, the main idea is to treat different groups differently according to the overall feature of each group. The mean value of its φk (namely ˜ µj) is calculated for each group (MSG(j)) firstly. The inter-layer optim loss (Linter) is defined as follows. (see in Eq. 10)

Here K(.) works as a kernel function to measure the distance between two representations, and here we realize it with L1 Norm.

For Llinterj , which denotes the inter-layer optim loss for l-th layer at j-scale (j ∈ {1, 2, ..., J}, l ∈ {1, 2, ..., L - 1}), it considers two cases: If ˜ µj is negative, we consider LAKD which makes the distance closer between node pairs, otherwise LDKD to keep the difference. The abstract value of ˜ µj reflects the extent of knowledge to be distilled.

For the last layer of MLP which also serves as the output layer (zo ), we calculate the optim function between it and the last layer of the teacher model (ho ). The formula is the same as Eq. 10 except for zl → zo and zl+1 → ho (denote as Lointer ). Each layer's inter-layer loss together makes the final Linter (shown in Eq. 11).

4.4 Hierarchical Weighted Sum as MLP Output

The receptive field of GNN may expand as the layer of the model increases and different layers will capture features to different extent [Li et al., 2020; Zhang et al., 2022; Chen et al., 2020b]. Thus, we utilize the weighted sum of each layer as the output of the final layer: zo = ∑Ll=1 wizi, where {wi}i are learnable parameters and ∑i wi = 1.

The final format of the optim loss is the combination of the two items: Loptim = Lintra+Linter. Here we utilize Loptim as an appendix item for the knowledge distillation loss, and the final loss function is the combination of the classification loss, distillation loss and our optim loss, as is shown in Eq. 12, where γ is another hyper-parameter.

5 Experiments

5.1 Experiment Setup

Datasets. We conduct our experiments on node-classification task and select eight benchmark datasets: (1) three citation networks: Cora, Pubmed, Citeseer [Bojchevski and G¨unnemann, 2017]; (2) two co-purchase networks: Amazon Computers and Amazon Photo; (3) two co-authorship networks: Coauthor CS and Coauthor Physics [Shchur et al., 2018] (4) one large graph from ogbn dataset: Ogbn-Arxiv [Hu et al., 2020].

Baselines. For teacher models, we select three classic spatial GNNs and one spectral GNN. For spatial GNNs, we select GCN [Kipf and Welling, 2016], GAT [Veliˇckovi´c et al., 2017] and GraphSAGE [Hamilton et al., 2017]. We set these models L = 3, d = 128 for layer and hidden dimension respectively. For spectral GNN, we select BernNet [He et al., 2021]. We also select two classic different distillation methods: GLNN [Zhang et al., 2021] and FF-G2M [Wu et al., 2023a]. Code implementations of all baseline methods are taken from their respective original papers and the best results of the mentioned distillation methods are recorded.

Experimental Settings. Our task follows the semisupervised node classification task [Kipf and Welling, 2016] and the splitting ratio of the datasets is 20%/20%/60% for train/val/test for all datasets except Arxiv, which follows standard ogbn dataset splitting. As the preprocess of the graph requires information of whole graph, so here all the experiments follow semi-supervised transductive setting, and the datasets are split randomly and run for 10 times. We calculate their mean accuracy as well as the standard deviation and record the best result. The maximum epochs of each training process is 2000 and we utilize an early stop mechanism with a patience of 300. For ogbn dataset, we rewrite the loss function in Eq. 12 as Eq. 13.

For student model, we utilize a 3-layer MLP with a dimension of 64 for all the datasets except Arxiv. For Arxiv the layer is 3 and the dimension is 512. We select the best model by performing a random search on three hyper-parameters in our method: α ∈ {0.0, 0.2, 0.4, 0.6, 0.8, 1.0}, J ∈ {2, 3, 4, 5, 6, 7}, γ ∈ {0.4, 0.5, 0.6, 0.7}. We optimize the model through Adam optimizer. The model is implemented through PYG and all the experiments are conducted on NVIDIA 3090s. Code can be available at https://github.com/KamonRiderDR/MSN-GDM.

5.2 Comparison Results on Different Teachers

The results on different teacher models and different distillation methods are shown in Tab. 1. In general, the proposed MSN-GDM significantly improves the performance of MLP, making it outperform the given teacher models, among which the maximum improvement in absolute accuracy is 9.71% in Citeseer with GrapgSAGE as teacher model. Moreover, compared to other Graph-MLP distillation methods, MSN-GDM achieves best or second in all the datasets. Specially, For Cora, Pubmed and CS, MSN-GDM outperforms all the other three models, and the maximum advantages over the second method (GLNN) reaches 2.61% for Cora where GraphSAGE works as the teacher model.

5.3 Comparison Results with Other Distillation Methods

We also conduct comparison with other SOTA distillation methods. We choose GraphSAGE (SAGE) as teacher model. Besides from the two distillation methods above, we also select several different distillation methods: KRD [Wu et al., 2023b], NOSMOG [Tian et al., 2022], VQGraph [Yang et al., 2024] and SimMLP [Wang et al., 2024]. The comparision experiments follow their own implementations and the results are listed in Tab. 2. The results show the effectiveness of MSN-GDM, no matter in small graphs (e.g. Cora, Pubmed, Citeseer), medium graphs (e.g. Photo, CS) and large graphs (e.g. Arxiv).

5.4 Ablation Studies

Effects of sub-modules. We first conduct our ablation studies on the two sub-modules of MSN-GDM: hierarchical weight sum (weight) and optim loss (optim). As is shown in left Fig. 5, the accuracy improves with the addition of each sub-module. Moreover, for most datasets the addition of optim sub-module witnesses a larger improvement in accuracy compared to weight sub-module. The histogram shows the effectiveness of each sub-module in MSN-GDM.

Effects of different optim loss. Further, as our optim loss consists of two parts namely Linter and Lintra, here we also explore the effects of the two types of loss within Loptim, which is shown in right Fig. 5. Discarding either type of loss may lead to a decline in the result, which demonstrates the effectiveness of both Lintra and Linter inside Loptim.

5.5 Sensitivity Analysis

Firstly, we conduct sensitivity analysis experiment on J (shown in Fig. 6 (a)), where we can see that for different datasets, the value J for best result remains different. For dataset Cora and GCN as teacher, the best value of is 6, while for datasets Computers it is 4. Moreover, for different datasets, the results fluctuate differently when J changes. For further analysis, we plot the heatmap of hyper-parameter J and α, as the final grouping result in preproess part is decided jointly by both two hyper-parameters. For different J, the corresponding optimal α is also different, and heatmap can reflect the relation of the two factors. The results in Fig.

6 (b) show that the best result lies in the middle part of α and J (α = 0.2, J = 6). We also conduct the sensitivity analysis experiment on the loss function coefficient γ (shown in Fig. 6 (c)). For Cora, MSN-GDM with different teacher models has best results with different γ, and the results have a smaller variation as γ changes.

Finally we conduct experiments on inference speed and parameters during inference time and results in Tab. 3 show MSN-GDM's effectiveness compared to teacher models.

5.6 Empirical Analysis

Finally, we conduct empirical analysis experiment on our method. Corresponding to the first subsection in Methods section, we aim to verify if the student model learned multiscale knowledge from teacher model. If the mean value ˜ µj is negative, then we plot MSM curve (left of Fig. 7) of its nodes pair group, otherwise MDM curve (right of Fig. 7).

From Fig. 7, we can see that for nodes assigned to "similar" groups (here the group is when (j = 2)), the similarity value increases and is larger than other methods, indicating that for similar node-set group, MSN-GDM enhances the model's ability to capture "similar" patterns; for nodes assigned to "different" groups (here j = 3), the value decreases gradually and remains the smallest among the four distillation methods, showing the model's advantage in capturing "different" patterns among corresponding group. The figures of 2 cases demonstrate that our MSN-GDM is better at capturing knowledge at different scales.

6 Conclusion

In this paper, we propose a GNN-MLP distillation method that captures muti-scale knowledge among different node pairs. We firstly design a multi-scale node pair grouping strategy based on our node pair similarity metric, then a multiscale optim function is proposed to capture knowledge at different scales. Extensive experiment results demonstrate the effectiveness of our proposed MSN-GDM. For future work, we may explore more adaptive and more efficient grouping methods for better distillation. We may also consider to apply it into inductive setting and more general graph tasks.

Acknowledgments

The work is supported in part by the National Natural Science Foundation of China under Grant No. 62471133, the Central University Basic Research Fund of China under Grant No. 2242024K40020, and the Big Data Computing Center of Southeast University.