HeteroFL: Computation and Communication Efficient Federated Learning for Heterogeneous Clients

1 INTRODUCTION . Mobile devices and the Internet of Things ( IoT ) devices are becoming the primary computing resource for billions of users worldwide ( Lim et al. , 2020 ) . These devices generate a significant amount of data that can be used to improve numerous existing applications ( Hard et al. , 2018 ) . From the privacy and economic point of view , due to these devices ’ growing computational capabilities , it becomes increasingly attractive to store data and train models locally . Federated learning ( FL ) ( Konečnỳ et al. , 2016 ; McMahan et al. , 2017 ) is a distributed machine learning framework that enables a number of clients to produce a global inference model without sharing local data by aggregating locally trained model parameters . A widely accepted assumption is that local models have to share the same architecture as the global model ( Li et al. , 2020b ) to produce a single global inference model . With this underlying assumption , we have to limit the global model complexity for the most indigent client to train its data . In practice , the computation and communication capabilities of each client may vary significantly and even dynamically . It is crucial to address heterogeneous clients equipped with very different computation and communication capabilities . In this work , we propose a new federated learning framework called HeteroFL to train heterogeneous local models with varying computation complexities and still produce a single global inference model . This model heterogeneity differs significantly from the classical distributed machine learning framework where local data are trained with the same model architecture ( Li et al. , 2020b ; Ben-Nun & Hoefler , 2019 ) . It is natural to adaptively distribute subnetworks according to clients ’ capabilities . However , to stably aggregate heterogeneous local models to a single global model under various heterogeneous settings is not apparent . Addressing these issues is thus a key component of our work . Our main contributions of this work are three-fold . • We identify the possibility of model heterogeneity and propose an easy-to-implement framework HeteroFL that can train heterogeneous local models and aggregate them stably and effectively into a single global inference model . Our approach outperforms state-ofthe-art results without introducing additional computation overhead . • Our proposed solution addresses various heterogeneous settings where different proportions of clients have distinct capabilities . Our results demonstrate that even when the model heterogeneity changes dynamically , the learning result from our framework is still stable and effective . • We introduce several strategies for improving FL training and demonstrate that our method is robust against the balanced non-IID statistical heterogeneity . Also , the proposed method can reduce the number of communication rounds needed to obtain state-of-the-art results . Experimental studies have been performed to evaluate the proposed approach . 2 RELATED WORK . Federated Learning aims to train massively distributed models at a large scale ( Bonawitz et al. , 2019 ) . FedAvg proposed by McMahan et al . ( 2017 ) is currently the most widely adopted FL baseline , which reduces communication cost by allowing clients to train multiple iterations locally . Major challenges involved in FL include communication efficiency , system heterogeneity , statistical heterogeneity , and privacy ( Li et al. , 2020b ) . To reduce communication costs in FL , some studies propose to use data compression techniques such as quantization and sketching ( Konečnỳ et al. , 2016 ; Alistarh et al. , 2017 ; Ivkin et al. , 2019 ) , and some propose to adopt split learning ( Thapa et al. , 2020 ) . To tackle system heterogeneity , techniques of asynchronous communication and active sampling of clients have been developed ( Bonawitz et al. , 2019 ; Nishio & Yonetani , 2019 ) . Statistical heterogeneity is the major battleground for current FL research . A research trend is to adapt the global model to accommodate personalized local models for non-IID data ( Liang et al. , 2020 ) , e.g. , by integrating FL with other frameworks such as assisted learning ( Xian et al. , 2020 ) , metalearning ( Jiang et al. , 2019 ; Khodak et al. , 2019 ) , multi-task learning ( Smith et al. , 2017 ) , transfer learning ( Wang et al. , 2019 ; Mansour et al. , 2020 ) , knowledge distillation ( Li & Wang , 2019 ) and lottery ticket hypothesis ( Li et al. , 2020a ) . Nevertheless , these personalization methods often introduce additional computation and communication overhead that may not be necessary . Another major concern of FL is data privacy ( Lyu et al. , 2020 ) , as model gradient updates can reveal sensitive information ( Melis et al. , 2019 ) and even local training data ( Zhu et al. , 2019 ; Zhao et al. , 2020 ) . To our best knowledge , what we present is the first work that allows local models to have different architectures from the global model . Heterogeneous local models can allow local clients to adaptively contribute to the training of global models . System heterogeneity and communication efficiency can be well addressed by our approach , where local clients can optimize low computation complexity models and therefore communicate a small number of model parameters . To address statistical heterogeneity , we propose a ” Masking Trick ” for balanced non-IID data partition in classification problems . We also propose a modification of Batch Normalization ( BN ) ( Ioffe & Szegedy , 2015 ) as privacy concern of running estimates hinders the usage of advanced deep learning models . 3 HETEROGENEOUS FEDERATED LEARNING . 3.1 HETEROGENEOUS MODELS . Federated Learning aims to train a global inference model from locally distributed data { X1 , . . . , Xm } across m clients . The local models are parameterized by model parameters { W1 , . . . , Wm } . The server will receive local model parameters and aggregate them into a global model Wg through model averaging . This process iterates multiple communication rounds and can be formulated as W tg = 1 m ∑m i=1W t i at iteration t. At the next iteration , W t g is transmitted to a subset of local clients and update their local models as W t+1i =W t g . In this work , we focus on the relaxation of the assumption that local models need to share the same architecture as the global model . Since our primary motivation is to reduce the computation and communication complexity of local clients , we consider local models to have similar architecture but can shrink their complexity within the same model class . To simplify global aggregation and local update , it is tempting to propose local model parameters to be a subset of global model parameters W t+1i ⊆ W tg . However , this raises several new challenges like the optimal way to select subsets of global model parameters , compatibility of the-state-of-art model architecture , and minimum modification from the existing FL framework . We develop Heterogeneous Federated Learning ( HeteroFL ) to address these issues in the context of deep learning models . A variety of works show that we can modulate the size of deep neural networks by varying the width and depth of networks ( Zagoruyko & Komodakis , 2016 ; Tan & Le , 2019 ) . Because we aim to reduce the computation complexity of local models , we choose to vary the width of hidden channels . In this way , we can significantly reduce the number of local model parameters , while the local and global model architectures are also within the same model class , which stabilizes global model aggregation . We demonstrate our method of selecting subsets of global model parameters Wl for a single hidden layer parameterized by Wg ∈ Rdg×kg in Fig . 1 , where dg and kg are the output and input channel size of this layer . It is possible to have multiple computation complexity levels W pl ⊂ W p−1 l · · · ⊂ W 1l as illustrated in Fig . 1 . Let r be the hidden channel shrinkage ratio such that d p l = r p−1dg and kpl = r p−1kg . It follows that the size of local model parameters |W pl | = r2 ( p−1 ) |Wg| and the model shrinkage ratio R = |W p l | |Wg| = r 2 ( p−1 ) . With this construction , we can adaptively allocate subsets of global model parameters according to the corresponding capabilities of local clients . Suppose that number of clients in each computation complexity level is { m1 , . . . , mp } . Specifically , we perform global aggregation in the following way . W pl = 1 m m∑ i=1 W pi , W p−1 l \W p l = 1 m−mp m−mp∑ i=1 W p−1i \W p i , . . . ( 1 ) W 1l \W 2l = 1 m−m2 : p m−m2 : p∑ i=1 W 1i \W 2i ( 2 ) Wg =W 1 l =W p l ∪ ( W p−1 l \W p l ) ∪ · · · ∪ ( W 1 l \W 2l ) ( 3 ) For notational convenience , we have dropped the iteration index t. We denote the W pi as a matrix/tensor . The W tg [ : dm , : km ] denotes the upper left submatrix with a size of dm × km . Also , W p−1 , t+1g \W p , t+1g denotes the set of elements included in W p−1 , t+1g but excluded in W p , t+1g . We exemplify the above equations using Fig . 1 . The first part of Equation ( 1 ) shows that the smallest part of model parameters ( blue , p = 3 ) is aggregated from all the local clients that contain it . In the second part of Equation ( 1 ) , the set difference between part p − 1 ( orange ) and p ( blue ) of model parameters is aggregated from local clients with computation complexity level smaller than p − 1 . In Equation ( 2 ) , the red part of model parameters can be similarly aggregated from m −m2 : p = m1 clients . In Equation ( 3 ) , the global model parameters W tg is constructed from the union of all disjoint sets of the partition . In summary , each parameter will be averaged from those clients whose allocated parameter matrix contains that parameter . Thus , a model of an intermediate complexity will have parameters fully averaged with all the other larger models but partially with smaller models ( according to the corresponding upper left submatrix ) . Several works show that wide neural networks can drop a tremendous number of parameters per layer and still produce acceptable results ( Han et al. , 2015 ; Frankle & Carbin , 2018 ) . The intuition is thus to perform global aggregation across all local models , at least on one subnetwork . To stabilize global model aggregation , we also allocate a fixed subnetwork for every computation complexity level . Our proposed inclusive subsets of global model parameters also guarantee that smaller local models will aggregate with more local models . Thus , small local models can benefit more from global aggregation by performing less global aggregation for part of larger local model parameters . We empirically found that this approach produces better results than uniformly sampled subnetworks for each client or computation complexity level . 3.2 STATIC BATCH NORMALIZATION . After global model parameters are distributed to active local clients , we can optimize local model parameters with private data . It is well-known that the latest deep learning models usually adopt Batch Normalization ( BN ) to facilitate and stabilize optimization . However , classical FedAvg and most recent works avoid BN . A major concern of BN is that it requires running estimates of representations at every hidden layer . Uploading these statistics to the server will cause higher communication costs and privacy issues Andreux et al . ( 2020 ) proposes to track running statistics locally . We highlight an adaptation of BN named as static Batch Normaliztion ( sBN ) for optimizing privacy constrained heterogeneous models . During the training phase , sBN does not track running estimates and simply normalize batch data . We do not track the local running statistics as the size of local models may also vary dynamically . This method is suitable for HeteroFL as every communication round is independent . After the training process finishes , the server sequentially query local clients and cumulatively update global BN statistics . There exist privacy concerns about calculating global statistics cumulatively and we hope to address those issues in the future work . We also empirically found this trick significantly outperforms other forms of normalization methods including the InstanceNorm ( Ulyanov et al. , 2016 ) , GroupNorm ( Wu & He , 2018 ) , and LayerNorm ( Ba et al. , 2016 ) as shown in Table 4 and Table 5 .

Summary: This work presents a novel FL algorithm named HeteroFL (the name might sounds weird to some peoples) and 3 different simple methods to improve FL in heterogeneous conditions (i.e. both in term of clients and data partitioning). These tricks are: 1. A revised batchnormalisation; 2. a pre-activity scaling; 3. a masked loss (i.e only consider local classes)  to help with non-IID datasets. All these modifications have been tested on 3 different datasets and 2 different tasks. From the results, we can see that the proposed approach works better. Although, it is not clear from where the benefit comes.