FedNAS: Federated Deep Learning via Neural Architecture Search

1 INTRODUCTION . Federated Learning ( FL ) is a promising approach for decentralized machine learning , which aims to avoid data sharing and lower the communication cost ( McMahan et al. , 2016 ) . As such , it has gained a lot of attention in various domains of machine learning such as computer vision , natural language processing , and data mining . Despite its widespread popularity , one of the key challenges of FL is data heterogeneity . Since users ’ data are not identically or independently distributed ( non-IID ) in nature , a globally learned model may not perform optimally on all user devices . When interweaving with data heterogeneity , data invisibility is another issue that has rarely been studied . For this reason , to find a better model architecture with higher accuracy , developers must design or choose multiple architectures , then tune hyperparameters remotely to fit the scattered data . This process is extremely expensive because attempting many rounds of training on edge devices results in a remarkably higher communication cost and on-device computational burden than the data center environment . To mitigate the challenge of data heterogeneity , researchers have proposed methods to train a global model , including FedProx Li et al . ( 2018 ) , FedNova Wang et al . ( 2020 ) , and FedOPT Reddi et al . ( 2020 ) . Additionally , personalized frameworks such as Ditto Li et al . ( 2020 ) , pFedMe Dinh et al . ( 2020 ) , and PerFedAvg Fallah et al . ( 2020 ) have been recently developed to optimize personalized models to adapt to individual user ’ s data . These prior works have made remarkable progress in designing optimization schemes for pre-defined model architectures operated at pure optimization . However , these algorithms all require lots of effort to tune hyperparameters ; this is attributed to their strong prior assumptions , which may not always match the unknown data distribution . For example , practitioners must tune the regularization parameter in Ditto Li et al . ( 2020 ) and pFedMe Dinh et al . ( 2020 ) to find a proper correlation between the aggregated global model and local model . Moreover , their design is only in optimization level and does not consider the efficacy of model selection and neural architecture design , leading to a suboptimal solution when using a pre-defined model . We aim to address data heterogeneity in FL via a different and complementary approach that is based on model personalization through neural architecture search ( NAS ) . NAS has recently gained much momentum to adapt heterogeneity in neural architecture design Smithson et al . ( 2016 ) ; Chen et al . ( 2020 ) , latency Wu et al . ( 2019 ) ; Tan et al . ( 2019 ) ; Cai et al . ( 2019 ) , memory footprint Cai et al . ( 2018 ) ; Marchisio et al . ( 2020 ) , energy consumption Hsu et al . ( 2018 ) ; Yang et al . ( 2020 ) for edge devices . NAS methods are often categorized into three types : gradient-based methods Liu et al . ( 2018 ) , evolutionary methods Liu et al . ( 2020 ) , and reinforcement learning ( RL ) -based methods Jaafra et al . ( 2019 ) . Among these , gradient-based methods are the most efficient as they can finish searching in only a few hours , compared to thousands of GPU days with other methods . In this work , to search for a personalized neural architecture for mitigating the data heterogeneity , we adopt an improved variant of the gradient-based method , MiLeNAS He et al . ( 2020c ) , which is computationally tractable and particularly suitable for resource-constrained edge devices . Particularly , we propose a new method named Federated NAS ( FedNAS ) to search model architectures among edge devices collaboratively . As shown in Figure 1 , FedNAS works in the following way . We first utilize the MiLeNAS He et al . ( 2020c ) as a local searcher on each client ’ s local data , which can be distributed easily and efficiently in search time ( Step 1 ) . Formally , it formulates NAS as a mixed-level problem : w = w − ηw∇wLtr ( w , α ) , α = α − ηα ( ∇αLtr ( w , α ) + λ∇αLval ( w , α ) ) , where w represents the network weight and α represents the neural architecture . Ltr ( w , α ) and Lval ( w , α ) denote the loss with respect to training data and validation data , respectively . After the local search , each client then transmits weights w and architecture α to the server ( Step 2 ) . The server then applies a weighted aggregation to obtain the server-side α and w ( Step 3 ) and sends the updated parameters back to each client for the next round of searching ( Step 4 ) . During the searching process , we can personalize the α and w parameters by alternative local adaptation . Such personalization method can either obtain a higher accuracy for various data distributions , or automate the training process with lightweight hyper-parameter searching efforts . We evaluate FedNAS comprehensively in curated non-I.I.D . datasets , including CIFAR-10 and GLD-23K . Our datasets cover both global model training and personalized model training . We also consider different training scenarios : cross-silo FL and cross-device FL , which has a different total number of clients and number of clients per round . We demonstrate that the personalized model architectures learned by the individual clients perform better than the fine-tuned FedAvg and other representative personalized FL methods such Ditto Li et al . ( 2020 ) and perFedAvg Fallah et al . ( 2020 ) with default hyper-parameters in most settings . In summary , our main contributions in this paper are three-fold . 1 . We propose the FedNAS method to search for both global model and personalized model architectures collaboratively among edge devices and show its satisfying performance in a variety of FL settings . 2 . We investigate the role of NAS to address the challenge of data-heterogeneity in FL and demonstrate via experimental results that it can adapt to users ’ data better than existing local adaptation and personalization schemes . 3 . We experimentally show that FedNAS can achieve state-of-the-art performance for both cross-silo and cross-device settings . 2 PROPOSED METHOD . 2.1 PROBLEM DEFINITION . In the federated learning setting , there are K nodes in the network . Each node has a dataset Dk : = { ( xki , yi ) } Nk i=1 which is non-IID . When collaboratively training a deep neural network ( DNN ) model with K nodes , the objective function is defined as : min w f ( w , α︸︷︷︸ fixed ) def = min w K∑ k=1 Nk N · 1 Nk ∑ i∈Dk ` ( xi , yi ; w , α︸︷︷︸ fixed ) ( 1 ) wherew represents the network weight , α determines the neural architecture , and ` is the loss function of the DNN model . To minimize the objective function above , previous works choose a fixed model architecture α then design variant optimization techniques to train the model w. We propose to optimize the federated learning problem from a completely different angle , optimizing w and α simultaneously . Formally , we can reformulate the objective function as : min w , α f ( w , α ) def = min w , α K∑ k=1 Nk N · 1 Nk ∑ i∈Dk ` ( xi , yi ; w , α ) ( 2 ) In other words , for the non-IID dataset scattered across many workers , our goal is to search for an optimal architecture α and related model parameters w to fit the dataset more effectively thus achieve better model performance . In this work , we consider searching for CNN architecture to improve the performance of the image classification task . 2.2 SEARCH SPACE . Normally , NAS includes three consecutive components : the search space definition , the search algorithm , and the performance estimation strategy Hutter et al . ( 2019 ) . Our search space follows the mixed-operation search space defined in DARTS Liu et al . ( 2018 ) and MiLeNAS He et al . ( 2020c ) , where we search in two shared convolutional cells and then build it up as an entire model architecture ( as shown in Figure 2 ) . Inside the cell , to relax the categorical candidate operations between two nodes ( e.g. , convolution , max pooling , skip connection , zero ) to a continuous search space , mixed operation using softmax over all possible operations is proposed : ō ( i , j ) ( x ) = d∑ k=1 exp ( α ( i , j ) k ) ∑d k′=1 exp ( α ( i , j ) k′ ) ︸ ︷︷ ︸ pk ok ( x ) ( 3 ) where the weight pk of the mixed operation ō ( i , j ) ( x ) for a pair of nodes ( i , j ) is parameterized by a vector αi , j . Thus , all architecture operation options inside a network ( model ) can be parameterized as α . More details are introduced in Appendix A.1.1 . 2.3 LOCAL SEARCH . Following the above-mentioned search space , each worker searches locally by utilizing the mixedlevel optimization technique MiLeNAS He et al . ( 2020c ) : w = w − ηw∇wLtr ( w , α ) α = α− ηα ( ∇αLtr ( w , α ) + λ∇αLval ( w , α ) ) ( 4 ) where Ltr ( w , α ) and Lval ( w , α ) denote the loss with respect to the local training data and validation data with w and α , respectively . 2.4 FEDNAS : FEDERATED NEURAL ARCHITECTURE SEARCH . Algorithm 1 FedNAS Algorithm . 1 : Initialization : initialize w0 and α0 ; K clients are selected and indexed by k ; E is the number of local epochs ; T is the number of rounds . 2 : Server executes : 3 : for each round t = 0 , 1 , 2 , ... , T − 1 do 4 : for each client k in parallel do 5 : wkt+1 , αkt+1 ← ClientLocalSearch ( k , wt , αt ) 6 : wt+1 ← ∑K k=1 Nk N wkt+1 7 : αt+1 ← ∑K k=1 Nk N αkt+1 8 : 9 : ClientLocalSearch ( k , w , α ) : // Run on client k 10 : for e in epoch do 11 : for minibatch in training and validation data do 12 : Update w = w − ηw∇wLtr ( w , α ) 13 : Update 14 : α = α− ηα ( ∇αLtr ( w , α ) + λ∇αLval ( w , α ) ) 15 : return w , α to server We propose FedNAS , a distributed neural architecture search algorithm that aims at optimizing the objective function in Equation 2 under the FL setting . We introduce FedNAS corresponding to four steps in Figure 1 : 1 ) The local searching process : each worker optimizes α and w simultaneously using Eq . 4 for several epochs ; 2 ) All clients send their α and w to the server ; 3 ) The central server aggregates these gradients as follows : wt+1 ← K∑ k=1 Nk N wkt+1 αt+1 ← K∑ k=1 Nk N αkt+1 ( 5 ) 4 ) The server sends back the updated α and w to clients , and each client updates its local α and w accordingly , before running the next round of searching . This process is summarized in Algorithm 1 . After searching , an additional evaluation stage is conducted by using a traditional federated optimization method such as FedAvg McMahan et al . ( 2016 ) .

Summary: This paper combines gradient-based NAS (DARTS-like algorithm: MileNAS) with Federated Learning (FL) setup, to improve both global and personalization performance with learned neural architecture. Since both NAS and FL learning are based on gradient, the extension to FL setup becomes intuitive and effective. Empirically FedNAS shows improved performance comparing to existing FL methods.