笔记整理：张津瑞，天津大学硕士，研究方向为知识图谱
链接：https://dl.acm.org/doi/10.1145/3488560.3498437
动机
知识图谱已被证明可用于各种 AI 任务，如语义搜索，信息提取和问答等。然而众所周知，知识图谱还远非完备，这进而也促进了许多关于知识图谱完备性的研究。其中比较常见且广泛使用的方法是知识图谱嵌入（KGE Knowledge Graph Embedding），如 TransE、ComplEx 和 RotatE 等。同时，为了获得更好的性能，通常首选训练具有更高维度的 KGE。但是模型大小（参数的数量）以及推理时间的成本通常随 embedding 维度的增加而快速增加，如图 1 所示：随着 embedding 维度的增大，性能增益越来越小，而模型大小和推理成本却仍几乎保持线性增长。此外，高维 KGE 在许多现实场景中是不切实际的，尤其是对于计算资源有限或者是在推理时间有限的应用中，低维的 KGE 是必不可少的。然而，直接训练一个小尺寸 KGE 通常表现不佳，我们进一步提出一个新的研究问题：是否有可能从预训练的高维 KGE 中获得低维 KGE，在更快成本更低的情况下取得良好的效果。
亮点
DualDE的亮点主要包括：
1. 提出了一种新颖的框架，能从高维 KGE 中提取低维 KGE，并取得良好的性能；
2. 在蒸馏过程中考虑了教师和学生之间的双重影响，并提出了一种软标签评估机制，为不同的三元组自适应地分配不同的软标签和硬标签权重，以及一种两阶段的蒸馏方法来提高学生对教师的接受程度。
概念及模型
模型整体框架如图 1 所示：
图1 DualDE总体架构图
蒸馏目标
准备一个预训练好的高维 KGE 模型（教师），并随机初始化一个低维 KGE 模型（学生）。在 DualDE 中，训练学生模型的硬标签损失是KGE方法的原始损失，通常是二元交叉熵损失。此外，我们让学生从目标三元组的整体可信度和嵌入结构这两个方面模仿教师。
首先，对一个三元组 (h, r, t)，教师和学生模型可以通过评分函数为其分配一个分数：1）学生模仿教师对三元组的整体可信度可以通过拟合两个模型输出的三元组分数完成；2）学生模仿教师的嵌入结构可以通过拟合两个模型中三元组头实体 embedding 和尾实体 embedding 的长度比和角度来反映。最后我们将教师和学生的三元组分数差异和嵌入结构差异之和作为软标签优化目标。
软标签评估机制
软标签评估机制可评估教师提供的软标签的质量，并自适应地为不同的三元组分配不同的软标签和硬标签权重，从而保留高质量软标签的积极作用，避免低质量软标签的负面影响。
理论上，KGE 模型会给正三元组更高的分数，给负三元组更低的分数，但对于一些 KGE 模型难以掌握的三元组则相反。具体来说，如果教师给一个负（正）三元组打高（低）分，这意味着教师倾向于将其判断为正（负）三元组，那么教师输出的这个三元组的软标签是不可靠的，甚至会误导学生。对于这个三元组，我们需要削弱软标签的权重，鼓励学生更多地从硬标签中学习。
两阶段蒸馏方式
上一部分介绍了如何让学生从 KGE 教师那里提取知识，其中学生用硬标签训练，软标签由固定教师生成。为了获得更好的学生，我们提出了一种两阶段蒸馏方法，通过解冻教师并让其在蒸馏的第二阶段向学生学习来提高学生对教师的接受度。
第一阶段。第一阶段类似于传统的知识蒸馏方法，其中教师在培训学生时保持不变。第二阶段。在第二阶段调整教师的同时，对于那些学生没有掌握好的三元组，我们也希望减少学生的输出对教师的负面影响，让教师更多从硬标签中学习，从而以保持教师的高准确性。因此，我们也将软标签评估机制应用于教师的调整。通过评估学生给每个三元组的分数，教师的硬标签和软标签的权重被自适应分配。在此阶段，教师和学生一起优化。
实验
我们在典型的KGE基准上评估了DualDE，并做实验探究了以下几个问题：
（1）DualDE是否能够从高维的教师中提炼出一个好的低维学生，并比从头开始训练的相同维度的模型表现更好？没有蒸馏或使用其他KD方法训练的相同维度模型？在没有蒸馏或使用其他KD方法的情况下？
（2）蒸馏后的推理时间有多大提高？
（3）软标签评估机制和两阶段蒸馏方法对我们的提议是否有贡献，贡献有多大？
DualDE 在常用的数据集 WN18RR 和 FB15K-237 上进行了实验，与几个当前先进的蒸馏方法相比，DualDE 体现了优越的性能。下面是实验结果：
表1 WN18RR的链接预测结果
表2 FB15k-237的链接预测结果
Q1：我们的方法是否成功地提炼出一个好学生？
首先，我们在表1中分析了WN18RR的结果。表1表明，随着嵌入维数的减小，“无DS”模型的性能显著下降。对于Simple，与512维教师相比，32维“无DS”模型在MRR、Hit@3和Hit@1上仅获得64.8%、66.1%和47.8%的结果。而对于ComplEx，MRR从0.433下降到0.268（38.1%）。这说明直接训练低维KGE产生差的结果。
与“no-DS”相比，DualDE大大提高了32维学生的成绩。TransE, SimplE,ComplEx 和 RotatE对WN18RR的MRR从0.164提高到0.21（28.0%），从0.273提高到0.384（40.7%），从0.268提高到0.397（48.1%），从0.421提高到0.468（11.2%）。在“no-DS”的基础上，我们的32个维度的学生在这四个KGE中的MRR、Hit@10、Hit@3和Hit@1上平均提高了32.0%、23.0%、33.9%和46.7%，最终达到了教师在MRR、Hit@10、Hit@2和Hit@3上92.9%、94.8%、93.1%和102.3%的水平。我们还可以在表2中观察到FB15k-237的类似结果。实验结果表明，DualDE在保留教师的大部分性能（90%以上）的同时，可以达到16倍（512：32）的嵌入压缩比，尽管有一定的性能损失，但仍然比直接训练低维模型好得多。
Q2：训练学生加快推理速度和多少？
为了测试推理速度,我们在93003条实验样本WN18RR和310116条FB15k237样本进行链路预测实验。由于推断速度不受预测模式（头部预测或尾部预测）的影响，因此统一比较尾部预测时间。推理在单个Tesla-V100 GPU上执行，测试批量大小设置为实体总数：WN18RR为40943，FB15k-237为14541。为了避免偶然因素，我们重复实验3次并报告平均时间。表3示出了推理时间成本的结果（以秒为单位）。
表3 推理时间（秒）
结果表明，经过训练的学生大大加快了推理速度.以ComplEx和RotatE为例，512维教师在WN18RR上的推理时间是32维学生的7.03倍和7.81倍。与教师相比，64维学生在两个数据集的TransE、SimplE、ComplEx和RotatE上的平均速度提高了2.25 ×、2.22 ×、3.66 ×和3.98 ×，32维学生的平均速度提高了3.11 ×、3.35 ×、5.90 ×和5.76 ×。
Q3：软标签评价机制和两阶段蒸馏法是否有贡献？贡献程度如何？
我们进行了一系列消融研究，以评价DualDE两种策略的影响：软标签评价机制和两阶段蒸馏方法。
首先,研究软标签评价机制的影响,我们比较我们的方法与删除软标签评价机制之后的方法。然后,研究影响二级蒸馏方法,我们比较DS和删除第一阶段(- s1)和删除第二阶段(s2)之后的模型。表4总结了WN18RR数据集上MRR 和Hit@10的结果。
表4:WN18RR消融研究。D指的是学生的维度和M指的是方法。
在去除SEM（见 -SEM）后，所有学生的表现与DS相比都有所下降。在这四个KGE中，64维学生的MRR和Hit@10平均下降了3.7%和2.8%，32维学生的MRR和Hit@10平均下降了7.9%和5.4%。实验结果表明，软标签评价模块对每个三元组的软标签质量进行评价，并对不同的三元组赋予不同的软标签和硬标签权重，确实有利于学生模型掌握较难的三元组，获得更好的性能。
删除S1，仅保留S2（见 -S1），总体性能低于DS。据推测，原因是教师和学生都将在S2适应对方。对于随机初始化的学生，学生向教师传递大部分无用的信息，这可能是误导性的并且将使教师崩溃。此外，“-S1”的性能非常不稳定。在“-S1”设置下，64维学生的成绩仅略差于DS，而32维学生的成绩明显非常差。对于SimplE 32维学生来说，“-S1”的MRR和Hit@10比DS下降了21.4%和10.6%。这甚至比使用最基本的蒸馏方法BKD更糟，表明第一阶段对于DualDE是必要的。
删除S2，仅保留S1（见 -S2），性能在几乎所有的指标上都降低了。与DS相比,64维和32维-s2的学生平均下降了2.4%和3.8%,表明第二阶段确实可以使老师和学生互相适应,并进一步改善结果。
这些结果支持我们的二级蒸馏的有效性,首先培养学生S1收敛到一定的性能,然后在S2共同优化老师和学生。
总结
太多的知识图谱的嵌入参数将带来巨大的存储和计算挑战实际应用场景。在这项工作中,我们提出一种新颖的KGE蒸馏法DualDE，将KGE压缩到低维空间,有效转移老师的知识到学生中。考虑教师和学生之间的双重影响,我们建议DualDE使用两种蒸馏策略分为: 软标签评估机制，用于自适应地为不同的三元组分配不同的软标签和硬标签权重;以及两阶段蒸馏方法，用于通过鼓励学生老师相互学习来提高学生对教师的接受度。我们已经在几个KGE和基准数据集上通过链接预测任务评估了DualDE。实验结果表明，该方法可以有效地减少嵌入参数，在性能损失很小或没有损失的情况下，极大地提高了高维KGE的推理速度。
OpenKG
OpenKG（中文开放知识图谱）旨在推动以中文为核心的知识图谱数据的开放、互联及众包，并促进知识图谱算法、工具及平台的开源开放。
点击阅读原文，进入 OpenKG 网站。