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 paper proposes a new federated learning framework called HeteroFL, which supports the training of different sizes of local models in heterogeneous clients. Clients with higher computation capability can train larger models while clients with less computation capability train smaller models, and all these model architectures belong to the same model class. This approach dramatically benefits clients with limited computation capability and fully exploits their computation power.