Efficient Split-Mix Federated Learning for On-Demand and In-Situ Customization
This addresses the problem of limited applicability of FL for participants with varying hardware and inference needs, offering an incremental improvement over existing heterogeneous-architecture FL methods.
The paper tackles the challenge of heterogeneous resources and dynamic inference requirements in federated learning by proposing a Split-Mix FL strategy that enables on-demand customization of model sizes and robustness after training, achieving high efficiency in communication, storage, and inference.
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.