ICML 2024 | 基于特征基匹配的图蒸馏方法

©作者 | 刘洋

单位 | 北邮GAMMA Lab

研究方向 | 图神经网络

随着图数据量的不断增加，高效训练图神经网络（graph neural networks, GNNs）的需求也在不断提升。新兴的图蒸馏（graph distillation, GD）技术为这一挑战提供了解决思路，图蒸馏通过生成一个小的合成图来代替真实的大规模图，并确保在合成图上训练的 GNNs 表现出与在真实图上训练相当的性能。然而，现有的方法依赖与 GNNs 相关的信息进行监督，例如梯度、表征和轨迹，这样做有两个局限性。

首先，GNNs 可能会影响真实图的谱（spectrum，即特征值），造成在合成图上产生谱偏差。其次，使用不同架构的 GNNs 蒸馏会生成不同的合成图，为了训练不同的 GNNs 都能获得最优性能，需要遍历各种架构的 GNNs 蒸馏出相应的合成图。

为了解决这些问题，我们提出了基于特征基匹配的图蒸馏方法（Graph Distillation with Eigenbasis Matching, GDEM），该方法在蒸馏过程中对齐了真实图和合成图的特征基和节点特征，并直接复制真实图的谱来构建合成图，以避免 GNNs 带来的影响。

除此之外，我们还设计了一个判别约束项，并通过调整各损失项的权重来平衡 GDEM 的有效性和泛化性。理论分析表明，GDEM 蒸馏出的合成图保留了真实图的谱相似性。大量实验表明，GDEM 的跨架构泛化能力和蒸馏效率远远超越了现有最先进的图蒸馏方法。

论文题目： 

Graph Distillation with Eigenbasis Matching

论文链接： 

https://arxiv.org/abs/2310.09202

代码链接： 

https://github.com/liuyang-tian/GDEM

背景与动机

图神经网络（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 方法。

符号定义

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

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

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

梯度匹配中的谱偏差

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

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

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

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

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

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

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

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

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

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

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

理论分析

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

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

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

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

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

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

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

实验

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 能够保留真实图的谱相似性。

总结

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

更多阅读

#投 稿 通 道#

 让你的文字被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。

📝 稿件基本要求：

• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注 

• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题

• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算

📬 投稿通道：

• 投稿邮箱：hr@paperweekly.site 

• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿

△长按添加PaperWeekly小编

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·

·

·