从对比链接中蒸馏自知识:非消息传递的图节点分类

这篇文章从图网络的边中蒸馏知识,使 MLP 在图节点分类任务上达到了 GNN 的表现。 不传递消息、不使用邻接矩阵、在边集上训练,因此空间占用低、分类速度快。

此外,作者还从随机点对中蒸馏“逆边”的知识,进一步提升了模型的效果。

方法

分叉 MLP

通常的图神经网络方法(GNN)基于消息传递模式,使用邻居节点传递来的消息帮助中心节点的分类。 那么反过来说,任一节点中应含有其邻居节点的部分信息。

因此构造分叉 MLP:从 MLP 的倒数第二层分叉出另一个输出层,为与原输出层相区别称为“推理层”。 分叉 MLP 接受一个节点的特征作为输入,输出该节点和它邻居节点的标签分布。

设对任一边 $(v_i, v_j)$,节点 $v_i$ 和 $v_j$ 的特征分别为 $x_i$ 和 $x_j$,通过分叉 MLP 得到 $z_i, s_i$ 和 $z_j, s_j$,节点标签为 $y_i$ 和 $y_j$。

则根据分叉 MLP 的设计,可让 $z_i$ 和 $s_j$ 逼近 $y_i$,让 $z_j$ 和 $s_i$ 逼近 $y_j$,最小化它们的交叉熵。 同时,让 $z_i$ 逼近 $s_j$,让 $z_j$ 逼近 $s_i$,最小化它们的误差(MSE)。 前者是常规的监督学习,而后者可以看作是在做自知识蒸馏。

边蒸馏

随机点对的对比蒸馏

如果随机采样一对节点,这对节点之间就(大概率)不能借助分叉 MLP 相互推导了。 因此用对比学习的思路,最大化一个节点的标签和另一个节点的推理层输出之间的交叉熵,同时最大化一个节点的输出层输出和另一个节点的推理层输出。

逆边的对比蒸馏

两种分类模式

分叉 MLP 训练好后,可以将推理层扔掉变回一个经典的 MLP 结构,输入节点特征对节点进行分类。 这种模型根据是否对逆边做对比蒸馏命名为 CoLinkDistMLP 和 LinkDistMLP。

另外,也可以保留推理层,将输出加到邻居节点上,提升邻居节点的分类准确率。 但这样一来,就相当于又用上消息传递模式了。 这种模型根据是否对逆边做对比蒸馏在文章中称为 CoLinkDist 和 LinkDist。

实验

做了归纳式、转导式和全监督三种设置下 8 个常用图数据集的节点分类实验,对比的基准除了 MLP 和 GCN 以外,还有 GCN2MLP 代表用训练好的 GCN 模型蒸馏得到的 MLP。

半监督、归纳式:

semi-supervised and inductive

半监督、转导式:

semi-supervised and transductive

全监督(转导式):

full-supervised

实验结果

  1. 即使仅仅是 LinkDistMLP,也接近甚至超过了 GCN 的表现,说明这种方式训练出的 MLP 确实能达到 GNN 的效果
  2. 用上消息传递和对比蒸馏能进一步提升模型效果
  3. GCN2MLP 在半监督时效果比 GCN 还要好,但在全监督时可用于蒸馏的知识不多,最后表现跟 MLP 差不多

此外,还用了 Open Graph Benchmark 的数据集进行节点分类实验

OGB

Page Not Found

Try to search through the entire repo.