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

Lan Zhang, Xiaoyong Yuan|arXiv (Cornell University)|Sep 8, 2021

Privacy-Preserving Technologies in Data参考文献 43被引用 6

一句话总结

FedZKT 提出了一种联邦学习框架，通过零样本蒸馏实现在异构设备模型之间的知识迁移，其中服务器利用针对多样化客户端模型训练的生成器执行计算密集型蒸馏，显著降低了设备端工作负载，同时在非独立同分布（non-IID）数据和延迟客户端（straggler）条件下保持了模型性能。

ABSTRACT

Federated learning enables distributed devices to collaboratively learn a shared prediction model without centralizing on-device training data. Most of the current algorithms require comparable individual efforts to train on-device models with the same structure and size, impeding participation from resource-constrained devices. Given the widespread yet heterogeneous devices nowadays, this paper proposes a new framework supporting federated learning across heterogeneous on-device models via Zero-shot Knowledge Transfer, named by FedZKT. Specifically, FedZKT allows participating devices to independently determine their on-device models. To transfer knowledge across on-device models, FedZKT develops a zero-shot distillation approach contrary to certain prior research based on a public dataset or a pre-trained data generator. To utmostly reduce on-device workload, the resource-intensive distillation task is assigned to the server, which constructs a generator to adversarially train with the ensemble of the received heterogeneous on-device models. The distilled central knowledge will then be sent back in the form of the corresponding on-device model parameters, which can be easily absorbed at the device side. Experimental studies demonstrate the effectiveness and the robustness of FedZKT towards heterogeneous on-device models and challenging federated learning scenarios, such as non-iid data distribution and straggler effects.

研究动机与目标

• 解决由于设备资源和模型架构差异，导致在联邦学习中训练异构设备端模型的挑战。
• 在无需对齐训练数据或共享模型结构的前提下，实现不同大小和结构模型之间的有效知识迁移。
• 通过将计算密集型的蒸馏过程卸载至服务器，降低设备端的计算负担。
• 在实际部署场景中，支持在非独立同分布（non-IID）数据分布和延迟客户端影响下的鲁棒联邦学习。

提出的方法

• 允许每个设备根据本地数据和资源约束，独立选择并训练其自身的设备端模型。
• 开发一种零样本蒸馏方法，实现知识迁移，而无需依赖共享的预训练模型或数据生成器。
• 将蒸馏任务分配给服务器，由其训练一个生成器，以对抗性方式与异构客户端模型的集成对齐。
• 通过服务器生成器与聚合客户端模型之间的对抗性训练构建中心化知识，随后将该知识以模型参数形式回蒸馏至客户端。
• 通过直接使用服务器提供的知识更新本地模型参数，实现设备端对蒸馏知识的适应。
• 利用服务器的计算能力处理知识蒸馏的复杂性，最大限度减少设备端计算量。

实验结果

研究问题

• RQ1 在联邦学习中，是否可以在不依赖模型结构对齐或共享数据的前提下，实现异构设备端模型之间的有效知识迁移？
• RQ2 在联邦学习中，如何最小化资源受限设备上知识蒸馏的计算成本？
• RQ3 FedZKT 在非独立同分布（non-IID）数据分布和延迟客户端影响下，其模型性能能保持多大程度的稳定？
• RQ4 基于服务器端生成器的对抗性训练的零样本蒸馏机制，能否在多样化模型之间实现鲁棒的知识迁移？

主要发现

• FedZKT 实现了在不同大小和架构的设备端模型之间有效进行知识迁移，且无需模型或数据对齐。
• 基于服务器的蒸馏机制显著降低了设备端计算负载，使联邦学习在资源受限设备上成为可能。
• FedZKT 在非独立同分布（non-IID）数据分布下表现出鲁棒性能，即使在客户端数据分布偏斜的情况下也能保持高模型准确率。
• 该框架对延迟客户端效应具有鲁棒性，即慢速或不可靠的客户端不会降低整体模型性能。
• 实验结果表明，FedZKT 在保持设备端计算最小化的同时，实现了与基线方法相当的准确率。
• 零样本蒸馏方法无需依赖公开数据集或预训练数据生成器，从而增强了隐私保护和可扩展性。