Personalized Neural Architecture Search for Federated Learning

1 INTRODUCTION . Federated Learning ( FL ) ( McMahan et al. , 2017 ) is a variant of distributed learning where the objective function can be decomposed into a linear combination of M local objective functions . Each function depends on its private data hosted by a local client and a set of shared parameters w , argmin w L ( w ) ≡ argmin w M∑ i=1 Li ( w | Di ) , ( 1 ) where Di denotes the ith local training dataset comprising input-output tuples ( x , y ) . In a standard supervised learning task where the predictive model is modeled as a fixed deep neural network ψ with learnable weights w , let ` ( x , y ) denote the loss incurred by predicting ψ ( x ; w ) when the true output is y . The expected loss of ψ ( x ; w ) on Di is given as Li ( w | Di ) = E ( x , y ) ∼Di [ ` ( x , y ; ψ ) ] . ( 2 ) This is not applicable to scenarios where local models are expected to solve different tasks which are similar in broad sense yet diverge in finer details . For example , consider the task of recognizing the outcome of a coin flip given images collected by two clients : one capture the coin from above , the other from below . This setting implies that when the same input image is provided by both clients , the correct classifications must be the opposite of one another . However , since existing FL methods converge on a single model architecture and weight , there can only be one predictive outcome which can not satisfy both tasks . To relax this constraint , the recent work of Fallah et al . ( 2020 ) extends FL by incorporating ideas from meta learning ( Finn et al. , 2017 ) which results in a new framework of personalized FL . The new framework can accommodate for such task heterogeneity but still requires all client models to agree on a single architecture beforehand , which is sub-optimal . To address this shortcoming , one naive idea is to adopt existing ideas in Neural Architecture Search ( NAS ) via Reinforcement Learning ( Zoph and Le , 2016 ; Pham et al. , 2018 ) which act as an outer loop to the existing FL routine . However , this simple approach does not allow client models to adapt to local tasks on an architecture level and is often not preferred due to the cost of repeated FL training . This paper proposes a novel personalized NAS algorithm for federated learning , which generalizes ideas in respective areas of NAS ( Zoph and Le , 2016 ; Pham et al. , 2018 ) originally developed for single-task scenarios , and FL ( Fallah et al. , 2020 ) under a unified len of federated personalized neural architecture search ( FEDPNAS ) . In particular , to customize the model architecture for each task in the FL workflow , FEDPNAS first represents the model architecture for each task as a sub-network sampled from a large , overparameterized network . The sampling distribution is ( collaboratively ) learned along with the parameters of the sampled network via a generalization of the recently proposed Discrete Stochastic NAS ( DSNAS ) method ( Hu et al. , 2020 ) . Unlike DSNAS , which lacks the ability to customize architecture for individual tasks , our generalized FEDPNAS incorporates model personalization on an architecture level . Our contributions include : 1 . A novel architecture that factorizes into a base component ( shared across tasks ) and a personalizable component , which respectively capture the task-agnostic and task-specific information ( Section 3.2 ) . 2 . A context-aware sampling distribution conditioned on specific task instance , which captures taskspecific information and naturally incorporates personalization into architecture search ( Section 3.4 ) . 3 . An FL algorithm that optimizes for a common architecture , followed by a personalization phase where each client subsequently adapts only the personalized component to fit its own task via finetuning with local data ( Section 3.1 ) . To ensure that the common architecture distribution converges at a vantage point that is relevant and beneficial to all clients , we generalize the vanilla FL objective in Eq . equation 1 such that local gradient steps directly optimize for expected improvement resulting from future fine-tuning ( Section 3.3 ) . 4 . A theoretical perspective on our FL objective ( Section 3.5 and thorough empirical analysis showing significant performance gain compared to state-of-the-art FL and NAS methods ( Section 4 ) . 2 RELATED WORKS . 2.1 TWO-STAGE NEURAL ARCHITECTURE SEARCH . Most existing NAS frameworks separately optimize for the optimal architecture and its parameters in two stages : searching and evaluation . The former stage usually employs evolutionary-based strategies ( Floreano et al. , 2008 ; Real et al. , 2019 ) , Bayesian optimization surrogates ( Bergstra et al. , 2013 ; Hu et al. , 2018 ) or Reinforcement Learning controllers ( Baker et al. , 2016 ; Zoph and Le , 2016 ; Pham et al. , 2018 ) to propose candidate architectures based on random mutations and/or observed experience ; while the latter optimizes the parameters of these architectures given task data and provide feedback to improve the search agent . Naturally , an extension of such methods to the FL setting is through distributing the evaluation workload over many clients , which does not require exposing private data . In practice , however , two-stage federated NAS frameworks are generally not suitable for the personalized FL setting for two reasons : ( a ) the clients often lack the computational capacity to repeatedly optimize the parameters for many candidate architectures ; and ( b ) the clients have to converge on a single architecture proposed by the central search agent . 2.2 DISCRETE STOCHASTIC NEURAL ARCHITECTURE SEARCH . Discrete stochastic neural architecture search ( DSNAS ) ( Hu et al. , 2020 ) addresses the computational issue of two-stage NAS by jointly optimizing the optimal architecture and its weight in an end-to-end fashion , which allows users to continually train a single network on demand over time as opposed to performing full parameter optimization for every candidate until a good architecture is discovered . The main idea of DSNAS is to combine weight training for an over-parameterized master architecture with discrete computational path sampling . DSNAS parameterizes the master architecture as a stack of modular cells : ψ ( x ) = ψC ◦ ψC−1 · · · ◦ ψ1 ( x ) 1 , where x is an arbitrary input , C is the number of cells , ψt denotes the tth cell in the stack , and ◦ denotes the compositional operator . The inner computation of ψt is in turn characterized by a directed acyclic graph ( DAG ) with V nodes { vi } |V |i=1 , where each node represents some intermediate feature map . For each directed edge ( vi , vj ) , 1Although each DSNAS cell receives the outputs of two previous cells as inputs , we simplify the number to one for ease of notation . there is an associated list of D possible network operations Oij = [ o1ij , o 2 ij . . . o D ij ] 2 where each operation okij transforms vi to vj . Here , v1 corresponds to the output of previous cell ψt−1 ( or input x when t = 1 ) . We recursively define intermediate nodes vj = ∑j−1 i=1 Z > ijOij ( vi ) , where Oij ( vi ) , [ o1ij ( vi ) , o 2 ij ( vi ) . . . o D ij ( vi ) ] and Zij is a one-hot vector sampled from the categorical distribution p ( Z | Π ) where the event probabilities Π = { π1 , π2 , . . . , πD | ∑D i=1 πi = 1 } are learnable . Essentially , learning the distribution p ( Z ) allows us to sample computational paths or sub-graphs of the original DAG that correspond to high-performing , compact architecture from the over-parameterized master network . Sampling discrete random variables from p ( Z ) , however , does not result in a gradient amenable to back-propagation . To sidestep this issue , DSNAS adopts the straight-through Gumbel-softmax trick ( Jang et al. , 2016 ) , which re-parameterizes the kth index of the one-hot variable as Zij [ k ] = I ( k , argmax t [ gt + log πt ] ) , where gt ∼ Gumbel ( 0 , 1 ) . While this forward computation does not have a gradient by itself , we can estimate the gradient through a proxy during the backward pass : ∇ Zij [ k ] ' ∇ Z̃ij [ k ] , ∇ ( exp ( ( gk + log πk ) /τ ) ∑D t=1 exp ( ( gt + log πt ) /τ ) ) ( 3 ) which is unbiased when converged as the temperature τ is steadily annealed to 0 ( Jang et al. , 2016 ) . This formulation , however , is not easily extended to the FL setting , especially when local tasks are not homogeneous . The key challenges in doing so are described in Section 3 , together with our proposed approaches . 3 PERSONALIZED NAS FOR FEDERATED LEARNING . 3.1 FEDERATED LEARNING OF DSNAS . Let W denote the concatenated weights of all network operations in the network architecture . The set up above of DSNAS ( Jang et al. , 2016 ) is then naïvely extendable to a FL setting via the following objective formulation : argmin W , Π L ( W , Π ) ≡ argmin W , Π 1 M M∑ i=1 Li ( W , Π | Di ) . ( 4 ) McMahan et al . ( 2017 ) optimizes this objective by alternating between ( a ) central agent broadcasting aggregated weights to local clients and ( b ) local clients sending gradient descent updated weights ( given local data ) to the central agent for aggregation . This , however , implies that after the last 2We drop the cell index in the definition operation for ease of notation . central aggregation step , all clients will follow the same architecture distribution induced by the final broadcasted copy of W and Π . As previously argued , this is not optimal in a heterogenous task setting which requires task-specific adaptation for local clients to achieve good performance . Furthermore , having the same sampling distribution p ( Z ) regardless of context ( i.e. , feature maps received as cell input ) limits the architecture discovery to those that perform reasonably on average over the entire dataset . However , we remark that restricting the architecture to be the same for every input samples is unnecessary and undermines the expressiveness of an over-parameterized search space . On the other hand , letting the architecture be determined on a per-sample basis makes better use of the search space and potentially improves the predictive performance . The focus of this work is therefore to incorporate both task-wise and context-wise personalization to federated neural architecture search in multitask scenarios , which is achieved through our proposed algorithm FEDPNAS . In general , FEDPNAS functions similarly to the vanilla FEDDSNAS algorithm described above , with an addition of a fine-tuning phase at each local client after the FL phase to adapt the aggregated common model for local task data , as shown in Fig . 1 . To make this work , however , we need to address the following key challenges : C1 . First , as previously argued in Section 1 , tasks across federated clients tend to share similarities in broad sense , and diverge in finer details . A good federated personalization search space , therefore , need to capture this fundamental observation through design and appropriate resource distribution . We address this challenge in Section 3.2 . C2 . Second , a major advantage of having an over-parameterized architecture search space is the flexibility of having specific computation paths for different samples , which is not exploited by DSNAS as reflected in its choice of context-independent sampling distribution p ( Z ) . To address this , Section 3.4 proposes a novel parameterization of p ( Z ) to incorporate context information into operator sampling . C3 . Last , while the fine-tuning phase is designed to incorporate task-personalization , there is no guarantee that the common model can be quickly adapted to client tasks ( Fallah et al. , 2020 ) . The common model may end up in a localization that favors one client over another , which makes it difficult for the latter to fine-tune . To address this concern , Section 3.3 proposes a new personalized federated NAS objective inspired by Finn et al . ( 2017 ) to optimize the common model in anticipation of further fine-tuning by the client models .

Summary: This paper proposes a personalized neural architecture search algorithm (FEDPNAS) for FL. FEDPNAS searches for an architecture with a base component (shared across clients) and a personalized component. It also uses a context-aware operator sampler to learn a sampling distribution for feature maps. It provides a theoretical analysis of the FL objective and empirically demonstrates that FEDPNAS outperforms FedAvg and FedDSNAS over image recognition tasks on CIFAR-10 and MNIST datasets.