Teaching MLP More Graph Information: A Three-stage Multitask Knowledge Distillation Framework

Junxian Li; Bin Shi; Erfei Cui; Hua Wei; Qinghua Zheng

cs.LGMar 2, 2024

Teaching MLP More Graph Information: A Three-stage Multitask Knowledge Distillation Framework

Authors: Junxian Li, Bin Shi, Erfei Cui, Hua Wei, Qinghua Zheng

Organizations: Xi’an Jiaotong University · Arizona State University, Tempe, USA · Tongji University, Shanghai, China

Abstract

We study the challenging problem for inference tasks on large-scale graph datasets of Graph Neural Networks: huge time and memory consumption, and try to overcome it by reducing reliance on graph structure. Even though distilling graph knowledge to student MLP is an excellent idea, it faces two major problems of positional information loss and low generalization. To solve the problems, we propose a new three-stage multitask distillation framework. In detail, we use Positional Encoding to capture positional information. Also, we introduce Neural Heat Kernels responsible for graph data processing in GNN and utilize hidden layer outputs matching for better performance of student MLP's hidden layers. To the best of our knowledge, it is the first work to include hidden layer distillation for student MLP on graphs and to combine graph Positional Encoding with MLP. We test its performance and robustness with several settings and draw the conclusion that our work can outperform well with good stability. Code implementation is in https://github.com/lijunxian111/KMP.

Figures & tables

Figure 1: Overall structure of our framework. It consists of three stages: GNN pretraining, distillation and inference. We transfer knowledge to student MLP during stage II. Best viewed in color.

Dataset	Models
Dataset	$SAGE$	$GCN$	$GAT$	$MLP$	$GLNN(SAGE)$	$KMP(SAGE)$	$GLNN(GCN)$	$KMP(GCN)$	$GLNN(GAT)$	$KMP(GAT)$
Cora	79.70 $\pm$ 1.90	79.58 $\pm$ 1.01	81.41 $\pm$ 1.15	58.66 $\pm$ 2.34	78.99 $\pm$ 1.90	79.48 $\pm$ 1.48	79.43 $\pm$ 0.82	80.03 $\pm$ 1.19	79.89 $\pm$ 1.44	80.44 $\pm$ 1.13
Citeseer	70.53 $\pm$ 2.16	70.09 $\pm$ 2.23	71.76 $\pm$ 1.54	58.01 $\pm$ 0.96	71.23 $\pm$ 2.13	71.96 $\pm$ 2.25	71.24 $\pm$ 2.98	72.18 $\pm$ 2.23	72.44 $\pm$ 1.63	72.72 $\pm$ 1.18
Pubmed	75.53 $\pm$ 2.39	77.96 $\pm$ 2.96	78.30 $\pm$ 2.44	67.72 $\pm$ 3.34	76.32 $\pm$ 2.67	76.80 $\pm$ 2.01	77.94 $\pm$ 2.95	78.48 $\pm$ 1.90	76.44 $\pm$ 2.28	76.86 $\pm$ 2.40
Amazon Photo	90.26 $\pm$ 1.14	87.18 $\pm$ 4.05	91.61 $\pm$ 0.94	78.39 $\pm$ 1.71	92.04 $\pm$ 0.91	92.23 $\pm$ 0.85	89.97 $\pm$ 3.08	90.48 $\pm$ 2.55	93.01 $\pm$ 0.94	93.19 $\pm$ 0.78
Amazon Computer	83.23 $\pm$ 1.49	80.19 $\pm$ 5.18	83.03 $\pm$ 2.30	69.01 $\pm$ 3.27	84.41 $\pm$ 1.51	84.75 $\pm$ 1.10	78.42 $\pm$ 4.26	78.96 $\pm$ 4.27	83.52 $\pm$ 2.17	82.94 $\pm$ 2.00

Table 1: Results of teachers, baseline models and KMP for trans setting. The best performance of student model is highlighted.

Dataset	Models
Dataset	$SAGE$	$GCN$	$GAT$	$MLP$	$GLNN(SAGE)$	$KMP(SAGE)$	$GLNN(GCN)$	$KMP(GCN)$	$GLNN(GAT)$	$KMP(GAT)$
Cora	78.64 $\pm$ 2.06	78.90 $\pm$ 1.69	81.10 $\pm$ 2.46	58.02 $\pm$ 2.72	71.94 $\pm$ 2.85	71.85 $\pm$ 2.40	72.39 $\pm$ 2.30	72.86 $\pm$ 2.37	72.41 $\pm$ 1.48	72.53 $\pm$ 1.38
Citeseer	70.83 $\pm$ 2.83	68.26 $\pm$ 3.65	70.75 $\pm$ 3.29	59.46 $\pm$ 4.55	70.80 $\pm$ 2.63	71.22 $\pm$ 2.04	67.15 $\pm$ 2.33	68.04 $\pm$ 2.52	68.92 $\pm$ 2.94	69.17 $\pm$ 3.01
Pubmed	74.86 $\pm$ 2.97	74.40 $\pm$ 2.94	76.64 $\pm$ 2.97	66.14 $\pm$ 5.03	74.76 $\pm$ 2.69	74.92 $\pm$ 3.03	74.89 $\pm$ 2.81	75.44 $\pm$ 2.75	74.63 $\pm$ 2.81	74.53 $\pm$ 3.06
Amazon Photo	91.33 $\pm$ 1.56	89.28 $\pm$ 2.12	92.27 $\pm$ 1.50	79.01 $\pm$ 1.80	89.75 $\pm$ 1.58	90.26 $\pm$ 2.43	88.54 $\pm$ 1.94	89.10 $\pm$ 1.69	90.02 $\pm$ 2.48	90.54 $\pm$ 1.37
Amazon Computer	82.73 $\pm$ 1.54	75.21 $\pm$ 3.56	83.38 $\pm$ 2.50	67.95 $\pm$ 3.88	81.18 $\pm$ 2.53	80.94 $\pm$ 2.33	74.01 $\pm$ 3.26	74.05 $\pm$ 4.01	81.08 $\pm$ 2.54	81.56 $\pm$ 2.57

Table 2: Results of teachers, baseline models and KMP for induc setting. The best performance of student model is highlighted.

Dataset	Models
Dataset	$SAGE$	$GCN$	$GAT$	$MLP$	$GLNN(SAGE)$	$KMP(SAGE)$	$GLNN(GCN)$	$KMP(GCN)$	$GLNN(GAT)$	$KMP(GAT)$
OGBN-Arxiv	71.54 $\pm$ 0.47	69.84 $\pm$ 0.41	73.45 $\pm$ 0.29	54.12 $\pm$ 2.97	64.48 $\pm$ 1.29	64.31 $\pm$ 0.50	60.07 $\pm$ 0.18	60.96 $\pm$ 1.11	63.45 $\pm$ 4.06	65.02 $\pm$ 2.23
OGBN-Products	78.07 $\pm$ 0.06	75.02 $\pm$ 0.23	78.35 $\pm$ 0.14	60.88 $\pm$ 0.17	67.44 $\pm$ 0.60	68.01 $\pm$ 0.52	66.92 $\pm$ 0.41	66.07 $\pm$ 0.64	67.53 $\pm$ 0.33	68.20 $\pm$ 0.69

Table 3: Results of teachers, baseline models and KMP for trans setting on large-scale graph datasets. The best performance of student model is highlighted.

Dataset	Models
Dataset	$SAGE$	$GCN$	$GAT$	$MLP$	$GLNN(SAGE)$	$KMP(SAGE)$	$GLNN(GCN)$	$KMP(GCN)$	$GLNN(GAT)$	$KMP(GAT)$
OGBN-Arxiv	70.53 $\pm$ 0.58	69.79 $\pm$ 0.41	—	54.46 $\pm$ 2.75	58.83 $\pm$ 0.33	59.31 $\pm$ 0.22	58.94 $\pm$ 0.71	60.01 $\pm$ 1.02	—	—
OGBN-Products	77.25 $\pm$ 0.22	74.93 $\pm$ 0.15	—	61.05 $\pm$ 0.36	65.24 $\pm$ 0.39	65.93 $\pm$ 0.53	64.37 $\pm$ 0.13	63.95 $\pm$ 0.17	—	—

Table 4: Results of teachers, baseline models and KMP for induc setting on large-scale graph datasets. The best performance of student model is highlighted. (’—’ means that the GNN shows poor performance and we do not use it as teacher.)

Dataset	Models(SAGE teacher, trans)			Models(SAGE teacher, induc)
Dataset	$KMP(ours)$	$KMP+PE(ours)$	$SA-MLP$	$KMP(ours)$	$KMP+PE(ours)$	$SA-MLP$
Cora	79.48 $\pm$ 1.48	80.97 $\pm$ 1.03	79.33 $\pm$ 1.55	71.85 $\pm$ 2.40	71.90 $\pm$ 1.81	71.54 $\pm$ 2.26
Citeseer	71.96 $\pm$ 2.25	72.07 $\pm$ 2.51	71.25 $\pm$ 1.57	71.22 $\pm$ 2.04	71.46 $\pm$ 2.65	70.95 $\pm$ 2.49
Pubmed	76.80 $\pm$ 2.01	76.58 $\pm$ 2.84	76.45 $\pm$ 2.50	74.92 $\pm$ 3.03	74.71 $\pm$ 2.09	74.99 $\pm$ 3.74
Amazon Photo	92.27 $\pm$ 1.50	92.38 $\pm$ 0.82	90.67 $\pm$ 0.93	90.26 $\pm$ 2.43	90.77 $\pm$ 0.84	90.27 $\pm$ 1.44
Amazon Computer	84.75 $\pm$ 1.10	85.10 $\pm$ 1.66	83.38 $\pm$ 2.50	80.94 $\pm$ 2.33	81.50 $\pm$ 2.58	81.01 $\pm$ 2.89

Table 5: Results of KMP+PE, KMP and SA-MLP for both trans and inductive setting. The best performance of the student model is highlighted.

Panel of Figure 2 — Figure 2: Results of Robustness study with different teachers on three datasets. The x-axis shows percentage of feature noise and y-axis means test accuracy(%).

Panel of Figure 3 — Figure 3: Results of sensitiveness study with different teachers on three datasets. Mean and variance of test accuracy reported. We study if changes of $\gamma$ can influence student MLP’s performance greatly. The x-axis shows different $\gamma$ s and y-axis means test accuracy(%) with them.

Dataset	Models
Dataset	$Teacher$	$MLP$	$GLNN$	$KMP$
MiniGC	91.90 $\pm$ 3.47	66.35 $\pm$ 5.36	75.44 $\pm$ 4.69	78.95 $\pm$ 5.01

Table 6: Results for graph classification on MiniGC dataset. The best performance of student model is highlighted.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Nodes	Edges	Feature dim	Classes
Cora	2708	5429	1433	7
Citeseer	3327	4732	3703	6
Pubmed	19717	44324	500	3
Amazon Photo	7487	119043	745	8
Amazon Computer	13381	245778	767	10
OGBN-Arxiv	169343	1166243	128	40

Appendix

Table 7: Details of the seven used datasets

	GCN	SAGE	GAT
number of layers	2	2	2
hidden layer dim	64	128	128
weight decay	0.001	5e-4	0.01
dropout ratio	0.8	0	0.6
attention dropout ratio	—	—	0.6
fan_out	—	5,5	—

Appendix

Table 8: Experiment settings for the 5 datasets.

$\gamma$	0.1,0.3,0.5,0.7,0.9,1,3,10,30
temperature( $t$ )	0.25,0.5,1,2,4,10
$\theta$ (percentage of hard labels)	0,0.2,0.4,0.6,0.8,1
learning rate	0.01,5e-3,1e-3,5e-4,1e-4

Appendix

Table 9: Grid Search of hyperparameters.

Panel of Figure 4 — Figure 4: Results of trainable kernel on three benchmark datasets, compared to kernels without self training

Percentage of noise	GLNN(SAGE teacher)	KMP(SAGE teacher)
10%	47.70 $\pm$ 2.91	48.91 $\pm$ 2.44
20%	40.75 $\pm$ 2.29	41.60 $\pm$ 2.47
30%	37.08 $\pm$ 3.96	37.41 $\pm$ 4.65
40%	33.96 $\pm$ 5.23	33.16 $\pm$ 5.88
50%	—	—

Appendix

Table 10: Results on Pubmed(with feature noise).

Explore similar work

Distilling Graph Geometry: Knowledge Gap from GNNs to MLPs
Oct 7, 2026Zhewei Chen, Hao Zhu, Jiaojiao Jiang +1Knowledge Distillation Cross-Architecture Knowledge Distillation
Australian National University · Data61, CSIRO · University of New South Wales
LLM-Guided Knowledge Distillation for Temporal Knowledge Graph Reasoning
Feb 16, 2026Wang Xing, Wei Song, Siyu Lin +2Knowledge Distillation Temporal KGs
School of Computer Science and Technology, Xidian University · School of Computing and Artificial Intelligence, Southwest Jiaotong University · School of Information Science and Engineering, Chongqing Jiaotong University
Heterophily-Aware Adaptive Knowledge Distillation for Hypergraph Neural Networks
Jun 8, 2026Joohee Cho, David Yoon Suk Kang, Yunyong KoEfficient Neural Network Inference Hypergraph Neural Networks
Chung-Ang University · Seoul, South Korea · Chungbuk National University +1

Oct 7, 2026J/K move · Enter open · S save

Distilling Graph Geometry: Knowledge Gap from GNNs to MLPs

Zhewei Chen, Hao Zhu, Jiaojiao Jiang, Ahad N. Zehmakan

Australian National University · Data61, CSIRO · University of New South Wales

Knowledge Distillation Cross-Architecture Knowledge Distillation Feature-Based Knowledge Distillation

Open paper AskPDF arXiv

GNN-to-MLP distillation aims to retain the predictive accuracy of a message-passing teacher while deploying a graph-free MLP at inference. Existing methods mainly transfer node-wise predictions or use confidence-based reweighting, but they do not specify where the student should preserve the teacher's graph-induced geometry. We show that this omission leads to two spectral failure modes in the student's representation space. On sparse graphs, the student suffers from spectral underfit, missing high-energy teacher directions concentrated near boundary regions. On dense graphs, it suffers from spectral overfit, retaining spurious directions that the teacher has collapsed through aggregation. Motivated by an energy-weighted teacher-student alignment objective, we propose Graph Geometry-aware MLP (G^2MLP), a training-time distillation framework guided by Ollivier-Ricci curvature. Curvature identifies where the two spectral errors concentrate and is used to allocate supervision between prediction-level and representation-level alignment. The deployed model remains a standard MLP and requires no graph access at inference. Across node-classification benchmarks, G^2MLP consistently improves over graph-free distillation baselines, reduces the teacher-student rank gap in both regimes, and transfers without architectural changes to Graph Transformer teachers and link prediction.

Abstract

Explore similar work

Distilling Graph Geometry: Knowledge Gap from GNNs to MLPs

LLM-Guided Knowledge Distillation for Temporal Knowledge Graph Reasoning

Heterophily-Aware Adaptive Knowledge Distillation for Hypergraph Neural Networks