FedMorph: Communication Efficient Federated Learning via Morphing Neural Network

1 INTRODUCTION . Federated Learning ( FL ) ( Li et al. , 2021 ; Bonawitz et al. , 2019 ; Kairouz et al. , 2019 ; Li et al. , 2019 ) ensures data privacy by decoupling the training on the dataset in local clients and model aggregation in the global server . The iterative process of the local updates optimized on each client and aggregating the updates by the global server assures the convergence of the global model and the excellent generalization on test datasets . ( A ) . Despite the benefits gained by the paradigm of local-train and global-aggregation of FL , the server-clients communication burden and limited local computing resources suspend the deployment of the State-of-The-Art ( SOTA ) large neural networks on a wide scale . Recently the lossy model compression techniques , such as gradient sparsification ( Wangni et al. , 2017 ; Zhou et al. , 2021 ) and quantization ( Courbariaux et al. , 2016 ; Hubara et al. , 2017 ) , have been studied to relieve the communication issues . Practical as they announced , but the limitations are also apparent . The gradient sparsification technique reduces the communication burden by selectively updating local gradients . Hence , they only work for client-to-server communication and have no compression benefits in the downstream direction . Meanwhile , the quantization technique compresses the communication and computation by changing the floating operations into low precision operations , and it works in both directions . However , its compression ability is at most 32× by reducing 32-bits floating operations to the 1-bit operations , not to mention the performance degradation . Absorbing the benefits from the former lossy compression techniques , FedDropout ( Caldas et al. , 2018 ) is another trial to reduce the downstream direction communication . It builds upon the basic idea of dropout ( Srivastava et al. , 2014 ) by randomly dropping some neurons of the global model at each round , resulting in a smaller sub-network with smaller weight matrices before broadcasting it to local clients . The updated sub-networks by local clients will map back to the global network . Distillation of the ‘ knowledge ’ from an extensive neural network ( Hinton et al. , 2015 ) is another branch to solve local devices ’ communication and computation burdens . The main idea is to train a small model to keep the output similar to the large one on a sizeable prepared dataset . It has significant performance on the recent developments ( Sanh et al. , 2019 ) ( Jiao et al. , 2019 ) . However , it requires a well pre-trained large model and a vast dataset for training the smaller network , which is not always satisfactory in the context of FL . ( B ) . Besides considering communication and computation budget , broadcasting a raw neural network may also cause a severe generalization problem . While it is almost common sense in the deep learning community that training a well-designed deep neural network performs better than a shallow one , the statement only establishes when the training dataset is sufficient to avoid overfitting ( Pitt & Myung , 2002 ; Hinton et al. , 2012 ) . In FL settings , each local client holds a dataset in a small ratio of the whole , which is universal when the number of clients is enormous . Hence , human labor ’ s well-designed deep neural network works well on the whole dataset , nevertheless easily overfit the local datasets . The overfitting will impede the number of local iterations , which is a crucial hyper-parameter in FL settings ( McMahan et al. , 2017 ) , resulting in a deteriorated model performance . Kairouz et al . ( 2019 ) has a similar statement in their Network Architecture Search ( NAS ) section that a well-designed network architecture in centralized settings may not work well in federated learning settings . This paper tries to solve the earlier issues by decoupling the network architectures on local clients and the global server . In detail , the global server maintains a large neural network suitable in a centralized setting , and keeps the architecture untouched throughout the entire process . During each communication round , the global server morphs a new sub-network from the maintained architecture and broadcasts it to selected clients for local optimization . The maintained network optimizes its weights at the morphing process by learning from the average aggregated network updated by local clients . Meanwhile , a newly morphed sub-network is constrained by keeping the knowledge similar to the average aggregated one while minimizing the number of its parameters . By morphing the shared neural network into a smaller size among all communication rounds , we ( 1 ) decrease the communication in both upload and download directions ; ( 2 ) relieve the computation overload of local clients ; ( 3 ) reduce the generalization error caused by overfitting on local datasets with a large network . 2 RELATED WORKS . Neural Architecture Search Our work is closely related to the NAS ( Zoph & Le , 2016 ) , whose purpose is automatically designing neural network architecture such as the number of layers and the number of neurons or filters in each layer . The search strategies include random search , evolutionary methods , Bayesian optimization , and gradient-based methods ( Elsken et al. , 2019 ) . Since NAS involves a vast search space , recent works focus on improving the speed and efficiency by attentive sampling ( Wang et al. , 2021 ) , untrained scheme ( Mellor et al. , 2021 ) , and block-wisely search with knowledge distillation ( Li et al. , 2020 ) . Besides , the network architecture search under the scenario of federated learning scenario was also explored in ( Zhu et al. , 2021 ) , aiming to reduce both computation and communication . Model Compression To reduce the complexity of deep neural networks , model compression was first proposed in ( Buciluǎ et al. , 2006 ) , followed by tremendous attention in both academia and industry . One of the most straightforward method to reduce model size is parameter pruning and sharing by removing redundant parameters that are not critical to model performance ( Han et al. , 2015 ; Blakeney et al. , 2020 ) . Similarly , the informative parameters can also be measured and selected by low-rank factorization ( Sainath et al. , 2013 ; Denton et al. , 2014 ) . To simultaneously reduce the computation and storage of a deep model , approaches based on transferred or compact convolutional filters were further proposed by designing special structural kernels ( Cohen & Welling , 2016 ; Wu et al. , 2016 ) and achieved benefits in domains with human prior . Besides , some other works are focusing on transferring the learned knowledge of a large teacher network to a small and lightweight student network , which yields the concept of knowledge distillation ( Hinton et al. , 2015 ; Mirzadeh et al. , 2020 ; Chen et al. , 2021 ; Gao et al. , 2021 ) , which is suitable for small- or medium-size data sets ( Cheng et al. , 2018 ) . 3 PROBLEM DEFINITION . 3.1 PRELIMINARY . In this work , we consider the following federated learning optimization problem : min w { F ( w ) , K∑ k=1 pkFk ( w ) } , ( 1 ) where K is the number of clients , and pk is the weight of the k-th client such that pk ≥ 0 and∑K k=1 pk = 1 . Suppose the k-th client holds the nk training data : xk,1 , xk,2 , · · · , xk , nk . The local objective Fk ( · ) is defined by Fk ( w ) , 1 nk nk∑ j=1 ` ( w ; xk , j ) , ( 2 ) where ` ( · ; · ) is a user-specified loss function . Specifically we consider a C class classification problem defined over a compact space X and a label space Y = [ C ] , where [ C ] = { 1 , . . . , C } . The data point { x , y } is a random sample over X ×Y . A function fθ : X → S maps x to the probability simplex S , where S = { z | ∑C c=1 zc = 1 , zc ≥ 0 , ∀c ∈ [ C ] } , and z is a C-dimensional column vector . fθ is parameterized over the hypothesis class w , i.e. , the weight of the neural network θ . We define the loss function ` ( w , x ) with the widely used cross-entropy loss as ` ( w , x ) = Ey [ log fθ ( w , x ) ] . 3.2 ALGORITHM DESCRIPTION . Definition 1 . Morphing Set : Given a neural network θ with its parameter weight w , its Morphing Set ( Θ , W ) is defined as a set that contains every pair of sub-network and the corresponding weight of the original network θ . We define the concept of Morphing Set for the convenience of explanation . Before the description of the proposed algorithm , we define θo and wo as the server maintained neural network and its parameter weights , and denote ( Θo , Wo ) as the Morphing Set of θo . Here , we describe one round ( say the t-th ) of our proposed algorithm . First , the central server broadcasts the latest model , ( θt , wt ) , to the selected clients ( say the K ) . Second , every client ( say the k-th ) begins with wkt , e=0 = wt , and then performsE ≥ 1 local updates with a randomly selected batch samples ξkt , e as follows , wkt , e+1 ←− wkt , e − ηt , e∇Fk ( wkt , e , ξ k t , e ) ( 3 ) Third , unlike Federated Average ( FedAvg ) ( McMahan et al. , 2017 ) or other conventional algorithms letting w̃t+1 = ∑ k∈K wkt , E |K| be the model weights for next round , our method , in order to compress the communication , morphs a small neural network from θo by minimizing the performance divergence of the morphed model wt+1 with w̃t+1 as follows , ( θt+1 , wt+1 ) = arg min ( θ , w ) ∈ ( Θo , Wo ) Lt ( w ; w̃t+1 ) = arg min ( θ , w ) ∈ ( Θo , Wo ) ∑ x∈Xv JS ( fθ ( w , x ) ||fθt ( w̃t+1 , x ) ) , ( 4 ) where w̃t+1 = ∑ k∈K wkt , E |K| follows the FedAvg algorithm to be the average aggregated value of local clients updated parameter weights . ( θ , w ) ∈ ( Θo , Wo ) denotes the newly morphed network and its weights , which is a sub-network of the server maintained one . fθ is the corresponding mapping function defined on the morphed neural network θ . The Jensen-Shannon ( JS ) divergence is a well-known metric for symmetrically measuring the distance of two probabilities . We apply the JS divergence as the loss function to minimize the output performance of the morphed sub-network and average aggregated network on the validation compact space . Here , we use the finite-sum on the validation dataset Xv ( without label information ) to evaluate the performance divergence . Finally , the newly morphed sub-network with the updated weights ( θt+1 , wt+1 ) are broadcasted to selected clients for next round local optimization . The optimization objective is to maintain as much knowledge as possible from the aggregated model in previous round , while keeping the morphed network as small as possible . Therefore , we consider the network architecture as a regularizer and reformat equation ( 4 ) as follows , ( θt+1 , wt+1 ) = arg min ( θ , w ) ∈ ( Θo , Wo ) { Lt ( w ; w̃t+1 ) + λLc ( θ ) } , ( 5 ) where Lc ( θ ) is a regularizer to measure the constraints on the number of neural network parameters ( PARAMs ) or the computation floating-point operations per second ( FLOPs ) or any other constraints as required . λ is an adjustable non-trained variable to balance the two losses .

Summary: The paper proposes FedMorph to address the communication and computation heterogeneity problem in federated learning. At every round, small morphed sub networks are send to clients for local training. The server updates the global model by distilling from the aggregated morphed subnets.