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: This paper proposes a customisation strategy named "Split-Max" for federated learning. The authors identify the heterogeneity of devices and data in FL scenarios. They present the importance of considering devices' budgets and dynamics when dispatching training models. Split-max can adjust the model size according to the devices' budget while maintaining good accuracy and robustness.  Split-max works in steps. First, multiple base models from different initialisations are trained to improve diversity. These base models are randomly given to clients to extract generalisable features. Then, base models are aggregated to the server. Secondly, to provide devices with models with different robustness, it trains two similar models together to capture both the standard-training accuracy and adversarial-training accuracy. Then layer-wise mixing is conducted to achieve both standard accuracy and adversarial accuracy.   Experiments show that Split-Mix achieves better accuracy than naive approaches. Moreover, with customisation, the models are smaller and more robust under budget constraints.