STRIDE: Structure and Embedding Distillation with Attention for Graph Neural Networks

📄 arXiv: 2310.15938v2 📥 PDF

作者: Anshul Ahluwalia, Payman Behnam, Rohit Das, Alind Khare, Biswadeep Chakraborty, Pan Li, Alexey Tumanov

分类: cs.LG

发布日期: 2023-10-24 (更新: 2025-08-16)

💡 一句话要点

提出STRIDE以解决图神经网络压缩中的知识蒸馏问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 图神经网络 知识蒸馏 模型压缩 注意力机制 中间层学习 图结构对齐 节点嵌入

📋 核心要点

  1. 现有的GNN压缩方法大多忽略中间层输出，导致在高压缩比时准确率显著下降。
  2. STRIDE通过注意力机制识别重要的中间层教师-学生对，专注于对齐图结构和节点嵌入。
  3. 在OGBN-Mag数据集上，STRIDE实现了2.13%的准确率提升，压缩比达到32.3倍，显示出其有效性。

📝 摘要（中文）

近年来，图神经网络（GNNs）的发展导致模型规模增大，以提高其能力和准确性。然而，这些大型模型的高内存使用、延迟和计算成本限制了其推理部署。现有的GNN压缩技术通过知识蒸馏（KD）将大型GNN压缩为较小的模型，但大多数KD方法仅考虑最后一层的输出，忽略了中间层的重要信息。为了解决这一问题，本文提出了一种新的KD方法——结构与嵌入蒸馏（STRIDE），利用注意力机制识别重要的中间层教师-学生对，并对齐图结构和节点嵌入。实验结果表明，STRIDE在多个数据集上表现出色，尤其是在OGBN-Mag数据集上实现了2.13%的准确率提升，同时压缩比达到32.3倍。

🔬 方法详解

问题定义：本文旨在解决图神经网络压缩中的知识蒸馏问题，现有方法主要关注最后一层输出，忽略中间层信息，导致准确率下降。

核心思路：STRIDE提出通过注意力机制识别重要的中间层教师-学生对，利用这些对齐图结构和节点嵌入，从而有效提升压缩模型的性能。

技术框架：STRIDE的整体架构包括教师模型和学生模型，利用注意力机制选择中间层对，并通过蒸馏损失函数进行训练，确保学生模型能够学习到重要的结构信息。

关键创新：STRIDE的创新在于引入了中间层的知识蒸馏，特别是通过注意力机制选择重要层对，与传统方法相比，能够更好地保留模型的性能。

关键设计：在设计中，STRIDE使用了特定的损失函数来平衡中间层和最终层的蒸馏，同时在不同模型架构（如GCNIIs、RGCNs和GraphSAGE）上进行了优化，以确保广泛适用性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在OGBN-Mag数据集上，STRIDE实现了2.13%的准确率提升，压缩比达到32.3倍，相较于最先进的方法表现出显著优势。在较小的数据集（如Pubmed）上，STRIDE也达到了141倍的压缩比，同时保持了与最先进方法相同的准确率。

🎯 应用场景

STRIDE的研究成果在图数据分析、社交网络分析和推荐系统等领域具有广泛的应用潜力。通过有效压缩GNN模型，能够在资源受限的环境中实现高效推理，推动智能系统的实际部署和应用。

📄 摘要（原文）

Recent advancements in Graph Neural Networks (GNNs) have led to increased model sizes to enhance their capacity and accuracy. Such large models incur high memory usage, latency, and computational costs, thereby restricting their inference deployment. GNN compression techniques compress large GNNs into smaller ones with negligible accuracy loss. One of the most promising compression techniques is knowledge distillation (KD). However, most KD approaches for GNNs only consider the outputs of the last layers and do not consider the outputs of the intermediate layers of the GNNs. The intermediate layers may contain important inductive biases indicated by the graph structure and embeddings. Ignoring these layers may lead to a high accuracy drop, especially when the compression ratio is high. To address these shortcomings, we propose a novel KD approach for GNN compression that we call Structure and Embedding Distillation with Attention (STRIDE). STRIDE utilizes attention to identify important intermediate teacher-student layer pairs and focuses on using those pairs to align graph structure and node embeddings. We evaluate STRIDE on several datasets, such as OGBN-Mag and OGBN-Arxiv, using different model architectures, including GCNIIs, RGCNs, and GraphSAGE. On average, STRIDE achieves a 2.13% increase in accuracy with a 32.3X compression ratio on OGBN-Mag, a large graph dataset, compared to state-of-the-art approaches. On smaller datasets (e.g., Pubmed), STRIDE achieves up to a 141X compression ratio with the same accuracy as state-of-the-art approaches. These results highlight the effectiveness of focusing on intermediate-layer knowledge to obtain compact, accurate, and practical GNN models.