Low-Dimensional Federated Knowledge Graph Embedding via Knowledge Distillation

📄 arXiv: 2408.05748v3 📥 PDF

作者: Xiaoxiong Zhang, Zhiwei Zeng, Xin Zhou, Chunyan Miao

分类: cs.AI, cs.LG

发布日期: 2024-08-11 (更新: 2026-01-10)

💡 一句话要点

提出FedKD，通过知识蒸馏实现低维联邦知识图谱嵌入，提升通信效率。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 联邦学习 知识图谱嵌入 知识蒸馏 模型压缩 通信效率

📋 核心要点

  1. FKGE面临高维嵌入带来的存储和通信挑战，现有压缩方法应用于FKGE时通信成本过高。
  2. 提出FedKD，利用知识蒸馏使低维学生模型模仿高维教师模型的得分分布，降低模型维度。
  3. FedKD自适应调整正负三元组的温度，并动态调整KD损失权重，实验证明其有效性。

📝 摘要（中文）

联邦知识图谱嵌入(FKGE)旨在促进多个客户端之间分布式知识图谱(KG)的实体和关系嵌入的协同学习，同时保护数据隐私。通常，训练更高维度的FKGE模型更有利，因为它们有可能实现卓越的性能。然而，高维嵌入在存储资源和推理速度方面提出了重大挑战。与传统的KG嵌入方法不同，FKGE涉及多个客户端-服务器通信轮次，其中通信效率至关重要。现有的传统KG嵌入压缩方法可能不直接适用于FKGE，因为它们通常需要多次模型训练，这可能会产生大量的通信成本。在本文中，我们提出了一种基于知识蒸馏(KD)的轻量级组件，名为FedKD，专门为FKGE方法量身定制。在客户端本地训练期间，FedKD利用KL散度损失，使低维学生模型模仿高维教师模型的三元组得分分布。与传统的KD方式不同，FedKD自适应地学习温度来缩放正三元组的得分，并使用预定义的温度分别调整相应负三元组的得分，从而缓解教师过度自信的问题。此外，我们动态调整KD损失的权重以优化训练过程。在三个数据集上的大量实验支持了FedKD的有效性。

🔬 方法详解

问题定义：联邦知识图谱嵌入(FKGE)旨在分布式环境下协同学习知识图谱嵌入，同时保护数据隐私。高维嵌入模型性能更好，但存储和通信成本高昂。现有知识图谱嵌入压缩方法应用于FKGE时，需要多次模型训练，导致通信开销巨大，不适用于通信敏感的联邦学习场景。

核心思路：利用知识蒸馏(KD)技术，训练一个低维的学生模型来模仿高维教师模型的输出，从而在不显著降低模型性能的前提下，降低模型维度，减少存储和通信开销。核心在于设计一种轻量级的KD方法，使其适用于FKGE场景，并尽可能减少额外的通信成本。

技术框架：FedKD框架主要包含客户端本地训练和服务器聚合两个阶段。在客户端，首先使用高维模型作为教师模型，低维模型作为学生模型。然后，学生模型通过最小化与教师模型输出的KL散度损失进行训练。服务器端负责聚合客户端上传的模型参数，并更新全局模型。

关键创新：FedKD的关键创新在于针对FKGE场景改进了传统的知识蒸馏方法。具体来说，FedKD自适应地学习温度参数来缩放正三元组的得分，并使用预定义的温度单独调整负三元组的得分，从而缓解教师模型过度自信的问题。此外，FedKD还动态调整KD损失的权重，以优化训练过程，平衡知识蒸馏和原始任务损失。

关键设计：FedKD使用KL散度作为知识蒸馏的损失函数，目标是使学生模型的输出分布尽可能接近教师模型的输出分布。为了缓解教师模型过度自信的问题，FedKD引入了温度参数，分别调整正负三元组的得分。正三元组的温度参数是自适应学习的，而负三元组的温度参数是预先设定的。此外，FedKD还动态调整KD损失的权重，以平衡知识蒸馏和原始任务损失。具体来说，KD损失的权重随着训练的进行而逐渐增加，以确保学生模型能够充分学习教师模型的知识。

📊 实验亮点

在三个数据集上的实验结果表明，FedKD能够有效地降低FKGE模型的维度，同时保持甚至提高模型性能。与现有的FKGE方法相比，FedKD在通信效率方面具有显著优势。例如，在某些数据集上，FedKD可以在降低模型维度50%的情况下，仍然保持与高维模型相当的性能。

🎯 应用场景

该研究成果可应用于各种需要联邦学习和知识图谱嵌入的场景，例如医疗知识图谱、金融知识图谱等。通过降低模型维度，可以有效减少存储和通信开销，提高联邦学习的效率和可扩展性。此外，该方法还可以应用于其他类型的联邦学习任务，例如联邦推荐系统、联邦自然语言处理等。

📄 摘要（原文）

Federated Knowledge Graph Embedding (FKGE) aims to facilitate collaborative learning of entity and relation embeddings from distributed Knowledge Graphs (KGs) across multiple clients, while preserving data privacy. Training FKGE models with higher dimensions is typically favored due to their potential for achieving superior performance. However, high-dimensional embeddings present significant challenges in terms of storage resource and inference speed. Unlike traditional KG embedding methods, FKGE involves multiple client-server communication rounds, where communication efficiency is critical. Existing embedding compression methods for traditional KGs may not be directly applicable to FKGE as they often require multiple model trainings which potentially incur substantial communication costs. In this paper, we propose a light-weight component based on Knowledge Distillation (KD) which is titled FedKD and tailored specifically for FKGE methods. During client-side local training, FedKD facilitates the low-dimensional student model to mimic the score distribution of triples from the high-dimensional teacher model using KL divergence loss. Unlike traditional KD way, FedKD adaptively learns a temperature to scale the score of positive triples and separately adjusts the scores of corresponding negative triples using a predefined temperature, thereby mitigating teacher over-confidence issue. Furthermore, we dynamically adjust the weight of KD loss to optimize the training process. Extensive experiments on three datasets support the effectiveness of FedKD.