Graph Distillation with Eigenbasis Matching

「作者」 Yang Liu, Deyu Bo, Chuan Shi

「单位」 北京邮电大学

「摘要」 随着图数据量的不断增加，高效训练图神经网络（graph neural networks, GNNs）的需求也在不断提升。新兴的图蒸馏（graph distillation, GD）技术为这一挑战提供了解决思路，图蒸馏通过生成一个小的合成图来代替真实的大规模图，并确保在合成图上训练的GNNs表现出与在真实图上训练相当的性能。然而，现有的方法依赖与GNNs相关的信息进行监督，例如梯度、表征和轨迹，这样做有两个局限性。首先，GNNs可能会影响真实图的谱（spectrum，即特征值），造成在合成图上产生谱偏差。其次，使用不同架构的GNNs蒸馏会生成不同的合成图，为了训练不同的GNNs都能获得最优性能，需要遍历各种架构的GNNs蒸馏出相应的合成图。为了解决这些问题，我们提出了基于特征基匹配的图蒸馏方法（Graph Distillation with Eigenbasis Matching, GDEM），该方法在蒸馏过程中对齐了真实图和合成图的特征基和节点特征，并直接复制真实图的谱来构建合成图，以避免GNNs带来的影响。除此之外，我们还设计了一个判别约束项，并通过调整各损失项的权重来平衡GDEM的有效性和泛化性。理论分析表明，GDEM蒸馏出的合成图保留了真实图的谱相似性。大量实验表明，GDEM的跨架构泛化能力和蒸馏效率远远超越了现有最先进的图蒸馏方法。

1. 背景与动机

图神经网络（Graph neural networks, GNNs）在各种图相关任务中已被证明是有效的。然而，图的非欧几里得性质为GNNs的训练效率和可扩展性带来了挑战。为了加速训练，一种以数据为中心的方法是将大规模图简化为一个更小的图。传统方法包括图稀疏化和图粗化，然而，这些方法通常旨在优化一些启发式指标，例如谱相似性和成对距离，这些指标可能与下游任务无关，从而导致次优的性能。

近年来，图蒸馏（GD），也称作图压缩，凭借其显著压缩比和无损性能在图简化领域引起了广泛的关注。GD的目标是合成一个小图，使得在该小图上训练的GNNs能够表现出与在真实大图上训练的GNNs相当的性能。为实现这一目标，现有方法通过匹配一些与GNNs相关的信息（如梯度、表征和训练轨迹）来优化合成图。合成图借此将分布对齐到了真实图，同时还结合了来自下游任务的信息。

尽管取得了相当大的进展，现有图蒸馏（GD）方法需要预先选择一个特定的GNN作为蒸馏模型，这造成了两方面限制：

「(1) 用于蒸馏的GNNs会影响真实谱，导致合成图产生谱偏差，即少数特征值主导了数据分布。」

Figure 1展示了真实图和合成图的总变差（Total Variation, TV），TV通常用于反映信号在图上的平滑度，小的TV值表示低频分布，反之亦然。我们可以观察到，由低通滤波器蒸馏的合成图中的TV值始终低于真实图的TV值，而高通滤波器的情况则相反，从而验证了谱偏差的存在。

「(2) 为获得最优性能，需要遍历各种GNNs架构，这导致了不可忽视的计算成本。」

Table 1展示了GCOND在六个常用GNNs（包括GCN，SGC，PPNP，ChebyNet，BernNet和GPRGNN）的跨架构结果。可以看到，每行中不同GNNs的性能差异很大。因此，现有GD方法需要遍历多种GNNs架构来进行蒸馏，以获得最优性能，这显著增加了时间开销。

根据现有方法的不足，一个自然的问题由此产生：「如何在不受GNNs影响的情况下进行图蒸馏？」因此，我们提出了基于特征基匹配的图蒸馏方法（GDEM）。具体来说，GDEM将图结构分解为特征值和特征基，在蒸馏过程中，GDEM只匹配真实图和合成图的特征基和节点特征，从而均衡地保留不同频率的信息，解决了谱偏差问题。此外，我们优化时还引入了一个判别损失项以提高GDEM的性能，并通过调整损失项之间的权重来平衡其有效性和泛化性。在完成优化后，GDEM利用真实图的谱和合成特征基构建完整的合成图，从而防止谱受到GNNs的影响，并确保了合成图的唯一性，从而避免了遍历需求，提高了蒸馏效率。

本文的主要贡献如下：(1) 我们系统分析了现有蒸馏方法的局限性，包括谱偏差和遍历需求。(2) 我们提出了GDEM，一个新的图蒸馏框架，通过匹配特征基而不是整个图结构来减轻对特定GNNs的依赖。此外，理论分析证明了GDEM在蒸馏过程中保留了谱的相似性。(3) 我们在七个图数据集上进行了大量的实验，验证了GDEM在有效性、泛化性和效率等方面超越了现有最先进的GD方法。

2. 符号定义

本文主要以节点分类的任务场景为例，给定图 ，其中  是节点集合且 ， 表示边集合， 是节点特征矩阵。图  的邻接矩阵为 ，如果节点  和  之间存在边，则 ，否则 。对应的归一化拉普拉斯矩阵为 ，其中  是单位矩阵， 是度矩阵，其中对于每个节点 ，，当  时，。

「特征基和特征值」 归一化图拉普拉斯矩阵可以分解为 ，其中  是特征值，且 ， 是特征基，由一组特征向量组成。每个特征向量  有一个对应的特征值 ，且 。

「图蒸馏」 图蒸馏（GD）旨在从真实的大图  中蒸馏出一个小的合成图 ，其中  且 ，使得在  和  上训练的GNNs将具有相似的性能，从而加速GNNs的训练。现有框架可以分为三类：梯度匹配、分布匹配和轨迹匹配。

3. 梯度匹配中的谱偏差

我们系统分析了梯度匹配的目标函数，并为GDEM的设计提供了灵感与指导。我们从一个简单的例子开始：当蒸馏模型为单层GCN，GNNs的目标函数为MSE Loss时：

其中， 是模型参数。真实图和合成图上的梯度分别为：

假设梯度匹配的目标是两个梯度之间的MSE Loss，即 ， 的上界可表示为：

其中， 和  是用于监督合成图更新的两个目标分布，根据以下分析，它们都会被少数特征值主导，导致谱偏差。

「引理 3.1」 当蒸馏模型GCN经过多层堆叠后，目标分布会由最小的特征值主导。

「证明」 此时目标分布可以表示为：

其中，是层数。当趋于无穷大时，只有最小的特征值  保留其系数，即 ，其他系数趋近于0，因此，目标分布  被  主导。 同理可得。

「引理 3.2」 假设蒸馏GNNs的滤波函数为 ，目标分布将由滤波值大于1（即 ）的特征值主导。

「证明」 此时蒸馏GNNs的目标函数为 ，目标分布变为  和 。因此，只有滤波值  的特征值保留了系数，并主导目标分布。

引理3.1和3.2表明，在蒸馏过程中使用GNNs相关的信息会在目标分布中引入谱偏差，导致合成图只能匹配真实图数据分布的一部分，导致其结构信息不完整。

4. 方法：GDEM

与先前方法（例如，梯度匹配Figure 2(a)和分布匹配Figure 2(b)）相比，GDEM（见Figure 2(c)）不依赖于特定的GNNs，其蒸馏过程分为两个步骤：(1) 匹配真实图和合成图之间的特征基和节点特征。(2) 使用合成的特征基和真实图的谱构建合成图。

4.1. 特征基匹配

图的特征基代表了其关键的结构信息，例如，较小特征值的特征向量反映了全局社区结构，而较大特征值的特征向量编码了局部细节。由于特征向量的数量与图中的节点数相同，我们无法在合成图中保留所有真实的特征基，因此，GDEM匹配了个最小特征值和个最大特征值的特征向量，其中和是超参数，且 。这种方法在图粗化和谱图神经网络中被普遍证明是有效的。我们通过优化矩阵  以匹配真实图的主要特征基 。为消除GNNs的影响，GDEM在蒸馏过程中不使用谱信息，因此，公式3中的第一个项变为：

其中， 和  是由真实图和合成图中的第个特征向量构成的子空间。

此外，作为图傅里叶变换的基，特征向量具有正交归一化的性质，然而，通过梯度下降直接优化不能保留这种性质，因此，我们使用了一个额外的正则化来约束表示空间：

4.2. 判别约束

我们在实践中发现，特征基匹配虽然提高了GDEM的跨架构泛化能力，但由于它仅保留全局分布（即 ），而没有考虑下游任务的信息，因此对节点分类性能的贡献较小。因此，我们需要近似公式3中的第二项，我们发现  可以看作是类别级表示，即每个类有一个维的表示。然而，MSE Loss仅强调真实图和合成图之间的类内相似性，而忽略了类间差异性。基于这一发现，我们设计了一个判别约束来有效保留类别级的信息，这也可以看作一种类别感知正则化技术。具体来说，我们首先学习真实图和合成图的类别级表示：

其中，是真实图拉普拉斯矩阵的第个特征值。然后我们约束和之间的余弦相似度：

虽然判别约束在蒸馏过程中引入了谱信息，这与特征基匹配相冲突，然而，我们发现调整特征基匹配和判别约束的权重可以有效平衡GDEM的性能和泛化性。

4.3. 优化目标函数和合成图构建

GDEM的整体损失函数为三个正则项的加权和：

其中，，和是超参数。GDEM的伪代码如Algorithm 1所示。

完成优化后，GDEM的输出是合成图的特征基和节点特征，由于缺乏图谱，此时数据还不完整。由于谱编码了图的全局结构，理想情况下，如果合成图保留了真实图的分布，它们的谱应该尽可能相似。因此，我们直接复制真实图的谱到合成图中，以构建其拉普拉斯矩阵或邻接矩阵：

5. 理论分析

我们进一步进行了理论分析，证明GDEM能够保留了真实图的谱相似性。

「定义 5.1」 「Restricted Spectral Similarity（RSS）」如果存在满足以下条件，则称合成图拉普拉斯矩阵保留了真实图拉普拉斯矩阵的RSS。

「命题 5.2」 GDEM蒸馏出的合成图是真实图的-谱近似。

「证明」 我们使用个主要特征值和特征向量作为真实图的截断表示：

其中，， 表示特征基  的非正交项，因此  是严格正交的。

结合公式11和12，我们有：

根据上述不等式，特征基匹配的目标函数是真实图和合成图之间谱差异的上界，优化  和  能够使上界更小，从而保留真实图的谱相似性。合成图是真实图的-谱近似，且 。

6. 实验

6.1. 节点分类

为了验证我们的方法的有效性，我们在七个数据集 (五个同配图数据集Citeseer，Pubmed，Ogbn-arxiv，Filckr和Reddit，以及两个异配图数据集Squirrel和Gamers）上与七个基线模型比较节点分类性能。

实验结果表明：

• GD方法始终优于传统方法（核心集和图粗化）。其原因主要有两方面：一方面，GD方法可以利用GNNs强大的表示学习能力来合成图数据。另一方面，蒸馏过程涉及下游任务的信息，相比之下，传统方法只能利用结构信息。
• GDEM在七个图数据集中有六个达到了最先进的性能，表明了其在保留真实图分布方面的有效性。现有GD方法严重依赖GNNs的信息来蒸馏合成图，GDEM的结果表明，匹配特征基也可以合成质量良好的图。
• GDEM在Ogbn-arxiv上的表现略逊一筹，但在其他大规模图上的性能良好。我们推测这是因为在0.05%-0.50%的压缩比下，只有几百个特征向量进行了特征基匹配，这不足以覆盖Ogbn-arxiv中所有有用的子空间。

6.2. 跨架构泛化性

我们评估了由GCOND，SFGC，SGDD和GDEM四种不同GD方法蒸馏出的合成图，在六个GNNs（GCN，SGC，PPNP，ChebyNet，BernNet和GPR-GNN）上的泛化能力。

我们发现：

• 除了Ogbn-arxiv，GDEM在所有数据集上的平均准确率均最高。这表明GDEM蒸馏出的合成图能够代替真实图训练各种GNNs，并表现出一致良好的性能。
• GDEM显著减少了不同GNNs之间的性能差距。GCOND的方差比GDEM高出2-6倍；SGDD将结构信息传播到合成图，表现出比GCOND更好的泛化能力，这表明保留图结构可以提高合成图的泛化能力；SFGC提出了无结构的蒸馏策略，然而由于缺乏明确的图结构，这种策略可能导致应用场景受限。

6.3. 最优性能和时间开销

我们在Pubmed数据集上，通过遍历各种GNNs架构来比较不同GD方法的最佳性能和时间开销。

在Table 4中可以看到，与Table 3中的结果相比，通过遍历不同的GNNs架构，GCOND和SGDD的性能有所提高，然而，这种策略也在蒸馏阶段引入了额外的计算成本。如Table 5所示，GCOND和SGDD的复杂度与蒸馏GNNs的复杂度相关，当选择复杂度高的GNNs（例如BernNet）时，它们的时间开销将显著增加。而GDEM在性能上仍然显著优于GCOND和SGDD的遍历结果，并且GDEM的复杂度不会受到GNNs的影响，同时不需要进行遍历。因此，GDEM的整体时间开销显著小于GCOND和SGDD，这验证了GDEM的效率。

6.4. 消融实验

我们在Pubmed和Gamers数据集上进行了消融实验，以验证不同正则项（即 ， 和 ）的有效性。

「模型分析」 Table 6验证了不同正则项的作用。 和  提升了GDEM的泛化能力，去掉它们中的任何一个都会显著增加GNNs的方差。 会损害GDEM的泛化性，原因是判别约束使用了图谱信息并引入了一定的低频谱偏差，但  提高了GDEM的性能。

「参数分析」 我们进一步对  和  权重超参数的影响进行了探究，如Figure 3所示。随着的增加，GDEM的方差逐渐减少，然而较高的值会导致性能下降。增加的值会持续增加GDEM的方差，并且当超过一定阈值后，准确率也会随之下降。

6.5. 可视化实验

Figure 4展示了由GCOND，SGDD和GDEM蒸馏出的合成图的TV值分布，可以观察到GDEM中的TV值最接近真实图。SGDD比GCOND更接近真实图的分布，这表明SGDD能够更好地保留结构信息，然而其性能仍不如GDEM，这验证了特征基匹配的有效性。

此外，我们在不同的训练周期中对由GDEM蒸馏出的合成图进行了可视化，如Figure 5所示。可以观察到，随着GDEM的优化，合成图中的TV值逐渐接近真实图（0.42 → 0.73 → 0.88），这验证了命题 5.2，即，GDEM能够保留真实图的谱相似性。

7. 总结

本文提出了一种新的基于特征基匹配的图蒸馏方法，该方法仅对真实图和合成图的特征基和节点特征进行对齐，从而缓解了现有方法的谱偏差问题，避免了遍历需求。我们从理论上证明了，GDEM能够保留真实图的-谱近似。大量实验验证了所提出方法的有效性、泛化性和高效性。