Personalized Neural Architecture Search for Federated Learning
Authors: Minh Hoang
Carl Kingsford
Venue: ICLR 2022 Submitted (2022)
Keywords: Personalized Learning
Federated Learning

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.

Under review as a conference paper at ICLR 2022
PERSONALIZED NEURAL ARCHITECTURE SEARCH FOR FEDERATED LEARNING
Anonymous authors Paper under double-bling review
ABSTRACT
Federated Learning (FL) is recently proposed learning paradigm for decentralized devices to collaboratively train 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 learn effectively. This motivates us to develop novel personalized neural architecture search (NAS) algorithm for FL. Our algorithm, FEDPNAS, learns 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.
INTRODUCTION
Federated Learning (FL) (MeMahan al. 2017) is variant of distributed learning where the objective function can be decomposed into a linear combination of M local objective functions. Each function depends its private data hosted by local client and set of shared parameters w,
M argmin L(w) = argmin Σ fww (D))
(1)
where D; denotes the 7th local training dataset comprising input-output tuples (x, v). In standard supervised learning task where the predictive model is modeled as fixed deep neural network eb with learnable weights let f(x, y) denote the loss incurred by predicting vxx; w) when the true output y. The expected loss of (1;;; w) on D; is given as
C D:) E( (z y: v(
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 provided by both clients, the correct classifications must be the opposite of one another. However, since existing FL methods converge on 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 al.) (2020) extends FL by incorporating ideas from meta learning (Finn et al. (017)) which results in new framework of personalized FI.. The new framework can accommodate for such task heterogeneity
still requires all client models to agree on single architecture beforchand, which is sub-optimal. To address this shortcoming, one naive idea is to adopt existing idcas in Neural Architecture Search (NAS) via Reinforcement Learning (Zoph and Le 2016 Pham et al. 2018) which act outer loop 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 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

---

Under review as a conference paper at ICLR 2022
FL (Fallah et al., 2020) under 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, over- parameterized network. The sampling distribution is (collaboratively) learned along with the pa- rameters 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 task- specific 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 fine- tuning with local data (Section 3.1 To ensure that the common architecture distribution converges at a vantage point that relevant and beneficial to all clients, we generalize the vanilla FL. objective in Eq. equation | 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
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-base based
strategies Florcano et al. 2008 Real et al. 2019),, Bayesian optimization surrogates (Bergstra et al.
2013. Hu et al. 2018) or keinforcement Learning controllers (Baker et al., 2016 Zoph and Le 2016 Pham et 20188 to propose candidate architectures based on random mutations and/or observed
experience; while the latter optimizes the parameters 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 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.
generally not suitable
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 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 good architecture is discovered.
The main idea of DSNAS is to combine weight training for an parameteaettri master architecture with discrete computational path sampling. DSNAS parameterizes the master architecture
stack of modular cells: vb(x) the 11-1 ... w (x) where x is an arbitrary input, C' is the number of cells, v/, denotes the (th cell in the stack, and denotes the compositional operator. The inner computation of dh is in turn characterized by a directed acyclic graph (DAG) with V nodes {u, [V] where each node represents some intermediate feature map. For each directed edge (v,, U2).
Althouh each DSNAS cell receives the outputs of two previous cells inputs, we simplify the number to for ease of notation.
2

---

Under review as a conference paper at ICLR 2022
Federated Phase
Server
Fine-tune Phase
Server
Client G_L.x
Client Ci
Weight
*B0/
1/D-
t
Euerry fecture map
Weight broedcast
VP
WPP
FedPNNAS Update
GD Update
Figure Our proposed method FEDPNAS consists of (1) federated learning phase, where each client updates both the base component (vl²) and the personalized component (1/,)) the architecture using the FEDPNAS update (Section 3.3) and sends its parameters to the central server for aggrega- tion; and (2) fine-tune phase, where each client updates only the personalized component of the architecture using standard gradient update.
there is associated list of D possible network operations Oij [o] ... of where each operation oi transforms v; to Uy. Here, U1 corresponds to the output of previous cell v/--1 (or input when t = 1). We recursively define intermediate nodes U, E Zr O.y(s,,), where (1)( & [oly (v:), (f::)....o) (f..)) distribution p(Z IT) where the event probabilities II { learnable. Essentially, learning the distribution p(Z) allows us sample computational paths or sub-graphs of the original DAG that correspond to high-pos seah-eeforrmngg. compact architecture from the over-parameterized master network. Sampling discrete random variables from p(Z). however, does not result in a gradient amenable to pack-propagation. To sidestep this issue, DSNAS adopts the Brraight-hrouuee Gumbel-softmax trick (Jang et al. 2016), which re- parameterizes the Ath index of
and Zij is a one-hot vector sampled from the categorical
TD Σ
are
where 9: Gumbel(O, 1). While this forward computation does not have gradient by itself, we can estimate the gradient through proxy during the backward pass:
the one-hot variable as Z [k] = Il arg max 9e + log क.
exp ((gx + log TT) /T) V
V Zij k
Z k
Σe exp ((g+ + log (T) /T)
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 setting, especially when local tasks are 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:
M
= argmin Σ C.W.II D.).
argmin W.n
(4)
W.n. M
McMahan et al. (017) 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
We drop the cell index in the definition operation for of notation.

---

Under review as a conference paper at ICLR 2022
central aggregation step, all clients will follow the same architecture distribution induced by the final broadcasted copy of W and II. 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.c., feature maps received 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-wiss wise personalization to federated neural architecture search in multitask scenarios, which achieved through our proposed algorithm FEDPNAS. In general, FEDPNAS functions similarly to the vanilla FEDDSNAS algo- rithm 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. I however, we need to address the following key challenges:
To make this work,
C1. First, as previously argued in Section [ tasks across federated clients tend to share similarities in broad sense, and diverge 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, major advantage of having wer-pparaeeterized architecture search space is the flexibility of having specific computation paths for different samples, which not exploited by DSNAS reflected in its choice of context- context-independent sampling distribution p(Z). To address this, Section 3.4 proposes a novel parameterization of pZ) to incorporate context information into operator sampling.
C3. Last, while the fine-tuning phase is designed to incorporate pask-personlliiation there is no guarantee that the common model can be quickly adapted to client tasks (Fallah et 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 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.
3.2 PERSONALIZABLA ARCHITECTURE SEARCH SPACE
CONY 181
CONY IXI
COSY LXI
Similar DSNAS (Hu al. 2020), our frame- work adopts a cell-based representation (Sec- tion 2.2 trade-off search space expressive- ness for efficiency, which is extremely suit- able for FL where clients tend have low-end computational capacity. Unlike the original de- sign which assumes similar role for every cell in the architecture stack (i.e., reflected by their choice of fully factorizable path sampling distribution p( (Z)), stack into two components with separate meta- roles catering to the federated personalization task: (a) a base stack th, = { (v4, vol &&, which aims to capture the broad commonali- ties of data samples across client tasks; and (b) personalized stack + {vi,, v/ v J.
x
(x)
instead split our cell
Figure 2: Feature mapping down the component stacks of our architecture space. Every base cell takes as inputs (a) the outputs from its immediate predecessor and (b) the one before it through skip-ahead connection. On the other hand, every personalization cell takes as input only the output from the previous cell.
which will be fine-tuned with local data cap- ture task-specific details.
We explain the main difference between these components to account for different level of expres- siveness requirements below:
Base stack. Every cell vi in the base stack takes as inputs the outputs of its previous two cells vif and w (replaced with raw input x when necessary for 1 2). The output of the skip-ahead cell
4

---

Under review as a conference paper at ICLR 2022
2i222 additionally passed through a 1 x convolution layer cost-effective way control the number channels. Additionally, the operators available to the base cell include large convolution layers with size 5 x 5 and 7 7. To compensate for the growing number of channels, 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 vif in the personal- ized stack has minimal expressiveness. That is, vi excludes large operators and only takes input the output of its immediate predecessor - (or Vo, when / = 1). There are two reasons for this choice. First, as the fine-tuning phase has access fewer data samples than the federated phase, having more compact fine-tuning space helps to improve the rate of convergence. Second, as we will discuss in Section 3.4 below, our personalized FI. objective requires the Hessian the personalized parameters, which is computationally expensive. As such, we only restrict the personalization to happen on the more compact personalized stack.
3.3 PERSONALIZED FEDERATED LEARNING OBJECTIVE
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 good initial point for every client model to improve upon via mne-tuning. To address this, we adopt the concept of training in anticipation of future adaptation introduced by MAML (Finn et al. (017). 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 small gradient step in the fine-tuning direction.
Algorithm 1 FEDPNAS FEDERATED PHASE
Algorithm 2 FEDPNAS EVAL
I: CENTRALAGGREGATION:
Input: 2, y,
2: 6g + -NITIALIZEPARAMTEER
2: SKIP. PREV +- Ir
3: for 1,2... do
3: W, + 0
4: for k == 1,2....... M in parallel do
4: for CELL & E 26 do
5:
05 t CLIENTUPDATE(K, 0(-1)
5: Z + SAMPLEOPS (I, PREV; II) PREV; 6: 1b < EXTRACTCHILDNET (Z, W)
6:
0₂ + -EE 1
7: CLIENTUPDATE(M; 0):
7:
if th E Th, then 8: OUTPUT + (PREV, SKIP) else if 1h E 1/ then
8: for 1 1,2....7. do
9:
for batch (z,y) v) E Dj do
9:
10:
Lx. VLx + EVAL(Z, yi 0b 0p)
10:
OUTPUT + (PREV)
11:
8p + GRADUPDATE(V v0 Cx)
SKIP + PREV
12:
Ž,, VÂn + EVAL(2: Jr, yi On, 0p)
12: PREV +- OUTPUT
C L-LOS(OUUTUT,)) y)
13:
Vo, Cx <- -Q.5
14:
VL ((-BAAKPROP(L)
9o. 0p + GRADUPDATE( V To.0, Cx)
15: return L, VL
15: return 0 to central server
For simplicity, let 0 = (0,, 0p) respectively denote all trainable parameters of the base stack and the personalized stack, i.e., 6, = { (,, I1.. (W p: IIp). The personalized FL objective at client i is then given by (.,. 0p) (.,, 0p) where 0, 0, Ve, (0,, 0p) adjusts the parameters of the personalized component account for small tne-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 Hessian term whose computation is expensive repeat over many epochs. To circumvent this problem, we use the first-order Taylor approximation to estimate the Hessian term by the outer product Jacobian, which results in gradient that requires exactly two forward/backward passes to compute:
Vo, E. = (2 (Vo,C.)
I -I-vVa Vo C.) (Vēg o0 C.)
I 1- V- L. Vo (Vo
(5)
5

---

Under review as a conference paper at ICLR 2022
where Cj and Ci are short-hande for (0, 8p) and (0,,, 02) respectively. The FL phase of our FEDPNAS framework is detailed in Alg. [ 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
3.4 CONTEXT-AWARE OPORATOR SAMPLER
The choice of fully factorizable sampling distribution p(Z) in DSNAS follows that of SNAS (Xie et al., 2018), which argues that the Markov assumption for p is not necessary because NAS has fully delayed rewards deterministic environment. However, this generally only holds for two-stage NAS (Section 2.1p and does not apply to end-to-end frameworks such as SNAS and DSNAS. We instead to take advantage of the paraammametere architecture via factorizing the conditional p(Z] x)) which takes into account the temporal dependency of structural decisions:
p(2 x))
p(z
II t=2
Z,-1 ... 21,
C p(z, x)
p(Z, (x) II
t=2 C
= pZz ())
II p pZ vf,z),
(6)
t=2
(i.i)
where Z. Zi and of respectively denote all the samples, the sample at edge (i.j) and the input at cell th. 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, approximate p(Ze Z Z1, x) = p(Z, vf,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(Z, uf,x)) is fully factorizable across edges in the same cell and parameterize p(Zi, x)
((55)) (1)))))) (v), where is deep classification network whose output dimension equal the number of edges in cell thy Samples of Z can then be generated using the straight-through Gumbel-softmax reparameterization similar to Jong et al. (2016).
3.5 THEORETICAL CONNECTION TO STANDARD GRADIENT UPDATE
Finally, we analyze the connection of our gradient update framework to the standard gradient update, and explain why it is critical in achieving vantage point that improves average objective value without compromising any local objective. First, we note that the gradient update Eq. in Section 3.3 at the t-th iteration can be written as:
M
M
0.+1 = P
0¹
12
Va, C, 0p,)
T1 M
Σ a. Vo, C,(,,
M
M
and 0.+1 = 0¹
12
V C. (0,,
M
where ph denotes the i-th local personalized parameters; 7/ and 172 are two separate learning rates and o; A VT, L(0g, Vo,, 0/,,) (See Appendix A for detailed derivation). This implics that our federated personalize update corresponds to federated update scheme with three gradient steps: (1) 0p takes a local gradient (w.rt. locally updated parameters) step of size mi (2) Both Og and 6p take a federated gradient (w.r.t. server-wide parameters averaging) step of size 72 (3) 8p takes a weighted federatted federated gradient step of size 7) 772, where the weight of client is given by og-
Explicitly, Step and 2 together comprise a special instance of FEDAVERAGING (MMMahan et al. 2017), where 0, take one gradient step for every two gradient steps taken by 0,. Step 3, on the other

---

Under review as a conference paper at ICLR 2022
hand, takes the information of the two gradient steps of 8g and adjust the magnitude of the local gradient step (whose direction is given by Ve, (8 (,)) 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 For a fixed instance of 0b, let 0p.5 0p.5 V, [(0,, 0p.)) denote the personalized parameters after a local update step (i.e. step above) at client i, then there exists a distribution S on matrix S R (l that satisfies
Vx IR",
||x||2 1: EsnS [[|Sx[]] 2 - 119] S E.,,
(8)
M
Pr SES
Ve, [(ss, 0.:) - - S
Σ M
Ve, [(,,
02x) (0. S
6
2 2-1
1-8
1:
where - O (C 0,112²) ( for some constant C and 02 denotes the optimal base parameters.
The above assumption implies that, as 9, 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 Ce,,, (p) from the federated gradient L J7 Σ Ve, (0t, 65), then the rank of this affine transformation would be inversely proportionate to |0, 0g 2-
Finally, Proposition| T below shows that when this assumption holds and 0, converges to the optimal parameter 0ī (i.e., the error term tends to 0), then with very high probability, the coefficient o, of the weighted federated gradient step (i.e., step above) accurately captures the cosine similarity between the local gradient (step 1) and the federated gradient (step 2).
Proposition 1 Suppose assumption holds, then with probability at least 1 - 28 and normalized gradients, we have:
0x Vo, (06, 0,)L O(10,, 0511/5)
(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, expect 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 in the minority. Proposition I however, implies that whenever the local gradient contradicts the federated gradient, ag will be close to the cosine similarity term, which is negative. This in turn results in dampening effect on the federated gradient and helps to preserve the client performance on its own local task.
4 EXPERIMENTS
This section describes our experiments 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.) 22099) which aims to predict image labels from 10 classes given a train/test set of 50000/10006 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 train/test set of 60000/ 10000/10000 10000 grayscale images of dimension 28 x 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.
7

---

Under review as a conference paper at ICLR 2022
On simulate heterogenous predictive tasks. To simulate this scenario, first distribute the data Li.d across clients (10000/2000 and 12000/2000 training/test images per client for CIFAR-10 and MNIST datasets respectively). Then, we independently apply 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 subjected to rotations of -30 -15,, 0%, 15 and 30" respectively, This data generation protocol reflects 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.
CA-PUEPNAS
CA-FUIDSNAS
FuiDSNNAS
0.6
0.4
0.-
6:
250 300
250
20
Epochs
(a)
(b)
(c)
Figure 3: Plotting average classification accuracy of various methods against no. training epochs on heterogeneous tasks derived from (a) MNIST dataset; and (b) CIFAR-10 dataset. Figure (c) compares cumulative running time of various methods against no. training epochs on CIFAR-10 dataset.
On the MNIST dataset (Fig. (b), all methods eventually converge to 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 sufficient for the MNIST task. On the other hand, we observe that both CA- PA-FEDPNAS and CA-FEDDSNAS converge much faster than their counterparts without the context-awars 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 10 dataset (Fig. 3)), 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 operation sampler, which confirms the intuition above. Most remarkably, our proposed framework CA-FEDPNAS AS achieves the best performance (0.8) and significantly out- performed both variants of federated DSNAS (0.71 for CA-FEDDSNAS and 0.63 for FEDDSNAS).
Lastly, Fig.Be shows the runtime comparison between three methods on the CIFAR- 10 experiment. In terms of sampling time, we observe that there negligible overhead incurred by using our context-aware sampler (CA-FEDDSNAS vs. FEDDSNAS). The time incurred by our update (CA-FEDPNAS) scales by constant factor compared to CA-FEDDSNAS since use exactly one extra forward/backwarr 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 [ from these two settings. We observe that our method CA-FEDPNA on most tasks and the performance gaps tasks with higher heterogeneity (i.e., up to 7% improvement on ROTATE 90 task). This clearly shows the importance of architecture personalization when the training tasks significantly different and justifies research goal.
show the respective result of each task achieves better performance more pronounced
8

---

Under review as a conference paper at ICLR 2022
TASK
HETEROGENBITY
FEDDSNAS
CA-FEDPNAS
DESCRIPTION
ROTATE -30
0.947
0.978
ROTATE -15
0.973
0.976
LOW
VANILLA
0.988
0.985
ROTATE 15
0.986
0.987
ROTATE 30
0.972
0.981
HUEJITTER -0.5
0.966
0.978
HUEJITTER 0.5
0.967
0.972
HIGH
VANILLA
0.988
0.989
ROTATE -90
0.892
0.932
ROTATE 90
0.866
0.932
Table I: Predictive accuracy of CA-FEDPNAS FEDDSNAS on tasks with varying heterogeneity levels. ROTATE X denotes rotation transformation of X on client data; VANILLA denotes the original MNIST images; and HUEJITTER X denotes hue jitter transformation of training images by factor of X. The best performance in each row is in bold font.
On knowledge transfer to completely new tasks. Finally, we investigate 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..,, 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 related but completely unseen transformations (i.e., 90 and 900 rotations).
Our results are summarized in Table 2 First, we measure the performance of CA-FEDPNAS and FEDDSNAS without any weight retraining. When received 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' 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.
UNSEEN TASK
FEDDSNAS
CA-FEDPNAS
FEDDSNAS
CA-FEDPNAS
DESCRIPTION
(RETRAINED)
(RETRAINED)
ROTATE -90
0.545 + 0.04
0.578 + 0.09
0.699 0.12
0.734 + 0.17
ROTATE 90
0.553 + 0.12
0.569 + 0.06
0.673 = 0.13
0.727 + 0.22
Table 2: Predictive accuracy (averaged over clients) and standard deviation CA-FEDPNAS and FEDDSNAS two unseen tasks (CIFAR-10), The best performance in each row is in bold font.
5 CONCLUSION
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) mask-personalization, which aims to select architectures best suited for specific learning objectives; and (2) context-personanl contex--peroonalizationn which aims to select archi- tectures 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-awend aware operator sampler which learns a sampling distribution for feature maps along master architecture; and (2) a personalized federated learning objective which anticipates client fine-tuning and guides the federated model update regions that tolerate future local updates.

---

Under review as a conference paper at ICLR 2022
REPRODUCIBILITY & ETHIC STATEMENT
This work contributes the literature of Federated Learning through improving the state-of-the- art performance. As such, could have significant broader impact by allowing users to more accurately solve practical problems. While applications of our work real data could result in ethical considerations, this an indirect (and unpredictable) side-effect of our work. Our ex- perimental work uses publicly available datasets evaluate the performance of our algorithms; no ethical considerations are raised. Our implementation code is published anonymously at https: //g thub. com/ i cm 12021 f fedpnas/ fedpnas All proofs and details of various architectures are included in the Appendix of this paper.
REFERENCES
B. Baker, O. Gupta, N. Naik, and R. Raskar. Designing neural network architectures using reinforce- ment learning. arXiv preprint arXi:: 2016.
J. Bergstra, D. Yamins, and D. Cox. Making science of model search: Hyperparameter optimization in hundreds of dimensions for vision architectures. In International conference on machine learning, pages 115-123. PMLR, 2013.
A. Fallah, A. Mokhtari, and A. Ozdaglar. Personalized federated learning: Modcl-agnostic meta- learning approach. In Proc. NeurIPS, 2020.
C. Finn, Abbeel, and Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In Proc. ICML, pages 1126-1135, 2017.
D. Floreano, P. Dür, and C. Mattiussi. Neuroevolution: from architectures to learning. Evolutionary Intelligence, 1(1):47-63) 2008.
Hu, L. Shen, and Sun. Squeeze-and-eeciiation networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 7132-7141, 2018.
S. Hu, S. Xie, H. Zheng, C. Liu, J. Shi, X. Liu, and D. Lin. DSNAS: Direct neural architecture search without parameter retraining. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 12084-12092, 2020.
E. Jang, S. Gu, and B. Poole. Categorical reparameterizzation with gumbel-softmax. arXiv preprint arXiv:1611.00 144, 2016.
A. Krizhevsky et Learning multiple layers of features from tiny images, Citeseer, 2009.
Y. LeCun, Cortes, and C. Burges. Mnist handwritten digit database. ATT Labs [Online], Available: http://yann, lecun. com/exdb/mmist,, 2, 2010.
H. B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. Arcas. Communication-eftiiicenn learning of deep networks from decentralized data. Proc. AISTATS, pages 1273-1282, 2017.
H. Pham, M. Y. Guan, B. Zoph, Q. V. Le, and J. Dean. Efficient neural architecture search via parameter sharing, arxiv preprint arXiv:1802.08: 1802.03268, 2018.
E. Real, A. Aggarwal, Y. Huang, and Q L.e Regularized evolution for image classifier architecture search. In Proceedings of the AAAI conference Artificial Intelligence, volume 33, pages 4780-4789, 2019.
S. Xie, H. Zheng, C. Liu, and L. Lin. SNAS: stochastic neural architecture search. arXiv preprint 182:::11...0922, 2018.
B. Zoph and Q V. Le. Neural architecture search with reinforcement learning. arXiv preprint 2016.
10


Reviews:
[{"review_id": "JnkTZOpo4gT", "reviewer": "Reviewer_MieJ", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n+ The authors target the personalization problem in federated learning from a combination of task and context points of view.\n+ The authors propose a novel algorithm heavily inspired by DSNAS and adopt it in a federated setting. The algorithm accounts for model heterogeneity by incorporating a base component(trained via federated learning) and a fine tuning component(trained on client data) which is the primary contribution.\n+ The paper propose a context aware operator sampler to incorporate context into operator sampling stage of NAS\n+ The paper backs up their claims with clear theoretical foundations.\n\nWeakness:\n- The paper does not show the variety of architectures that were created due to the added personalization. It would be helpful to see how the personalized architectures differ from each other to judge how well it is really helping accuracy.\n- A comparison of the base component of the model learned from federated setting to the model trained on regular NAS with the full data should be present in the paper to get a better idea of how the algorithm performs.\n- The search space contains 2^40 possible architecture candidates. The authors should show the search cost as well.\n\n", "summary_of_the_paper": "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.", "summary_of_the_review": "This paper proposes a novel NAS algorithm in a federated setting and backs their claims with theoretical rigor. However, the lack of ablation studies diminishes the merits of the proposed technique.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "5WdrGIohYE", "reviewer": "Reviewer_1UyE", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n1. The algorithm of learning architecture with a base component and a personalized component in FL is novel and interesting.\n2. Experiment results are promising, showing that utilizing the context-aware operator sampler significantly improves the performance of both FedPNAS and FedDSNAS.\n\nWeaknesses:\n1. Presentation needs to be improved. Also, there are many typos in the manuscript. Fig. 2 is not referred to in the paper.\n2. The concept of 'task-personalization' is confusing. What is the difference between 'task-personalization' in this paper and distribution shift across clients? Or client data heterogeneity? \n3. The experiment setting is far from a realistic scenario. First, clients' data are not iid distributed. I am curious whether non-iid data distribution will make any difference in the resulting model performance. There are natural federated EMNIST dataset and synthetic federated Cifar10 datasets. Why not use these federated datasets for evaluation? Also, the number of clients participating in the training is 5, which is too small. \n4. According to the experiment results in Fig. 3, FedPNAS performs worse than FedDSNAS. Only after adding the context-aware operator sampler (CA), CA-FedPNS outperforms CA-FedDSNAS and FedDSNAS on Cifar-10. Therefore, the experiment results cannot demonstrate the usefulness of the personalized component.\n\nOther comments and questions:\n1. It would be interesting to see the accuracy for each client and see if the proposed method is suitable for personalization overall (e.g., all clients' models have high accuracy).\n2. According to experiment results, CA is the critical reason for CA-FedPNAS outperforming FedDSNAS. It would be good to see if this conclusion is valid on a more complicated/larger dataset.", "summary_of_the_paper": "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.", "summary_of_the_review": "The idea of personalized NAS in FL is interesting. The proposed algorithm to learn architecture with a base component and a personalized component is novel. However, the definition of task-personalization is confusing, and the experiment used for evaluation is not sufficient to fully demonstrate the effectiveness of FEDPNAS. Therefore, I think this paper is marginally below the acceptance threshold.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "2q2DRKhb2Sp", "reviewer": "Reviewer_RES6", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Positive points:\n\n1.\tThe paper is well-organized with good motivation on conducting personalized NAS for FL.\n2.\tThe idea of adopting NAS for personalized FL learning seems to be novel and worth explored for researchers and practitioners in both communities. \n\nConcerns:\n\n1.\tLack of experimental results on real dataset to provide the effectiveness of the proposed method. The significance of the experiments is not convincing enough by only using the two very cliché image benchmarks with limited clients.\n2.\tSome literatures (baselines) are missed such as FedNAS (https://arxiv.org/abs/2004.08546). \n3.\tThe method seems to be a combination of several existing well-known techniques without deliberate narrative and in-depth discussion of why we need this specific design and how they works under diversified settings. For example, I doubt about whether the NAS method could handle the non-iid data situation or stateless setting, which are a common case in FL setting. As even NAS on a single dataset itself is a very difficult problem, more experiments should be conducted along this direction to really show how NAS method applied in the proposed framework performs and affected under different FL settings. Similar discussions should also be considered for the meta-learning algorithms adopted. \n\n\n", "summary_of_the_paper": "This paper presents a personalized federated learning framework, which leverage a stacked search space of base and personalized architectures along with the context-aware search strategy to learn personalized neural network in the FL scenario. The proposed method is validated on two benchmark image datasets compared to its variants and the FedAvg baseline.", "summary_of_the_review": "Based on the above concern described above, I currently tend to weakly reject the paper but would like to hear more feedback from the author and discussion from other reviewers towards the final decision.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "00KwoHi99-U", "reviewer": "Reviewer_EuZf", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength:\n- The combination of NAS and FL is mostly novel, there are a few works done but so far all are very recent and unpublished.\n\nWeaknesses:\n- Main issue is the experiments. The authors do not compare to any method besides the versions they propose here and fedAvg baseline. Experiments are all very limited as well. The result is that one cannot say anything about the abilities of these models in practice.\n- Missing many references to PFL works, and should have compared to at least some of them. A partial list:  Shamsian et al \"Personalized federated learning using hypernetworks\", Dinh et al \"Personalized federated learning with Moreau envelopes\", \tZhang et al \"Personalized Federated Learning with First Order Model Optimization\", and Collins et al \"Exploiting Shared Representations for Personalized Federated Learning\".\n- We note that in Shamsian et al the authors do investigate having a few models of different sizes to fit varying computational resources. While not the same as NAS it should be referenced as related work.\n- The theoretical part is not clear at all. Assumption 1 seems strong and unintuitive, while the meaning of the result in Proposition 1 isn't clear from the text.\n- The algorithm isn't clear and so is the discussion about the Hessian that follows. I am pretty sure one can avoid computing the Hessian with efficient Jacobian-vector multiplications that should be easy to compute via backprob, but I am not clear from the text where exactly does the Hessian come from.", "summary_of_the_paper": "The paper proposes a way to combine neural architecture search (NAS) with federated learning (FL) and personalized federated learning (PFL). The authors describe the basic extension of DSNAS to federated learning and propose some modifications to better suit the federated framework.", "summary_of_the_review": "The experiments are very lacking, theory and writing is unclear.  ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers had remarkably consistent feedback about this paper. They appreciated the formulation of the federated learning problem with architectures having both shared and private (personalized) components. On the other hand, they felt the experiments were insufficient to prove the effectiveness of the method, and had several suggestions in terms of tasks and datasets. They also felt that it's hard to assess whether the existence of private/personalized components is warranted without visualizing the difference between architectures. Overall, the reviewers had good feedback that could strengthen the paper.
Decision: Reject