[论文解读] FedZKT: Zero-Shot Knowledge Transfer towards Resource-Constrained Federated Learning with Heterogeneous On-Device Models

Lan Zhang, Dapeng Wu+1 more|arXiv (Cornell University)|Sep 8, 2021

Privacy-Preserving Technologies in Data被引用 8

一句话总结

FedZKT 提出了一种无需数据、零样本的知识迁移框架，用于联邦学习，使资源受限且本地模型异构的设备能够在不共享原始数据的情况下协同训练全局模型。通过将计算密集型蒸馏任务卸载至服务器，并利用生成器从聚合的设备模型中学习中心知识，FedZKT 在模型异构性、非独立同分布（non-IID）数据以及延迟设备（straggler）影响下均表现出稳健性能，同时设备端计算量极低，且无需依赖私有数据。

ABSTRACT

Federated learning enables multiple distributed devices to collaboratively learn a shared prediction model without centralizing their on-device data. Most of the current algorithms require comparable individual efforts for local training with the same structure and size of on-device models, which, however, impedes participation from resource-constrained devices. Given the widespread yet heterogeneous devices nowadays, in this paper, we propose an innovative federated learning framework with heterogeneous on-device models through Zero-shot Knowledge Transfer, named by FedZKT. Specifically, FedZKT allows devices to independently determine the on-device models upon their local resources. To achieve knowledge transfer across these heterogeneous on-device models, a zero-shot distillation approach is designed without any prerequisites for private on-device data, which is contrary to certain prior research based on a public dataset or a pre-trained data generator. Moreover, this compute-intensive distillation task is assigned to the server to allow the participation of resource-constrained devices, where a generator is adversarially learned with the ensemble of collected on-device models. The distilled central knowledge is then sent back in the form of the corresponding on-device model parameters, which can be easily absorbed on the device side. Extensive experimental studies demonstrate the effectiveness and robustness of FedZKT towards on-device knowledge agnostic, on-device model heterogeneity, and other challenging federated learning scenarios, such as heterogeneous on-device data and straggler effects.

研究动机与目标

• 为解决在联邦学习中整合资源受限且本地模型异构的设备所面临的挑战，其中经典方法要求统一的模型结构。
• 消除对私有设备数据或公开数据集的依赖，实现无需数据的知识迁移，从而支持无数据协作。
• 通过将计算密集型的蒸馏过程卸载至服务器，降低设备端计算负载。
• 确保在真实联邦学习环境中对模型异构性、非独立同分布（non-IID）数据分布以及延迟设备效应的鲁棒性。
• 在保持全局模型收敛性和性能的前提下，实现设备端模型设计的独立性。

提出的方法

• 设备根据其本地硬件约束独立训练紧凑且资源高效的模型。
• 通过与收集到的设备模型集成进行对抗训练，由服务器学习的生成器在无需私有数据的情况下蒸馏出中心知识。
• 引入一种新颖的SL（软标签）损失函数，实现在结构异构模型之间的零样本知识蒸馏。
• 将蒸馏得到的知识以设备端模型参数的形式回传，使设备能够轻松吸收并微调全局知识。
• 该框架将计算密集型的蒸馏任务分配给服务器，从而保持设备端的资源效率。
• 在设备端训练过程中应用L2正则化，以提升在非独立同分布数据分布下的鲁棒性。

Figure 1 : Overview of FedZKT: zero-shot knowledge transfer for federated learning with heterogeneous on-device models. Most compute-intensive workloads of FedZKT are done at the server. Participating devices only need to update the on-device models based on their local data.

实验结果

研究问题

• RQ1 在无需共享数据或模型结构对齐的前提下，能否在本地模型高度异构的设备之间有效实现联邦学习？
• RQ2 如何在不依赖公开数据集或预训练数据生成器的情况下，实现无数据依赖的异构模型间知识迁移？
• RQ3 当蒸馏任务被卸载至服务器时，资源受限设备在联邦学习中的参与程度在多大程度上得以提升？
• RQ4 该框架在真实部署中对非独立同分布（non-IID）数据分布和延迟设备效应的鲁棒性如何？
• RQ5 设备端模型架构和设备数量对全局模型收敛性和准确率的影响是什么？

主要发现

• FedZKT 在多种设备端模型上均表现出强大性能：在CIFAR-10（IID）数据集上，设备3使用模型C达到88.63%准确率，设备8使用模型C达到88.34%，表明其对模型异构性的高度适应能力。
• 即使每轮仅20%的设备活跃（p=0.2），框架仍能保持稳定训练，展现出对延迟设备效应的强韧性。
• L2正则化显著提升非IID数据下的性能：在C=5的不平衡设置下，准确率从56.58%提升至63.89%；在β=0.5的分布不平衡下，从66.17%提升至69.39%。
• 设备数量（K ∈ {5,10,15,20}）对性能影响较小，不同K值下的平均准确率仅波动±2%。
• FedZKT 使设备即使在资源极少的情况下也能有效贡献，因为蒸馏过程由服务器优化，且知识通过轻量级模型参数更新进行传输。
• 该框架在无需任何私有数据或公开数据依赖的情况下实现了有效的知识迁移，证实了在异构联邦学习环境中实现无数据、零样本蒸馏的可行性。

Figure 2 : Norm of gradients w.r.t input data (MNIST, IID). The gradients for the KL-divergence loss tend to vanish, while the gradients for the $\ell_{1}$ norm loss are much larger and unstable during the learning process. The proposed SL loss overcomes both problems in the federated learning.