[论文解读] FedKD: Communication Efficient Federated Learning via Knowledge Distillation

Chuhan Wu, Fangzhao Wu+3 more|arXiv (Cornell University)|Aug 30, 2021

Privacy-Preserving Technologies in Data参考文献 33被引用 18

一句话总结

FedKD 提出了一种基于自适应知识蒸馏和基于 SVD 的动态梯度压缩的通信高效联邦学习框架。它在客户端协作训练轻量化学生模型的同时，为每个客户端保留本地教师模型，通过仅共享小型学生模型并以可变精度压缩梯度，显著降低了通信成本，实现了与现有方法相当的性能，同时大幅减少了带宽使用。

ABSTRACT

Federated learning is widely used to learn intelligent models from decentralized data. In federated learning, clients need to communicate their local model updates in each iteration of model learning. However, model updates are large in size if the model contains numerous parameters, and there usually needs many rounds of communication until model converges. Thus, the communication cost in federated learning can be quite heavy. In this paper, we propose a communication efficient federated learning method based on knowledge distillation. Instead of directly communicating the large models between clients and server, we propose an adaptive mutual distillation framework to reciprocally learn a student and a teacher model on each client, where only the student model is shared by different clients and updated collaboratively to reduce the communication cost. Both the teacher and student on each client are learned on its local data and the knowledge distilled from each other, where their distillation intensities are controlled by their prediction quality. To further reduce the communication cost, we propose a dynamic gradient approximation method based on singular value decomposition to approximate the exchanged gradients with dynamic precision. Extensive experiments on benchmark datasets in different tasks show that our approach can effectively reduce the communication cost and achieve competitive results.

研究动机与目标

• 降低大型模型（如 BERT）在联邦学习中因大模型更新带来的高通信成本。
• 在具有隐私敏感数据的去中心化学习环境中，最小化带宽使用，同时保持模型性能。
• 开发一种通信高效的框架，通过客户端本地学生与教师模型之间的自适应知识蒸馏，维持模型准确率。
• 引入基于奇异值分解（SVD）的动态梯度近似方法，以可变精度进一步压缩模型更新。
• 在多种基准数据集和任务上验证该方法，以确保其鲁棒性和可扩展性。

提出的方法

• 提出一种自适应互知识蒸馏框架，每个客户端训练一个本地教师模型和一个共享学生模型，蒸馏强度根据预测正确性动态调整。
• 利用知识蒸馏在本地教师模型与学生模型之间传递知识，提升泛化能力并降低过拟合风险。
• 仅在客户端之间共享小型学生模型以实现协同训练，避免传输大型教师模型参数。
• 采用基于 SVD 的动态梯度近似方法，以可变精度压缩模型更新，降低通信开销。
• 应用基于能量的阈值策略，仅保留梯度矩阵中最显著的奇异值，以平衡准确率与压缩率。
• 采用两阶段阈值策略：训练初期 $T_{start} = 0.95$，训练末期 $T_{end} = 0.98$，以自适应调整压缩精度。

实验结果

研究问题

• RQ1 在联邦学习中，客户端本地学生与教师模型之间的自适应互知识蒸馏是否能在不牺牲模型性能的前提下降低通信成本？
• RQ2 通过 SVD 实现的动态梯度压缩如何影响联邦训练中的通信效率与模型准确率？
• RQ3 在所提出的 FedKD 框架中，客户端数量的变化对模型收敛性和性能有何影响？
• RQ4 梯度矩阵中奇异值的能量分布如何随训练过程演变？这对压缩效率有何启示？
• RQ5 与标准 FedAvg 相比，该方法是否能在显著减少通信量的同时，实现在基准数据集上的竞争力表现？

主要发现

• FedKD 通过仅共享小型学生模型而非完整的大模型，显著降低了通信成本，在极低带宽下实现了具有竞争力的性能。
• 自适应互知识蒸馏机制根据预测正确性动态调整蒸馏强度，相比标准蒸馏方法，提升了模型准确率并减轻了过拟合风险。
• 基于 SVD 的动态梯度近似实现了显著压缩，奇异值能量高度集中于低秩分量中，尤其在前馈网络参数中表现明显。
• 随着客户端数量的增加，性能保持稳定甚至略有提升，这是由于来自多个客户端本地教师模型的知识聚合更加丰富。
• 在 $T_{start} = 0.95$ 和 $T_{end} = 0.98$ 时，模型在通信成本与准确率之间达到最优权衡，兼顾了压缩效率与精度。
• 训练过程中所需奇异值的数量逐渐增加，这为采用动态阈值策略以在训练过程中保持准确率提供了合理依据。