Year: 2021

Abstract:
Federated Learning (FL) is a recently proposed learning paradigm for decentralized devices to collaboratively train a predictive model without exchanging private data. Existing FL frameworks, however, assume a one-size-fit-all model architecture to be collectively trained by local devices, which is determined prior to observing their data. Even with good engineering acumen, this often falls apart when local tasks are different and require diverging choices of architecture modelling to learn effectively. This motivates us to develop a novel personalized neural architecture search (NAS) algorithm for FL. Our algorithm, FEDPNAS, learns a base architecture that can be structurally personalized for quick adaptation to each local task. We empirically show that FEDPNAS significantly outperforms other NAS and FL benchmarks on several real-world datasets.

heading: 1 INTRODUCTION
text: 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 cannot 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).

heading: 2 RELATED WORKS

heading: 2.1 TWO-STAGE NEURAL ARCHITECTURE SEARCH
text: 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.

heading: 2.2 DISCRETE STOCHASTIC NEURAL ARCHITECTURE SEARCH
text: 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.

heading: 3 PERSONALIZED NAS FOR FEDERATED LEARNING

heading: 3.1 FEDERATED LEARNING OF DSNAS
text: 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.

heading: 3.2 PERSONALIZABLE ARCHITECTURE SEARCH SPACE
text: Similar to DSNAS (Hu et al., 2020), our framework adopts a cell-based representation (Section 2.2) to trade-off search space expressiveness for efficiency, which is extremely suitable for FL where clients tend to have low-end computational capacity. Unlike the original design which assumes similar role for every cell in the architecture stack (i.e., as reflected by their choice of fully factorizable path sampling distribution p(Z)), we instead split our cell stack into two components with separate metaroles catering to the federated personalization task: (a) a base stack ψb = {ψb1, ψb2 . . . ψbCb} which aims to capture the broad commonalities of data samples across client tasks; and (b) personalized stack ψp = {ψp1 , ψ p 2 . . . ψ p Cp }, which will be fine-tuned with local data to capture task-specific details.
We explain the main difference between these components to account for different level of expressiveness requirements below:
Base stack. Every cell ψbt in the base stack takes as inputs the outputs of its previous two cells ψbt−1 and ψbt−2 (replaced with raw input x when necessary for t ≤ 2). The output of the skip-ahead cell
ψbt−2 is additionally passed through a 1× 1 convolution layer as a cost-effective way to control the number of channels. Additionally, the operators available to the base cell include large convolution layers with size 5× 5 and 7× 7. To compensate for the growing number of channels, we periodically employ a reduction convolution (with stride larger than 1) similar to DSNAS (Hu et al., 2020) to reduce the feature dimension down the stack.
Personalized stack. As opposed to the design of the base cells above, every cell ψpt in the personalized stack has minimal expressiveness. That is, ψpt excludes large operators and only takes as input the output of its immediate predecessor ψpt−1 (or ψ b Cb
when t = 1). There are two reasons for this choice. First, as the fine-tuning phase has access to fewer data samples than the federated phase, having a more compact fine-tuning space helps to improve the rate of convergence. Second, as we will discuss in Section 3.4 below, our personalized FL objective requires the Hessian of the personalized parameters, which is computationally expensive. As such, we only restrict the personalization to happen on the more compact personalized stack.

heading: 3.3 PERSONALIZED FEDERATED LEARNING OBJECTIVE
text: Unlike FEDAVERAGING (McMahan et al., 2017), which assumes the clients will follow the consensus base model obtained after the federated phase, FEDPNAS expects clients to further personalize the base model with local task data. That said, while the base model is trained to work well in the expected sense over the task distribution, there is no guarantee that it is a good initial point for every client model to improve upon via fine-tuning. To address this, we adopt the concept of training in anticipation of future adaptation introduced by MAML (Finn et al., 2017). That is, during client update, instead of optimizing the loss with respect to the same consensus weight, each client will instead optimize the weight perturbed by a small gradient step in the fine-tuning direction.
Algorithm 1 FEDPNAS - FEDERATED PHASE 1: CENTRALAGGREGATION: 2: θ0 ← INITIALIZEPARAMETER 3: for t = 1, 2 . . . Ts do 4: for k = 1, 2 . . .M in parallel do 5: θkt ← CLIENTUPDATE(k, θt−1) 6: θt ← ∑M k=1 1 M θ k t
7: CLIENTUPDATE(k, θ): 8: for t = 1, 2 . . . Tc do 9: for batch (x, y) ∈ Dk do
10: Lk,∇Lk ← EVAL(x, y; θb, θp) 11: θ̃p ← GRADUPDATE(∇θ̃pLk) 12: L̃k,∇L̃k ← EVAL(x, y; θb, θ̃p) 13: ∇θpL̃k ← EQ. 5 14: θb, θp ← GRADUPDATE(∇θb,θpL̃k) 15: return θ to central server
Algorithm 2 FEDPNAS - EVAL 1: Input: x, y, θ 2: SKIP, PREV ← x, x 3: W,Π← θ 4: for CELL ψ ∈ ψ do 5: Z← SAMPLEOPS (x, PREV;Π) 6: ψ ← EXTRACTCHILDNET (Z,W)
7: if ψ ∈ ψb then 8: OUTPUT ← ψ(PREV, SKIP) 9: else if ψ ∈ ψp then
10: OUTPUT ← ψ(PREV) 11: SKIP ← PREV 12: PREV ← OUTPUT 13: L ← LOSS(OUTPUT, y) 14: ∇L ← BACKPROP(L) 15: return L,∇L
For simplicity, let θ = {θb, θp} respectively denote all trainable parameters of the base stack and the personalized stack, i.e., θb = {Wb,Πb}, θp = {Wp,Πp}. The personalized FL objective at client i is then given by L̃i(θb, θp) , Li(θb, θ̃p) where θ̃p , θ̃p − η∇θpLi(θb, θp) adjusts the parameters of the personalized component to account for a small fine-tuning gradient step. The adjusted local loss only depends on the respective client data and is amenable to federated learning. The local update gradient, however, involves a Hessian term whose computation is expensive to repeat over many epochs. To circumvent this problem, we use the first-order Taylor approximation to estimate the Hessian term by the outer product of Jacobian, which results in a gradient that requires exactly two forward/backward passes to compute:
∇θpL̃i = ( ∇θp θ̃p )( ∇θ̃pL̃i ) = ( I− η∇2θpLi )( ∇θ̃pL̃i
) ' ( I− η∇>θpLi∇θpLi )( ∇θ̃pL̃i ) (5)
where Li and L̃i are short-hands for Li(θb, θp) and L̃i(θb, θp) respectively. The FL phase of our FEDPNAS framework is detailed in Alg. 1. An instance of FEDPNAS’s forward and backward pass which sequentially unrolls down the component stacks, alternating between sampling and evaluation, is in turn given in Alg. 2.

heading: 3.4 CONTEXT-AWARE OPERATOR SAMPLER
text: The choice of a fully factorizable sampling distribution p(Z) in DSNAS follows that of SNAS (Xie et al., 2018), which argues that the Markov assumption for p(Z) is not necessary because NAS has fully delayed rewards in a deterministic environment. However, this generally only holds for two-stage NAS (Section 2.1) and does not apply to end-to-end frameworks such as SNAS and DSNAS. We instead to take advantage of the over-parameterized architecture via factorizing the conditional p(Z | x), which takes into account the temporal dependency of structural decisions:
p(Z | x) = p(Z1 | x) C∏ t=2 p(Zt | Zt−1 . . .Z1, x)
' p(Z1 | x) C∏ t=2 p(Zt | vt1, x)
= p(Z1 | x) C∏ t=2 ∏ (i,j) p(Ztij | vt1, x) , (6)
where Zt, Ztij and v t 1 respectively denote all the samples, the sample at edge (i, j) and the input at cell ψt. We have also assumed a single stack setting since the parameterization of p(Z) does not differ between base and personalized cells.
To reduce computational complexity, instead of conditioning the samples of subsequent cells on previous Z samples, we approximate p(Zt | Zt−1 . . .Z1, x) ' p(Zt | vt1, x) by the assumption that the cell contents are conditionally independent given the immediate cell input and the original input. Finally, we assume that p(Zt | vt1, x) is fully factorizable across edges in the same cell and parameterize p(Ztij | vt1, x) = φ(i,j)(vt1, x) where φ is a deep classification network whose output dimension equal the number of edges in cell ψt. Samples of Ztij can then be generated using the straight-through Gumbel-softmax reparameterization similar to Jang et al. (2016).

heading: 3.5 THEORETICAL CONNECTION TO STANDARD GRADIENT UPDATE
text: Finally, we analyze the connection of our gradient update framework to the standard gradient update, and explain why it is critical in achieving a vantage point that improves average objective value without compromising any local objective. First, we note that the gradient update Eq. 5 in Section 3.3 at the t-th iteration can be written as:
θt+1p = θ t p − η2 M M∑ i=1 ∇θpLi(θtb, θ̃tp,i) + η1η2 M M∑ i=1 αi∇θpLi(θtb, θtp) ,
and θt+1b = θ t b − η2 M M∑ i=1 ∇θbLi(θtb, θ̃tp,i) , (7)
where θ̃tp,i denotes the i-th local personalized parameters; η1 and η2 are two separate learning rates and αi , ∇>θpL(θ t b, θ t p,i)∇θp,iL(θtb, θ̃tp,i) (See Appendix A for detailed derivation). This implies that our federated personalize update corresponds to a federated update scheme with three gradient steps: (1) θp takes a local gradient (w.r.t. locally updated parameters) step of size η1; (2) Both θb and θp take a federated gradient (w.r.t. server-wide parameters averaging) step of size η2; (3) θp takes a weighted federated gradient step of size η1η2, where the weight of client i is given by αi.
Explicitly, Step 1 and 2 together comprise a special instance of FEDAVERAGING (McMahan et al., 2017), where θb take one gradient step for every two gradient steps taken by θp. Step 3, on the other
hand, takes the information of the two gradient steps of θp and adjust the magnitude of the local gradient step (whose direction is given by∇θpLi(θtb, θtp)) accordingly to trade-off between preserving local objective value and improving average objective value. We then theorize the scenario in which such an update is beneficial and state the following assumption to lay the foundation of our analysis:
Assumption 1 For a fixed instance of θb, let θ̃p,i = θp,i − η∇θpL(θb, θp,i) denote the personalized parameters after a local update step (i.e., step 1 above) at client i, then there exists a distribution S on matrix S ∈ Rk×|θp| that satisfies
∀x ∈ Rn, ‖x‖2 = 1 : ES∼S [ |‖Sx‖22 − 1|` ] ≤ ` · δ , (8)
Pr S∈S (∣∣∣∇θpL(θb, θ̃p,i)− S>S( 1M M∑ i=1 ∇θpL(θb, θ̃p,i) )∣∣∣ ≤√ 6 kδ ) ≥ 1− δ (9)
where k = O ( C‖θb − θ∗b‖ −2 2 ) for some constant C > 0 and θ∗b denotes the optimal base parameters.
The above assumption implies that, as θb improves and better captures the broad similarity across tasks, the local personalized components will diverge to capture the differences. It then becomes less likely for the FEDAVG personalized gradient to capture all these differences simultaneously. That is, suppose there exists an affine transformation to reconstruct the local component ∇θpLi(θtb, θ̃tp) from the federated gradient L , 1M ∑M i=1∇θpLi(θtb, θ̃tp), then the rank of this affine transformation would be inversely proportionate to ‖θb − θ∗b‖2. Finally, Proposition 1 below shows that when this assumption holds and θb converges to the optimal parameter θ∗b (i.e., the error term tends to 0), then with very high probability, the coefficient αi of the weighted federated gradient step (i.e., step 3 above) accurately captures the cosine similarity between the local gradient (step 1) and the federated gradient (step 2).
Proposition 1 Suppose assumption 1 holds, then with probability at least 1 − 2δ and normalized gradients, we have: ∣∣∣αi −∇>θpLi(θtb, θtp)L∣∣∣ = O(‖θb − θ∗b‖/δ) (10) Proof. See Appendix B
This result has strong implication with respect to the scenario with multiple heterogeneous tasks, whose local gradients contradict in directions. Per this setting, we expect a standard federated gradient update scheme to encourage parameters drifting in the general direction of the majority (i.e., captured by the federated gradient), thus worsening the performance of tasks that are in the minority. Proposition 1, however, implies that whenever the local gradient contradicts the federated gradient, αi will be close to the cosine similarity term, which is negative. This in turn results in a dampening effect on the federated gradient and helps to preserve the client performance on its own local task.

heading: 4 EXPERIMENTS
text: This section describes our experiments to showcase the performance of FEDPNAS compared to different NAS and FL benchmarks on various scenarios. All of our empirical studies are conducted on two image recognition datasets: (a) the CIFAR-10 dataset (Krizhevsky et al., 2009) which aims to predict image labels from 10 classes given a train/test set of 50000/10000 colour images of dimension 32× 32 pixels; and (b) the MNIST dataset (LeCun et al., 2010) which aims to predict handwritten digits (i.e. 0 to 9) given a train/test set of 60000/10000 grayscale images of dimension 28 × 28 pixels. Our search space entails 240 possible architectures, which is detailed in Appendix D. We compare two variants of our framework, CA-FEDPNAS (with context-aware operation sampler) and FEDPNAS (without the operation sampler), against: (a) FEDAVERAGING of a fixed architecture to justify the need for NAS in FL; (b) FEDDSNAS - the federated extension of DSNAS (Section 3.1) to show the effectiveness of our proposed context-aware sampler on NAS performance; and finally (c) CA-FEDDSNAS, which extends FEDDSNAS with our context-aware sampler.
On simulate heterogenous predictive tasks. To simulate this scenario, we first distribute the data i.i.d across clients (10000/2000 and 12000/2000 training/test images per client for CIFAR-10 and MNIST datasets respectively). Then, we independently apply a different transformation to each partitioned dataset. Input images within the same train/test set is subject to the same transformation. In both our experiments, the client datasets are subjected to rotations of −30◦,−15◦, 0◦, 15◦ and 30◦ respectively. This data generation protocol reflects a realistic and frequently seen scenario where independently collected data of the same phenomenon might contain systematic bias due to measurement errors and/or different collection protocols. Fig. 3 below shows the performance of all the methods in comparison, plotted against number of search epochs and averaged over the above rotated variants of CIFAR-10 and MNIST datasets.
On the MNIST dataset (Fig. 3b), all methods eventually converge to a similar performance. Among the NAS benchmarks, FEDPNAS and FEDDSNAS both converge slower than FEDAVG and start off with worse performance in early iterations, which is expected since FEDAVG does not have to search for the architecture and it is likely that the default architecture is sufficient for the MNIST task. On the other hand, we observe that both CA-FEDPNAS and CA-FEDDSNAS converge much faster than their counterparts without the context-aware operation sampler component. This shows that making use of contextual information helps to quickly locate regions of high-performing architectures, especially on similar inputs.
On the CIFAR-10 dataset (Fig. 3a), we instead observe significant gaps between the worst performing FEDAVG and other NAS methods. This is likely because the default architecture does not have sufficient learning capability, which confirms the need for customizing solutions. Among the NAS benchmarks, we again observe that both CA-FEDPNAS and CA-FEDDSNAS outperform their counterparts without our operation sampler, which confirms the intuition above. Most remarkably, our proposed framework CA-FEDPNAS achieves the best performance (0.8) and significantly outperformed both variants of federated DSNAS (0.71 for CA-FEDDSNAS and 0.63 for FEDDSNAS).
Lastly, Fig. 3c shows the runtime comparison between three methods on the CIFAR-10 experiment. In terms of sampling time, we observe that there is negligible overhead incurred by using our context-aware sampler (CA-FEDDSNAS vs. FEDDSNAS). The time incurred by our update (CA-FEDPNAS) scales by a constant factor compared to CA-FEDDSNAS since we use exactly one extra forward/backward pass per update.
On objectives with varying heterogeneity. We expand the above study by investigating respective performance of CA-FEDPNAS and FEDDSNAS on tasks with varying levels of heterogeneity. At low level of heterogeneity, we deploy these methods on 5 sets of slightly rotated MNIST images. At high level of heterogeneity, we employ a more diverse set of transformations on MNIST images, such as hue jitter and large angle rotations of 90◦ and−90◦. Table 1 show the respective result of each task from these two settings. We observe that our method CA-FEDPNAS achieves better performance on most tasks and the performance gaps on tasks with higher heterogeneity are more pronounced (i.e., up to 7% improvement on ROTATE 90 task). This clearly shows the importance of architecture personalization when the training tasks are significantly different and justifies our research goal.
On knowledge transfer to completely new tasks. Finally, we investigate a scenario where the architecture distributions discovered by CA-FEDPNAS and FEDDSNAS are required to generalize to completely unseen tasks. Particularly, we train both methods on five clients whose local data consist of 12000 slightly rotated CIFAR-10 images (i.e., in the range of ±30◦), similar to the setting of the first experiment. During testing, however, we supply each local client with 2000 test images subjected to related but completely unseen transformations (i.e., 90◦ and −90◦ rotations). Our results are summarized in Table 2. First, we measure the performance of CA-FEDPNAS and FEDDSNAS without any weight retraining. When received no additional information from the unseen tasks, both methods perform poorly as expected. While CA-FEDPNAS achieves better predictive accuracy, the performance gap in this scenario is negligible. To provide additional clues for adaptation, albeit minimal, we retrain the weights of each local model with 200 images rotated according to respective unseen task description. Here, the parameters of our operator sampler component, (and respectively, FEDDSNAS’s categorical distribution parameters), are frozen to gauge the quality of the learned architecture distributions. Our results show that, with only 100 retraining iterations on limited data, CA-FEDPSNAS already outperforms FEDDSNAS (5% and 8% improvement respectively on two unseen tasks). This implies that CA-FEDPNAS has more accurately capture the broad similarity of the task spectrum through the personalized architecture distribution, which requires minimal additional information to successfully adapt to unseen tasks.

heading: 5 CONCLUSION
text: We demonstrate that federated learning for multi-task scenarios requires extensive personalization on the architecture level to obtain good predictive performance. This paper identifies two potential sources of model personalization: (1) task-personalization, which aims to select architectures best suited for specific learning objectives; and (2) context-personalization, which aims to select architectures best suited for specific input samples. To incorporate these aspects of personalization into Federated NAS, we propose FEDPNAS which consists of two main components: (1) a context-aware operator sampler which learns a sampling distribution for feature maps along a master architecture; and (2) a personalized federated learning objective which anticipates client fine-tuning and guides the federated model update to regions that tolerate future local updates.

heading: 6 REPRODUCIBILITY & ETHIC STATEMENT
text: This work contributes to the literature of Federated Learning through improving the state-of-theart performance. As such, it could have significant broader impact by allowing users to more accurately solve practical problems. While applications of our work to real data could result in ethical considerations, this is an indirect (and unpredictable) side-effect of our work. Our experimental work uses publicly available datasets to evaluate the performance of our algorithms; no ethical considerations are raised. Our implementation code is published anonymously at https://github.com/icml2021fedpnas/fedpnas. All proofs and details of various architectures are included in the Appendix of this paper.

Review summaries:
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.

The paper proposes a personalized neural architecture search technique for federated learning. The paper incorporates both task-personalization and context-personalization. Experimental results on both CIFAR-10 and MNIST datasets demonstrate the promise of the proposed method.