[论文解读] DualDE: Dually Distilling Knowledge Graph Embedding for Faster and Cheaper Reasoning

Yushan Zhu, Wen Zhang+5 more|arXiv (Cornell University)|Sep 13, 2020

Advanced Graph Neural Networks被引用 8

一句话总结

DualDE 提出了一种知识蒸馏方法，将高维知识图谱嵌入（KGEs）压缩为低维学生模型，以实现更快、更经济的推理。它采用软标签评估机制，自适应地加权教师模型提供的分数，并通过两阶段蒸馏过程增强学生模型的学习效果，在基准数据集上实现7×–15×的参数减少和2×–6×的加速，同时性能损失极小。

ABSTRACT

Knowledge Graph Embedding (KGE) is a popular method for KG reasoning and training KGEs with higher dimension are usually preferred since they have better reasoning capability. However, high-dimensional KGEs pose huge challenges to storage and computing resources and are not suitable for resource-limited or time-constrained applications, for which faster and cheaper reasoning is necessary. To address this problem, we propose DualDE, a knowledge distillation method to build low-dimensional student KGE from pre-trained high-dimensional teacher KGE. DualDE considers the dual-influence between the teacher and the student. In DualDE, we propose a soft label evaluation mechanism to adaptively assign different soft label and hard label weights to different triples, and a two-stage distillation approach to improve the student's acceptance of the teacher. Our DualDE is general enough to be applied to various KGEs. Experimental results show that our method can successfully reduce the embedding parameters of a high-dimensional KGE by 7 times - 15 times and increase the inference speed by 2 times - 6 times while retaining a high performance. We also experimentally prove the effectiveness of our soft label evaluation mechanism and two-stage distillation approach via ablation study.

研究动机与目标

• 解决在资源受限环境中高维KGE带来的高存储和计算成本问题。
• 在不牺牲推理性能的前提下，实现快速且低成本的推理。
• 开发一种可泛化的蒸馏方法，将高维教师KGE的知识迁移至低维学生KGE。
• 通过考虑教师模型提供的软标签质量，提升学生模型的性能。
• 通过两阶段蒸馏过程，实现教师与学生之间的相互适应。

提出的方法

• 引入一种软标签评估机制，用于评估教师模型对每个三元组预测的可靠性，并为软标签和硬标签分配自适应权重。
• 采用两阶段蒸馏方法：第一阶段，使用软标签和硬标签训练学生模型；第二阶段，联合优化教师和学生模型，实现相互适应。
• 将教师模型的输出分数作为软标签，通过置信度加权监督来降低预测质量差的三元组带来的噪声影响。
• 应用温度缩放的知识蒸馏，以提升学生模型的泛化能力和知识迁移效果。
• 通过架构无关性，支持多种KGE模型（如TransE、ComplEx、RotatE），实现跨模型的通用性。
• 采用结合软标签和真实标签交叉熵的蒸馏损失，并根据标签质量动态调整权重。

实验结果

研究问题

• RQ1 知识蒸馏能否在保持推理性能的前提下，有效将高维KGE压缩为低维模型？
• RQ2 如何评估教师模型提供的软标签质量，以避免因不可靠预测导致的负面知识迁移？
• RQ3 与单阶段蒸馏相比，包含初始学生训练后联合优化的两阶段蒸馏过程是否能提升学生模型性能？
• RQ4 自适应软标签加权在难以预测的三元组上能带来多大性能提升？
• RQ5 所提出的方法是否能在多种KGE架构上实现显著的加速和参数压缩？

主要发现

• DualDE 在保持高推理性能的同时，将KGE嵌入参数量减少了7×至15×，在WN18RR等基准数据集上表现优异。
• 该方法使推理速度提升2×至6×，适用于实时或边缘设备部署。
• 软标签评估机制使64维学生模型的MRR平均提升3.7%，Hit@10提升2.8%；32维学生模型则分别提升7.9%和5.4%。
• 若移除第一阶段蒸馏（S1），32维SimplE模型的MRR性能最高下降21.4%，表明S1在稳定训练中起关键作用。
• 第二阶段联合优化（S2）使64D学生模型的MRR平均提升2.4%，32D学生模型提升3.8%，证实其在相互适应中的有效性。
• 消融实验证实，软标签评估和两阶段蒸馏均对最优性能至关重要，显著提升了准确率和鲁棒性。