Teaching MLP More Graph Information: A Three-stage Multitask Knowledge Distillation Framework.^††thanks: Supported by organization x.

To improve knowledge transferring to student MLPs, it’s advisable to use distillation methods that go beyond just aligning soft label distributions. By defining the graph signal $X$, we consider two species of knowledge: local node representation and global positional information. Here we introduce approach called Neural Heat Kernel(NHK) Based Distillation for effective node representation generation from student MLPs. Moreover, we study Positional Encoding (PE) for better global positional information transfer.

To introduce NHK as a teaching tool for MLPs, it’s essential to include the following definition:

This transformation can facilitate the extraction of local structural features. We resort to the NHK to capture the influence exerted on the nodes by the message passing process during the graph convolution process, for the process of information exchange between two points on the manifold is similar to that of heat transform. Consequently, feature embedding of nodes after certain hidden layers is added with information from connected nodes with the mapping from kernel functions $f_{k}$. This can be regarded as a prior knowledge also named Graph Smoothness [29]. Layer-to-layer kernel mapping are proposed as:

where node-level smoothness is passed between $node_{\alpha}$ and $node_{\beta}$, which are arbitrary nodes from graph. $H_{\alpha}$ and $H_{\beta}$ are node representations after certain hidden layer. By characterizing fuzzy geometric information as a specific function, we adopt an easily comprehensible approach that simplifies local distillation, which will be discussed further below.

To improve the capturing of global information, we draw inspiration from graph transformer frameworks. Specially, Dwivedi et al. [32] utilize Laplacian Encoding as initial PE, calculated before training. By providing PE as prior knowledge to student model, its ability to perceive positional and structural information can be enhanced. In this work, we use Laplacian eigenvectors as PE. The eigenvectors are calculated by defining a Laplacian matrix:

where $A$ is the adjacency matrix, $D$ is the diagonal matrix for matrix standardization and $U$, $\Lambda$ represents eigenvectors and eigenvalues of the Laplacian matrix. Then we choose k-smallest non-trivial eigenvectors in real space for n nodes as an extra input $U_{selected}\in R^{n\times k}$. A theorem related to Spectral Clustering [30] demonstrates benefits of PE for classification tasks:

To quantify the degree of difference in the soft target distribution between student and teacher models, we introduce Kullback-Leible Divergence [31] (KL divergence). Consider two random variables $P$, $Q$, with discrete probability distributions(like the real situations) $p(\theta)$ and $q(\theta)$, respectively. We have KL divergence from P to Q as:

The smaller the difference between the soft logits of the student and teacher models, the smaller the KL divergence will be. This shows that we can use minimizing KL divergence as a way to simplify the utilization of soft target information.

Our proposed distillation process consists of three stages: GNN pre-training stage, distillation stage and inference stage. During the first stage, we perform simple message-passing training and save the trained graph model $GNN_{pre}$ and the soft target distribution $T_{t}$ of the model output.

During the second stage, we first add PE as a prior knowledge for student MLPs. Consider the $d$-dimensional input feature $X_{s_{i}}^{(0)}\in R^{d}$ of a certain node. We embed the $k$-dimensional Laplacian PE into a feature vector with the same dimension as $X_{s_{i}}^{(0)}$. Then we concatenate them or do a simple dot product between them. As the following equation shows:

In order to capture the $message-passing\ control\ information$, we make sure that nodes for training contain the same ones as being input towards teacher GNNs. Thus, the process of hidden mechanism control by Neural Heat Kernels can work. Our target is to make hidden layers of students able to predict the outputs of a guiding hidden layer of teachers [19]. Layer-to-layer mapping matrices as Sec. 3.2 on hidden outputs of $m$ training nodes are adopted as:

where $\mathbb{M}[\bullet]$ denotes for operation to obtain mapping matrices. We select four non-linear kernel functions as $f_{k}^{(l\ to\ l+a)}$ to do the calculation:

where $\sigma$ denotes for non-linear operations like $Sigmod$, $Tanh$ or $Relu$ ,$f_{proj}$ denotes for a linear transformation in Sigmod Kernel and $T$ denotes for a constant time interval in Gaussian Kernel. And $M_{r}$ is a randomized matrix whose vectors obey a Gaussian distribution in Randomized Kernel. We load the pre-trained $GNN_{pre}$ in the first stage to generate teacher mapping matrices and aim to minimize the distance between those of teachers and students. A simple L2 Norm is used to measure distance as:

Also, we adopt a process of soft logits matching, which is defined in Section 4.2. During stage III, we do inference tasks on either the existing graph or unseen new in the same graph dataset. We will discuss this further in Sec. 5.5.

During the entire training process, we consider two main parts of the distillation loss: the loss from model prediction matching and the loss from distance measurement between mapping matrices. Therefore, we can express the total multi-task distillation loss $L_{total}$ as:

where $\hat{Y_{s}}$ denotes for student prediction, and $Y_{t}$ can be true labels or soft targets from teacher. $\gamma$ is a constant to show the proportion of $\mathcal{L}_{dis}$. And that is calculated by a K-layer matching process between $Mat_{s}^{(l)}$ and $Mat_{t}^{(l)}$ included above. It is noteworthy that we provide two kinds of training nodes: with true labels or with soft targets from teacher GNNs. The student predictions of them can be represented as $\hat{Y_{sL}}$ and $\hat{Y_{sO}}$. So $\mathcal{L}_{out}$ can be specified as:

where $y_{t}$ are true labels of selected labeled nodes while $T_{t}$ are soft target distributions from teacher outputs, as mentioned in Sec. 4.1. They are two parts of $Y_{t}$. By minimizing $\mathcal{L}_{total}$, we can get a well-trained MLP to do inference tasks.

Getting inspired by Variational Autoencoder (VAE) [40] and GNN-level topology distillation [18], We have tried to optimize the heat kernel itself. A trainable matrix $W_{k}$ is introduced for transformation whose parameters can be optimized and shared by both student and teacher. The trainable reverse kernel can be represented as:

Like VAE does, we also adopt reconstruction loss between original inputs($X^{(0)}$) and outputs of the last hidden layer ($H^{(l+a)}$) and utilize it to optimize $W_{k}$ for data processing:

Consequently, we need to do a two-step optimization both on the student model and the kernel. That helps to explore proper kernels for distillation.

Our proposed methodology is mainly for solving two major problems of MLP classifier: $1.$ positional information loss $2.$ low generalization. To address these issues, we utilize two efficient approaches: PE and NHK based distillation. These approaches enable the transfer of knowledge from non-Euclidean domains into simpler domains in a explainable manner. Overall, our methodology offers a promising solution to enhance the performance of MLP classifiers on graph datasets. Also, this transfer is done in a pretty explainable way.

In this part, we do our research on five typical datasets [33]. Three common benchmark datasets: Cora, Citeseer, Pubmed and two larger ones: Amazon photos, Amazon Computers. Also, we introduce two OGB datasets [34]: Arxiv and Products, to see if our model perform well on large-scale graphs. The details of the datasets are in Appendix. 7. For each dataset, we label 20 nodes per class and select 30 nodes randomly for validation. We adopt both true or soft labels here. Other nodes are provided for testing.

We choose frequently used teacher models: SAGE, GCN and GAT. It’s easy for us to transfer knowledge from them into student MLPs. We use single MLP without distillation and GLNN as baseline model to check if our framework perform better.

In all experiments, we report accuracy($\%$)on test set. 10 experiments with random seeds are done, and mean value and standard deviation of test accuracy are reported.

During stage I and II, we set the max epochs during training to 1000, and an early stop is adopted when accuracy on validation set doesn’t increase for 50 epochs. We take 512 as our batch-size, $Adam$ as our optimizer and cuda as our training device. In order to report exact results, we follow the settings of teacher GNNs for 5 small datasets in Appendix. 0.A.2. A grid-research is also adopted to find the most proper hyperparameters for loss computing, as introduced in Appendix. 0.A.3.

Transductive and inductive settings. We study two categories of node classification tasks. The most common setting for node classification is a transductive (trans) setting, where all nodes are split into training set, validation set and test set. Under this setting we can see the whole graph structure (adjacency matrix) while training. However, under real-life scenarios, an inductive (induc) setting is often more applicable, for we cannot always see all users and relations in real networks. Unobserved nodes are first taken out and all information of them can’t be seen during stage II. Then a train-valid-test split is adopted. Something to note while transductive or inductive distillation is as follows:

We apply experiments to see performance of our framework on large-scale graphs on Arxiv and Products. Results are shown in Table. 3 and Table. 4.

Results. We can see that we improve baseline GLNN’s performance by about 0.5 to 1.5$\%$ in 7 of 10 total cases via our method, with transductive or inductive setting. It demonstrates that our method really works on large-scale graphs.

As mentioned in Sec. 4.3, we study trainable reverse kernel’s performance. We do our study on three benchmark datasets, Cora, Citeseer and Pubmed. Note that we use default hyperparameters for distillation, so we don’t really report the ’best’ results of our framework here but only a comparison. We show the results in Fig. 4 in Appendix. Teaching MLP More Graph Information: A Three-stage Multitask Knowledge Distillation Framework.^††thanks: Supported by organization x..

Now we consider our KMP’s results when utilizing PE as an initial node feature. We test it on the seven datasets mentioned above. We only adopt Laplacian PE, for it’s easy to calculate and proved efficient [32]. Note that we also introduce a method SA-MLP [35] here. It tries to deal with problems of student MLP by simply mapping the adjacency matrix and node features to the same dimension and then concatenating them together. The complete source code of SA-MLP hasn’t been provided yet, so we reproduce a training process for it by ourselves. We provide results of SAGE as teacher GNN. Results are listed in Table. 5.

Results.

Now we see that our framework can reach the best performance on all datasets with transductive setting. Moreover, it performs best with inductive setting on 5 datasets. The conclusion is that Positional Encoding can really help KMP work better with little extra consumption for doing calculations. Only a tiny linear layer projection is necessary for change PE into node features(much tinier than SA-MLP). It is a great guide to real-life inference tasks.

Dataset. We do our robustness and sensitiveness study on benchmark datasets: Cora, Citeseer, Pubmed. As they are easy for us to compute student predictions. This is an intuitive guide for large-scale graphs.

Problem definition. We test our framework’s robustness by adding noise attack to initial node features $X_{init}$. Consider that a random matrix $X_{rand}$ is added to $X_{init}$ by a certain percentage, and our doubt on robustness is whether KMP can work well with interference of noise features. We choose a percentage in $\{10\%,20\%,30\%,40\%,50\%\}$. Note that we also provide GLNN’s performance for comparison. Next, we do sensitiveness test by modifying the constant $\gamma$, mentioned in Sec. 4.2, as our framework is much sensitive to only $\gamma$ revealed by previous experiments. We choose $\gamma\in\{0.1,0.3,0.5,0.7,1,3,10,30\}$. Note that we choose the default hyperparameters and try to get the test curve for comparison but not the best results.

Results. The results of this robustness study is shown in Fig. 3. (We do not report results on Pubmed, for its test accuracy bitterly declined to very poor values with noise. This can be seen in Appendix. 0.A.5. Our frameworks ability to reduce harm of feature noise is stronger than baseline GLNN by about 2% in most cases. Next, we lay out results of sensitiveness study in Fig. 3. We discover the fact that our method works stably when $\gamma$ is between 0 and 3, and a sharp decline in test accuracy occurs with $\gamma$ larger than 3. Usually, we choose $\gamma\in(0,1]$ and our framework can work normally.

Since our framework can achieve good results with inductive setting for node classification, we consider another inductive task: graph classification.

Dataset. We utilize MiniGC [36] for this additional study. It contains 8 kinds of graphs in total with variable number of nodes. Here we generate different graphs with 10-20 nodes, and choose 240 graphs for training, 60 for validation and 60 for testing. Our target is to train a student MLP excellent enough for classification.

Results. Our graph classification results are listed in Table. 6. We can see that our KMP increases test accuracy of GLNN by about 3.5$\%$. This means that our approach may also have positive results for larger graph classification datasets.