arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2403.01079cs.LGcs.AI
收藏

教给 MLP 更多图信息:一种三阶段多任务知识蒸馏框架

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

Xi’an Jiaotong University(西安交通大学);Arizona State University(亚利桑那州立大学);Tongji University(同济大学)

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

AI总结:

针对 GNN 在大规模图推理中的高时空开销及图知识蒸馏到 MLP 时的位置信息丢失、泛化不足问题,提出融合位置编码、Neural Heat Kernels 与隐藏层匹配的三阶段多任务蒸馏框架。

摘要:

我们研究图神经网络在大规模图数据集上执行推理任务时面临的挑战性问题:巨大的时间与内存消耗,并尝试通过降低对图结构的依赖来克服它。尽管将图知识蒸馏到学生 MLP 是一个出色的想法,但它面临位置信息丢失和泛化能力低这两个主要问题。为解决这些问题,我们提出一种新的三阶段多任务蒸馏框架。具体而言,我们使用 Positional Encoding(位置编码)捕获位置信息。此外,我们引入负责 GNN 中图数据处理的 Neural Heat Kernels,并利用隐藏层输出匹配,以提升学生 MLP 隐藏层的性能。据我们所知,这是首个在图上为学生 MLP 纳入隐藏层蒸馏、并将图 Positional Encoding 与 MLP 结合的工作。我们在多种设置下测试其性能和鲁棒性,并得出结论:我们的工作能够表现优异,并具有良好稳定性。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑