摘要：

图神经网络(GNN)在节点分类任务上表现优异, 但消息传递机制导致的邻域获取延迟问题限制了其在延迟敏感应用中的部署。尽管多层感知机(MLP)在节点分类任务上的准确性不及GNN, 但由于具有理想的推理效率, 因此其仍是实际工业应用的主要工具。鉴于GNN和MLP在各自优势与劣势上的互补性, 提出基于对抗训练与对比表示蒸馏的GNN推理优化方法, 以将GNN教师模型学到的知识传递给更高效的MLP学生模型。通过将快速梯度符号法(FGSM)生成的特征扰动与节点内容特征结合并作为学生模型的输入, 然后在真实标签与教师模型的Softmax概率分布的指导下, 对学生模型进行对抗训练, 以降低其对节点特征噪声的敏感性。此外, 设计对比表示蒸馏模块, 将学生和教师模型对同一节点输出的嵌入视为正样本对, 对不同节点输出的嵌入视为负样本对, 通过缩小正样本对的距离并扩大负样本对的距离, 使学生模型能够捕捉教师模型输出的节点嵌入之间的关系, 从而保留GNN的全局拓扑结构。在公共数据集上的实验结果表明, 当教师模型为GraphSAGE时, 采用该方法训练得到的MLP学生模型推理速度比GraphSAGE快89倍, 准确率相较于普通MLP和GraphSAGE分别平均提高14.12和2.02百分点, 且优于两种基线方法。

关键词: 图神经网络, 知识蒸馏, 推理加速, 对比学习, 对抗训练

Abstract:

Graph Neural Network (GNN) excels in node classification tasks, but its message-passing mechanism causes neighbor-fetching latency, limiting deployment in latency-sensitive applications. Despite being less accurate than GNN in node classification tasks, Multi-Layer Perceptron (MLP) is preferred in practical industrial applications owing to its efficient inference. Given the complementary advantages and disadvantages of GNN and MLP, this paper proposes an optimized inference method for GNN based on adversarial training and contrastive representation distillation. This method aims to transfer the knowledge learned from a GNN teacher model to a more efficient MLP student model. This method uses the Fast Gradient Sign Method (FGSM) to generate feature perturbations and combines them with node content features as input for the student model. Adversarial training is conducted under the guidance of real labels and the teacher model's Softmax probability distribution to reduce the student model's sensitivity to node feature noise. The contrastive representation distillation module treats embeddings of the student and teacher models on the same node's output as positive sample pairs and embeddings of different nodes' outputs as negative sample pairs. By minimizing the distance between positive sample pairs and maximizing the distance between negative sample pairs, the student model can capture the relationships between node embeddings output by the teacher model, thereby preserving the global topological structure of GNN. Experiment results on public datasets demonstrate that, when using GraphSAGE as the teacher model, an MLP student model trained by this method achieves an inference speed that is 89 times that of GraphSAGE. Additionally, its accuracy improves by 14.12 and 2.02 percentage points on average compared to those of vanilla MLP and GraphSAGE, respectively, outperforming the two baseline methods.

Key words: Graph Neural Network (GNN), knowledge distillation, inference acceleration, contrastive learning, adversarial training