Federated Learning with Partial Model Personalization
Authors: Krishna Pillutla
Kshitiz Malik
Abdelrahman Mohamed
Michael Rabbat
Maziar Sanjabi
Lin Xiao
Venue: ICLR 2022 Submitted (2022)
Keywords: personalization
federated learning
partial personalization
adapter modules
nonconvex minimization

Abstract: We propose and analyze a general framework of federated learning with  partial model personalization. Compared with full model personalization, partial model personalization relies on domain knowledge to select a small portion of the model to personalize, thus imposing a much smaller on-device memory footprint. We propose two federated optimization algorithms for training partially personalized models, where the shared and personal parameters are updated either simultaneously or alternately on each device, but only the shared parameters are communicated and aggregated at the server. We give convergence analyses of both algorithms for minimizing smooth nonconvex functions, providing theoretical support of them for training deep learning models. Our experiments on real-world image and text datasets demonstrate that (a) partial model personalization can obtain most of the benefit of full model personalization with a small fraction of personalized parameters, and, (b) the alternating update algorithm often outperforms the simultaneous update algorithm.

Under review as a conference paper at ICLR 2022
FEDERATED LEARNING WITH PARTIAL MODEL PERSONALIZATION
Anonymous authors Paper under double-bling review
ABSTRACT
We propose and analyze general framework of federated learning with partial model personalization. Compared with full model personalization, partial model personalization relies domain knowledge select a small portion of the model personalize, thus im- posing a much smaller on-device memory footprint. We propose two federated optimization algorithms for training partially personalized models, where the shared and personal param- eters are updated either simultaneously or alternately on each device, but only the shared parameters are communicated and aggregated at the server. We give convergence analyses of both algorithms for minimizing smooth nonconvex functions, providing theoretical support of them for training deep learning models. Our experiments on real-world image and text datasets demonstrate that (a) partial model personalization can obtain most of the benefit of full model personalization with a small fraction personalized parameters, and, (b) the alternating update algorithm often outperforms the simultancous update algorithm.
INTRODUCTION
Federated Learning (McMahan et al. 2017) has emerged powerful paradigm for distributed and privacy-preservinn machine learning over a large number of edge devices (see Kairouz et al. 2021 and references therein). We consider typical setting of Federated Learning (FL) with devices (also called clients), where each device has a training dataset of N; samples 24,1, N, Let w E Rd represent the parameters of a (supervised) learning model and Se(w, be the loss of the model the training example Zi.j- Then the loss function associated with device is Fi(w) = (1/N,) E fe(w, 21,)). A common objective FL is to find model parameters that minimize the weighted average loss across all devices (without transferring the datasets):
minimize
og aFF((w).
(1)
where weights a are nonnegative and satisfy Σラ a; = A common practice is to choose the weights as a; N₁/N where N = Eh Nk. which corresponds to minimizing the unweighted average loss across all samples from the n devices: (1/N) Σ N. fi(w, (.,))
The main motivation for minimizing the average loss over all devices is to leverage their collective statistical power for better generalization, because the amount of data each device be very limited. This is especially important for training modern deep learning models with large number of parameters. However, this argument assumes that the datasets from different devices are sampled from the same, or at least very similar, distributions. Given the diverse characteristics of the users and increasing trend of personalized on-device services, such an i.i.d. assumption may not hold in practice. Thus, the one-model-fits-all formulation in (1) can be less effective and even undesirable.
Several approaches have been proposed for personalized FL, including ones based on multi-task learning (Smith et al. (0077), meta learning (Fallah et al., (020),, and proximal methods (Dinh et al. 2020 Li et al. 2021). simple formulation that captures their main idea is
Σ as F((w:)) + A 1| - woll?).
minimize wa
(2)
2

---

Under review as a conference paper at ICLR 2022
u)(
F((u,1))
1/
A
1/
input
input
input 1
imput 2
(a) Personalized input layer(s). (b) Personalized output layer(s). (c) Personalized split input layer(s).
Figure 1: Three simple examples of partitioning deep learning modcls.
where wi for i = -. 7. are personalized model parameters at the devices, Wo is reference model maintained by the server, and the X;'s are regularization weights that control the extent of personal- ization. A major disadvantage of the formulation (2). which we call full model personalization, is that it requires twice the memory footprint of the model, wy and Wp at each device, which severely limits the size of trainable models. On the other hand, the flexibility of full model personalization can be unnecessary. Modern deep learning models are composed of many simple functional umits and are typically organized into layers more general interconnected architecture. Personalizing the "right" components, selected with domain knowledge, may result in a substantial benefit with only a small increase in memory footprint. In addition, partial model personalization can be less susceptible to "catastrophic forgetting" (McCloskey & Cohen 1989), where a large model finetuned on a small local dataset forgets the original (non-personn personalized) task, leading to a degradation of test performance.
We propose framework for FL with partial model personalization. Specifically, we partition the model parameters into two groups: the shared parameters
I8 and the personal parameters V E IR4: for i = I,..., n. The full model on device is denoted as w; = (u, u.), and the local loss
function is Fu v.) (1/N,) EN fi ((u, 24.)) Our goal is to solve the optimization problem
minimize M, (oe)
,o.F(uu
Notice that the dimensions of U; can be differrnt across the devices, allowing the personal components of the model to have different number of parameters or even different architecture.
We investigate two FL algorithms for solving problem B): FedSim, simultaneous update algorithm and FedAlt, an alternating update algorithm. Both algorithms follow the standard FL protocol. During each round, the server randomly selects subset of the devices for update and broadcasts the current global version of the shared parameters to devices in the subset. Each selected device then performs one or more steps of (stochastic) gradient descent to update both the shared parameters and the personal parameters, and sends the updated shared parameters to the server for aggregation. The updated personal parameters are kept local at the device to serve as the initial states when the device is selected for another update. In FedSim, the shared and personal parameters are updated simultaneously during each local iteration. In FedAlt, the devices first update the personal parameters with the received shared parameters fixed and then update the shared parameters with the new personal parameters fixed. We provide convergence analysis and empirical evaluation of both methods.
The main contributions of this paper are summarized as follows:
We propose a general framework of FL. with partial model personalization, which relies on domain knowledge to select small portion of the model to personalize, thus imposing much smaller memory footprint on the devices than full model personalization. This framework unifies existing work on personalized FL and allows arbitrary partitioning of deep learning models.
We provide convergence guarantees for the FedSim and FedAlt methods the general (smooth) nonconvex setting. While both methods have appeared in the literature previously, they are either used without convergence analysis or with results on limited settings (assuming convexity or full participation) Our analysis provides theoretical support for the general nonconvex setting with partial participation. The analysis of FedAlt with partial participation is especially challenging and we develop novel technique of virtual full participation to overcome the difficulties.
2

---

Under review as a conference paper at ICLR 2022
add & norm
Fl(u, m.)
adapter
adlatee
feedforward
add & norm
adapter
attention
imput
(a) Transformer layer with two adapters.
(b) Generalized additive model.
Figure 2: More structured partial model personalization. (a) The adapter has a skip connection, thus collapses to the identity mapping if e; O; in addition, it has a bottleneck in the middle (Houlsby et al. 2019)) (b) The generalized additive model can be further augmented with shared input layer for representation learning.
We conduct extensive experiments on image classification and text prediction tasks, exploring different model personalization strategies for each task, and comparing with several strong baselines. Our results demonstrate that partial model personalization can obtain most of the benefit of full model personalization with a small fraction of personalized parameters, and FedAlt often outperforms FedSim.
Our experiments also reveal that personalization (full or partial) may lead to worse performance for some devices, despite improving the average. Typical forms of regularization such as weight decay and dropout do not mitigate this issue. This phenomenon has been overlooked in previous work and calls for future research to improve both performance and fairness.
Related work. Specific forms of partial model personalization have been considered in previous works. Liang et al. (2019) propose to personalize the input layers to learn personalized representation per-device (Figure lab. ArivazhArvaz et al. (2019) and Collins et al. (202T) propose to personalize the output layer while learning shared representation with the input layers (Figure lb). Both FedSim and FedAlt have appeared the literature before, but the scope of their convergence analysis is limited. Specifically, Liang et al. (2019), AAnvazhagan et al. (2019) and Hanzely et al. (021) use FedSim, while Collins al. 2021) and Singhal et al. (2021) proposed variants of FedAlt. Notably, Hanzely et al. (2021) establish convergence of FedSim with full device participation in the convex and non-convex cases, while [Collins et al. (002I) prove the linear convergence of FedAlt for a two-layer linear network where Fi(-, 0.) and (u, -) are both convex for fixed v; and respectively, We analyze both FedSim and FedAlt in the general nonconvex case with partial device participation, hence addressing more general and practical setting.
While we primarily consider the problem (3) in the context of partial model personalization, it can serve general formulation that covers many other problems. [Hanzely et al. (2021) demonstrate that various full model personalization formulations based on regularization (Dinh et al. 2020 Li
et al. 2021), including (2). well as interpolation (Deng al. 2020: Mansour al. 2020)
special cases of this problem. The rates of convergence we prove in $ are competitive with or better than those in previous works for full model personalization methods the non-convex case.
2 PARTIALLY PERSONALIZED MODELS
Modern deep learning models all have multi-layer architecture. While complete understanding of why they work so well is still out of reach, a general insight that the lower layers (close to the input) mostly responsible for feature extraction and the upper layers (close to the output) focus complex pattern recognition. Depending on the application scenarios and domain knowledge, we may personalize either the input layer(s) the output layer(s) of the model; see Figure
In Figure [c] the input layers are split horizontally into two parts, one shared and the other personal. They process different chunks of the input vector and their outputs are concatenated before feeding
3

---

Under review as a conference paper at ICLR 2022
Algorithm 1 Federated Learning with Partial Model Personalization TredSim//eddai
Input: initial states (9))) {oz (0)
J number of rounds T,, number of devices per round m
1: for = 0, 1,... T
do
2: server randomly samples m devices as S(t) {l,...,n]
3: server broadcasts 2 to each device in S(+)
4: for each device i E 5(4) in parallel, do
5:
(u. (u(+t)) 2j ( LocalSim/Locallee
(1),, )2()) ((+1)
((+1)
=0!- if S()
6: send u;
(++1) back to server
7: server updates (t++)) 1 m Eieson u;
to the upper layers of the model. As demonstrated in (Bui et al. 2019). this partitioning help protect user-specific private features (input in Figure Ic) as the corresponding feature embedding (through U2) personalized and kept local the device. Similar architectures have also been proposed in context-dependent language models (e.g., Mikolov & Zweig, 2012)..
more structured partitioning is illustrated in Figure Za| where typical transformer layer (Vaswani let al.] (2017) is augmented with two adapters. This architecture is proposed by Houlsby et al.) ( (2019) for finetuning large language models. Similar residual adapter modules are proposed by Rebuffi et al. (2017) for image classification models in the context of muulti-task learning. In the context of FL, we treat the adapter parameters as personal and the rest of the model parameters as shared.
Figure 2b shows a generalized additive model, where the outputs of two separate models, one shared and the other personalized, are fused to generate prediction. Suppose the shared model is h(u,) and the personal model is -1),,.)) -). For regression tasks with samples zi.j = (Ii.j. Vig), where vi the input and Vi.j IRP is the output, we let (u, v4) (1/N,) N. 1=1 fi(uu, v;), Ziz) with
f ((u, 24), 21,) = llue, - hu, z.,)) ha x2,))l²
In this special case, the personal model fits the residual of the shared model and vice-verss Agarwal et al., 2020). For classification tasks, l(u,,)) and (x(r;,)) -) produce probability distributions over muluple classes. We can use the cross-entropy loss between yt.j and a convex combination of the two model outputs: 0h( Fi.j) + (1 - 0)),, ... where E (0, 1) is a learnable parameter.
Finally, we can cast the formulation 2 of full model personalization as special case of 3 by letting
1/4 + wi Fi v;) + F((1:) + (A:/2) - uII?.
u +
Many other formulations of full personalization can be reduced to see Hanzely et al. 2021
3 ALGORIHTMMS AND CONVERGENCE ANALYSIS
In this section, we present and analyze two FL algorithms for solving problem (3) To simplify presentation, we denote V (U,,,..."9 Un) ) IR d and focus on the case of or; 1/n, i.e.,
minimize,,
F(u, V) = Fu vz).
This is equivalent to 33 if we scale F; by na,, thus does not lose generality. Moreover, we consider more general local functions Fl,u, v4) E [f, v.), 2)], where D; is the local distribution.
The FedSim and FedAlt algorithms share a common outer-loop description given in Algorithm | They differ only in the local update procedures LocalSim and LocalAIt, which are given in Algonthm and Algorithm 3 respectively.
the two local update procedures, V., and V, represent stochastic gradicnts with respect to and respectively. In LocalSim (Algorithm 2). the personal variables i/ and local version of the shared parameters u, are updated simultaneously. with their (stochastic) partial gradients evaluated at the same point. In LocalAlt (Algorithm B), the personal parameters are updated first with the received shared parameters fixed, then the shared parameters are updated with the new personal parameters fixed. They are analogous the classical Jacobi update and Gauss-Seidel update in numerical linear algebra (e.g., Demmel 1997 $6.5).
In order to analyze the convergence of the algorithms, we make the following assumptions.
4

---

Under review as a conference paper at ICLR 2022
Algorithm 2 LocalSim v4)
Algorithm 3 LLoalAAl (u,
Input: number of steps step sizes Y/ and Ya
Input: number of steps Tux Tux step sizes Yeu Yu 1: initialize V1.0 V/
I: initialize 0; = 0;
2: initialize 44 = u
2: for k = 0, 1, ... To do
3: for k 0, 1....
do
3: U1,k+1 Vi,k To V Fi (u,,k)
4: 01,k+1 = Vi,k - 7oo V.F. (u,,,, vi.k)
4: update U, = Vt and initialize 21,0 = 2
5: U1,++1 44
Tu V.Fi (u V(,k)
5: for k = 0, 1, 1, Tu do
6: update v," Vi,T
6:
Ui,k+1 Ui Yn V Fi (U u,
7: update u, = UiT
7: update u, = 11,Ta
8: return (u, t, V
8: return 27")
Assumption (Smoothness). The function F; is continuously differentiable for each 1,..., and there exist constants Lu, Ly, Luu and Lvu such that for each i 1...., n. it holds that
Vu,F, 0z) is L, LupLishhhizz with respect and L -ipsccht with respect to vg (.,,(u, L, Lipschhtt with respect to v, and L.,. Lipschitz with respect to u.
Due to the definition of F(u, V) in 4 it is easy to verify that V ((u, V) has Lipschitz constant L. with respect to u, / /v with respect V, and Trwe/ n with respect any u;- We also define
X max Luv,, Lra)/ Lr,Le.
(5)
which measures the relative cross- sensitivity of F,F, with respect to e and V,F, with respect to u. Assumption 2 (Bounded Variance). The stochastic gradients in Algorithm 2 Znd Algorithm 3 are unbiased and have bounded variance. That is, for all u and Us.
E V ((,, = V.,Fi(u, vz).
(i,,F((,,e)) vi)] = V,FF(u,
Furthermore, there exist constants on and on such that
E V.lluu - V F. [x)]a o
E[I v.) F. ut)l S so.
We can view F.,,, us), when is randomly sampled from (I....,n). as a stochastic partial gradient of F(u, V) with respect to u. The following assumption imposes a variance bound.
Assumption (Partial Gradient Diversity). There exist 5 2 0 and p 2 0 such that for all u and V,
1 Σ V Fi v)) V Fu, 12 < 6 + p² V.,,,
With / = 0, this assumption is similar to constant variance bound on the stochastic gradient V, (u,, v:); with > 0, it allows the variance to grow with the norm of the full gradient.
Throughout this paper, we assume F is bounded below by F* and denote A Fu F ((0) Further, we use shorthand v(a) = (v (() and
s = V.F (w()) V( and
For smooth and nonconvex loss functions F, we obtain convergence in expectation to stationary point of F if the expected values of these two sequences converge to zero.
We first present main result for FedSim (Algorithm I with LocalSim), proved in Appendix A.2 Theorem (Convergence of FedSim). Suppose Assumptions and hold and the learning rates in FedSim are chosen as Tu = (I.TT) and 7o /1/ (1/(TT) with
1n/n n min [21++ 196(1 x1+x 111+
Then, ignoring absolute constants, we have
+
1 E La
m E[ RL.
<
AFe
+1(11 + x2) 7/
(oi+09(1- m)
+
mo?
mbo
nue
- - (1 + xx2)
2200
+
L
o2
(6)
5

---

Under review as a conference paper at ICLR 2022
Table 1: Convergence rates of FedSim and FedAlt in different regimes along with optimal learning rate via 7). We only show the dominant terms and hide the lower order terms in ''' for simplicity.
Condition
FedSim
FedAlt
General
T T-2/2
AF.
1
x"
+
mo?
AFo VT
od +
m)
of
m +
x²(n
m)
VT
vT
m/y
alw
mL.VT
LeVT
1
of
0
T
1
AFb
5'
1
AFD +
5 x.
7
72/3
73/3
o2 0
APo
A
0(1)
T
The left-hand side of (6) is the average over time of a weighted sum of E 2|A and E | The right-hand side contains three terms of order ((1/(177) ((/T)), O(1) and O respectively. We can minimize the right-hand side by optimizing over 17. By considering special cases such as oi = o = 0 and m = n, some terms on the right-hand side disappear and we can obtain improved rates. Table shows the results in several different regimes along with the optimal choices of fj.
Challenge in Analyzing FedAIt. We now turn to FedAlt. Note that the personal parameters updated only for the selected devices in ( in each round t. Specifically,
(8)
7th Ei F.
if St),
(
if 8(1)
Consequently, the vector (1+1) of personal parameters depends on the random variable 5(1). This makes it challenging to analyze the u-update steps of FedAlt because they are performed after ((++1) is generated (as opposed to simultaneously in FedSim). When we take expectations with respect to the sampling of (()) in analyzing the u-updates, ((++) becomes dependent random variable, which prevents standard proof techniques from going through (see details in Appendix A..).
We develop novel technique called virtual full participation to overcome this challenge. Specifically, we define a virtual vector ((++),, which is the result if every device were perform local u-updates. It is independent of the sampling of ((t) and can derive convergence rate for related quantities. We carefully translate this rate from the virtual ((++1) the actual (()) get the following result.
Theorem (Convergence of FedAlt). Suppose Assumptions 7 in FedAlt are chosen as Tu /1/ Tu) and 7ou (LoT)) with
2
and 3
hold and the learning rates
1 5 min 24(1+)2)) (28x=(n-m) V
Xan
Then, ignoring absolute constants, we have
+T=
E + mL. E Luu
S
AFo
+
c2+ m)
m+ x (n-m) ai
7/f'
mbo
Lx n
22 1/2 Lu
02m
22
(1
Tu
+
Zen
Lo
The proof of Theorem 2 is given in Appendix A.3 Similar to the results for FedSim, we can choose to minimize the above upper bound to obtain the best convergence rate, as summarized in Table
Comparing FedSim and FedAlt. Table | shows that both FedSim and FedAlt exhibit the standard O(1/VT) rate in the general case. Comparing the constants their rates, we identify two regimes in terms of problem parameters. The regime where FedAlt dominates FedSim is characterized by
2 (1
2m 12 < (1-
A practically relevant scenario where this is true is of x O and of from using large or full batch on a small number of samples per device. Here, the rate of FedAlt is better than FedSim by a factor of (1 + x²), indicating that the rate of FedAlt is less affected by the coupling between the personal and shared parameters. Our experiments in $ corroborate the practical relevance of this regime.
The rates from Table T also apply for full personalization schemes without convergence guarantees in the nonconvex case (Agarwal et al. 2020 Mansour et al. 20200] Li et al. 2021). Our rates are better than those of (Dinh al. 2020) for their pFedMe objective.

---

Under review as a conference paper at ICLR 2022
Table Summary of datasets and models. A histogram of data per device given in Figure (Appendix
Task
Dataset
#Classes
Model
# Model Params
#Devices
#Data device
Mcan
Max
Next-word prediction
StackOverflow
10000
1-layer transformer
6.M
1000
1964
15520
Landmark recognition
GL.Dv2
ResNet-18 ResNet-18
12M
Charucter recogmition
EMNIST
11M
1114
298
418
4 EXPERIMENTS
In this section, experimentally compare different model personalization schemes using FedAlt and FedSim as well as no model personalization. Details about the experiments, hyperparameters and additional results are provided in the appendices. The code to reproduce the experimental results will be publicly released.
Datasets, Tasks and Models. We consider three learning tasks; they are summarized in Table 2
(a) Next-Word Prediction: We use the StackOverflow dataset, where each device corresponds to the questions and answers of one user on stackoverflow com This task representative of mobile keyboard predictions. We use a 4-layer transformer model (Vaswani et al. 2017))
(b) Visual Landmark Recognition: We use the GLDv2 dataset (Weyand et al. 2020 Hsu al. 2020), lcage-scale dataset with real images of global landmarks. Each device corresponds to Wikipedia contributor who uploaded images. This task resembles scenario where smartphone users capture images of landmarks while traveling. We use ResNet-18 (He et al. 2016 model with group norm instead of batch norm (Hsieh et al. 2020) and images are reshaped to 224 224.
(c) Character Recognition: We use the EMNIST dataset (Cohen et al. 2017), where the input is 28 x 28 grayscale image of handwritten character and the output its label (0-9, a-z, A-Z)). Each device corresponds to a writer of the character. We use ResNet-18 18 model, with input and output layers modified to accommodate the smaller image size and number of classes.
All models are trained with the cross entropy loss and evaluated with top-1 accuracy of classification.
Model Partitioning for Partial Personalization. We consider three partitioning schemes.
(a) Input layer personalization: This architecture learns a personalized representation per-device by personalizing the input layer, while the rest of the model is shared (Figure Tap. For the transformer, we use the first transformer layer in place of the embedding layer.
(b) Output layer personalization: This architecture learns a shared representation but personalizes the prediction layer (Figure 1b For the transformer model, we use the last transformer layer instead of the output layer.
(c) Adapter personalization: In this architecture, each device adds lightweight personalized adapter modules between specific layers of a shared model (Figure 2a). We use the transformer adapters of Houlssy et al. (2019) and for ResNet-18, 18, the residual adapters of Rebuffi et al. (2017).
Algorithms and Experimental Pipeline. For full model personalization, consider three baselines: (i) Finetune, where each device finetunes (using SGD locally) its personal full model starting from learned common model, (ii) Ditto (Li et al. 2021), which is finetuning with & regularization, and, (iii) PFedMe (Dinh et al., 2020) which minimizes the objective (2). All methods, including FedSim, FedAlt and the baselines are initialized with global model trained with FedAvg.
4.1 EXPERIMENTAL RESULTS
Partial personalization nearly matches full personalization and can sometimes outperform
Table 3 shows the average test accuracy across all devices of different FL algorithms. We see that on the StackOverflow dataset, output layer personalization (25.05%) makes up nearly 90% of the gap between the non-person non-personalized baseline (23.8%%) and full personalization (25.2 (5.21%) On EMNIST, adapter personalization exactly matches full personalization. Most surprisingly, on GLDv2, adapter personalization outperforms full personalization by 3.5pp (percentage points).
7

---

Under review as a conference paper at ICLR 2022
Table Comparison of partial model personalization with full model personalization terms of average test accuracy % across all devices. The subscript denotes the standard deviation over runs with different random seeds. The boldfaced/highhightee numbers denote the accuracies within one standard deviation of the maximum in each row. For partial personalization, show the accuracy of FedAlt; see Table A for FedSim.
Non-pers.
Full Model Personalization
Partial Model Personalization
FedAvg
Finetune
Ditto
PFedMe
Input Layer
Output Layer
Adapter
StackOverflow
23.82
25. 20
25.20 25.20
25.2 25.21g.or
24.44p.0
25.05g.ur 01
24. 24.82
GLDv2
51.43
02.88
62. 05
62.92,.u
53.94o.m 07
56 56.64p.0.
66.41u 06.01
EMNIST
93.18
p4.13p.01
94. o4.13
94.13p.o1
93.62p.o4 04
93.57p.o5
p3..3
StackOverflow
EMNIST
A Accuracy
Accuracy
Adazir
Accuracy
0.2
1.0
Acuptar
to.
0.
Partial
0.5
Leyar
Oetput
Full
10
# Personalized Params.
Personallzed Params.
Personal Ized Params,
Figure 3: Absolute change in accuracy (percentage points) due personalization plotted against number of personal parameters (i.e., dimensionality v.). Note that the z-axis is in log scale.
This success of adapter personalization can be explained partly by the nature of GLDv2. On average, the training data on each device contains 25 classes out of possible 2028 while the testing data contains 10 classes not seen in its own training data. These unseen classes account for nearly 23% of all testing data. Personalizing the full model is susceptible to "forgetting" the original task (Kirkpatrick et al. 2017), making harder to get these unseen classes right. Such catastrophic forgetting is worse when finetuning very small local dataset, as we often have in FL, On the other hand, personalizing the adapters does not suffer as much from this issue (Rebuffi et al. 2017).
Partial personalization only requires a fraction of the parameters to be personalized. Figure shows that the number of personalized parameters required to compete with full model personalization is rather small. On StackOverflow, personalizing 1.2% of the parameters with adapters captures 72% of the accuracy boost from personalizing all 5.7M parameters; this can be improved nearly 90% by personalizing 14% of the parameters (output layer). Likewise, we match full personalization on EMNIST and exceed it on GLDv2 with adapters, personalizing 11.5-12.1 5-12.5% of parameters.
The best personalized architecture is model and task dependent. Table that personalizing the final transformer layer (denoted as "Output Layer") achieves the best performance for StackOver- flow, while the residual adapter achieves the best performance for GLDv2 and EMNIST, This shows that the approach of personalizing fixed model part, in several past works, is suboptimal. Our framework allows for the use of domain knowledge to determine customized personalization.
Finetuning is competitive with other full personalization methods. Full finetuning matches the performance of pFedMe and Ditto on StackOverflow and EMNIST, On GLDv2, however, pFeddM outperforms finetuning by 0.07pp, but is still 3.5pp worse than adapter personalization.
FedAlt outperforms FedSim for partial personalization. If the optimization problem convex, we would expect similar performance from FedAlt and FedSim. However, with nonconvex optimization problems such as the ones considered here, the choice of the optimization algorithm often affects the quality of the solution found. We see from Table 4 that FedAlt almost always better than FedSim by small margin, e.g., 0.08pp for StackOverflo//Adaape FedSim in turn yields a higher accuracy than simply finetuning the personalized part of the model, by large margin, e.g., 0. 0.pppp for StackOverflow/OOrppt Layer and 2.55pp for GL GI.Dy2/Adapter
3 were
and 0.3pp for GLDv2/Input Layer.
8

---

Under review as a conference paper at ICLR 2022
Table Comparing FedAlt and FedSim partial model personalization. "FT (part.)" means finetuning personal parameters v, while fixing the shared parameters u from non personnand training. The numbers are averaged over 5 runs with different random seeds. The standard deviations are given in Table T (Appendix C
StackOverflow
GLDv2
EMNIST
FT (part.)
FedAlt
FedSim
(part.)
FedAlt
FedSim
FT (part.)
FedAlt
FedSim
Input Layer
24.96
24.44
24.81
61.97
53.94
53.64
93.29
93.62
93.55
Output Layer
24.93
25.05
25.02
63.21
56.64
56.24
93.37
93.57
93 55
Adapter 24.71
24.82
24.74
63.86
66.41
66.35
93.66
94.13
94.07
4.2 EFFECTS OF PERSONALIZATION ON PER-DEVICE GENERALIZATION
Personalization hurts the test accuracy some devices. Figure 4 shows the change in training and test accuracy of each device, compared with a non non-personee personallzee model trained by FedAvg. We see that personalization leads to improvement in training accuracy across all devices, but reduction in test accuracy on some the devices over the non-persee non-personalized baseline. In particular, devices whose testing performance hurt by personalization are mostly the left side of the plot, meaning that they have relatively small number of training samples. On the other hand, many devices with the most improved accuracy also appear on the left side, signaling the benefit of personalization. Therefore, there is large variation of results for devices with few samples.
Additional experiments (see Appendix C show that using l regularization, as in (2). or weight decay does not mitigate this issue. In particular, increasing regularization strength (less personalization) can reduce the spread of per-device accuracy, but only leads a worse average accuracy that close to using a common model. Other simple strategies such as dropout also do not fix this issue.
An ideal personalized method would boost performance on most of the devices without causing reduction in (test) accuracy on any device. Realizing this goal calls for sound statistical analysis for personalized FL. and may require sophisticated methods for local performance diagnosis and more structured regularization. These are very promising directions for future research.
Full (train)
Full (test)
Partial (train)
Partial (test)
mulos
Acccurayyy
2000 1200 Intes 1250e 1506a # Data per device
pal 700o J0dee Laiee 35000
Z:30 Sale 20000 12500 12020
2510 5060 7506 225:0 19a02
# Data per device
Data per device
* Data per device
Figure StackOverflow task: Scatter plot change in training and test accuracy (percentage points) per-device versus the number of training samples on the device for (a) Left: full personalization with finetuning, and, (b) Right: partial personalization with the output layer.
5 DISCUSSION
In addition to a much smaller memory footprint than full model personalization and being less susceptible catastrophic forgetting, partial model personalization has other advantages. For example, it reduces the amount communication between the server and the devices because only the shared parameters are transmitted. While the communication saving may not be significant (especially when the personal parameters are only a small fraction of the full model), communicating only the shared parameters may have significant implications for privacy. Intuitively, can be harder to infer private information from partial model information. This is especially the case if the more sensitive features of the data are processed through personal components of the model that are kept local the devices. For example, parameters in order to satisfy differential privacy requirements (Abadi et al. 2016). This a very promising direction for future research.
speculate that less noise needs to be added to the communicated

---

Under review as a conference paper at ICLR 2022
REPRODUCIBILITY STATEMENT
For theoretical results, state and discuss the assumptions in Appendix The full proofs of all theoretical statements are also given there.
For our numerical results, we take multiple steps for reproducibility. First. we run each numerical experiment for five random seeds, and report both the mean and standard deviation over these runs. Second, we only use publicly available datasets and report the preprocessing at length in Appendix B Third, we give the full list of hyperparameters used in our experiments in Table in Appendix B Finally, we will publicly release the code to reproduce the our experimental results.
ETHICS STATEMENT
The proposed framework for partial model personalization is immediately applicable for a range of practical federated learning applications in edge devices such text prediction and speech recognition. One of key considerations federated learning is privacy. Partial model personalization maintains the all the privacy benefits of current non-personalized federated learning systems. Indeed, approach is compatible with techniques to enhance privacy such as differential privacy and secure aggregation. We also speculate that partial personalization has the potential for further reducing the privacy footprint - an investigation of this subject is beyond the scope of this work and is an interesting direction for future work.
On the flip side, we also observed in experiments that personalization (both full or partial) leads to reduction in test performance on some of the devices. This has important implications for fairness, and calls for further research into the statistical aspects of personalization, performance diagnostics as well as more nuanced definitions of fairness in federated learning.
REFERENCES
Martín Abadi, Andy Chu, Ian J. Goodfellow, H. Brendan McMahan, Ilya Mironov, Kunal Talwar, and Li Zhang. Deep Learning with Differential Privacy. In Proc. of ACM SIGSAC, pp. 308-318. ACM, 2016.
Alekh Agarwal, John Langford, and Chen- Yu Wei. Federated Residual Learning. arXiv Preprint, 2020.
Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary. Feder- ated Learning with Personalization Layers. arXiv Preprint, 2019.
Duc Bui, Kshitiz Malik, Jack Goetz, Honglei Liu, Seungwhan Moon, Anuj Kumar, and Kang G. Shin. Federated User Representation Learning. arXiv Preprint, 2019.
Gregory Cohen, Saced Afshar, Jonathan Tapson, and André van Schaik. EMNIST: an extension of MNIST to handwritten letters. arXiv Preprint, 2017.
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting Shared Repre- sentations for Personalized Federated Learning. Proc. of ICML, volume 139, pp. 2089-2099, 2021.
James W. Demmel. Applied Numerical Linear Algebra. SIAM, Philadelphia, 1997.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Fei-Fei Li. ImageNet: A large-scale hierarchical image database. In Proc. of CVPR, pp. 248-255, 2009.
Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Adaptive Personalized Federated Learning, arXiv Preprint, 2020.
Canh T. Dinh, Nguyen Tran, and Josh Nguyen. Personalized Federated Learning with Moreau Envelopes. In Proc. of NeurlPS, volume 33, pp. 21394-21405, 2020.
Alireza Fallah, Aryan Mokhtari, and Asuman E. Ozdaglar. Personalized Federated Learning with Theoretical Guarantees: A Model-Agnostic Meta-Learning Approach. In Proc. of NeurlPS, 2020.
10

---

Under review as a conference paper at ICLR 2022
Filip Hanzely, Boxin Zhao, and Mladen Kolar. Personalized Federated Learning: A Unified Frame- work and Universal Optimization Techniques. arXiv Preprint, 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep Residual Learning for Image Recognition. In Proc. of CVPR, pp. 770-778, 2016.
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin de Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-Efficient Transfer Learning for NLP. In Proc. of ICML,, volume 97, pp. 2090-279999999 2019.
Kevin Hsieh, Amar Phanishayee, Onur Mutlu, and Phillip B. Gibbons. The Non-IID Data Quagmire of Decentralized Machine Learning. In Proc. of ICML, volume 119, pp. 4387-4398. PMLR, 2020.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Federated Visual Classification with Real-World Data Distribution. In Proc. of ECCV, volume 12355, pp. 76 (6-92,224 2020.
Peter Kairouz, H. Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista A. Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, Rafael L. Hubert Eichner, Salim El Rouayheb, David Evans, Josh Gardner, Zachary Garrett, Adrià Gascón, Badih Ghazi, Phillip B. Gibbons, Marco Gruteser, Zaïd Harchaoui, Chaoyang He, Lie He. Zhouyuan Huo, Ben Hutchinson, Justin Hsu, Martin Jaggi. Tara Javidi, Gauri Joshi, Mikhail Khodak, Jakub Konečny, Aleksandra Korolova, Farinaz Koushanfar, Sanmi Koyejo, Tanerède Lepoint, Yang Liu, Prateek Mittal, Mehryar Mohri, Richard Nock, Ayfer Özgüir, Rasmus Pagh, Hang Qi. Daniel Ramage, Ramesh Raskar, Mariana Raykova, Dawn Song. Weikang Song, Sehastian Stich, Ziteng Sun, Ananda Theertha Suresh, Florian Tramèr, Praneeth Vepakomma, Jianyu Wang. Li Xiong, Zheng Xu, Qiang Yang, Felix X. Yu, Han Yu, and Sen Zhao. Advances and Open Problems in Federated Learning. Found. Trends Mach. Learn., 14(1-2):1-210, 2021.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sehastian Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic controlled averaging for federated learning. In Proc. of ICML, 2020.
James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka BarbskkaBarrinska, Demis Hassabis, Claudia Clopath, Dharshan Kumaran, and Raia Hadsell. Overcoming catastrophic forgetting in neural networks. Proceedings of the National Academy of Sciences, 114(13):3521-332 2017.
Anastasia Koloskova, Nicolas Loizou, Sadra Boreiri, Martin Jaggi, and Sebastian Stich. A Unified Theory of Decentralized SGD with Changing Topology and Local Updates. In Proc. of ICML, 2020.
Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and Robust Federated Learning Through Personalization. In Proc. of ICML, volume 139, pp. 6357-6368, 2021.
Paul Pu Liang, Terrance Liu, Ziyin Liu, Ruslan Salakhutdinov, and Louis-Philippe Morency. Think Locally, Act Globally: Federated Learning with Local and Global Representations. In NeurIPS Workshop on Federated Learning, 2019.
Yishay Mansour, Mehryar Mohri, Jac Ro, and Ananda Theertha Suresh. Three Approaches for Personalization with Applications Federated Learning. arXiv Preprint, 2020.
Michael McCloskey and Neal J. Cohen. Catastrophic Interference in Connectionist Networks: The Sequential Learning Problem. volume 24 of Psychology of Learning and Motivation, pp. 109-165. Academic Press, 1989.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera Arcas. Communication-Efiee Effiicient Learning of Deep Networks from Decentralized Data. In Proc. of AISTATS, pp. 1273-1282, 2017.
Tomas Mikolov and Geoffrey Zweig. Context dependent recurrent neural network language model. In JEEE SLT, pp. 234-239, 2012.
Sylvestre-Alvise Rebuffi, Hakan Bilen, and Andrea Vedaldi. Learning multiple visual domains with residual adapters. In Proc. of NeurlPS, pp. 506-516, 2017.
11

---

Under review as a conference paper at ICLR 2022
Sashank J. Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Hugh Brendan McMahan. Adaptive Federated Optimization. In Proc. of ICLR, 2021.
Karan Singhal, Hakim Sidahmed, Zachary Garrett, Shanshan Wu, Keith Rush, and Sushant Prakash. Federated reconstruction: Partially local federated learning. In Proc. of NeurIPS, 2021.
Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet S Talwalkar. Federated Multi-Tasse Learning. Proc. of NeurIPS, pp. 4424 4434, 2017.
TensorFlow Federated. Ittps::/ /twps::/www... /www. www.oe/fow...rgeeeeeeeeee..ee.e.
Iulia Ture, Ming- Chang, Kenton Lee, and Kristina Toutanova. Well-read students learn better: On the importance of pre-training compact models. arXiv Preprint, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is All you Need. In Proc. of NeurIPS, pp. 5998 6008, 2017.
Tobias Weyand, André Araujo, Bingyi Cao, and Jack Sim. Google Landmarks Dataset v2 A Large-Scale Benchmark for Instance-Level Recognition and Retrieval. In Proc. of CVPR, pp. 2572-2581, 2020.
12

---

Under review as a conference paper at ICLR 2022
Appendix
Table of Contents
Convergence Analysis: Full Proofs
14
A.1 Review of Setup and Assumptions
14
A.2 Convergence Analysis of FedSim
15
A.3 Convergence Analysis of FedAli
22
Technical Lemmas
B
Experiments: Detailed Setup and Hyperparameters
31
B.
Datasets, Tasks and Models
................................
32
B.2
Experimental Pipeline line and Baselines
................................
34
B.3
Hyperparameters and Evaluation Details
35
Experiments: Additional Results
36
Ablation: Final Finetuning FedAlt and FedSim
36
Effect of Personalization on Per Device Gencralization
................................
37
Partial Personalization for Stateless Devices
38
13

---

Under review as a conference paper at ICLR 2022
A
CONVERGENCE ANALYSIS: FULL PROOFS
We give the full convergence proofs here. The outline of this section is:
A.I
Review of setup and assumptions;
A.2
Convergence analysis of FedSim and the full proof of Theorem
A.3
Convergence analysis of FedAlt and the full proof of Theorem 2
A.4
Technical lemmas used in the analysis.
A.1
REVIEW OF SETUP AND ASSUMPTIONS
We consider federated learning system with n devices. Let the loss function on device be (u, vi), where u E Rd denotes the shared parameters across all devices and 8; E Rdd denotes the personal parameters device i. We aim minimize the function
1 Σ Fr(u,;;)
F(u, V) =-
where V (vi,... Un) of (3) with the equal per-device weights, i.e., a; 1/n. Recall that we assume that F is bounded from below by F*.
is a concatenation of all the personalized parameters. This is a special case
For convenience, we reiterate Assumptions 1 and 3 from the main-pan paper Assumptions
and below respectively, with some additional comments and discussion.
Assumption I (Smoothness). For each device i 1...., the objective F, is smooth, i.e., it is continuously differentiable and.
(a) Fu,F v:) is Lip-Lipschitz.fo for all uo. (b) U; + V.F.(u. Us) is L,pLipschitz for all u,
(c) U; + F. v.) is Lie-Liish for all u, and,
(d) - V.F.uu, vz) is Amp-hiipchitz for all Us.
Further, we assume for some X 0 that
max (uuv Lon S X L.l...
The smoothness assumption is standard one. We can assume without loss of generality that the cross-Lipschitz coefficients Luws Lew are equal. Indeed, if F, twice continuously differentiable, we show that Luv, Lou are both equal to the operator norm Fu,, ea)llop of the mixed second derivative matrix. Further, X denotes the extent to which u impacts the gradient of o; and vice-versa.
Our next assumption is about the variance of the stochastic gradients, and is standard in literature. Compared to the main paper, we adopt a more precise notation about stochastic gradients.
Assumption 2 (Bounded Variance). Let DD denote a probability distribution over the data space Z on device i. There exist functions Gi and Gi,y which are unbiased estimates of V F, and V F respectively. That is, for all u, Vg:
E [Gr,auu, v, 2]] Fu,, v4), and EEND, [Gi, v, 2]] (u,,(u,
Furthermore, the variance of these estimators is at most o and oi respectively. That is.
EgwD, (G..uuu, 2, - F, us 11)1 < oi E Gi, v, - V.F(u,, uz)l < o².
In practice, one usually has G (u, Ui V.,i((u, vz), 2), which is the gradient of the loss on datapoint z ~ D; under the model (u, Uz), and similarly for Gi
Finally, make a gradient diversity assumption.
Assumption 3 (Partial Gradient Diversity). There exist 8 2 0 and 2 0 such that for all u and V,
IV. v.F(u, < 88 + (u,
(8)
72
14

---

Under review as a conference paper at ICLR 2022
Algorithm 4 FedSim: Simultaneous update of shared and personal parameters
Input: Initial iterates (0)), v Number of communication rounds T, Number of devices per round m, Number of local updates т, Local step sizes Yux Tv-
do 2: Sample devices from [n] without replacement in S(1)
1: for t 0, 1,
3: for each selected device i S in parallel do
Initialize 0000 = vi (t) and u = u(t)
4:
5: for 0, .... 7
Update all parameters jointly
6:
Sample data (t)
i, ~ D;
(1)
No Gr (t) (1) (1)
7:
i,++1 i,k
i,k i,k
8:
11,+1 Hikh Ya (ui,k. ki 2 (i
(1) (t)
u
(8) (1) (0)
9:
Update vi
((+1) = =// and u;
((+1) Uir
(0)
Update (141) Licoty a.f((ti)
10:
at the server with secure aggregation
Licort) a
return (T)) (T) (T)
11:
un
This assumption is analogous to the bounded variance assumption (Assumption 2 but with the stochasticity coming from the sampling of devices. It characterizes how much local steps on one device help or hurt convergence globally. Similar gradient diversity assumptions are often used for analyzing non-personalized federated learning (Koloskova et al.] 2020, Karimireddy al. 2020). Finally, it suffices for the partial gradient diversity assumption to only hold at the iterates (u V() generated by either FedSim or FedAlt.
A.2 CONVERGENCE ANALYSIS OF FEDSIM
We give the full form of FedSim in Algorithm 4 for the general case of unequal 's but focus on
a; = 1 i/n for the analysis. In order to simplify presentation, we denote v(t) (v] (t) define the following shorthand for gradient terms
....,, u and
A (c) = V.F (
)( and A (1)
1
V.F 400) (00 )
1=1
For convenience, restate Theorem from the main paper.
Theorem (Convergence of FedSim). Suppose Assumptions
and in FedSim are chosen Yu = -/ and No -n/[[T,p with
hold and the learning rates
m/n S min x10+)) 196(1
Then, ignoring absolute constants, have
A
E
m nL.
E(A
<
o7 + ((1+x2))
(
(11t0(1-- m²)
mo? ntre
Tru
mlr
+ 11 + +x3)
(12 2
o2
(6)
Ly
Before proving the theorem, we give the following corollary with optimized learning rates. Corollary 3. Consider the setting of Theorem and let 0 be given. Suppose the learning rates 7/u = 1// LT)) and To = n/ (T Le). where (ignoring absolute constants).
1/2
E
17
62
L. (1+
82
o2 Io (1 7-1 x2)
71
m/n
1/2
p2) A
(1+
(1- p2) x2)
15

---

Under review as a conference paper at ICLR 2022
We have,
V.
(w.), v())
+ m T
V.
(u (
Se
T
La
after communication rounds, where, ignoring absolute constants,
AFo(1+
0² + +2 m
mo nho
T
62
mbu
Aov (1 1-1 + xx²
Ja 1o
+ Ov
c3/2
VTr
ving
AFo
1+ 1+/(1= g 1-
(1-
+
(1 + x (1 + p²)
Proof. The choice of the constant 7) ensures that each of the constant terms in the bound of Theorem is OEE) The final rate is now
plugging in the value of r7 completes the proof.
We now prove Theorem
The proof mainly applies the smoothness upper bound write out a descent condition with suitably small noise terms. We start with some notation.
Proof of Theorem 7
Notation. Let ((()) denote the o-algebra generated by (1(1), ((t) and denote E. E[-[F(()) For
((t) all devices, including those not selected in each round, we define virtual sequences 2 o the SGD updates in Algorithm for all devices regardless of whether they are selected. For the selected
(8) (1(0 ((t)) devices E s((), we have "2 2,k)
1,ke 2, Note now that the random variables "i i
(())
independent of the device selection S(D), The updates for the devices i S(O) are given by
(++1]
(t)
Vo
Σ
Gi
(4) 1,k 21,
k=
and the server update is given by
(++1)
(t) Yo
(00) (1(0 Gi
=
i,6
m
k=0
Proof Outline. We use the smoothness of Fs, more precisely Lemma 16 obtain ((1)) v(+t)) F F(ul), v()) V V.F(! v(0), (1+ (1)))
(1 22 )
(1)
(++1)
T2.0
Ti.v
(10)
L.(1 + x²) (C-1) (0)
L (1 + 2)
(+++) (1)
2
1=1
T2.5
T2.0
Our goal will be to bound each these terms to get a descent condition from each step of the form
E, ((1(+1)) v((+1)) v(1+1))) F(u(1) v(t)]
s 5- 7/u v F (1 w(o)||=
mm
Σ V
(1)))
+
8n2
16

---

Under review as a conference paper at ICLR 2022
where the 0 + 7() terms are controlled using the bounded variance and gradient diversity assumptions. Telescoping this descent condition gives the final bound.
Main Proof. Towards this end, we prove non-asymptotic bounds on each of the terms Ti Ti T2 and T2 in Claims to respectively. We then invoke them get the bound
Es ((( F(ul), -(ti]]
- TuT A
Torm
A
4n
La + x
128
Lo + X))rrroo
(1 - m/nl) +
2 0u +
2n
Ey
Mink
(1)
wcoll (Tri + m
=xx (2ITTtoo Toot
(11)
2
MΣΣ
E
(2)
Note that simplified some constants appearing the gradient norm terms using
Ta - (1
and S (6))
Our next step is to bound the last two lines of 11 with Lemma assumption (Assumption as
8 and invoke the gradient diversity
1
V.
+ ( (()))
This gives, after plugging in the learning rates and further simplifying the constants,
E, ((((+))) v(a+1)) ((ultt),vot1)) - FF (w v(o)]
S 8L.
ca
am (1)
+c
(1 (x))
0
mo?
652
(1
m71
8L
2L
nIv
L,m
248
40
4
d (1 x11xx)111 (1
+
Ly
In
In
Taking full expectation, telescoping the series over t = 0,... , and rearranging the resulting terms give the desired bound Theorem [
Claim 4 (Bounding TI, Let TI,0 be defined as in 10 We have,
[T, S
Totm
V F. Fi
(0)
2n
+ m
E
x2 L Le
(1)) ii
- (1))
+ 0 (())
2 (1)
19]
km 0
Proof. Define T1,0, to be contribution of the ith term to Ti For i & S., we have that Ti, = 0,
since 0ū uz On the other hand, for i E s(t) use the unbiasedness of the gradient estimator
((-1) (1)
17

---

Under review as a conference paper at ICLR 2022
Gi and the independence of Fizk from Wigk: Vi to get
(6) (8) (0)
Er [Ti = 7o Σ E
V Fi
(w/),, 1 VoFi (())
k=0
Ep V VuFi (2('), o.))) v.li (uh 511))
k=0 0
V Fi m m1 }1
7o7
%
Σ
E,
F. (u(),,(()) _r(")) v.F.
ui o1x)
voF. (ulo) .2(10)
k=0
7o?
V.
(u(t) 2641) ) ||
+ To Σ Es
vuFi (u.f..1.1) 0(2) v.Fi (u(), o/ )|² (12)
2
2
For the second term, we add and subtract V.Fi (u(1) o14 and use smoothness to get
V.F (1 o(2)) voFs (1))
(ul),, )/0)|²
(ul)),()))(" s Ln 1u we0l
+ 21 000 (
(13)
Since the right hand side of this bound is independent of St, we get,
E,
[T..]
m n
Eg
T.e.s
m n²
Σ Eg TT,
m
(ES())
and plugging in (12 and
13
completes the proof.
Claim 5 (Bounding T1,w). Consider T. defined in
10 We have the bound,
E, [T,.
TaT
(1)
)
2
Ya
Er
L Uik
(1)
1| 2 + x2 L Lo with v. 1
1=1
k=0
Proof. Due to the independence of (11) from (0) (1)
i,k. we have,
E, (t+1) - -x(()]
- E,
Σ
V
Bi
10 (t)
iE 516) k=
E
Va
Fs
Ui 2((1)
(it)
((((()) k==
E,
F.F:
(1))
18

---

Under review as a conference paper at ICLR 2022
where the last equality took an expectation over s(0, which is independent of (0) (0) the same sequence of arguments as Claim we have,
iiki Vi.k- Now, using
4
Es
(v.(((() v(1), ((++1) - u(()))
r-1
Σ E
1 v(t) V.F (()))
(1(0) Fv
010)
n
k=0
TuT
VuF (u()) x/()))
7a
(1) (00)
vaF (u(O) v(t))
E
V Fi
Ui 0(6)
2
2
K=
(4)
You
VuF (u()), r()) 1|²
E.
V F
(1)
v(())
vuFi (1)t) o1.1)) )|²
2
2n
YaT
2
VuF (0 v( )|²
+
7a
π
Σ
--1
Es
L2
Ugk
w00) 1|
2
+
L
uv
i,k - 0z 1991
i=1
km
where the inequality (+) follows from Jensen's 's inequality as
V F. (2(0 o22 vuF (ul), m(o))
1 n
VuFi
(0))
- -vaF (ufh
(0) b(t)) )|?
n
n2
Claim (Bounding T2, Consider T2. as defined in 110..
We have the bound,
E, [T,00] S 3Lv + 2 m
V. (0 2s
(1)
2
Lo + x² x2)22rmm0?
+
2n
2n
(1, + x² 22 mm
Es
12
2:1 1i - 1.001²
+ x2 Ton Inp
(0100
u
2n
Proof. We start with
Eg
((t)
-100
= EE
Σ Gi (2% 002 (00)
530r7 Σ
--1
Ee Gi.
(00))
|I ||
k=0
T-1 Σ
(1 x ||
5127207 + 10
Eg
V.F.
k=
< 2000203 0 + 3 22
Vofi (0) (.(() ||
2
|
+ 2 Lu Lu [1(01 - (0) 27
+ 3 2
Σ Ex
k=
Using (a) u. (t+1)
= =
(t) U,
for E S(t), and, (b) stt) is independent from u.. (1)) we get,
Lo + +m²)m
1 1o - - ,.00 ||?
Ep [2,u
E,
2n
m
< (1 + x2
ym
(1)
Eg
2n2
(() 1, - 600)
2
Plugging in the bound E,
completes the proof.
19

---

Under review as a conference paper at ICLR 2022
Claim 7 (Bounding T2,u). Consider T2,u as defined in (10) We have,
La (1 ?
[T S
(0 + 128² (1- m
m --) ))
2m
3 (1 + X 2) Ya 2 (1 + p²) VoFi (ul)) v(0) )(2
mL (1 + X Mar
ΣΣ Eu 1 - u w(0'| + x Lu ( 5.
T-1
|/]
2n
Proof. We proceed with the first two inequalities as in the proof of Claim 6 to get
E, (++1) -
i ((()) Fu
S
E
2, 1.k)
m
VES()
Ta.
For Ta.j (a) we add and subtract V. V.,((((,, v(9) and V, v,F((u!)), -.ft), (b) invoke the squared triangle inequality, and, (c) use smoothness to get
Tajj 66 E,
1
Σ
Vali
(((), (o()) Val Vaf(ulo, v10)
+6|
+ 6 val (wl), vit) |
m iES())
+ 3E,
1
Σ
(000 w001/ x I (000 -o.00l()))) [(01(() ?
m
SES())
For the first term, we use the fact that 5(1) is obtained by sampling without replacement to apply Lemma 17 together with the gradient diversity assumption to get
E
Σ
V.P.
au0t,,f0 V.F (n()) p(tt)
m
iES(1)
1
n m
VuFi (2)))
(2((),,(()) of(t) - vot (w)0, vl0)
m
i=1
1
n m
82 + 0² V.F (o )1
<
m
Therefore,
1282
m
+ 6(1 6 p²) V.F (w v(o) /
Ta.j
3
1 Likk - w00ll² -ucol
+ x2 Ta T.m 1::0 -(() 101
+
E,
n
where we also used the independence between g(t) and (ut2 (00)) Plugging this into the expression for E, []((((+1) 112 completes the proof.
Lemma 8. Let F; satisfy Assumptions and consider the iterates
U = ub Ta Gi Uk, 2k), and, Uk+1 = U2 To Gi Uk: 21)
for k = ..... , 1, where 2ạ ~ D.. Suppose the learning rates satisfy You = CT/(TLu) and
Th = c,,(TL,) with ca. S max(1x Further, define,
A = 7u 12, + Ix and, B 8% 12g + 27m, Toy Tep:
20

---

Under review as a conference paper at ICLR 2022
where f E (0, 1] is given. Then, we have the bound,
Σ O
E Allua +lllu+Blluk - - 4r² 11) +8,0
+
20)l 7,2 uo)ll3)
Proof. T 1, there is nothing to prove, so we assume 1. Let A. := Allux uo uol12 + Blu voll and denote by F& the sigma-algebra generated by (WR, uk). Further, let Exl:] = E[.]Fa]. We use the inequality 20 I a² /82 + 8² 8²,3² for reals a, B, to get,
1
Eg [ullux+1 - uoll
uu
<
+
| uoll The E Ex:ll Gi,u (uk, Uk, =x)1
1
14
|| uoll + 7 + T (u v2))
7
1
1
I|u
uol 2l 1ol22 + + 20 37 V Pi vo)l
+ +37 uu uoll 3 Luv Uk mull2
where the last inequality followed from the squared triangle inequality (from adding and subtracting V.F.(uo. 0x) and (uuo vn)) followed by smoothness. Together with the analogous inequality for the v-update, get,
1+ r ( Ae + A'Iluu -nol? + B'llee voll c
Ea(A++1] S
where we have
A' +2x and, B' B'= - -2xx LI, and, C' -2=
3 A (uo, vo) 1 3 uo)
Next, we apply Lemma 20 to get that A' < A/T and B' < B/T under the assumed conditions on the learning rates; this allows us to write the right hand side completely in terms of A& and unroll the recurrence. The intuition behind Lemma 20 is as follows. Ignoring the dependence for a moment, if Yhu and % are both (1)). then A', B' are both (m²²), while A and B are O(n). Thus, making 77 small enough should suffice to get A' < O(A) and B' < O(B).
T, Lus Le,
Concretely, Lemma 20 gives
(1+--
[
S
1+ T
+C, E[A,]+C,
and unrolling this recurrence gives for l: ---1
2 j
0 7-1 1 (
2 [ c
E[AN]
<
1+
1 --1
T-1
2
j=0
1
(
1+
-
2
r-1
0
2 (T
11),
2
2
where used (1 + 1/a)a < for all a 0. Summing over k and using the numerical bound
e² completes the proof.
Remark 9. We only invoked the partial gradient diversity assumption (Assumption 3 iterates (u(*), V(t)); therefore, it suffices if the assumption only holds at iterates (ult), v(()) generated by FedSim, rather than at all (u, V).
21

---

(1) ((+1) (0) Under review as a conference paper at ICLR 2022
Algorithm 5 FedAlt: Alternating updates of shared and personalized parameters
Input: Initial iterates (0)), v(0), Number of communication rounds T, Number of devices per round Number of local updates Tux Tux Local step sizes Yuu Yus
1: for t 0, 1, ...
do
2: Sample devices from [n] without replacement in S(1)
3: for each selected device S0)) in parallel do
4: Initialize 0000 = vi
5: for 0, .. Tv do
Update personalized parameters
6:
Sample data (t)
i ~ D;
1,++1 1++ i,k Gr, (0) (0)
(1)
7:
(t)
8:
Update v/
k.
Initialize (1) = (t)
9:
41,0
10:
for k = 0,... Tu - do
Update shared parameters
(0) Uik %a Ga Uike 0i =iik) (t)
(1)
U1,4+1 ((+1) Ui
12:
Update u,
Update u (1+1) Σ Eneoet) ((41)
at the server with secure aggregation
return (2)) (T) , Un
A.3 CONVERGENCE ANALYSIS OF FEDALT
We give the full form of FedAlt in Algorithms for the general of unequal 's but focus ag = 1/n n for the analysis. For convenience, we reiterate Theorem below. Recall the definitions
(0)
VuF (ut) vit and, A
F. (u"),,()''"
Theorem (Convergence of FedAlt). Suppose Assumptions in FedAlt are chosen as Tu /1/ (InuTu) and 7o n/ To) with
and 3 hold and the learning rates
n S min 24(1+p"1 (n-
x
Then, ignoring absolute constants, we have
T-1 +
La LEAA
+ -T, E[A
AFo 02+ (1- 1
o² m+ 2 (n-mm)
T/ +1 mbn..
(
1= (1
Ty 1) o.m (1 1o
X202
L
Before proving the theorem, we have the corollary with optimized learning rates.
Corollary 10. Consider the setting of Theorem 2 and fix some > 0. Suppose we set Thu (T()) and 7v =1/(TLn such that, ignoring absolute constants,
0. m EL, (= +x2(1 (mmmm/mm)))))))))))))) A
O + (1mm1mm)
0 + 82
-1/2
+x((1-m/n))
(1 - T
mbnE
LiGE
3.m
-1/2
N
Line (1 -oo1))
m
m
1+
2 m)
x
Then, we have,
b(())) (1) SE
E
V F
m
E
V
P.
(1)
T t=
La
Lnn
==1
22

---

Under review as a conference paper at ICLR 2022
after T communication rounds, where, ignoring absolute constants,
A
au +
(1
m
0
m
m
T <
X
1
g²
mLn
Lu
n1
71
Fo
Ou + V1
-1
+
Jv
V1
To
3/2
VLu
VLu
+
A Aoo
1 0²
x²
n2
+
Proof. We get the bound by balancing terms from the bound of Theorem The choice of ensures that all the (π)) and (1)²) terms are most O(e). Finally, the smallest number communication rounds to make the left hand side of the bound of Theorem 2 smaller than is A (ne).
We are now ready prove Theorem 2
Proof of Theorem The proof mainly applies the smoothness upper bound to write out a descent condition with suitably small noise terms. We start with some notation.
We introduce the notation as the analogue of A (t) with the virtual variable (1++1);
A 1
V (0) v(
Notation. Let FF(t) denote the o-algebra generated by (z(1), v(t) and denote E,[-] E[-]F(()). For
all devices, including those not selected in each round, we define virtual sequences úi (() SGD updates in Algorithm for all devices regardless of whether they are selected. For the selected
V, as the
devices E 5((), we have (t) ({t) ú,A fi and Wi,k ii Note now that the random variables =(t) =(0) " Vik are independent of the device selection 5(t), Finally, we have that the updates for the selected devices i ((!) are given by
(t) ((t)
(t+1)
Σ
(.()) (1) Gi (2(),((() 2i,k) 21
%
k=
and the server update is given by
u (t+1)
(1)
7a
Gi Wi 2100)
iES())
k==
Proof Outline and the Challenge of Dependent Random Variables. We start with
F (1))))) vioti) F v()) =F (u'), (u(",, v(ott)) F (u(), v(t) + F (+())), v((+1)) - (u(t), v((t)).
(14)
The first line corresponds to the effect of the e-step and the second line to the u-step. The former is easy to handle with standard techniques that rely on the smoothness of F (u(t),.), The latter is more challenging. particular, the smoothness bound for the u-step gives us
( (++1)) v(tt))) F (u()), v(+tt)
s (V (w!), vitti) ; (+11)
w(1) _u(tt)) | -
Lu ((+1))
The standard proofs of convergence of stochastic gradient methods rely on the fact that we can take expectation w.r.t. the sampling (()) of devices for the first order term. However, both ((++1) and
23

---

Under review as a conference paper at ICLR 2022
((++1) depend the sampling ( of devices. Therefore, cannot directly take an expectation with respect to the sampling of devices in S!O.
Virtual Full Participation to Circumvent Dependent Random Variables, The crux of the proof
lies in replacing V(++1) in the analysis of the u-step with the virtual iterate ((++) the dependence of the u-step ( to the u(++)) term. This allows to take an expectation;
move all
remains to carefully bound the resulting error terms.
Finally, we will arrive at bound of the form
T
TuTu B[A[]+ 8
ToToT 16n
E[A!!)
Ao T +70)+
Next, we translate this bound from gradient E gradient computed at the actual iterate V A careful analysis shows that we only incur lower
(1) of the virtual ((++1) E[A[), which the
order term of O(7in)) in this translation. Choosing 7a and 7n small enough will give us the final
result.
Analysis of the u-Step with Virtual Full Participation. We introduce the virtual iterates (++1) into the analysis of the u-step follows:
F (((+)) v(tt)) - F (1)) v(tt)
< VuF (100,(at)) ((tt)) (1(2))
- (0))
+ Lu
u ((+)) (0
2
(vrF 1
1 ((ott) (
) -
I.u
(c+1) - (00 1
+ 2
11
(00 (iits) v.r (n())) (((+1)) (2+1) m(())
A (u11),
(1(11) v(ct))
- (10)) + Lu
| - -2(0)/ 1/
2T. Vol (n!), v(tt) V. F (u(t), F((ttt) )1²
< VuF (1)),, )(c+1) (1+1) - ult)) + Ly [(++)) (0) w
2
2 L
(t+1)
(++1
2n
T2.u
Ta.
The last two inequalities follow from Young's inequality and Lipschitzness of V v.F(u, V) respectively.
We have now successfully eliminated the dependence the first-order term Ti,, on V(++), The virtual iterates ((++1) are now independent of (((),, This allows us to take an expectation w.r.t. the sampling 5(t) of the devices.
We bound each of these terms in Claims 11 to 13 below to get
E,
F (e+1),
v((+1)) F (u()) v((t))
Ju
E, [A
2% L2
n Tu
E
11
(0)
w01l
m/nn)
4
n
i=1 k=
T
T 22 m
o 352 (1- n
)) +
Note that we used the fact that 24L (1, (1+ p²) < to simply the coefficients of some of the terms above. The second term has also been referred to client drift in the literature; we bound it with
24

---

Under review as a conference paper at ICLR 2022
Lemma 18
and invoke the assumption on gradient diversity (Assumption 3 get
16 I (Tu 1)
Ex Vuli (u()) )(( L 1)0
T, S
16 L Tu Tu 1)
82 + p2 v.F (u') ((ot)) + 8 - 1)0
n
Plugging this back in, we get,
F [1(++) v((tt) - F (u()), (ul)),Vttt))
(ultt)))((tt))
E,
S Ju E, [00 + Tu (o + +22 (1- + 4.0 (1/mn))
8
m
+ L 1, x2 m/n) 1 - -
Note that we used 128 (T. 1)1² V 1, which is implied by 24L. The (1 + p²) S.
Bound with the Virual Iterates. We plug this analysis of the u-step and Claim 14 for the u-step into 14 next. We also simplify some coefficients using LoX 1) S 1. This gives us
F (1)+1) v(ott))
- (u(),, (l))
Eg
S
Yota Ee( (0 Tomm
E, E [ + 47,
m n
+ x² - (m/m)
16n
1 Lu (o + 28 (1 - mm/n)) + m
- () + +8 + 472
1)0,m
n
Taking unconditional expectation, summming it over ( = 0 to T 1 and rearranging this gives
Tu B + TeTm EIA (0)
(15)
8
16n
1=0
A Fo 2o700
(mh + x1(11-m/
(011mm/ -m/mm))) + Th (o 25 m/n))
+
T
m
1) + 282) + 4 2 To Tv 1)0,m
n
This is bound in terms of the virtual iterates ((+1),, However, we wish to show bound in terms of the actual iterate ((t).
Obtaining the Final Bound. It remains now to relate A with ( Using the inequality and smoothness, we have,
E, V. F (n()), (1(0)/1))) v())
F (10)) (0, ||
1
E, V.F, (1(0),.))) (0) V.Fi (n()) (u(0,0(tt))
) ||
n
x" Lu
(1+1) 0/10//
11 Eg
x" 2 Ln Lo
16
VeFi (u(),u(0) )/0))( + +8r21o0?
2 2 To ex L, + 16:3-0x x L.
where the last inequality followed from Lemma 19 Using
V. (1 )o)) 2 v.F (w(0) v(0) - F (1))) p(oti) || +2| V.F (n()), (.9,,o+t) |
25

---

Under review as a conference paper at ICLR 2022
get,
E S 2
32 2 E ) L
Therefore, we get,
Yu E A 0 + To
E
16
32n
Tu (
ToTum
1
x2n2x²m
E[A!! o Lubio
8
16n
2
T
Tu
E
(0)
(9)
ToTum
+ Tor oo Tru T
8
16n
where we used 324 x'
S 1/2, which is one of the conditions we assume on f/.
Summing this up and plugging in 15
gives
Tu
+ To E[A!)])
16
32n
T-1
Tu
E
(0)
Yo
()) + 7utunon 2.2 22200 Lnubo
8
16n
0=
<
AFo
4 2 Lu
m
x2 mm/mm/)))
+ PubuT (o 20 - m//n)
T
1) + +25 47 2 (Tu 1)
8 3
Tu - 2 In
Plugging in 7h = / (Lu and To / Tx) completes the proof.
The analysis of each of the terms in the u-step is given in the following claims.
Claim 11 (Bounding Ti We have,
E, [T.al Tu Eg VuF (1t), p(tn)))
Yu Tu Tu-
w0ol?
E.
n
(1) Proof. For i 8(0, we have that 2.122 u.k Therefore, we have,
[T,Ti = E. V.F (n v(tt)
1 Σ Σ m (ES())
Tu
VuFi U 0y s((12)
k=0
Using that a!()) independent s(t), we get,
E, [T,Ti. = -r E, V.P (1 v(+11)
V.
(100
((22)(tt))
K=0
YutuEe
ViF (1))), )(
V.F (1( v(tt))
Σ V.P.
(1 - vuFi (uM,,
7/
E,
K=
Invoking (x. < 12 ||/l1 /2 for vectors x, y followed by smoothness completes the proof.
0
Claim 12 (Bounding T2 We have,
2p2 (1 (1- m/n) E,
V.F V. (ul), F(ott)
E, [T,.a] S 3
1+ m
3 tu Tu
Tu
-r00 6I
Eg
(t)
(1 - m/n).
1=1 k=0
26

---

Under review as a conference paper at ICLR 2022
Proof. We use E| |||2]|| EE
E[2]1 for random vector 2 to get
Eg T2, S
Lu 2u
+ I.
E
1
Σ
F
(((+1))
m iES()
k=0
We break the term To
Tk 33
((0) VuFi
Fs
o(tt1) - VuFi (ul") o((
(ES(1)
+3
1 m
Σ
Vuli
(()) oo )((-11) v.F (n!),
(ul),((43)) F(41))
+3 V (ul)) p(41)) |
(ES())
For the first term, we Jensen's inequality take the squared inside the sum, then use smoothness and take an expectation over the sampling of devices to get
Σ (Vafi
(2 of o((+)))
VoFi (2 -f
n E
w00
E
m
n
iESIO
For the second term, we use the fact that s(() was sampled without replacement (cf. Lemma 17 and invoke the gradient diversity assumption (Assumption to get,
VuFi (1(() (.((((())))))
V F (1 p((t))
(E)())
m
V Fs (n()) . a(((t)) V.F (u, ((,t))))))'
n
mn
2
1 (1-m)
m)
62 p² E
V.F (u()) ((+t)))
(ul), )
m
To complete the proof, we plug these terms back into the definition of Th and E, T2,ua complete the proof.
Claim 13 (Bounding Ts,u). We have,
E, [Ts.ul 8 (1- n
A
4x
2n2 Io (1
7
Proof. Since us (t+1) = =0 (((+1)
for i E 5(t), we have that
Ts. 2
(t++)
1/ 1
2n
igs())
Since o ((((41) 100 is independent of s(t), we can take an expectation to get
[T(Tu
L²
Σ
P(i & s(t) E,
0
(t+1)
(t) 1
2n
x²
(1
n
((+1)
v. 10)2
2n
Plugging in Lemma 19 completes the proof.
27

---

Under review as a conference paper at ICLR 2022
The analysis of the v-step is given in the next result. Claim 14. Consider the setting of Theorem 2 and assume that VoTolo S 1/8. We have,
- Wot.mA!
Tot Lo 4T2 - 1)
E. I (ul), p(tti) -F (u()) v(o)
8n
2n
Proof. From smoothness, we get,
Fi (u("),((t))) o((+t) - Fs (u(0,,!0) 1 <Vo.s (ul), ((()) (oftt1)
- 21
+
Le
(((+1) 0000)
2
Ta.
We bound the first term as
[T,,] = 7o E, V.F. (u()),(0 (1)
Σ
Fo (u))) a2(2)
k=
Vooii (1()), (2(0),,()) ((0)) || 2
-60
Σ Er Vorli (1() -,!")) Vofi (u(t) (u²))0)))
V.P. (u'0) .o(0) )(
k=
To-1
To
V Pi (u()){(()) )1²
Σ
V.Fi (u(0,=,?)) _.! - v (,10,,070 )|²
+
E
2
2
k=
<
To
VuFi
(2(1), o6 )1²
+
Th
L
2550 |
2
2
Next, we observe that
Ez E u, =)]1 = IV. (M, u;)1 1 EE, us, 2) V
We invoke this inequality to handle the second term as
E. [T S
Lo
To 1
E
GiD (w (w0)0(0) 2,kt 71,k
((0
2
k=
To
2 Lo
E
V F; (10) 200,012)
2
2
k=0
1 L.
2
2
PePi (n( ; -p(") ||
+ No Inp To
2
To
V Fi (n()) [ v.(22) - F (n.), (0 )| ||
+ ve Ln To
Σ
Ex
h=0
2 T.0
+ 3 LoT
V.l (u0o,,)))
)3( + LLLT Ee.lo
- e.(1
2
k=
Plugging these bounds for Tw and T2, into the initial smoothness bound and using Yo To $ 114 gives
E, Fi ((0,,((42)) - Fg (ui9), of (
To V Fi (w((), .2.0)) 1| + To E
1
201p
22 2. +
4
2
We invoke Lemma 18 to bound the Ee Eg
i.k 0/
(1) ||2 term, which is also known as client drift.
We simplify some coefficients using PoTToLo
5 to get
E, F. (v()) 5 a((+1))
(w()))
(1)
F.
2
To 8
V.Fi
(1
1
2.00 2
+ 4- 2 - -1)03
2
28

---

Under review as a conference paper at ICLR 2022
It remains invoke that StO is uniformly random sample of m devices from {1, .n) and that i (++1) is independent of ((O, To this end, note that
Es F (u'), vltti) F (u(), v(oo
Et
Σ
Fd
( , offtt)) Fb ef0))
n
m
iES(
m7
(u'",,,00 01 offtt)))
- Fi (ul)) (1(0),())) nn(t]]]
E
Fi
Plugging in the previous bound completes the proof.
Remark 15. We only invoked the partial gradient diversity assumption (Assumption) 3 at (virtual)
iterates (u(0) V(++))); therefore, it suffices if the assumption only holds at iterates (ult) V(++))
generated by FedAlt, rather than at all (u, V).
A.4 TECHNICAL LEMMAS
The first lemma involves smoothness of two blocks of variables; we use this in the proof of FedSim. Lemma 16 (Block Smoothness). Suppose Fi IRd Rd, satisfy Assumption !|
Then, holds that
Fi v)) - Fi(w, v²) S (V Fiw, ui), w' - w) + V Fi vz), v/g vz)
+ 2 (1 [1+ xx+x))llu' w|1 + L
L
(1 + x3)1(e wil1
2
Proof. Using the Smooothnsss of F(,, vf) and the smosohhhssss of F(w, +), we have
L Fl , Fi(w, v < (V w Fl (w, 21), w/ w) + I| w|l
L F((w, (()))()= Fr vx) < (V m. (w, us), v/4 - -wi) + - l|u! vill?
2
Summing the above two inequalities together gives F, o Fi,w, e4) s (V Ww,FFW, u/), u/ - w) + (V FWW of - vz)
L
I|u w|12 + La ||u 1/112.
(16)
2
2
We can bound the first inner product term on the right- side of the above incquality as (V v/), w/ - w)
(V n Fil, uz), w/ w) + (V Fiw v/) V V.FF(w, vz), w/ w)
V ~ FiW Uz). w' w) + V Fww v/) Filw, w:lllllw - w|l
V ~ Fw, vz), w' w) + L - 0; I|||ww - wil
V n. (w,, 22). w/ w) + XV L L, 10| U/ vell|ww w|l
Lu 10 vill + x² Lu
V w. Fw, vz), w/ w) +
2 l|w' w|l
where the first inequality is due to Sauchy-Schwarz the second inequality is due to L Lipschitz property of F., the third inequality is due the definition of x in (5), and the last inequality is due to Young's inequality. Substiutiing the above inequality into (16) yields the desired result. r
Next, we have the variance of sampling without replacement. Note the correction factor of (n m)/(n 1) over sampling with replacement. We include the elementary proof for completeness.
Lemma 17 (Sampling Without Replacement). Let a1, an E Rdd be given. Let S be a uniformly random sample of size m from this collection, where the sampling is without replacement. Denoting the mean Em= di/n1 f1, we have,
Eg
a
0
<
m
Σ 7
a.
m
n2
m
IES
29

---

Under review as a conference paper at ICLR 2022
Proof. The statement is trivially true if m 1 or nt. Therefore, assume now that 2 I I n 1. Further, without loss of generality, we assume that ā = 0. Finally, let S denote the set of all subsets of [n] of size Note that |S| = (m)) We now have,
Es
Σ a
1a.12+ +
Σ sJES:izj
(an,))
m iE
n m
m
SES
EES
For the first term, we have,
ΣΣ SES SE
ag 2
||a:||=
11 1 i=1
|a:1 10
i=1 SES :
Likewise, for the second term, we use Ejji Ug = - to get,
Σ (a;, ax) = ΣΣ Σ
(a;
12 2
(oi,az) = - (n-2
Σ
||a:l'
i=1 iitt SES:iJEs
1=1 jzh
i=1
Therefore, we get,
Es
1
mn
n-1) n-2
m
|0:11
12-2 m
| ||
n m
|a,112.
m 2
2
mn
1)
iES
m
m
i=1
The next two lemmas are about the effect of the local updates in the local SGD literature. The first lemma has also appeared in (Karimireddy et al. 2020); we give the proof for completeness.
Lemma 18. Consider : RRd + R which is L-smooth and fix a w(() E RR, Define the sequence (w!") of ilerates produced by stochastic gradient descent with fixed learning rate 7y starting from (0):
(t+1) (t)) -gg
where glt) is an unbiased (and independent of w) estimator of V f(w) with bounded variance o². Fix a number T of steps. If S (V2rL)-1. we have the bound
(t) - 10 < 87 1) +f7 -1)02
Proof. If 1, we have nothing to prove. Assume now that + 2 2. Let (+( be the sigma-algebra generated by wit)) and denote E,[.] = El- [(]] We will use the inequality
Eg 9 E, g {(ww()) ||² + [)( 00² + |V( (17)
We now successively deduce,
Eg (e-1) - m. 10) = [1) - (0) - 112
WV
1
1
|l (0 12 + 2 - it
7
1+
2 +2² |1 - w 10) +
/|| |\(w0)) + 2
+ 220 ?
7
SV
1+
1
+ +,222 1.2 [| (0 11 + 27 [ (0) | 272702
7 --1
1
(4)
2 [100(0) 1012 +
2
1+ 7
1 10 + ~00
30

---

Under review as a conference paper at ICLR 2022
Above, we used (a) the inequality 2a,3 Σ a2/6² 8 for reals a,,,6, (b) Eq. 17 (c) L- smoothness of f. and, (d) the condition on the learning rate.
Let C 27 V f( [12 + T0 Unrolling the inequality and summing up the series gives for all t S
) - 2
2 i C
((r-1) (1+ -- -
2
| - - 101 SC
1 1+
--1
j=0
C $ 2
T-1
C (T (T- 1)e²
[T
|+ 7 --1
-
where we used the bound (1 + 1/a) I e for all a > 0. Summing over and using the numerical bound e² completes the proof.
Lemma 19. Consider the setting of Lemma
18
If S (2TL.-1 ! we have the bound
|| - w < 16
V ((w (0) (() )|² -872700²
Proof. Proceeding similar to the last proof (expect using ò = T) gives us
- - + 2 ||=(()))
(0) || + 7 V (0) || 2y2r0²
(0)
Unrolling and summing up the sequence completes the proof, similar to that of Lemma 18
Finally, the last lemma about constants.
Lemma 20. Let "/a Yos Twwy Teu X. ʃ E R2 and natural number T be given. Denote
A TuLa + f Lubuu and, : fruLl + 7ox L
Suppose 7a Cu/ ca/(T Twu) and 9o c c (T with Cus Cr satisfying
Cu. Cv S
max
V6
Then, we have that
R2x² L. + ?.L A/(379), and, Rux Ln.b B/(37?)
Proof. Note that it suffices to show
37 3 L. & A/2, and, 37 x2 Tu Lub A B/2.
Plugging in Tue To. these are equivalent to 6x²fc² fct + 2x S x² tu and, 6x²c + 6x* fce fc S fCo + xx²c..
The assumption c. implies that 6x fc. < fc. and 6x* 2&,cc. Cu. Therefore, the first condition holds. Similarly, the second condition holds too.
B EXPERIMENTS: DETAILED SETUP AND HYPERPARAMETERS
We conduct our experiments on three datasets from two modalities, namely images and text. The datasets contain natural, non-iid split of data which is reflective of data heterogeneity encountered federated learning. We describe in detail the experimental setup and hyperparameters. The code to reproduce the experimental results will be publicly released.
The outline of this section is:
BB.
describes the tasks and their associated datasets and metrics.
B.2
describes the experimental pipeline well the baselines compare to.
BB3
presents the hyperparameters of all the algorithms.
As discussed in I we take the weight ag to be proportional to the number datapoints available on the device.
31

---

Under review as a conference paper at ICLR 2022
StackOverflow
GLDv2
EMNIST
250
100
Count
203-
Count
200
Count
80-
150
150
80
100
40
50-
50
20
25000
12500 15000
50
150
Data per device (words)
per device (images)
per (images)
Figure 5: Distribution of number of training samples per device for each of the tasks considered in the experiments. GLDv2, we not show the long right tail, where the maximum number of data points per device is 1000 (cf. Table
B. DATASETS, TASKS AND MODELS
We consider three tasks motivated by real-world applications of federated learning. The tasks are summarized in Table the main paper and the distribution of data across the clients visualized
in Figure 5
For each model, we consider three partial personalization architectures:
(a) Input layer personalization: Motivated by Liang et al. (2019), this architecture places the first layer on-device to learn a personalized representation per-client, while the rest of the model is shared. For the transformer model, we use the first transformer layer in place of the word embedding layer owing to its large size.
Output layer personalization: Motivated by Collins et al. (2021). this architecture learns shared global representation but personalizes the prediction layer. For the transformer model, we use the last transformer layer place the last prediction layer owing its large size.
(c) Adapter personalization: We also consider novel partial personalization architecture, where the full model is shared among clients, while each client adds personalized adapter modules, which are lightweight modules added between layers of the shared model. We use the transformer adapters proposed by Hooulsy et al. (2019) and residual adapters proposed by Rebuffi et al. (007).
B. STACKOVERFLOW FOR NEXT WORD PREDICTION
Dataset. The StackOverflow dataset comprises of questions and answers from the programming question-answer website stackoverllow.com The goal of the next word prediction task is to predict the next word given a partial sequence of words in a question or answer. This task is good open- source benchmark for next word predictions in mobile keyboards. We use the StackOverflow dataset provided by TensorFlow Federated
Client Distributions. Each client corresponds to one user on Stack Overflow; the data on the client corresponds to the questions and answers posted by this user. We only consider clients with least 100 training sequences and 10 testing sequences, where a sequence refers to either question or an answer. We use fixed subsample of 1000 of them. Following Reddi et al. (202T), we restrict the vocabulary to the top 10000 most frequently occurring words in the dataset. We pad and truncate each sequence of each client length 20 and consider at most 1000 training sequences on each client.
Model. We use a transformer model (Vaswani et al. 2017) commensurate size with BERT Mini (Tur al.) 2019). It has with 4 transformer blocks and 4 attention heads in each self- attention layer with a transformer hidden dimension of 256 and fully-connected hidden dimension 1024. The output layer causal language modeling head, i.e., fully connected layer which assigns
32

---

Under review as a conference paper at ICLR 2022
Table Summary of partial personalization architectures the transformer model for next word prediction.
# Personalized Params.
Shared Params.
Personalization Type
Layer on-device
Input Layer
1st transformer block
0.8M
4.9.M
Output Layer
Last transformer block
0.8.M
4.9.M
Adapter
Adapter modules
0.07
5.7M
Table 6: Summary of partial personalization architectures for ResNet-18 model visual landmark recognition.
# Personalized Params.
Shared Params.
Personalization Type
Layer on-device
Input Layer
1st conv. layer
0.01M
12.2M
Output Layer
Last fully connected layer Residual adapter modules
1M
11.2M
Adapter
1.4M
12.2M
score for each possible vocabulary item, including the special tokens. The model has 6 million parameters, which require around 23 megabytes of memory.
Partial Personalization Architecture. The partial personalization architectures used are summa- rized in Table
Loss Function and Evaluation Metric. We train the model with the causal language modeling objective. That is, for each partial sequence, we treat the prediction of the next word as multiclass classification problem to minimize the multinomial logistic loss, also known as cross entropy loss. For evaluation, use the top-1 accuracy of predicting words in the proper 10000-word vocabulary (i.e., ignoring special tokens such as padding, out-of-vocabulary. and beginning/end of sequence).
B.1.2 GLDv2 FOR VISUAL LANDMARK RECOGNITION
Dataset. GLDv2 stands for Google Landmarks Dataset v2 /Weyand et al. 2020). which is a large- scale image dataset. It contains images of popular landmarks from around the world taken and uploaded by Wikipedia contributors. While the images vary in size, the most common image size is 800 600 pixels.
The goal of the visual landmark recognition task is to identify the landmark from its image. This task resembles scenario where smartphone users take photos natural and architectural landmarks while traveling. We use the federated version of the GLDv2 dataset introduced by Hsu et al. (020 with 2028 landmarks and provided by TensorFlow Federated
Client Distributions. Each client corresponds to one Wikipedia user and contains all the images contributed by that user. We only all 823 clients with at least 50 datapoints. We do not use original test set from GLDv2 from evaluation as it comes from different clients. Instead, we take 50% of the data each client testing set.
Model. We use ResNet-18 18 (He et al. 2016) model pretrained on ImageNet (Deng al. 2009). with group normalization instead of batch normalization (Hsich et al. (2020). We resize all images to 224 x 224. We use two data augmentations for training: random crop from 256 x 256 and random horizontal flip. The model has 12 million parameters, which require around 49 megabytes of storage.
Partial Personalization Architecture. The partial personalization architectures used are summa- rized in 6
Loss Function and Evaluation Metric. We use the multinomial logistic loss, also known as cross entropy loss. We evaluate the performance of the model using its classification accuracy.
33

---

Under review as a conference paper at ICLR 2022
Table 7: Summary partial personalization architectures for the ResNet- model character recognition.
Personalized Params.
Shared Params.
Personalization Type
Layer on-device
Input Layer Output Layer
1st conv. layer Last fully connected layer Residual adapter modules
0.7K
11.2M
0.03M
11.2M
Adapter
1.4M
11.2M
B.1.3 EMNIST FOR CHARACTER RECOGNITION
Dataset. EMNIST Cohen et al. 2017 is character recognition dataset. The goal is to identify images of handwritten digits or letters; there are 62 possible options (a-z.A-Z, 0-9). The images are grey-scaled pictures of 28 28 784 pixels. We use the EMNIST dataset provided by TensorFlow Federated!
Client Distributions. Each client corresponds to one "writer", i.e., the human subject who hand- wrote the digit/letter during the data collection process. We only use those clients with at least 100 training points and 25 testing points: there are 1114 of such clients.
Model. We use ResNet-18 (He et al., 2016) model with group normalization instead of batch normalization (Hsieh et al. 2020). We make two modifications to handle the smaller image size (28 28 1 as opposed the 224 x 224 x 3 which the original ResNet was designed to accept): (a) we use convolutional kernel of size x 3 rather than the original 7 x 7 in the first convolution layer, and, (b) we drop the first pooling layer. The model has 11 million parameters, which require around 45 megabytes. Note that the number of parameters in this ResNet is smaller than the one for GI.vv2 due the architectural modifications we make for smaller images as well the smaller number of classes.
Partial Personalization Architecture. The partial personalization architectures used are summa- rized in Table
Loss Function and Evaluation Metric. We use the multinomial logistic loss, also known as cross entropy loss. We evaluate the performance of the model using its classification accuracy.
B.2 EXPERIMENTAL PIPELINE AND BASELINES
There three components in the training pipeline for all experiments:
(a) Non-personal conppersonalized federated training: The first step involves training a global model wy using the
one- mne-model-fitsall approach of ( with FedAvg variants.
(b) Personalized federated training: This optional second step involves training the shared parameters w together with the personalized parameters 1y using personalized federated learning approach. We warm-start w, Uk from the non-personalized model wy from the previous step.
(c) Final finetuning: The last step involves only finetuning the personalized parameters us while the shared parameters w remain unchanged.
initialize Uk for each i to be the appropriate part of wy for input/output layer personalization. On the other hand, for adapters, we initialize eg to be equal to the same set of randomly initialized weights for each device
For step b
We consider the following baselines:
Non-personalized: This denotes the performance of step (a) of the pipeline above, i.e., non- personalized federated training with FedAvg variants.
Full model personalization: We consider three baselines of personalization the full model:
(i) Finetune: The non-personalized model from step (a) of the pipeline above is finetuned locally each client (step of the pipeline). Step Step(h) skipped for this baseline.
34

---

Under review as a conference paper at ICLR 2022
Table 8: Hyperparameters for each dataset/task.
Hyperparameter
StackOverflow
GLDv2
EMNIST
Batch size
64
64
Devices per round
50
50
10
Local epochs
1
1
Server Optimizer
FedAdam SGD
FedAdam
FedAAg
Common
Client Optimizer
SGD
SGD
Global Scheduler
Linear
Linear
Exponential
Warm up
10% of rounds
10% of rounds
N/A
I.R decay rounds
N/A
N/A
500
Max. grad. norm.
0.1
N/A
N/A
# Rounds
1000
2500
2000
Non-penn Non-personalized training (step (a) of the pipeline)
Server learning rate Client learning rate
5 10 44 1
10
10-2
0.5
Rounds
500
600
500
Personalized training (step (b) the pipeline)
Server learning rate
x 10-b
2 10-5
1.0
Client learning rate
10
10-3
10-2
#Epochs
5 SGD 10
5 SGD 10-3
Local finetuning
Optimizer
SGD
(step (e) of the pipeline)
Client learning rate
10
(ii) Ditto (Li ct al. 2021): The non non- non-personalized model from step (a) of the pipeline above is finetuned locally each client (step (f) the pipeline) with & regularization ||v wy 1
(b) skipped for this baseline.
(iii) MeedMe (Dinh et al. 2020): The non-personalized baseline model from step (a) is trained
further in step to to optimize (2) using the pFedMe algorithm oi et al. (000 Finally
the resulting model w is finetuned locally step [c)
Partial Model Personalization: We consider partial model personalization with three different
architectures, as defined in $B.I] For each personalization approach, we start with the non-
personalized model in step (a) continue personalization in step using either FedAlt FedSim
as the algorithm, and finally run (c)p for the local finetuning.
B.3 HYPERPARAMETERS AND EVALUATION DETAIL.S
All the tuning hyperparamettrss was performed on validation data, formed by holding out 20% of the training data each device. Once the tuning was complete, we reran the experiments on the full training data, including those held out for validation.
Evaluation Metric. Our primary evaluation metric is the weighted average the test accuracy on each client, weighted by the number of test examples (the details of how the accuracy is computed on each dataset is given in [B. in the paragraph on "Loss Function and Evaluation Metric").. This corresponds to the unweighted accuracy obtained by pooling all the data locally, similar to the loss as discussed in I The same metric is used for hyperparameter tuning and is reported in all the tables and plots, unless explicitly noted otherwise.
The final hyperparameters we use are given in Table g
Rounds. We start with the number of communication rounds (i.e., the number of calls to secure aggregation routine for the shared parameters), which used measure the progress of each algorithm. For the non-personalized training. we use 1000 rounds for StackOverflow, 2500 rounds for GLDv2 and 2000 rounds for EMNIST. For the personalized training, we warm-start the model from the non-personalized one, and run the training for 500 rounds for StackOverflow and EMNIST and 600 rounds for GLDv2.
Devices per Round. All devices are assumed to be available and selections are made uniformly at random. Following (Reddi et al. 2021 Weyand al. 2020), select 50 devices per round for
35

---

Under review as a conference paper at ICLR 2022
Table 9: This table shows Table of the main paper along with corresponding standard deviations as subscripts. We compare FedAlt and FedSim for partial model personalization in this table. "FT (part.)" corresponds to finetuning the personal parameters vi locally while fixing the shared parameters from a non- personalized training. The numbers averaged over 5 random seeds; the boldfaced numhers denote those within standard deviation of the highest accuracy in each
StackOverllow
GLDv2
EMNIST
FT (part.)
FedAh
FedSim
FT (part.)
FedAIt
FedSim
FT (part.)
FedAh
FedSim
Inpur Layer
24. 24.96o.0 o1
24.4-o.o1
24.8
51.97.o0 02
53.
53.64o.o
03.299..m
93.55o.0
Output Layer
21.93..o
25.05%.0r
25.02s.av
53.21o.01
56 64a.0s
56.24o.04
93.37e.01
93.57
93 55.50..00
Adapter
21.71a.00
24.82a.0
24.74b.01
63.866.00
66.
66.35o.00
93.066o.o
91.113g.0s
91.076.08
StackOverflow/GLDD22 and 10 per round for EMNIST, for both the pon-personalized well as the personalized training.
Local Updates and Minibatch Size. Each selected device locally runs 1 epoch of mini-batch stochastic gradient descent locally for pon-eeesoe as well as personalized federated training. The final finetuning at the end of personalized training is performed for 5 epochs. We use minibatch size of 64 for StackOverllowGLL.DD and 32 for EMNIST for all settings.
Server and Client Optimizer Details. We use FedAvg for EMNIST and FedAdam (Reddi et al. 202T) for StackOverflow and GLDv2. We also use global scheduler. which applies schedule on the client learning rates across rounds, while the client learning rate within each round is held constant. We use either linear scheduler or an exponential scheduler (also called "stepLR" in PyTorch). A linear scheduler applies a linear warmup. if applicable, until the maximum learning rate followed by a linear decay to exponential scheduler halves the client learning rate once every fixed number of rounds. Both the client and server learning rates are tuned using the validation set.
Regularization Coefficient for pFedMe and Ditto. We tune the regularization coefficient X for PFedMe and Ditto using the validation data from the set {10- 10-3, ...,, ,,,)) of possible values. The tuned values
StackOverflow: 10 for Ditto and 10 4 for PFedMe. GLDv2: 10-1 for both Ditto and pFedMe, EMNIST: 10-1 for both Ditto and PFedMe.
Random Seed. We report numbers averaged over 5 random seeds.
C EXPERIMENTS: ADDITIONAL RESULTS
We now present the detailed experimental results.
We start with Table which shows the standard deviations of Table 4 We see that the numbers are fairly consistent across runs so the standard deviation quite small (0.01 to 0.05 percentage points).
C.I ABLATION: FINAL FINETUNING FOR FEDALT AND FEDSIM
We now study the effect of the final finetuning (step (c) of the experimental pipeline; ef. B.2 for FedAlt and FedSim.
The final finetuning has minimal impact on partial personalization. We see from Table that the effect of the final finetuning is much smaller than the improvements from personalization. For instance, the improvements from finetuning are close to for FedAlt on the StackOverflow dataset. For GL.v2, the finetuning accounts for < 0.5pp of improvement, whereas personalization overall accounts for 5 to 15pp.
The final finetuning is more important to FedSim than FedAlt. Table 10 also shows that the final finetuning helps FedSim more than FedAlt. However, FedAlt still outperforms FedSim, in
36

---

Under review as a conference paper at ICLR 2022
Table 10:
The change accuracy (percentage points) from the final finetuning for FedAlt and FedSim with stateful devices. The subscript denotes the standard deviation over random seeds.
StackOverflow
GL.vv2
EMNIST
FedAlt
FedSim
FedAlt
FedSim
FedAIt
FedSim
Input Layer
0.06u.01
0.04p.02
0.1 0.120.02
0.170.08
0. 0.12p.01
0.12o.ox
Output Layer
0.00g.01
0.25p.02
0.490.02
0.57o.06
0.090.01
0.09o.or
Adapter
0.0lg.o
0. 0.40p.os
0.14 0.140.02
0.17o.01
0.27p.02
0.33p.03
Pen-clinnt Statitics (train]
Pon-clle: Statistics (best)
Effect of Regularization (test)
Effect of Dropout (test)
Accuracy IC
Accuracy
@FodMe
PFodMM
Best Reg. Large Reg.
Large
Figure 6: Left two: Distribution of change in the per-device train (left most) and test (center left) accuracy due personalization the StackOverflw dataset. Right two: Distribution of change in the per-device test accuracy of partial personalization under regularization the StackOverflow dataset: center right: adapter personalization under f2 regularization, and, (h) rightmost: output layer personalization under dropout. Note that the "No Reg." and "No d/o" plots on the right two are different because they personalize different model parts. Interpretation: The white dot in inside the violin denotes the median, while the black box enclosing this white dot marks the interquartile range (i.e., 25th and 75th percentiles). The body of the violin is kernel density estimate of the distribution of accuracies. The lines extend to minimum and maximum accuracy each case.
Table 9 Overall, this shows that FedAlt is better algorithm than FedSim. The final finetuning helps FedSim make up some percentage points in accuracy, but not enough to make its gap with FedAlt.
C.2 EFFECT OF PERSONALIZATION ON PER-DEVICE GENERALIZATION
Summary of all scatter plots. All the scatter plots shown in the main paper are summarized in the violin plot of Figure We see from the leftmost figure that the training accuracies on all devices improve with personalization. From the second figure, we see that the test accuracy of some of the devices reduces with personalization; this is true both partial and full personalization.
From the third plot of Figure we see that regularization does not mitigate this overfitting. In fact, the regularization tuned for best average accuracy leads to a nearly identical distribution of test accuracies. A larger regularization reduces the spread of accuracies, but does smaller median (white dot). The fourth plot of Figure 6 shows that the effect of dropout is similar. The best dropout improves the median accuracy, but it does not mitigate the issue of some devices being hurt by personalization.
the expense of
Train Accuracy plots for devices. From Figure 7 we see that personalization leads to reduction in test accuracy on some of the devices beyond the initial non-personalized model. The corresponding train accuracy plot is given in Figure | We observe that the personalization always leads to an improvement in the training accuracy but not in the test accuracy. The analogous plots for GLDv2 are in Figure g where the trends are similar.
Whether personalization helps device or not depends on the random seed. We see in Figure I0 that the shaded region for some of the devices intersects the dotted line at In other words, personalization sometimes helps this device and sometimes hurts it, depending the random seed. This indicates that the best fix in practice to use A/B testing on the deployed model to choose whether use the personalized model or the non-personalized one.
37

---

Under review as a conference paper at ICLR 2022
Fineture
Ditto
pfedMee
Partial
Accuracy
1200 per
35000
2303
1566 12500
12500
50000
33000
per
Fincture
Ditto
pFedMe
Portial
Accuracs
Pers Tark)
2307 3000 1306 19000 12:00 33000 Data per device
2903 5000 1506 19090 12500 13000
2300 se00 1500 12000 13500 15060 Data per device
2590 5000 1500 10030 12300 33600 Data per device
# Data per device
Figure 7: Scatter plot of change accuracy (pp) per-device device StackOverflow. Top: Training accuracy. Bottom: Test accuracy. This the version Figure from the main paper.
the number of training samples the
Fineture
Dio
pfedMe
Partial
mean
A.Acuracyy
Piin: Fumh
300
200
200
420
200
326
201
2i0
4ic
202
316
202
203
166
200
200
46c
502
#
Data per device
Data per device
Data per device
*
Data per device
Anetune
Ditto
PFedMe
Partial
Accuracy
Tarts
Data
device
Data device
Data
device
Data device
Figure 8: Scatter plot of change in accuracy (pp) per-device versus the number of training samples on the device for GLDv2. Top: Training accuracy. Bottom: Test accuracy.
Regularization and dropout do not mitigate this issue. From the first row of Figure 9 we see that the weight decay with best mean accuracy exactly matches the unreguarlized case in terms of per-device statistics. Increasing the regularization weight can reduce the spread of per-device accuracy. However, this only leads to worse mean accuracy and does not mitigate the issue of personalization hurting individual devices.
From the second of Figure 9 in average accuracy (0.18 pp). It also reduces the number of devices hurt by personalization from 256 of 1000 to 193, but it does not fix this issue. Increasing dropout further only leads degradation of per-device statistics.
that the best dropout (0.3 in this case) leads slight increase
C.3 PARTIAL PERSONALIZATIONN FOR STATELESS DEVICES
The algorithmss we considered in this paper, namely FedAlt and FedSim, require the devices to maintain the personalized parameters 's state across rounds. In cross-devics federated learning settings, it is also interesting to consider stateless devices, which are not allowed to maintain state between training rounds.
38

---

Under review as a conference paper at ICLR 2022
Param,
Param. best
Reg. large
Telo
mean
Park
2008 4600 6000 9283 200002000000000000
2003 1066 1000012608 16000
1000 1600 12300 3460e 16600
Data device
Data device
Data device
Dropout
Dropout
Drapout large
Fars Telp
mean
Accuracy
per device
Data per device
Figure 9: Scatter plot of change in accuracy (pp) per-device versus the number of training samples on the device with the effect of regularization. Top: & regularization a.k.a. weight decay. Bottom: dropout. The "best" values of the & regularization parameter and dropout are chosen maximize the average accuracy across all devices.
Finetune
Ditto
PFedMe
Partial
Pera ndica
Accuracy
& Accuracy
Accuracy
4. Accuracy
0-
200
000
B02
2000
200
400
200
1003
450
620
1020
200
402 000
a2o
2000
Device Rank
Device Rank
Device Rank
Device Rank
Figure 10: Change per-device accuracy (pp) due personalization. The solid line the mean nuns and the shaded area denotes the max/min across these runs. The devices are sorted in ascending order of accuracy change. The points in orange depict two example devices who might either be helped or harmed by personalization depending on the random seed.
random
We give preliminary experiments in this setting. We modify the FedAlt and FedSim algorithms from the main paper so that the personalized parameters for participation. We warm-start v; from the appropriate part of the non non-personalized model trained in step of the pipeline. For adapters, we fix random initialization once, and reuse it.
reinitialized each time device i is chosen
FedAlt is better than FedSim for stateless devices, although the improvement is smaller. We see from Table 11 that all algorithms perform similarly for the stateless setting. Nevertheless, we see that FedAlt obtains mild improvements over both FedSim and finetuning for GLDv2, e.g., 0.24pp with adapters.
The final finetuning is crucial for stateless devices. We see from Table 12 that the final finetuning accounts for most of improvements in the stateless case. For instance, for GLDy2, the final finetuning accounts for 11.68 and 10. 42.42pp out of total of 12.67 and 11.76pp for FedAlt and FedSim respec- tively. However, the personalized federated training (step (b) of the pipeline; ef. (B.2) still leads to an increase in accuracy of to 1.34pp.
39

---

Under review as a conference paper at ICLR 2022
Table 11: This is the counterpart of Table to stateless devices. We compare FedAIt and FedSim for partial model personalization with stateless devices. "FT (part.)" corresponds finetuning the personal parameters vs locally while fixing the shared parameters from non-personalized training. The numbers are averaged over random seeds; the boldfaced numbers denote the highest accuracy in each row.
StackOverllow
GLDv2
EMNIST
FT (part.)
FedAIt
FedSim
FT (part.)
FodAIt
FodSim
FT (part.)
FedAIt
FedSim
Input Layer
24.966.01
2188...01
24.896.01
51.97o.02
76..06
52.74..02
99..29.00
83.484.04
Output Layer
24.93p.o
24.94p.0
24.946.01
63.21o.
o3
53.30p.os
93.37o.m
93.53. oa
90.5lo.o
Adapter
24.71..00
24.69%.o1
24.71o.o1
63.86%.oe
10
63.19 19
66.66..
93.97p.0s
93.89..0
Table 12: The change in accuracy (percentage points) from the final finetuning for FedAlt and FedSim with stateless devices. The subscript denotes the standard deviation over 5 random seeds.
StackOverflow
GLDv2
EMNIST
FedAIt
FedSim
FedAlt
FedSim
FedAlt
FedSim
Input Layer
0.86p.o3
1.00p.02 02
0.44g.03
0.42p.os
0.11o.02
0.10g.04
Output Layer
1.08p.03
1.10o.02
1.470.04
1.466.05
0.150.02
0.110.02
Adapter
0.84p.04
0.88a.o2
11.68
10.42p.09
0.466.02
0.42a.04
40


Reviews:
[{"review_id": "rg7mwFhTH3T", "reviewer": "Reviewer_8i9S", "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": "My main concern about the paper is the limited novelty. The idea of partial personalization has been proposed in previous literature. The objective function (3) has been detailly studied in [1]. Moreover, the convergence of FedSim for nonconvex case has also been studied in [1] under very similar assumptions (See Section 3.1 of [1], where the name of FedSim is LSGD-PFL). Thus, the only new result left is the analysis of FedAlt under nonconvex setting. In order to derive Theorem 2, the authors claim that they proposed a new technique called virtual full participation; however, this technique is actually not new and has been used widely in convergence analysis for FedAvg/Local-SGD. This way, the whole contribution of the paper relies on the analysis of FedAlt in nonconvex case based on well-known techniques, which seems quite weak contribution to me. I would recommend the authors to give more detailed explanation on what the key theoretical challenge in deriving Theorem 2 is. \n\nThe extensive experiments of using different model structures on different datastes is appreciated. And the observation personalization can hurt the test accuracy on some devices is interesting. But given the idea of partial personalization is proposed elsewhere, the gain from the experiments is limited.\n\nThe writing is clear. I can easily understand the paper. Thanks for that.\n\nIn a word, given the ideas used in the paper are now new, I think this paper is maybe better written in a pure theoretical one. The key would be to stress the theoretical challenge and novelty. Based on the current version, I believe the contribution is not novel enough or at least it is not stated clearly.\n\n[1] Filip Hanzely, Boxin Zhao, and Mladen Kolar. Personalized Federated Learning: A Unified Framework and Universal Optimization Techniques. arXiv Preprint, 2021.", "summary_of_the_paper": "The paper discusses using partial personalization objective function defined in equation (3) to achieve personalized models in federated learning. While the idea of partial personalization and the objective of (3) have been widely studied in previous literatures, the main contributions of the paper are the convergence analysis of two proposed algorithms FedSim and FedAlt in nonconvex case. The authors also did extensive experiments to compare partial personalized models with fully personalized models by using various model structures and on different datasets.", "summary_of_the_review": "The contribution is unclear, which constitutes strong reason for not accepting the paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and 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": "NDADjCqM2LJ", "reviewer": "Reviewer_447D", "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": "Pros:\n1. The authors proposed a unified framework for personalized federated learning. To my best knowledge, they give the first analysis of this framework and algorithm on smooth nonconvex function.\n2. The proposed algorithm can outperform existing personalization algorithms on NLP or Vision tasks.\n\nCons:\n1. The personalized layer idea is not new, as authors mentioned, it has been proposed in prior works [Arivazhagan et al. (2019) and Collins et al. (2021)].\n2. I think the established theory does not perfectly fit with experiment model, since they assume smoothness in theory, but in practice the motivated example is usually non-smooth ReLU network. Hence I think it will be great if they can perform analysis on ReLU network, perhaps very simple one layer setting.\n\n", "summary_of_the_paper": "This paper proposed a personalized FL framework with partial model personalization. It separates the model parameters into two parts, shared model and personalized model, and optimize them in an interleaving manner. The authors proposed two optimization algorithm, named as FedSim and FedAlt, with partial clients participation. They also analyzed the algorithms' convergence rate on general smooth nonconvex function. In experiments, they consider 3 different model split methods: 1. input layer as personalized model, the rest as shared model; 2. output layer as personalized model, the rest as shared model; 3. adding adapter as personalized model.  The experiments conduct on NLP or vision tasks also demonstrate that the proposed algorithm outperforms other personalization method.", "summary_of_the_review": "Overall, this framework is interesting and theory is new. However, the novelty for me is somewhat limited. If they can come up with theory that can explain their empirical observations (ReLU network), I believe it will be a strong work.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and 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": "NA2VchyA2c_", "reviewer": "Reviewer_Nawz", "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": "This is a well written paper and easy to follow. However, I have the following concerns. \n\nAn important related work is missing, and I strongly suggest the authors discuss\n- [Singhal et al. 2021 Federated Reconstruction: Partially Local Federated Learning https://arxiv.org/abs/2102.03448]. If I understand correctly, the proposed FedAlt algorithm is very similar to the paper, except that FedAlt requires a local state v_i. \n\nClarification on contributions\nIf I understand correctly, FedSim has been used in [Liang et al. 2019, Arivazhagan et al. 2019, Collins et al. 2021, Li et al. 2021 ], FedAlt is relatively new but closely related to [Singhal et al. 2021]. This paper makes the following contributions\n- Convergence guarantees of FedSim and FedAlt: standard convergence rate under reasonable assumptions.  Li et al. 2021 presented a convergence rate of FedSim, could the authors compare?  Regarding the novelty of FedAlt proof, I believe virtual proxy variable is a common technique in federated optimization, (e.g., eq(13) in [Wang et al. 2021 A Field Guide to Federated Optimization https://arxiv.org/pdf/2107.06917.pdf]), and I would appreciate it if the authors elaborate more on the novelty. \n- FedAlt algorithm, which is a client state variant of Singhal et al. 2021, and achieves marginal improvement over FedSim [Liang et al. 2019, Arivazhagan et al. 2019, Collins et al. 2021, Li et al. 2021] (always < 0.5%.)\n- Theoretical insights on when FedAlt is better than FedSim. I found this to be interesting, and I would encourage the authors to provide more discussion and connections to empirical results.\n\n\nSavings of memory footprint\n- It would be good to show either analytically or empirically how much memory the partial model training can reduce compared to full model training as this is the main motivation. \n\nClarification on some details. \n-  Authors mentioned “All methods, including FedSim, FedAlt and the baselines are initialized with a global model trained with FedAvg. ”: how is the global model trained? And is it trained on the same set of clients as personalization?\n- For table 3, does the full model in baselines include adapter parameters? \n- How do the authors split the dataset for training, personalization and testing?\n", "summary_of_the_paper": "This paper studies personalization in federated setting, i.e., instead of collaboratively training a global model, personalizing a model for each client.  This paper proposed to personalize only part of the model parameters instead of the full model, and studied two algorithms FedSim and FedAlt. In local client updates, FedSim will simultaneously train the shared and personalized parameters, while FedAlt will train the personalized model first, then train the shared parameter. ", "summary_of_the_review": "I strongly encourage the authors clarify contributions explicitly acknowledging previous works [Liang et al. 2019, Arivazhagan et al. 2019, Collins et al. 2021, Li et al. 2021 ], as well as  [Singhal et al. 2021 Federated Reconstruction: Partially Local Federated Learning https://arxiv.org/abs/2102.03448]. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and 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": "94UFgpJ2oBk", "reviewer": "Reviewer_mEwK", "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": "This paper studies an interesting question and proposed idea looks interesting with corroborating empirical results, but there are few issues that prevents me from giving it a high score:\n\n- The key motivation for the paper is to utilize a small footprint in clients but the proposal is somehow mis-leading. I understand the learned personalized model is significantly smaller than the shared model (e.g. 1-2% as observed in experiments for some applications) and requires a small footprint after deployment (inference stage), but during the training the clients need enough memory for both models to participate in collaborative training. Also, I was left wondering how large the coupling between personalized and shared parameters captured by $\\chi$ would be in presence of heterogeneity and significant difference in model sizes. \n\n- While the proof of theoretical results looks sound as far as I checked, the obtained rates are hard to interpret and poorly elaborated. For example, the gradient diversity assumption (Assumption 3) is only defined over the shared model. Also, the role of gradient diversity term, \\rho, which appears in tuning the learning rate $\\eta$ is ignored in discussions (e.g. Table 1) which might dominate other terms if incorporated properly. So, it seems hard to put the obtained rates in the context of known results.\n", "summary_of_the_paper": "To overcome statistical heterogeneity among data shards in collaborative federated learning, this paper proposes a novel personalization schema which only requires smaller memory footprint in clients and possibly is less susceptible to catastrophic forgetting. The main idea is to split the trainable parameters into  shared and personalized parameters where, unlike existing personalization schema, only the shared model is exchanged with the server in communication rounds to be aggregated- thus partial personalization. In regression, this corresponds to learning the residual error of shared model via personalized model and in the classification setting corresponds to output averaging (unlike interpolation based personalization methods that do parameter mixing). The authors propose two algorithms FedSim (simultaneous updating of shared and personalized models locally) and FedAlt (alternative updating of shared and personalized models locally), to learn in this setting, theoretically analyze the convergence rates in non-convex settings and conduct empirical studies on image classification and next-word prediction to evaluate the proposed methods.\n", "summary_of_the_review": "This paper proposes a solution for personalization in federated learning where clients only exchange a small portion of parameters with a server for collaborative learning. The problem and proposal looks interesting, but it would be great to hear from authors on some issues raised above before final evaluation.\n", "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: Dear authors,

I have read the reviewers and your careful rebuttals. I would have liked to see much more engagement from the reviewers. However, even after your rebuttal, no reviewer suggested acceptance, with two reviewers proposing reject (3) and two proposing weak reject (5).

The reviewers found the paper well written. I concur. The reviewers also notice that the contributions are very marginal compared to prior literature. Personalized FL formulation studied here was in a simpler form first proposed by Hanzely and Richtarik (Federated learning of a mixture of global and local models, 2020) and later generalized by Hanzely et al - a paper the authors cite. That work performed an in-depth analysis, also including the nonconvex case, which the authors (claim) did not notice. Compared to that work, the authors perform an analysis in the partial participation regime. However, partial participation is by now a standard technique which can usually be combined with other techniques without much difficulty. The authors tried to argue that their analysis approach is unique, but the reviewers remained unconvinced. 

In summary, I think this is a solid piece of work which is perhaps judged, looking at the raw scores, a bit too harshly. However, most verbal comments are indeed fair. I am also of the opinion that the paper in its current form does not reach the necessary bar for acceptance. I would encourage the authors to carefully revise the manuscript, taking into account all feedback that they find useful. I think the paper can be improved, with not too much effort perhaps, to a state in which the bar could be reached. 

Kind regards,

Area chair
Decision: Reject