将GNN逐层蒸馏为MLP,提升推理速度同时保持性能
TINED: GNNs-to-MLPs by Teacher Injection and Dirichlet Energy Distillation
- 分层蒸馏:用教师注入和狄利克雷能量匹配GNN的特征变换与图传播
- 在7个数据集上超越GNN和现有蒸馏方法,最高提升12.3%准确率
- 适合需要低延迟推理的图学习场景,如实时推荐系统
图神经网络(GNN)在节点分类等图学习任务中表现优异,但其推理依赖多跳信息,难以满足低延迟需求。现有将GNN蒸馏为多层感知机(MLP)的方法常忽略层次结构信息。本文提出TINED,通过教师注入与狄利克雷能量蒸馏,实现逐层蒸馏。我们识别出GNN中的特征变换(FT)在计算上等价于MLP的全连接层,因此直接迁移参数并微调。MLP中的全连接层复现GNN中FT与图传播(GP)的顺序。我们建立GP近似的理论边界,并发现FT与GP具有相反的平滑效应:GP激进,FT保守。基于狄利克雷能量,我们设计DE比率量化该特性,并提出能量蒸馏机制传递此信息。大量实验表明,TINED在七种数据集上均优于原始GNN及主流蒸馏方法,显著提升推理效率。
原文摘要 · Abstract (English)
Graph Neural Networks (GNNs) are pivotal in graph-based learning, particularly excelling in node classification. However, their scalability is hindered by the need for multi-hop data during inference, limiting their application in latency-sensitive scenarios. Recent efforts to distill GNNs into multi-layer perceptrons (MLPs) for faster inference often underutilize the layer-level insights of GNNs. In this paper, we present TINED, a novel approach that distills GNNs to MLPs on a layer-by-layer basis using Teacher Injection and Dirichlet Energy Distillation techniques. We focus on two key operations in GNN layers: feature transformation (FT) and graph propagation (GP). We recognize that FT is computationally equivalent to a fully-connected (FC) layer in MLPs. Thus, we propose directly transferring teacher parameters from an FT in a GNN to an FC layer in the student MLP, enhanced by fine-tuning. In TINED, the FC layers in an MLP replicate the sequence of FTs and GPs in the GNN. We also establish a theoretical bound for GP approximation. Furthermore, we note that FT and GP operations in GNN layers often exhibit opposing smoothing effects: GP is aggressive, while FT is conservative. Using Dirichlet energy, we develop a DE ratio to measure these effects and propose Dirichlet Energy Distillation to convey these characteristics from GNN layers to MLP layers. Extensive experiments show that TINED outperforms GNNs and leading distillation methods across various settings and seven datasets. Source code are available at https://github.com/scottjiao/TINED_ICML25/.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。