Efficient Split-Mix Federated Learning for On-Demand and In-Situ Customization

TL;DR

This paper introduces a Split-Mix federated learning approach that enables efficient, on-demand customization of model size and robustness for heterogeneous participants, addressing resource disparities and dynamic inference needs.

Contribution

It proposes a novel Split-Mix FL strategy that learns base sub-networks for flexible, in-situ model customization, improving efficiency and adaptability over existing methods.

Findings

Outperforms existing heterogeneous FL methods in customization quality

Reduces communication, storage, and inference costs

Enables flexible model adaptation for diverse resource constraints

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.