解耦反向Kullback-Leibler：重新审视知识蒸馏损失以桥接GNNs与MLPs

《Pattern Recognition》：Decoupled reverse Kullback–Leibler: Rethinking knowledge distillation loss for bridging GNNs and MLPs

　　

• •
  反向KL散度的优化特性比前向KL散度更适合GKD。
• •
  传统反向KL将拟合与正则化耦合在一起，且权重固定相等。
• •
  DRKL通过可调节的超参数将两个项解耦，从而实现灵活的拟合。
• •
  DRKL在多种设置下均一致地优于前向KL和反向KL。
• •
  DRKL是一种即插即用的损失函数，可用于基于输出logits进行蒸馏的GKD框架。

引言

图神经网络（GNNs）在处理非欧几里得结构化数据方面表现出色[1]，广泛应用于社交网络[2]、[3]、[4]、生物信息学[5]、[6]、[7]和交通运输[8]、[9]等领域。然而，它们的工业部署受到“邻域获取”机制的阻碍：虽然该机制通过邻域聚合捕获关系信息，但同时也引入了大量的计算和延迟，使GNN不适合轻量化、低延迟的应用场景。

为了弥合这一差距，图知识蒸馏（GKD）[10]采用Kullback-Leibler（KL）散度将高成本教师GNN的图知识转移到低成本学生MLP中，并且仅部署MLP。尽管MLP的图拓扑表达能力有限，但它们具有高效率、快速推理和易于部署的优势[11]。因此，GKD将具有图感知能力的GNN与推理高效的MLP相结合，实现了高速下的竞争性性能。

在GKD中，KL散度用于对齐GNN和MLP的预测分布[12]。由于KL散度是不对称的，它被分为前向KL（FKL）和反向KL（RKL）[13]。大多数现有的GKD方法采用FKL，因为经验证据表明FKL优于RKL[14]。然而，当我们分析FKL和RKL的优化目标时，我们发现RKL的两个项表现出不同的行为，它们的组合对于GKD似乎具有互补性。

最小化FKL会产生一种分散拟合模式：如图1(a)所示，学生分布倾向于覆盖所有类别，包括低概率类别，因为FKL不会对教师低质量区域的高概率质量进行惩罚。相比之下，最小化RKL涉及两种模式：一个集中拟合项，专注于高概率类别；以及一个分散正则化项，鼓励更宽的分布以避免坍缩。这两个项对学生分布施加了相反的压力，这表明它们的平衡（而非单个项）决定了拟合的质量（图1(b)）。那么为什么RKL在实践中表现不佳呢？

更仔细的观察揭示了RKL损失高度耦合的特性。集中拟合项和分散正则化项具有固定且相等的权重。过高的正则化比例会抑制学生准确学习教师的预测，从而抑制分布的对齐。这一观察为分析结果与实际性能之间的差距提供了一个合理的解释，同时也为GKD开辟了一个新方向。

因此，我们提出了解耦反向Kullback-Leibler（DRKL）。通过将固定常数替换为可调节的超参数，DRKL独立调整集中拟合项和分散正则化项的权重，从而实现对多样化预测场景的灵活适应，并提高蒸馏性能。此外，DRKL作为即插即用的GKD损失函数，可以直接应用于现有的蒸馏框架，无需更改其训练流程。

我们的贡献包括：

• •
  通过从优化视角研究KL散度，我们提供了分析表明，传统RKL公式的耦合（而非RKL本身）是GKD的制约因素，这与普遍偏向FKL的经验偏好形成对比。
• •
  我们消除了传统RKL中的耦合；DRKL提供了平衡拟合项和正则化项的灵活性，以适应不同场景，从而缩小了我们分析所建议的优化行为与实际观察之间的差距。
• •
  在公共基准上的大量实验证明了我们的方法在广泛的基准数据集、教师GNN架构和GKD框架下的有效性、鲁棒性和效率，并且DRKL可以无额外成本地集成到现有的GKD方法中。

章节摘要

GNN的模型压缩

早期的图神经网络模型压缩技术主要集中在剪枝[15]、[16]、[17]和量化[18]、[19]、[20]。尽管这些方法可以通过构建稀疏图或减少乘加（MAC）操作在一定程度上加快GNN推理，但由于底层GNN架构未改变，它们的改进有限。固有的图依赖性以及推理阶段持续摄取子图使得GNN难以

方法论

在本节中，我们首先介绍图、图神经网络和知识蒸馏的相关定义。然后我们分析图知识蒸馏的机制，以检验前向和反向KL散度优化目标之间的差异。最后，我们通过解耦传统反向KL提出我们的GKD损失函数。

实验

在本节中，我们通过回答以下问题来验证我们方法的有效性、鲁棒性和效率：(1) 与GNN、MLP和现有GKD方法相比，DRKL的表现如何？(2) DRKL能否在实际生产环境中有效运行？(3) DRKL能否进一步提升最先进GKD方法的性能？(4) DRKL的鲁棒性如何？(5) DRKL在推理速度方面的表现如何？(6) DRKL的泛化能力

结论与讨论

在本文中，我们首先从优化目标的角度对知识蒸馏损失（即前向KL散度和反向KL散度）进行了深入的研究。我们的分析表明，与前向KL仅依赖分散拟合模式不同，反向KL结合了两种互补的拟合行为，与GKD目标更为匹配。然而在实际中，反向KL通常被发现表现不如前向KL。

CRediT作者贡献声明

Shaokai Xu：写作-初稿，软件，方法论，调研，资金获取，数据整理。Muyang Li：可视化，软件。Qiyan Huang：写作-审阅与编辑，方法论。Yingjie Tian：写作-审阅与编辑，监督，资源。

利益冲突声明

作者声明不存在可能影响本文报告工作的已知竞争经济利益或个人关系。

致谢

作者感谢中央高校基本科研业务费的支持。

Shaokai Xu|Muyang Li|Qiyan Huang|Yingjie Tian