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

1 INTRODUCTION . Federated learning ( Konečný et al. , 2015 ) is a distributed learning paradigm that leverages data from remote participants and aggregates their knowledge without requiring their raw data to be transferred to a central server , thereby largely reducing the concerns from data security and privacy . FedAvg ( McMahan et al. , 2017 ) is among the most popular federated instantiations , which aggregates knowledge by averaging homogeneous model parameters . In many practical federated learning scenarios , the participants usually exhibit a variety of heterogeneity . The training datasets from participants can be non-i.i.d . and therefore lead to data heterogeneity , which has received extensive attention in the community ( Li et al. , 2020b ; Fallah et al. , 2020 ; Hong et al. , 2021 ; Li et al. , 2021 ) . In addition , the participants can also have resource heterogeneity , i.e. , participants have different hardware resources available , such as memory , computing power , and network bandwidth ( Ignatov et al. , 2018 ) . For example , in a learning task for face recognition , clients may use different types of devices ( e.g. , desktop computers , different generations of tablets or smartphones ) to participate in learning , which thereby leads to drastic differences in computational resources for training models , and in sizes of models that can be loaded to the memory and transferred through the network . The disparity of hardware resources results in varying budgets in both training and inference stages ( Xu et al. , 2021 ) . To accommodate different model sizes , one can turn to more flexible architectures trained by distillation from ensemble ( Lin et al. , 2020 ) , partial model averaging ( Diao et al. , 2021 ) , or directly combining predictions ( Shi et al. , 2021 ) . Yet besides heterogeneity , the deployment of models usually has a great amount of dynamics . One type of dynamics is directly related to the resources of participants ’ devices . For example , memory consumptions and computational capabilities may vary drastically between training and inference stages , and the same dynamics in hardware can also happen during inference ( Yu et al. , 2018 ; Xu et al. , 2021 ) , depending on how the resource allocation of the running programs is prioritized on a participant ’ s device . Another example of dynamics is the requirements of properties of the model , such as adversarial robustness . In autonomous driving , for instance , cloudy and rainy weather can ( a ) Illustration of FedAvg ( McMahan et al. , 2017 ) with a device-incompatible model and a heterogenousarchitecture variant ( HeteroFL ) ( Diao et al. , 2021 ) with under-trained wide models . ( b ) The proposed Split-Mix framework provides in-situ customization of widths and adversarial robustness to address heterogeneity and dynamics , enabling efficient training and inference . In this example , we use a subnet with 1/4 channels ( or widths ) per layer as a base model for model-width customization . For simplicity , we denote it ×0.25 net , and a ×1 net can be split into 4 ×0.25 base models . Figure 1 : Comparison of traditional and proposed methods . lower the quality of camera and other sensors , and therefore one would like to quickly switch to a more robust prediction model , as compared to the one used on sunny days . The dynamics along with the low latency requirements for the inference , altogether demand a model that can be customized in situ , meaning that we can instantly tailor our predictive model to a proper size that fits our available budget , and promptly adjust its robustness . Most existing federate learning models used a rather static model architecture that prevents such customization . A naive solution is to simultaneously train multiple models with wanted architectures and properties , which however may lead to prohibitive training cost and latency of switching models in a budget-limited device . Moreover , the trade-off between model adversarial robustness and generalization performance can hardly be reconcilable ( Tsipras et al. , 2019 ; Merchant et al. , 2020 ) , which further complicates the customization . To address the aforementioned challenges from heterogeneity and dynamics , we study a novel SplitMix federated learning approach for in-situ customization , which follows a communication-efficient paradigm and yields a model whose size and robustness can be efficiently customized , without additional training . Specifically , we first split the complete knowledge in a large model into several small base sub-networks ( shards ) according to model widths and robustness levels . To complete the knowledge , we let the base models be fully trained on all clients . To provide customized models , we mix selected base models to construct desired widths and required robustness . We illustrate the main idea in Fig . 1b . Overall , our contributions can be summarized in three folds : • Within the domain of heterogeneous federated learning , we are the first to study training a model with the capability of in-situ customization with heterogeneous local computation budgets . Illustrated in Fig . 1a , we show that conventional static-model solutions , such as FedAvg , can not provide customized models for dynamic client environments , and that its heterogeneous-architecture variant ( HeteroFL ) suffers from under-training in its large models due to local budget constraints . • To address the challenge , we propose a novel Split-Mix framework that aggregates knowledge from heterogeneous clients into a width- and robustness-adjustable model structure . Remarkably , due to fewer parameters and modular nature , our framework is not only efficient in federated communication and flexibly adaptable to various client budgets during training , but also efficient and flexible in storage , model loading and execution during inference . • Empirically , we demonstrate that the performance of the proposed method is better than other FL baselines under heterogeneous budget constraints . Moreover , we show its effectiveness when facing the challenge of data heterogeneity , i.e. , clients ’ training datasets suffer from distributional shifts . 2 RELATED WORK . Heterogeneous Federated Learning . Federated learning with budget-insufficient devices has attracted a great amount of attention . For example , FedDistill ( Lin et al. , 2020 ) used logit averages to distill a network of the same size or several prototypes , which will be communicated with users . FedDistill made an assumption that the central server has access to a public dataset of the same learning task , which is impractical for many applications . He et al . ( 2020 ) introduced a distillation-based method after aggregating private representations from all participants . The method closely resembles centralized learning because all encoded samples are gathered , and however it is less efficient when local clients have large data dimensions or sample sizes . Importantly , the method may not transfer adversarial robustness knowledge through the intermediate representations due to the decoupling of the input and prediction . On the other hand , HeteroFL ( Diao et al. , 2021 ) avoids distillation , allowing participants to train different prototypes and sharing parameters among prototypes to reduce redundant parameters . However , HeteroFL also reduces the samples available for training each prototype , which leads to degraded performance . FedEnsemble ( Shi et al. , 2021 ) is technically related to the proposed approach , which uses ensemble of diverse models to accommodate non-i.i.d . heterogeneity . The authors showed that combining multiple base models trained simultaneously in FL can outperform a single base model in testing . A critical difference between the proposed approach and FedEnsemble comes from the challenging problem setting of constrained heterogeneous computation budgets . For the first time , we show that base models can be trained adaptively under budget constraints and used to customize efficient inference networks that can outperform a single model of the same width but trained in a heterogeneous-architecture way . Customizable Models . To our best knowledge , this is the first paper discussing in-situ customization in federated learning and here we review similar concepts in central learning . First , customization of robustness and accuracy was discussed by Wang et al . ( 2020 ) , where an adjustable-conditional layer together with decoupled batch-normalization were used . The conditional layer enables continuous trade-off but also brings in over 50 % more parameters . In comparison , a simple weighted combination without additional parameters is used in our method and is very efficient in both communication and inference . In terms of model complexity , a series of work on dynamic neural networks were proposed to provide instant , adaptive and efficient trade-off between accuracy and latency ( mainly related to model complexity ) of neural networks at the inference stage . Typically , sub-path networks of different complexity ( Liu & Deng , 2018 ) or sub-depth networks ( Huang et al. , 2018 ; Wu et al. , 2019 ; Wang et al. , 2018 ; Zhang et al. , 2021 ) are trained together . However , due to the large memory footprint brought by a constant number of channels per layer , the memory footprint at inference is barely reduced . To address the challenge , slimmable neural network ( SNN ) ( Yu et al. , 2018 ) was proposed to train networks with adaptive layer widths . Distinct from SNN , we consider a more challenging scenario with distributed and non-sharable data and heterogeneous client capabilities . 3 PROBLEM SETTING . The goal of this work is to develop a federated learning ( FL ) strategy , that yields a global model , which can be instantly customized for dynamic needs . Formally , FL minimizes the objective 1∑ k |Dk| ∑K k=1 ∑ ( x , y ) ∈Dk L ( f ( x ; W ) , y ) with distributed data { Dk } K k=1 across K participants/clients , where L is loss function , e.g. , cross-entropy loss , f is a model parameterized by W , and x , y are a sample and a label , respectively . In this paper , we consider two kinds of customization : adjustable model sizes to trade off between execution complexity and accuracy , adjustable robustness for dynamic security sensitivity . Though other types of customization could be considered , for example , neural-architecture search or fairness , we believe the two types considered are quite representative for both architecture and non-architecture knowledge ( Wang et al. , 2020 ) . When training customizable/adjustable models , significant challengs arise from heterogeneity among clients . In this paper , we consider the following : 1 ) Heterogeneous computational budgets during training of clients constrain the maximal complexity of local models . The complexity of deep neural networks can be defined in multiple dimensions , like depths or widths ( Zagoruyko & Komodakis , 2017 ; Tan & Le , 2019 ) . In this paper , we consider the width of hidden channels , because it can significantly reduce not only the number of model parameters but also the layer output cache . Thus , we assume clients have confined width capabilities { Rk ∈ ( 0 , 1 ] } Kk=1 , defined by width ratios w.r.t . the original model , i.e. , ×Rk net as presented in Fig . 1b . Similar to FedAvg and HeteroFL , the same architecture is used by clients and therefore the model width can be tailored according to local budgets . 2 ) Imbalance computation budgets . In many applications , there are usually a significant number of devices with insufficient computational budgets ( Enge , 2021 ) . Without loss of generality , we assume exponentially distributed budgets in uniformly divided client groups : Rk = ( 1/2 ) ⌈4k/K⌉ , i.e. , the first group with 1/4 clients is capable for a 1-width and the rest are for 0.5 , 0.25 , 0.125-widths respectively . Such a budget distribution could result in additional challenges in the presence of data heterogeneity , e.g. , non-i.i.d . features , Dk ∼ Di for ordered domain i , since training a single model on multiple domains ( Li et al. , 2020b ) or models in a single domain ( due to budget constraints ) ( Hong et al. , 2021 ) are known to be suboptimal on average performance . Finally , following the standard FL setting ( McMahan et al. , 2017 ) , only model parameters can be shared to protect privacy . In summary , the problem setting has challenges from 1 ) resource heterogeneity , where disparity in computation budgets will significantly affect the training of large models because they can only be trained scarce budget-sufficient clients ; 2 ) data heterogeneity , where the under-trained large models may perform poorly on unseen in-domain or out-of-domain samples .

Summary: Split-Mix is a Federated Learning strategy aimed at easing the problems that arise when a heterogeneous pool of devices/clients (i.e. some with more compute/memory capabilities than others, different data distributions) collaboratively train a global model. Split-Mix trains a model that can later on be customized in terms of model size and robustness. As the name suggests, Split-Mix has two stages: During `split`, a large model is split into smaller base models. Base models are constructed by discarding channels but maintaining the the number and type of layers in the network. During `mix`, the server samples a fraction of the base models (depending on a given client's compute capabilities) and fuses them into a single one that is send to a client to train. All base models are trained on all clients (should these meet the compute requirements of a sub-model) in a federated manner. This means that the more capable devices, train all base models. This is envisioned to happen in a parallel fashion in a given device. Once FL training is completed, a customized model can be deployed to a device/client in hardware-aware and robustness-ware fashion. The Authors refer to this as _in-situ customization_