HYDRA-FL: Hybrid Knowledge Distillation for Robust and Accurate Federated Learning

📄 arXiv: 2409.19912v3 📥 PDF

作者: Momin Ahmad Khan, Yasra Chandio, Fatima Muhammad Anwar

分类: cs.LG, cs.CR

发布日期: 2024-09-30 (更新: 2024-10-07)

备注: Annual Conference on Neural Information Processing Systems (NeurIPS), 2024

💡 一句话要点

提出HYDRA-FL，通过混合知识蒸馏提升联邦学习在异构数据和攻击下的鲁棒性和准确性

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

关键词: 联邦学习 知识蒸馏 模型中毒攻击 数据异构性 鲁棒性 混合蒸馏 辅助分类器

📋 核心要点

  1. 联邦学习中知识蒸馏方法在数据异构性下表现良好，但易受模型中毒攻击，存在攻击放大效应。
  2. HYDRA-FL通过混合知识蒸馏，将部分蒸馏损失转移到浅层，降低攻击对全局模型的影响。
  3. 实验表明，HYDRA-FL在攻击场景下优于现有方法，并在良性环境下保持了相当的性能。

📝 摘要（中文）

联邦学习(FL)中用户间的数据异构性是一个重大挑战，会导致全局模型性能下降。社区设计了各种技术来解决这个问题，其中基于知识蒸馏(KD)的技术很常见。虽然这些技术有效地提高了高异构性下的性能，但它们无意中导致了模型中毒攻击下更高的精度下降（称为攻击放大）。本文通过一个案例研究揭示了KD-based FL系统中这一关键漏洞。我们通过经验证据展示了KD导致此问题的原因，并以此为动机设计了一种混合蒸馏技术。我们提出了一种新的算法，即用于鲁棒和精确FL的混合知识蒸馏(HYDRA-FL)，它通过辅助分类器将一些KD损失转移到浅层，从而减少了攻击场景中攻击的影响。我们将HYDRA-FL建模为一个通用框架，并将其应用于两种基于KD的FL算法，FedNTD和MOON。通过将这两种算法作为案例研究，我们证明了我们的技术在攻击设置中优于基线，同时在良性设置中保持了相当的性能。

🔬 方法详解

问题定义：联邦学习在数据异构性环境下，利用知识蒸馏能够提升模型性能。然而，现有的基于知识蒸馏的联邦学习方法在面对模型中毒攻击时，会放大攻击效果，导致模型精度大幅下降。因此，需要设计一种既能应对数据异构性，又能抵抗模型中毒攻击的联邦学习方法。

核心思路：HYDRA-FL的核心思路是采用混合知识蒸馏策略，将知识蒸馏损失分散到模型的不同层。具体来说，除了在输出层进行知识蒸馏外，还在模型的浅层添加一个辅助分类器，并对该辅助分类器也进行知识蒸馏。这样可以降低攻击者对输出层的影响，从而提高模型的鲁棒性。

技术框架：HYDRA-FL是一个通用的框架，可以应用于各种基于知识蒸馏的联邦学习算法。其主要包括以下几个模块：1）客户端本地训练：每个客户端使用本地数据训练模型，并计算知识蒸馏损失和辅助分类器损失。2）服务器聚合：服务器收集客户端的模型参数和损失信息，并进行聚合。3）全局模型更新：服务器使用聚合后的信息更新全局模型。4）模型分发：服务器将更新后的全局模型分发给客户端。

关键创新：HYDRA-FL的关键创新在于引入了混合知识蒸馏策略，通过在浅层添加辅助分类器并进行知识蒸馏，降低了攻击者对输出层的影响，从而提高了模型的鲁棒性。与传统的知识蒸馏方法相比，HYDRA-FL能够更好地平衡模型性能和鲁棒性。

关键设计：HYDRA-FL的关键设计包括：1）辅助分类器的网络结构：辅助分类器通常是一个简单的线性分类器或小型神经网络。2）知识蒸馏损失函数：可以使用KL散度或均方误差等损失函数。3）辅助分类器损失权重：需要仔细调整辅助分类器损失的权重，以平衡模型性能和鲁棒性。4）适配不同的KD算法：HYDRA-FL可以适配不同的KD算法，例如FedNTD和MOON。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文在FedNTD和MOON两种算法上验证了HYDRA-FL的有效性。实验结果表明，在模型中毒攻击下，HYDRA-FL显著优于基线方法，降低了攻击带来的性能损失，同时在良性环境下保持了与基线方法相当的性能。具体性能提升幅度取决于攻击强度和数据集，但总体趋势是HYDRA-FL能够有效提升联邦学习系统的鲁棒性。

🎯 应用场景

HYDRA-FL可应用于各种对数据安全和模型鲁棒性有较高要求的联邦学习场景，例如金融风控、医疗诊断、自动驾驶等。该方法能够有效提升模型在异构数据和恶意攻击下的性能，保障联邦学习系统的安全可靠运行，具有重要的实际应用价值和广阔的应用前景。

📄 摘要（原文）

Data heterogeneity among Federated Learning (FL) users poses a significant challenge, resulting in reduced global model performance. The community has designed various techniques to tackle this issue, among which Knowledge Distillation (KD)-based techniques are common. While these techniques effectively improve performance under high heterogeneity, they inadvertently cause higher accuracy degradation under model poisoning attacks (known as attack amplification). This paper presents a case study to reveal this critical vulnerability in KD-based FL systems. We show why KD causes this issue through empirical evidence and use it as motivation to design a hybrid distillation technique. We introduce a novel algorithm, Hybrid Knowledge Distillation for Robust and Accurate FL (HYDRA-FL), which reduces the impact of attacks in attack scenarios by offloading some of the KD loss to a shallow layer via an auxiliary classifier. We model HYDRA-FL as a generic framework and adapt it to two KD-based FL algorithms, FedNTD and MOON. Using these two as case studies, we demonstrate that our technique outperforms baselines in attack settings while maintaining comparable performance in benign settings.