[论文解读] Efficient Split-Mix Federated Learning for On-Demand and In-Situ Customization

Junyuan Hong, Haotao Wang+2 more|arXiv (Cornell University)|Mar 18, 2022

Privacy-Preserving Technologies in Data被引用 6

一句话总结

本文提出了一种名为 Split-Mix 联邦学习的新框架，可在异构联邦学习环境中实现按需且现场的模型尺寸与鲁棒性定制。通过在每个客户端独立训练并聚合多个宽度和鲁棒性水平各异的基础子网络，该方法可在推理时动态组合模型，相较于现有方法在通信、存储和推理效率方面表现更优，且在资源受限条件下仍能保持极低的精度下降。

ABSTRACT

Federated learning (FL) provides a distributed learning framework for multiple participants to collaborate learning without sharing raw data. In many practical FL scenarios, participants have heterogeneous resources due to disparities in hardware and inference dynamics that require quickly loading models of different sizes and levels of robustness. The heterogeneity and dynamics together impose significant challenges to existing FL approaches and thus greatly limit FL's applicability. In this paper, we propose a novel Split-Mix FL strategy for heterogeneous participants that, once training is done, provides in-situ customization of model sizes and robustness. Specifically, we achieve customization by learning a set of base sub-networks of different sizes and robustness levels, which are later aggregated on-demand according to inference requirements. This split-mix strategy achieves customization with high efficiency in communication, storage, and inference. Extensive experiments demonstrate that our method provides better in-situ customization than the existing heterogeneous-architecture FL methods. Codes and pre-trained models are available: https://github.com/illidanlab/SplitMix.

研究动机与目标

• 解决由于设备异构性和环境变化导致的联邦学习中运行时资源与鲁棒性需求动态变化的挑战。
• 实现在现场且按需定制模型尺寸与对抗鲁棒性，且无需承担高昂的训练、存储或推理成本。
• 克服现有异构架构联邦学习方法的局限性，如大模型中的训练不足问题以及在非独立同分布（non-IID）数据下的泛化能力差。
• 开发一种可扩展、模块化的框架，支持在硬件和预算各异的多样化客户端之间高效的知识迁移与模型组合。

提出的方法

• 该方法将全宽模型拆分为多个不同宽度（例如 ×0.25、×0.5）和鲁棒性水平的基础子网络，并在每个客户端上独立训练。
• 每个基础模型在其本地数据上完成完整训练，确保知识的完全利用，缓解大模型中常见的训练不足问题。
• 在推理时通过加权聚合选择的基础模型组合形成定制化模型，实现对设备资源约束的动态适应。
• 引入一种新型可微架构搜索（基于 DBN）方法，在训练过程中联合优化鲁棒性与宽度，提升泛化能力。
• 客户端使用一种采样策略（算法 2）确保所有基础模型在所有设备上均被训练，从而增强知识多样性与模型稳定性。
• 该框架可同时支持宽度与鲁棒性定制，通过单一条件变量平滑调节二者之间的权衡。

实验结果

研究问题

• RQ1 联邦学习框架是否能在不重新训练的前提下，实现对模型尺寸与鲁棒性的按需且现场定制？
• RQ2 如何在资源预算各异的异构客户端之间高效共享与聚合知识？
• RQ3 独立训练的基础子网络是否能在资源受限条件下实现比联合训练的大模型更好的泛化性能与更低的精度损失？
• RQ4 与现有异构架构联邦学习方法相比，所提出的 Split-Mix 策略在精度、通信效率与推理效率方面表现如何？
• RQ5 当动态调整模型宽度与对抗鲁棒性时，该方法在多大程度上能维持精度与鲁棒性？

主要发现

• 在资源受限条件下，Split-Mix 在现场定制方面显著优于 HeteroFL 与 FedAvg，精度下降更小。
• 在非独立同分布（non-IID）数据设置下，与 HeteroFL 相比，该方法将精度损失降低了最多 3.5%，尤其在 clipart 与 painting 等模型训练不足严重的领域表现更优。
• Split-Mix 确保所有基础模型在各客户端上均被完全训练，实现每个领域 100% 的参数利用率，而 HeteroFL 则存在高达 60% 的参数未被充分训练。
• 所提出的 Split-Mix+DAT 扩展在鲁棒性与标准精度之间的权衡上比 FedAvg+OAT 更为平滑，无收敛问题，且性能优于 FedAvg+AT 基线。
• 在可调节宽度的场景中，Split-Mix+DAT 随着宽度增加，同时提升了鲁棒与标准精度，展现出高效且稳定的联合定制能力。
• 该框架通过避免存储或重新训练多个完整模型，显著降低了通信与存储成本，实现了高效的设备端模型切换。