[论文解读] Geometric Knowledge Distillation: Topology Compression for Graph Neural Networks

Chenxiao Yang, Qitian Wu+1 more|arXiv (Cornell University)|Oct 24, 2022

Neural Networks and Applications被引用 5

一句话总结

该论文提出几何知识蒸馏（GKD），一种新颖的框架，通过对齐编码底层图流形几何特性的神经热核（NHK），将教师图神经网络（GNN）在完整图上学习到的拓扑知识，迁移至在部分或稀疏图上运行的学生GNN。该方法在性能上达到与最优模型相当的水平，并在模型压缩、自蒸馏和在线蒸馏等多种设置下，超越了学生和教师模型的表现。

ABSTRACT

We study a new paradigm of knowledge transfer that aims at encoding graph topological information into graph neural networks (GNNs) by distilling knowledge from a teacher GNN model trained on a complete graph to a student GNN model operating on a smaller or sparser graph. To this end, we revisit the connection between thermodynamics and the behavior of GNN, based on which we propose Neural Heat Kernel (NHK) to encapsulate the geometric property of the underlying manifold concerning the architecture of GNNs. A fundamental and principled solution is derived by aligning NHKs on teacher and student models, dubbed as Geometric Knowledge Distillation. We develop non- and parametric instantiations and demonstrate their efficacy in various experimental settings for knowledge distillation regarding different types of privileged topological information and teacher-student schemes.

研究动机与目标

• 解决从完整图向在部分或稀疏图上运行的GNN迁移几何知识的挑战。
• 形式化一种系统性方法，将图拓扑编码为GNN中的几何先验，即使在无法获取完整拓扑信息时亦可实现。
• 开发一种灵活的蒸馏框架，实现基于潜在几何结构的知识迁移，而不仅依赖于特征。
• 在多种设置下验证方法的有效性：模型压缩、自蒸馏和在线蒸馏。
• 基于图上的热方程推导出的神经热核（NHK），提供一种理论基础坚实的几何知识表征。

提出的方法

• 提出神经热核（NHK）作为潜在图流形几何特性的数学表征，源自图上的热方程。
• 利用NHK矩阵封装并从在完整图上训练的教师GNN向在部分图上训练的学生GNN迁移几何知识。
• 提出两种实现方式：非参数化GKD，在对潜在空间做假设下显式计算NHK；参数化GKD，通过可微模块学习NHK。
• 通过蒸馏损失对齐教师与学生模型的NHK，促使学生模仿教师在底层流形上的信息流动态。
• 为参数化版本采用类似EM的优化策略，联合训练GNN与NHK近似模块。
• 将该框架应用于多种设置：模型压缩、自蒸馏和在线蒸馏，均在基线方法上实现一致性能提升。

实验结果

研究问题

• RQ1 即使模型仅能访问部分图，能否将图拓扑作为几何先验蒸馏进GNN模型？
• RQ2 如何形式化并实现教师与学生GNN之间由热核编码的图几何结构的转移？
• RQ3 能否通过一种基于热力学启发的系统性框架，在缺乏完整拓扑信息的情况下，提升GNN在稀疏图上的性能？
• RQ4 与传统知识蒸馏相比，所提出的几何知识蒸馏框架在性能和泛化能力方面表现如何？
• RQ5 在使用参数化与非参数化NHK近似时，训练效率与模型容量之间存在何种权衡？

主要发现

• 在在线蒸馏设置下，GKD-G、GKD-S和GKD-R在Cora数据集上的测试准确率分别为88.48±0.59、88.50±0.33和88.41±0.62，优于教师模型（84.61±0.37），并接近最优模型（88.63±0.48）。
• 在模型压缩设置中，GKD将SGC学生模型的准确率从85.93±0.17提升至87.15±0.85，将GCN-8从84.52±1.34提升至88.30±0.46，超越了学生和教师模型的表现。
• 在自蒸馏设置中，GKD将GCN-32的准确率从88.63±0.48提升至89.23±0.52，证明了其在性能优化方面的有效性。
• 参数化GKD（PGKD）在200个周期内收敛，最佳验证准确率在前100个周期内达到，训练开销可接受。
• GKD变体的收敛速度与标准知识蒸馏相当，Sigmoid和Gaussian NHK的收敛速度与标准KD一致。
• 在所有评估设置中，该方法始终优于学生模型，并与最优模型表现相当，验证了其在几何知识迁移中的有效性。