WSDM23 | 学习蒸馏图神经网络

今天介绍的是由北京邮电大学、昆士兰大学等高校机构共同发表在WSDM 2023会议上的一篇论文：

Learning to Distill Graph Neural Networks

01 前言

图神经网络(GNNs)可以有效地获取图的拓扑和属性信息，在许多领域得到了广泛的研究。近年来，为了提高GNN的效率和有效性，将GNN与知识蒸馏相结合成为一种新趋势。然而，现有的应用于GNN的知识蒸馏方法都采用了预定义的蒸馏过程，这些过程由几个超参数控制，而不受蒸馏模型性能的监督。蒸馏和评价之间的这种隔离会导致次优结果。在这项工作中，作者提出一个通用的知识蒸馏框架，可以应用于任何预训练的GNN模型，以进一步提高它们的性能。为了解决分离问题，作者提出了参数化和学习适合GNN蒸馏的蒸馏过程。具体地说，本文不像以前的大多数工作那样引入统一的温度超参数，本文通过学习节点的特定蒸馏温度来获得更好的蒸馏模型性能。作者首先通过其邻域的编码和预测函数来参数化每个节点的温度，然后设计了一种新的迭代学习过程用于模型蒸馏和温度学习。作者还引入了一种可扩展的方法来加速模型训练。在5个基准数据集上的实验结果表明，本文提出的框架可以应用于5个流行的GNN模型，并持续提高其预测精度，平均相对增强3.12%。此外，可扩展的变体模型以1%的预测精度为代价，使训练速度提高8倍。

图神经网络(GNNs)已经成为在图上进行半监督学习的最先进的技术，在过去的五年中受到了广泛的关注。数百个GNN模型已经被提出并成功应用于各个领域，例如计算机视觉、自然语言处理和数据挖掘。近年来，为GNN配置知识蒸馏来加速模型推理，并利用无标记数据从而提高其效率或有效性是一种新兴趋势。具体来说，在典型的知识蒸馏中，轻量级模型(即学生模型)通过训练来模仿已训练模型(即教师模型)的软预测来学习知识。从效率的角度来看，可以利用知识蒸馏将一个较深的GCN模型压缩为较浅的模型，从而更快地推理。从有效性的角度来看，知识蒸馏可以提取GNN模型(教师模型)的知识，并将其注入到非GNN模型(学生模型)中，以便使用更多的先验知识和未标记的数据进行更精确的预测。

除了教师和学生的选择外，决定教师和学生模型的软预测在损失函数中如何匹配的蒸馏过程对蒸馏学生在下游任务上的预测性能也至关重要。例如，全局超参数“温度”在知识蒸馏中被广泛采用，它软化了教师模型和学生模型的预测，以促进知识的迁移。然而，现有的GNN模型上应用的知识蒸馏方法均采用预定义的蒸馏过程，即只有超参数但没有任何可学习参数。换句话说，蒸馏过程是启发式或经验式设计的，没有任何来自蒸馏学生表现的监督，这将蒸馏过程与结果评估隔离，从而导致次优结果。

在这项工作中，作者的目标是提出一个通用的知识蒸馏框架，可以应用于任何预训练的GNN模型，以进一步提高它们的性能。请注意，本文关注的是蒸馏过程，而不是学生模型的选择，因此简单地让学生模型具有与BAN所建议的老师相同的神经结构。为了克服蒸馏和评估之间的隔离问题，作者没有将全局温度作为超参数引入，而是创新地提出了由蒸馏GNN学生的性能监督的节点特定温度。具体来说，通过其邻域编码和预测的函数来参数化每个节点的温度。由于传统知识蒸馏框架中存在隔离问题，蒸馏学生的成绩对节点温度的偏导数不存在，使得在温度参数化中学习参数变得非平凡。因此，本文设计了一种新的迭代学习过程，交替执行准备、蒸馏和学习步骤来进行参数训练。

在准备步骤中，作者将利用当前参数计算每个节点的温度，并建立基于节点特定温度的知识蒸馏损失；在蒸馏步骤中，学生模型的参数将根据蒸馏损失进行更新;在学习步骤中，温度建模中的参数将更新，以提高蒸馏GNN学生的分类精度。此外，作者引入了一种可扩展的方法，通过启发式地更新节点特定的温度。

02 预备知识

2.1 节点分类

节点分类是在给定有标记节点和图结构的情况下，对图中未标记节点进行分类的一种典型的图上半监督学习任务，被广泛应用于许多GNN模型的评估中。

形式上，给定一个连通图 ， 是顶点集, 是边集，节点分类的任务是基于图结构 有标签的节点集 和节点特征 来预测没有标签的节点集 中每个节点 的标签。其中矩阵 的每一行 示节点 的 维特征。设 为节点标签的集合，则每个节点的真实标签可以表示为一个 维的独热向量 。

2.2 图神经网络

图神经网络可以通过迭代聚合邻居信息，即消息传递机制，将每个节点 编码为 维logit向量 。在本文中，作者提出的算法不是针对特定的图神经网络模型设计的，而是可以应用于任何图神经网络。因此，简单地将图神经网络编码器以黑盒形式描述化为：

其中 是图神经网络中的可学习参数， 是在softmax函数归一化之后的预测标签分布。然后图神经网络会对每个有标签的节点最小化该节点的真实标签与预测标签之间的距离，通常采用交叉熵损失来训练参数 :

2.3 知识蒸馏

在本工作中，本文关注的是蒸馏过程的研究，而不是学生模型的选择。因此，只需让教师模型和学生模型具有相同的神经结构，并分别表示为GNNT 和GNNS参数分别为 和 。给定教师模型的预训练参数，通过对GNNT 和GNNS之间的软预测进行对齐，训练学生模型的参数。从形式上讲，知识蒸馏框架旨在优化： 其中第一项是学生预测和教师预测之间的交叉熵，第二项是 中节点的学生预测与真实标签的交叉熵。许多知识蒸馏方法会引入额外的温度超参数来软化教师和学生的预测:

其中是温度超参数。温度等于1时对应原始的softmax操作。温度越高，预测就越软(趋向均匀分布)，而温度越低，预测就越硬(趋向独热分布)。在最流行的蒸馏框架[中，所有的温度都被设置为相同的超参数，即为每个节点𝑣设置。通过调整全局温度超参数，然后对经过蒸馏的学生模型进行评估，并期望其性能优于教师。

03 方法

图1 (a)用于蒸馏GNN的典型框架;(b)本文提出的蒸馏框架

3.1 参数化温度

直接为每个节点指定一个自由参数作为节点特定温度将导致严重的过拟合问题。因此，假设具有相似编码和邻域预测的节点应该具有相似的蒸馏温度。在实际应用中，每个节点 的温度可以通过一个函数来参数化，该函数需要用到以下特征:(1)学生的logit向量，它直接表征了学生模型当前的预测状态;(2) Logits向量的L2范数，由于softmax函数中的指数算子，较大的范数通常表示较硬的预测分布;(3)中心节点邻居的预测熵，描述了节点邻居的标签多样性。直观上，上述所有特性都会影响模型预测的置信度，因此在温度参数化中都应考虑。形式上，作者将所有学生模型的温度设置为1，以实现更精确的预测，并将教师温度参数化：

作者使用教师而不是学生来建模邻居的预测熵，以获得更好的数值稳定性。在实验部分将研究每个连接组件的影响，并讨论学习的温度。此外，为了避免梯度爆炸或消失的问题，作者还通过一个基于sigmoid操作的函数 ，将温度限制在 范围内。注意，用来建模温度的其他特征也可能存在，作者发现使用的三个特征足以提高性能，并且这三个特征都是有用的。

3.2 迭代学习过程

作者将标记的节点集划分为两个不相交的节点集 和 。 在损失函数的第二项中用于蒸馏， 于评估蒸馏的学生模型和学习节点温度。蒸馏部分的损失可以写成：

评估蒸馏的学生和监督温度的损失是：

由于蒸馏和评估之间的差距，评估损失只与学生模型的参数有关，评估损失对温度参数的偏导数不存在，这使它不可能通过反向传播来学习温度。为了解决这个问题，作者提出了一个新的迭代学习过程，交替执行以下准备、提炼和学习步骤: 准备步骤:首先计算每个节点 的温度，然后设定好蒸馏损失。蒸馏步骤:对于模型蒸馏，作者通过单步反向传播来更新学生模型参数:

其中 是蒸馏过程的学习率。学习步骤:作者用更新之后的学生参数计算评估损失，然后通过链式规则在温度参数上进行反向传播:

其中 是蒸馏过程的学习率。这里作者将评估损失对温度参数的偏导分解为评估损失对学生参数的偏导和学生参数对温度参数的偏导的乘积，这两项都可以分别由上述式子的偏导计算得到。通过迭代执行准备、蒸馏和学习步骤，可以训练参数化的节点特定温度，从而提高蒸馏学生的预测性能。整个算法的伪代码如下：

表1 DPK整体伪代码

04 实验

作者在图像分类和目标检测方面进行了广泛的实验。此外，还对所提出的方法进行了各种消融实验和数据分析。此外作者的训练数据与代码是公开提供的以供再现。

3.1 图像分类

作者在CIFAR-100和ImageNet数据集上评估DPK方法进行图像分类的性能。作者通过将DPK与其他多种基线模型对比，证明了模型的有效性，实验结果如表2,3所示：

表2 CIFAR-100验证集上的结果

表3 ImageNet验证集上的结果

表2显示了CIFAR-100验证集上，同构师生对方法的top-1准确率。用粗体突出显示最佳模型结果。同时表3在大规模的图像数据集ImageNet进行了实验，显示了由ResNet-34指导的ResNet-18的top-1和top-5的准确率。结果表明，DPK在性能上优于其他基线。

3.2 更好的老师，更好的学生

上述实验表明，DPK在普遍的学生老师模型对中表现良好。在此，作者展示了DPK可以进一步改进与更好的教师。如图2所示：

图2 更好的老师，更好的学生

如图2所示：通过逐步替换较大的教师模型，DPK训练的学生模型的准确性不断提高，而其他算法训练的模型性能不断波动。同时，DPK在每个阶段都超过了它对应的方法，并逐步扩大了性能差距。需要注意的是，给定学生的表现逐渐饱和是一个合理的现象，但表现的波动会在实际应用中造成许多困难。

3.3 消融实验

在本节中，作者提供了广泛的消融研究来分析DPK的每个成分的影响。实验在ImageNet上进行分类任务，采用ResNet34和ResNet18作为教师和学生。

掩码比：特征掩码是DPK的一个关键组成部分，表4报告了不同掩码比下不同DPK变体的结果：

表4 不同掩码比例消融实验

实验结果表明从15%到95%的广泛的掩码比可以为学生实现性能的提高。这意味着教师提供的先验知识非常有利于学生的网络学习。此外不同师生对下的最优掩模比并不一致，例如55%的ResNet-18和ResNet-34表现最好，而ResNet-18和ResNet-101的最优掩模比为75%。表4还显示了学生使用所提出的动态掩码策略获得了最好的准确性，这表明了所提出的动态掩码策略对于自动选择掩码比例的必要性和有效性。

先验知识：表5阐述了整合教师先验知识在构建混合学生特征中的重要性：

表5 先验知识的消融实验

具体来说，使用零填充和可学习的掩码来发挥教师的特征在掩码位置上的作用。结果表明，没有教师提供的先验知识会导致较差的表现，因为这加重了从学生到教师的特征模仿的负担。这有力地证明了通过特征掩码来证明提供来自教师的先验知识的有效性。

05 结论

在本文中，作者提出了一种新的知识蒸馏框架LTD，可以应用于任何预训练的GNN模型，以进一步提高其预测性能。作者没有像以前的大多数工作那样引入一个全局温度超参数，而是创新地提出通过蒸馏学生的表现来学习节点特定的蒸馏温度。具体而言，作者通过邻域编码和预测情况的函数来参数化每个节点的温度，并设计了一种新的迭代学习过程，用于模型提取和参数学习。作为一种成本有效的选择，LTD的可扩展变体提出了启发式更新节点特定温度。作者在五个基准数据集上进行了实验，并表明本文提出的框架可以成功地应用于五个流行的GNN模型。大量的研究进一步证明了该方法的有效性。

论文链接：

https://dl.acm.org/doi/abs/10.1145/3539597.3570480

代码链接：

https://github.com/BUPT-GAMMA/LTD