【论文笔记_知识蒸馏_2022】Compressing Deep Graph Neural Networks via Adversarial Knowledge Distillation

通过对立知识蒸馏压缩深度图神经网络

摘要

深度图神经网络(GNNs)已被证明对图结构数据的建模具有表达能力。然而，深度图模型的过度堆积的体系结构使得在移动或嵌入式系统上部署和快速测试变得困难。为了压缩过度堆积的广义神经网络，通过师生架构进行知识提炼是一种有效的技术，其中关键步骤是用预定义的距离函数来度量教师和学生网络之间的差异。然而，对各种结构的图使用相同的距离可能是不合适的，并且最佳距离公式很难确定。为了解决这些问题，我们提出了一个新的图模型的对抗性知识蒸馏框架GraphAKD，它对抗性地训练一个鉴别器和一个生成器来自适应地检测和减少差异。具体来说，注意到良好捕获的节点间和类间相关性有利于深度GNNs的成功，我们建议用可训练的鉴别器从节点级和类级的角度来批评继承的知识。鉴别者区分老师的知识和学生继承的知识，而学生GNN作为一个生成器工作，旨在愚弄鉴别者。据我们所知，GraphAKD是第一个将对抗性训练引入图领域知识提取的软件。在节点级和图级分类基准上的实验表明，GraphAKD大幅度提高了学生的成绩。结果表明，GraphAKD能够准确地将知识从复杂的教师GNN转移到简洁的学生GNN。

图1:节点分类精度与图形大小的关系。每个气泡的面积与模型的参数数量成比例。带的型号名称表示变体。这些统计数据来自OGB排行榜。*

1.介绍

近年来，图形神经网络(GNNs)已经成为图形相关应用的标准工具包，包括推荐系统[20，57]，社交网络[19，29]和生物化学[10，13]。然而，在大规模图上表现出强大表达能力的gnn往往是过度参数化的[22]。随着包括微软学术图形(MAG) [45]和开放图形基准(OGB) [23]在内的大规模图形基准的涌现，复杂和过度堆叠的GNNs，30]已经被开发出来以实现最先进的性能。图1展示了模型性能与图形大小(即图形中的节点数量)的关系。我们注意到，在大规模图上，深度和复杂的GNNs明显优于浅层模型，这意味着过参数化GNNs的强大表达能力。然而，过度堆叠的架构频繁且不可避免地降低了GNNs [30]的参数效率和时间效率，这使得它们不适用于计算受限的平台，如移动或嵌入式系统。为了压缩深层gnn并保持其表达能力，我们探索了近年来越来越受到关注的图域知识提取技术。

在视觉学习和语言建模任务[2，25]中，知识提炼已被证明对压缩巨大的神经网络非常有效，尤其是在师生架构中。主要思想是，学生网络模仿教师网络的行为，以获得有竞争力的甚至更好的表现[14，37]，而教师网络将软目标、隐藏的特征图或层对之间的关系作为提取的知识转移到浅层学生网络。然而，使知识提炼适应于图域的现有算法[53，55，59]主要提出专门设计的固定距离函数来测量教师和学生图模型之间的差异，这导致以下两个固有的限制：
1.他们用手工制作的距离函数迫使学生网络模仿教师网络，其中的最佳公式很难确定[48]。更糟糕的是，王等人[46，47]已经指出，以这种方式训练的学生的表现总是次优的，因为很难从老师那里了解准确的分布。
2.预定义的固定距离不适于测量教师和学生在不同特征空间中的分布差异。例如，由于文本数据和视觉数据之间的内在差异，引用网络和图像网络具有不同的特征空间。4.2节的实验也证实了这一说法。

在本文中，我们提出了一个新的对抗性知识提取框架GraphAKD来解决上述问题。具体来说，我们开发了一个可训练的鉴别器，从节点表示和逻辑的角度区分学生和教师，而不是用手工设计的距离函数强迫学生网络精确地模仿教师网络。鉴别器将师生架构修改为生成性对抗网络(GANs) [17]，其中学生模型作为生成器工作。两个标识符构成鉴别符，即代表标识符和logit标识符。表示识别器通过批评连接节点的局部相似性和(补丁-摘要)对的全局相似性来区分学生和教师节点表示，而逻辑识别器利用剩余多层感知器(MLP)来区分教师和学生逻辑。我们认为提出的鉴别器是拓扑感知的，因为它考虑了图形结构。生成器，即学生网络，被训练来产生类似于教师分布的节点表示和逻辑，使得鉴别器不能区分。通过交替优化鉴别器和生成器，GraphAKD能够将节点间和类间相关性从复杂的教师GNN转移到紧凑的学生GNN。我们还注意到，鉴别器比诸如Kullback-Leibler (KL)散度和欧几里德距离等预定义的距离公式更宽容。我们可以将可训练的鉴别器视为教学助手，帮助弥合教师和学生模型之间的能力差距。此外，提出的拓扑感知鉴别器可以降低过拟合的风险。

为了评估我们提出的GraphAKD的有效性，我们在八个节点分类基准和两个图分类基准上进行了大量的实验。实验表明，GraphAKD能够使紧凑的学生gnn获得与深度教师模型相当甚至更好的结果，而只需要相应教师的10%~40%的参数。
表1：符号词汇表

2.背景

在这一部分，我们回顾了GNNs的基本概念。表1总结了主要符号。gnn被设计为非欧几里得数据的卷积的扩展[4]。在本文中，我们主要选择基于消息传递的GNNs [16]作为教师和学生模型，其中消息在节点之间交换，并用神经网络进行更新[16]。设G = (V，ε)表示具有节点v ∈ V的特征向量X_v的图。我们对两个任务感兴趣，即(1)节点分类，其中每个节点v ∈ V具有相关联的标签y_v，并且目标是学习v的表示向量h_v，其有助于预测v’s标签为ˆy_v=f(h_v)；以及(2)图分类，其中给定一组图{G1，…，GN }以及相应的标签{yG1，…，yGN }，目标是学习概括向量sG，使得整个图的标签可以被预测为’ yG = f (sG)。因此，我们将一般的GNN分解成一个编码器和一个分类器。编码器遵循消息传递方案，其中节点v ∈ V的隐藏嵌入hv(K+1)根据在第k次迭代中从v’s图邻域N (v)聚集的信息来更新。这种消息传递更新[18]可以表示为：
注意，如果k = 0，则初始嵌入被设置为所有节点的输入特征，即h_v(0) = X_v，∀v ∈ V。在运行GNN消息传递的K轮迭代之后，我们导出每个节点v ∈ V的最终表示hv = hv^(K)。对于图级分类，读出函数聚集最终节点嵌入以获得整个图的汇总向量sG，即：

其中读出函数可以是简单的排列不变函数，例如最大汇集和平均汇集[19，50]。然后，分类器读入节点或图的最终表示，用于节点级或图级分类，即：

这里我们通常把zv(或zG)和ˆyv(或ˇyg)分别解释为一个节点(或一个图)的logit和prediction。

3.方法

在这一部分中，我们首先在3.1节中介绍我们的对抗性知识提炼框架。接下来，我们在第3.2节中详述了建议的表示标识符，在第3.3节中详述了logit标识符。
3.1图形神经网络的对抗性知识蒸馏框架
我们注意到，一个强大的图域知识提取方法应该能够1)自适应地检测教师和学生在各个领域的图上的差异；2)尽可能继承老师知识；以及3)以拓扑感知的方式传递教师知识。
之前的工作[8，46，52]已经证明，距离函数，如ℓ_p距离和KL散度，对于容量较小的学生模型来说过于强大。更糟糕的是，Wang等人[48]宣称很难确定哪个距离公式是最优的。因此，我们开发了第一个对抗的图域提取框架来满足第一个要求。此外，受中间表示可以为知识转移提供hint的事实的启发[37]，我们利用从深度教师模型导出的节点表示和逻辑来改进紧凑学生网络的训练。最后，为了满足第三个要求，我们开发了一个拓扑感知鉴别器，它刺激学生网络模仿教师，并产生相似的连接节点的局部相似性和补丁-摘要对的全局相似性。
图2展示了GraphAKD的整体架构，它在一个双人极大极小游戏中针对拓扑感知鉴别器对学生模型进行对抗性训练。学生GNN作为一个生成器，产生类似于教师输出的节点嵌入和逻辑，而鉴别器旨在区分教师输出和生成器产生的内容。极小极大博弈确保学生网络通过对抗损失完美地模拟均衡状态下教师知识的概率分布[46，47]。

除了教师和学生gnn之外，GraphAKD还包括一个新颖的鉴别器，它可以被分解为一个表示标识符和一个logit标识符。作为节点分类任务在现有的蒸馏文献[53，55，59]中很流行，我们以节点分类为例，描述算法1中GraphAKD对节点级分类的训练过程。至于图级算法，请参考附录a，设De和Dl是分别作用于节点嵌入和逻辑的两个标识符。假设我们有一个预先训练的老师GNN HT和伴随的知识，即节点表示HT和逻辑ZT。目标是用更少的参数训练学生GNN GS，同时保留老师GT的表达能力。也就是说，我们期望学生GS可以生成高质量的节点表示HS和logits ZS，以实现与老师的竞争表现。由于可训练鉴别器比包括KL-divergence和ℓ푝距离[8，52]在内的大多数特定距离更加灵活和宽容，GraphAKD使学生gnn能够捕捉节点间和类间的相关性，而不是模仿教师知识的精确分布。

图2:建议的对抗性知识提炼框架GraphAKD的图示。
3.2表示标识符
在本节中，我们将介绍我们的拓扑感知表示标识符。请注意，在第2节中，我们将GNN分解为编码器和分类器。表示标识符集中于GNN编码器的输出，即最终的节点嵌入。我们不是直接匹配教师和学生的特征图[55]，而是从局部和全局视图中提取教师模型的节点表示。
…见原文

3.3 logit标识符
在本节中，我们将介绍所提议的框架的第二个标识符，它对输出逻辑进行操作。为了便于表示，我们在节点级分类的上下文中引入了logit标识符。基于GNN的分类器的输出是类别的概率分布。该概率是通过对最后一个完全连接的层的输出应用softmax函数(也称为logits)而得到的。通过利用对抗性训练，我们旨在将课堂间的相关性[54]从复杂的教师gnn转移到紧凑的学生gnn。

我们不是通过最小化KL-divergence [21]或其他预定义的距离来强迫学生精确地模仿教师，而是通过logit标识符来传递隐藏在教师逻辑中的类间相关性。受视觉表征学习[48，52]中的对抗性训练的启发，我们的逻辑标识符被训练来区分学生逻辑和教师逻辑，而生成器(即学生GNN)被对抗性训练来欺骗标识符。也就是说，我们期望紧凑的学生GNN输出与教师逻辑相似的逻辑，使得标识符不能区分。

由于剩余学习可以缩小教师和学生之间的差距[15]，我们使用具有剩余连接的MLP作为我们的logit标识符Dℓ.每层中隐藏单元的数量与logit的维数相同，logit的维数等于类别C的数量。一个简单的标识符读入每个节点的logit，并预测一个二进制值“真/假”,该值指示logit是由教师还是学生导出的。将教师和学生导出的节点푣的对数分别表示为z_vT和z_vS。简单标识符D旨在最大化对数似然，因为：

正如Xu等人[52]指出的，普通版本速度慢且不稳定，我们遵循[36，52]并修改目标以预测特定的节点标签。因此，Dℓ的输出是C + 1维向量，第一个C用于标签预测，最后一个用于真/假(即老师/学生)指示符。因此，我们最大限度地为Dℓ培训以下目标：

至于生成器的训练，即学生GNN GS，我们遵循[24，52]并且除了类别级对齐之外，还在教师和学生逻辑之间引入实例级对齐。因此，我们最小化训练GS的损失函数：

与现有的GNNs知识提炼模型[53，59]相比，logit标识符放松了学生和教师之间的严格耦合。此外，对抗性训练方法减轻了失败的痛苦。

表2:八个节点分类基准的统计数据。

4.实验

…

5.相关工作

在这一部分中，我们首先在第5.1节中介绍了将知识提炼应用于图领域的现有工作。接下来，我们在5.2节中回顾对抗训练在图域中的应用。

5.1图模型知识蒸馏
知识蒸馏在视觉学习和语言建模任务的网络压缩方面取得了巨大成功[2，25，37，39]。然而，将视觉学习和语言建模中已建立的方法直接应用于图域是不适用的，因为图包含特征和拓扑结构。

最近在GNNs中的成功推动了GNNs知识提炼的出现。在现有的工作中，[53，55，59]是与本文最相关的，因为它们都使用教师-学生架构，并专注于节点分类。然而，自适应蒸馏策略使我们的工作不同于现有的研究[53，55，59]。具体来说，LSP [55]将教师和学生的节点表示与基于核函数的距离对齐；张等人[59]和杨等人[53]使用欧氏距离来匹配教师和学生的概率分布。与此相反，我们提出的GraphAKD对抗性地训练一个鉴别器和一个生成器来自适应地检测和减少教师和学生之间的差异。此外，[53，55，59]仅进行节点级分类，并专注于少于100K个节点的图，而提出的GraphAKD广泛适用于节点级和图级分类任务，并在节点数从2K到2M变化的图上表现良好。

5.2图模型的对抗性训练
对抗训练的思想源于生成对抗网络(GANs) [17]，其中生成器和鉴别器相互竞争以提高它们的性能.近年来，对抗性训练在不同目标的图域中显示出优越的性能。具体而言，Wang等人[44]和Feng等人[11]利用对抗性架构来学习通用和健壮的图表示，而Dai等人[9]则探索了对图结构数据的对抗性攻击。Alam等人[1]使用图模型执行对立的域适应，而Suresh等人[42]开发了对立的图增强以提高自我监督学习的性能。与他们不同的是，这项工作旨在对图模型进行对立的知识提炼，这导致了贡献。

6.结论

在大规模图形数据上，过度堆叠的gnn通常具有很强的表现力和强大的功能。为了压缩深层GNNs，我们提出了一种新的图领域对抗性知识提取框架GraphAKD，它首次将对抗性训练引入到拓扑软件知识转移中。通过对抗性地训练鉴别器和生成器，GraphAKD能够将节点间和类间相关性从复杂的教师GNN转移到紧凑的学生GNN(即生成器)。实验表明，在10个基准数据集上，GraphAKD在节点级和图级任务上产生了一致且显著的改进。以这种方式训练的学生GNNs实现了与他们的教师图模型竞争甚至更好的结果，而只需要很小比例的参数。在未来的工作中，我们计划探索GraphAKD在分类之外的图任务上的潜在应用。