VQGraph:重新思考图表示空间以桥接 GNN 与 MLP
机器学习
2024-03-07 v3 人工智能
计算机视觉与模式识别
摘要
GNN 到 MLP 的知识蒸馏旨在利用知识蒸馏(KD)通过学习计算高效的多层感知机(学生 MLP)在图数据上模仿教师 GNN 的输出表示。现有方法主要让 MLP 模仿 GNN 在少数类标签上的预测。然而,类空间可能不足以表达众多多样的局部图结构,从而限制了从 GNN 到 MLP 的知识迁移性能。为解决此问题,我们提出通过直接标注节点多样的局部结构来学习一个新的强大图表示空间,用于 GNN 到 MLP 的蒸馏。具体而言,我们提出 VQ-VAE 的一个变体,在图数据上学习一个结构感知的 tokenizer,可将每个节点的局部子结构编码为离散码。这些离散码构成一个码本,作为一种新的图表示空间,能够通过相应的码索引识别节点的不同局部图结构。然后,基于学习到的码本,我们提出一个新的蒸馏目标,即软码分配(soft code assignments),直接将每个节点的结构知识从 GNN 迁移到 MLP。所得框架 VQGraph 在七个图数据集的转导与归纳设定下,于 GNN 到 MLP 蒸馏中取得了新的 SOTA 性能。我们展示了 VQGraph 在取得更优性能的同时,推理比 GNN 快 828 倍,并且相对于 GNN 和独立 MLP 平均准确率分别提升 3.90% 和 28.05%。代码:https://github.com/YangLing0818/VQGraph。
引用
@article{arxiv.2308.02117,
title = {VQGraph: Rethinking Graph Representation Space for Bridging GNNs and MLPs},
author = {Ling Yang and Ye Tian and Minkai Xu and Zhongyi Liu and Shenda Hong and Wei Qu and Wentao Zhang and Bin Cui and Muhan Zhang and Jure Leskovec},
journal= {arXiv preprint arXiv:2308.02117},
year = {2024}
}
备注
ICLR 2024. Code: https://github.com/YangLing0818/VQGraph