ICLR'24 | VQGraph: 重新审视图表示学习，将Graph空间Token化

ICLR 2024 | VQGraph: 重新审视图表示学习，将Graph空间Token化

论文题目： VQGraph: Rethinking Graph Representation Space for Bridging GNNs and MLPs

作者单位：北京大学，蚂蚁集团，斯坦福大学

代码链接：https://github.com/YangLing0818/VQGraph

介绍：

GNN-to-MLP distillation旨在利用知识蒸馏（KD）通过模拟teacher GNN的输出表示来学习计算高效的多层感知器（student MLP）用于图数据。现有方法主要使MLP模仿GNN对少量类别标签的预测。然而，类别空间可能不足以覆盖众多不同的局部图结构，从而限制了从GNN到MLP的知识迁移的性能。为解决这一问题，我们提出通过直接标记节点的多样局部结构来学习新的强大图表示空间用于GNN-to-MLP蒸馏。具体来说，我们提出了VQ-VAE的变体在图数据上学习一个结构感知的标记器，该标记器可以将每个节点的局部子结构编码为离散化语义token。如下图所示

这些token构成了一个代码簿（codebook），作为一个新的图表示空间，能够识别具有相应代码索引的节点的不同局部图结构。然后，基于学习的codebook，我们提出了一种新的蒸馏目标，即soft code assignments，以直接将每个节点的结构知识从GNN传输到MLP。由此产生的框架VQGraph在跨七个图数据集的传导和归纳设置中实现了新的最先进性能。我们展示了VQGraph比GNN快828倍，并且分别比GNN和独立的MLP平均提高了3.90%和28.05%的准确率。

VQGraph的训练流程：

VQGraph的训练流程包含两个阶段：

• 首先是Graph Tokenizer Training，该阶段是为了在graph空间里学习一个包含不同局部字结构信息的codebook，即码本。码本中的每一个token是一个向量，并有唯一的index去代表它。码本学习的同时，teacher GNN也会同时进行学习。该阶段训练好坏的判断标准是将学到的token分配到不同graph node上之后，能否解码重构原有的图结构信息（所有node和edge）。我们在decoder部分的设计非常简单，就是普通的MLP结构。这样设计的初衷是，如果解码器设计过于强大，那么模型的重构就会过分依赖解码器，那样就得不到很好的具有表达力的码本空间。

第一阶段训练目标函数：

• 第二部分是基于学习得到的teacher GNN和codebook，进行GNN-to-MLP训练。相比较之前的蒸馏方法，我们不仅让MLP去模仿GNN的类别预测或者特征，我们还让MLP去模仿GNN上每个node的token分布情况。这样的蒸馏过程能更好地让MLP学习到GNN中的局部结构信息，分类也会更准确。

第二阶段训练目标函数：

整体框架图：

实验分析：

本文通过一系列实验分析验证了VQGraph模型的有效性和泛化性。在节点分类场景下，VQGraph蒸馏得到的MLP分类效果在所有数据集上（包含大规模图数据）超越了之前的蒸馏方法，包括teacher GNN的表现，这是因为我们新提出的基于codebook的图表征空间比之前更加有效。

在工业级场景下，VQGraph也是相比之前的方法取得了全面的领先

我们对VQGraph的inference-accuracy之间的trade-off以及特征空间可视化都进行了探索。可以发现，我们蒸馏得到的graph MLP在速度、精度以及表达性上都要优于之前的方法。

最后我们还把VQGraph扩展到异构图数据，发现我们新提出的图表征空间依然能提升现有方法的效果，说明局部结构信息的显式化表达确实能够增强不同类型图数据分类场景下的表现。

总结

在这篇论文中，我们通过为节点的多样化局部子结构学习新的图表征空间-codebook，以及利用codebook促进结构感知的 GNN 到 MLP 蒸馏，改进了现有图表示空间的表达能力。对七个数据集进行的广泛实验表明，我们的VQGraph可以显著提高 GNN 的性能，平均准确率提升 3.90%，提高 MLP 的性能达到 28.05%，提高了最先进的 GNN 到 MLP 蒸馏方法的性能达到 1.39%，同时与 GNN 相比保持了快速的推理速度，为 828 倍。此外，我们还提供了额外的可视化和统计分析，以及消融实验，以证明所提出模型的优越性。在拓展到异构图上的实验中，我们还发现了VQGraph强大的泛化性以及在更多复杂图数据集上进行应用的可能性。