知识蒸馏再升级：用大语言模型赋能图神经网络

一、背景介绍

知识蒸馏（Knowledge Distillation）作为一种模型压缩方法，通过将复杂教师模型（Teacher Model）的知识传递给较轻量的学生模型（Student Model），在提高推理效率的同时保持模型性能。在图神经网络（GNN）和大语言模型（LLM）的结合领域，知识蒸馏已经展现出巨大的潜力，特别是在处理图结构数据和多模态数据时。本专题聚焦于两篇论文，分别探讨了LLM和GNN结合的知识蒸馏方法在不同任务中的表现。

二、GALLON: LLM与GNN的多模态知识蒸馏框架

GALLON（Graph Learning from Large Language Model Distillation）框架旨在将大语言模型（LLM）的表征能力与图神经网络（GNN）的结构分析能力结合，通过知识蒸馏将这两者的优势整合到一个轻量级的多层感知机（MLP）中，显著提高分子属性预测任务的性能和效率。

框架结构

分子数据通常以多模态形式存在，包括文本格式（SMILES字符串）、视觉表示（分子图像）、图结构（节点与边的关系）。GNN擅长捕获图结构信息，但在处理文本和视觉数据方面表现有限。LLM对文本和视觉数据具有强大的表征能力，但在处理图结构数据时效率低下。为了解决这些问题，GALLON提出将LLM和GNN的知识蒸馏到一个MLP中，以实现高效的推理。

多模态输入： 提取SMILES字符串、分子图像和图结构，作为模型的多模态输入。

教师模型：

• GNN处理图结构数据，生成图表征。
• LLM结合文本和视觉信息生成多模态表征。

学生模型（MLP）： 通过知识蒸馏从GNN和LLM中学习，将多模态知识压缩到轻量级的MLP中。

知识蒸馏： 包括标签蒸馏和表示蒸馏，通过对齐教师模型与学生模型的表征和预测分布来传递知识。

知识蒸馏目标： GALLON使用标签蒸馏（Label Distillation）和表示蒸馏（Representation Distillation）来训练学生模型：

• 标签蒸馏：

其中，、、 分别是学生模型、LLM 和 GNN 的输出预测分布， 和 是权重超参数。

• 表示蒸馏：

其中，、、 为学生和教师模型的表征， 是映射到统一空间的变换函数。

最终损失： 将预测损失与知识蒸馏损失结合：

GALLON的实验结果

模型性能： 在多个分子属性预测数据集（如BACE、BBBP、HIV）上，GALLON显著超越传统的GNN和LLM模型，同时减少了模型参数量和推理时间。

效率对比： 展示了ROCAUC与推理时间的对数和模型大小之间的关系，结果显示，GALLON在最短的推理时间和最小的模型大小下，达到了最佳性能。这些结果表明，我们的方法有效地将GNN和LLM的知识蒸馏到MLP中，在效率和准确性方面都达到了最先进的水平。

三、MuGSI: 通过多粒度结构信息蒸馏的图分类框架

框架介绍

MuGSI（Multi-Granularity Structural Information for Graph Classification）提出了一种新的图神经网络（GNN）到多层感知机（MLP）的知识蒸馏（KD）框架，解决了图分类任务中知识蒸馏面临的两个主要挑战：

• 稀疏的学习信号： 图分类任务中，学习信号稀疏，因为标签是图级别的。
• 学生模型表达能力不足： MLP作为学生模型，虽然推理速度较快，但在表达图结构时通常不如GNN强大。

为了应对这些问题，MuGSI框架提出了多粒度结构信息蒸馏，结合了图级、子图级和节点级的蒸馏信息，确保学生模型能够充分学习教师模型的多层次结构信息。

MuGSI的关键组成部分

图1：MuGSI框架的整体结构图，展示了从教师模型（GNN）到学生模型（MLP）蒸馏过程中的多粒度结构信息。

• 图级蒸馏（Graph-Level Distillation）：在图级别进行蒸馏，通过对比学生和教师模型的全图表示，使学生模型学习到更全面的全图信息。图级蒸馏的损失函数如下所示：

其中， 和 分别是教师和学生模型在图级别的表示，损失函数通过计算图级别的 L2 范数差异来引导学生模型对齐教师模型的表示。

• 子图级蒸馏（Subgraph-Level Distillation）：通过将图划分为多个子图并对齐子图级的表示，MuGSI能够捕捉到更细粒度的结构信息。子图级蒸馏通过最大化集群间的相似性来进行，使用一个核矩阵（）来衡量集群之间的相似性，计算得到的损失为：

这里 和 分别是学生和教师模型从集群级别获得的核矩阵。

• 节点级蒸馏（Node-Level Distillation）：通过学习节点级的局部子结构，MuGSI加强了学生模型在节点特征上的表达能力。路径一致性损失（Path-Consistency Loss）用于确保学生和教师模型对节点间的相似性具有一致的看法，计算公式如下：

  其中， 和 分别是教师和学生模型在路径 上的条件概率。

MuGSI通过引入多粒度的结构信息蒸馏，成功解决了图分类任务中面临的稀疏学习信号和学生模型表达能力不足的问题。通过图级、子图级和节点级的蒸馏，MuGSI有效提升了学生模型的性能

MuGSI的实验结果

MuGSI在多个数据集上的实验表明，相比于其他知识蒸馏方法（如GLNN），MuGSI在图分类任务中表现出了更优的效果。例如，在使用GNN作为教师模型时，MuGSI在PROTEINS和REDDIT-BINARY等数据集上的表现优于其他方法，尤其是在速度和推理效率方面。

四、总结

通过解读GALLON和MuGSI两篇论文，可以发现LLM与GNN结合的知识蒸馏在以下方面具有显著潜力：

1. 多模态与多粒度： GALLON侧重多模态数据的知识融合，MuGSI聚焦多粒度图结构信息的蒸馏，分别在不同任务场景中展现了优势。
2. 高效性与实用性： 两者均通过蒸馏技术简化了复杂模型，显著提升了推理效率，满足工业界对实时性和低资源消耗的需求。
3. 未来展望： 将多模态与多粒度蒸馏相结合，进一步提升学生模型的适应性；探索更强大的学生模型结构，如混合型GNN-MLP框架。

本专题揭示了知识蒸馏在大模型与图学习结合中的关键作用，为相关研究和实际应用提供了新的方向和灵感。

五、参考文献

【1】Xu J, Wu Z, Lin M, et al. LLM and GNN are Complementary: Distilling LLM for Multimodal Graph Learning[J]. arXiv preprint arXiv:2406.01032, 2024.

【2】Yao T, Sun J, Cao D, et al. MuGSI: Distilling GNNs with Multi-Granularity Structural Information for Graph Classification[C]//Proceedings of the ACM on Web Conference 2024. 2024: 709-720.