F EDERATED L EARNING WITH P ARTIAL M ODEL P ERSONALIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
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 beneﬁt of full model personalization with a small fraction of personalized parameters, and, (b) the alternating update algorithm often outperforms the simultaneous update algorithm.
I NTRODUCTION
Federated Learning (McMahan et al., 2017) has emerged as a powerful paradigm for distributed and privacy-preserving machine learning over a large number of edge devices (see Kairouz et al., 2021, and references therein).
We consider a typical setting of Federated Learning (FL) with n devices (also called clients), where each device i has a training dataset of N i samples z i, 1 , · · · , z i,N i .
Let w ∈ R
d represent the parameters of a (supervised) learning model and f i ( w, z i,j ) be the loss of the model on the training example
z i,j .
Then the loss function associated with device i is
F i ( w ) =
(1 /N i ) (cid:80) N i j =
f
i ( w, z i,j ) .
A common objective of FL is to ﬁnd model parameters that minimize the weighted average loss across all devices (without transferring the datasets): minimize
[Equation]
where weights α
i are nonnegative and satisfy (cid:80) ni
=
α i = 1 .
A common practice is to choose the weights as α i
= N
i /N where N = (cid:80) nk =1 N k , which corresponds to minimizing the unweighted average loss across all samples from the n devices: (1 /N )
(cid:80) ni
=1 (cid:80) N i j
=1 f i ( w, z i,j ) .
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 on each device can 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-ﬁts-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., 2017), meta learning (Fallah et al., 2020), and proximal methods (Dinh et al., 2020;
Li et al., 2021).
A simple formulation that captures their main idea is minimize w
[Equation]
i
[Equation]
[Figure]
( u
[Figure]
i ) u v i (
[Figure]
input 1 input
F i ( u,v i ) u
v i (c) Personalized
[Figure]
input layer(
[Figure]
where w i for i = 1 , . .
. , n are personalized model parameters at the devices, w 0 is a reference model maintained by the server, and the λ i ’s are regularization weights that control the extent of personalization.
A major disadvantage of the formulation (2), which we call full model personalization , is that it requires twice the memory footprint of the model, w i and w 0 at each device, which severely limits the size of trainable models.
On the other hand, the ﬂexibility of full model personalization can be unnecessary.
Modern deep learning models are composed of many simple functional units and are typically organized into layers or a more general interconnected architecture.
Personalizing the “right” components, selected with domain knowledge, may result in a substantial beneﬁt 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 ﬁnetuned on a small local dataset forgets the original (non-personalized) task, leading to a degradation of test performance.
We propose a framework for FL with partial model personalization .
Speciﬁcally, we partition the model parameters into two groups: the shared parameters u ∈ R d 0 and the personal parameters v i ∈
R
d i for i = 1 , . . . , n .
The full model on device i is denoted as w i =
( u, v i ) , and the local loss function is F i ( u, v i ) =
(1 /N i ) (cid:80) N
i
i =
1 f i (cid:0) ( u, v i ) , z i,j (cid:1) .
Our goal is to solve the optimization problem minimize
[Equation]
Notice that the dimensions of v i can be different 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 (3): FedSim , a simultaneous update algorithm and FedAlt , an alternating update algorithm.
Both algorithms follow the standard FL protocol.
During each round, the server randomly selects a 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 ﬁrst update the personal parameters with the received shared parameters ﬁxed and then update the shared parameters with the new personal parameters ﬁxed.
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 a small portion of the model to personalize, thus imposing a much smaller memory footprint on the devices than full model personalization.
This framework uniﬁes existing work on personalized FL and allows arbitrary partitioning of deep learning models.
• We provide convergence guarantees for the FedSim and FedAlt methods in the general (smooth) nonconvex setting.
While both methods have appeared in the literature previously, they are either used without convergence analysis or only analyzed for convex models.
Our analysis in the nonconvex setting provides the ﬁrst theoretical support of them for training general deep learning models.
The analysis of FedAlt with partial participation is especially challenging and we develop a novel technique of virtual full participation to overcome the difﬁculties.
adapter adapter [Figure] add & norm add & norm feedforward attention [Figure]
i
[Figure]
input
F i ( u,v i ) u
v i (b) Generalized additive model.
• We conduct extensive experiments on image classiﬁcation 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 beneﬁt 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.
Speciﬁc forms of partial model personalization have been considered in previous works.
Liang et al.
(2019) propose to personalize the input layers to learn a personalized representation per-device (Figure 1a), while Arivazhagan et al.
and Collins et al.
(2021) propose to personalize the output layer while learning a shared representation with the input layers (Figure 1b).
Both FedSim and FedAlt have appeared in previous literature, although their convergence analysis is limited. Liang et al.
, Arivazhagan et al.
(2019) and Hanzely et al. (2021) use FedSim, while
Collins et al.
(2021) and Singhal et al. (2021) use 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.
(2021) prove the linear convergence of FedAlt for a two-layer linear network where F i ( · , v i ) and F i ( u, · ) are both convex for ﬁxed v
i and u
respectively.
We analyze both FedSim and
FedAlt in the general nonconvex case with partial device participation, hence addressing a more practical setting.
While we primarily consider the problem (3) in the context of partial model personalization, it can serve as a 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), as well as interpolation (Deng et al., 2020; Mansour et al., 2020) are special cases of this problem.
The rates of convergence we prove in §3 are competitive with or better than those in previous works for full model personalization methods in the non-convex case.
P ARTIALLY P ERSONALIZED M ODELS
Modern deep learning models all have a multi-layer architecture.
While a complete understanding of why they work so well is still out of reach, a general insight is that the lower layers (close to the input) are mostly responsible for feature extraction and the upper layers (close to the output) focus on complex pattern recognition.
Depending on the application scenarios and domain knowledge, we may personalize either the input layer(s) or the output layer(s) of the model; see Figure 1.
In Figure 1c, 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 to the upper layers of the model.
As demonstrated in (Bui et al., 2019), this partitioning can help protect user-speciﬁc private features (input 2 in Figure 1c) as the corresponding feature embedding (through v i ) are personalized and kept local at the device.
Similar architectures have also been proposed in context-dependent language models (e.g., Mikolov & Zweig, 2012).
A more structured partitioning is illustrated in Figure 2a, where a typical transformer layer (Vaswani et al., 2017) is augmented with two adapters.
This architecture is proposed by Houlsby et al.
for ﬁnetuning large language models.
Similar residual adapter modules are proposed by Rebufﬁ et al.
(2017) for image classiﬁcation models in the context of multi-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 a prediction.
Suppose the shared model is h ( u, · ) and the personal model is h
i ( v i , · ) .
For regression tasks with samples z i,j = ( x i,j , y i,j ) , where x i,j is the input and y i,j ∈ R p is the output, we let F i ( u, v i ) =
(1 /N i ) (cid:80) N i j =
f
i (cid:0) ( u, v i ) , z i,j (cid:1) with f i (cid:0) ( u, v i ) , z i,j (cid:1) =
(cid:107) y
i,j − h ( u, x i,j ) − h i ( v i , x i,j ) (cid:107) 2 .
In this special case, the personal model ﬁts the residual of the shared model and vice-versa (Agarwal et al., 2020).
For classiﬁcation tasks, h ( u, · ) and h i ( v i , · ) produce probability distributions over multiple classes.
We can use the cross-entropy loss between y i,j and a convex combination of the two model outputs: θh ( u, x i,j ) + (1 − θ )
h i ( v i , x i,j ) , where θ ∈ (0 , 1) is a learnable parameter.
Finally, we can cast the formulation (2) of full model personalization as a special case of (3) by letting
[Equation]
Many other formulations of full personalization can be reduced to (3); see Hanzely et al.
A LGORITHMS AND C ONVERGENCE A NALYSIS
In this section, we present and analyze two FL algorithms for solving problem (3).
To simplify presentation, we denote V =
( v 1 , . .
. , v n ) ∈ R
d
+ ... + d n and focus on the case of α i = 1 /n , i.e., minimize
[Equation]
This is equivalent to (3) if we scale F i by nα i , thus does not lose generality.
Moreover, we consider more general local functions F i ( u, v i ) =
E z
∼D
i
[ f i (( u, v i ) , z )] , where D i is the local distribution.
The FedSim and FedAlt algorithms share a common outer-loop description given in Algorithm 1.
They differ only in the local update procedures LocalSim and LocalAlt, which are given in Algorithmand Algorithm 3 respectively.
In the two local update procedures, (cid:101) ∇ u and (cid:101) ∇ v represent stochastic gradients with respect to w and v i respectively.
In LocalSim (Algorithm 2), the personal variables v i and local version of the shared parameters u i are updated simultaneously, with their (stochastic) partial gradients evaluated at the same point.
In LocalAlt (Algorithm 3), the personal parameters are updated ﬁrst with the received shared parameters ﬁxed, then the shared parameters are updated with the new personal parameters ﬁxed.
They are analogous to the classical Jacobi update and
Gauss-Seidel update in numerical linear algebra (e.g., Demmel, 1997, §6.5).
1: initialize v i, 0 = v i 2: initialize u i, 0 = u 3: for k = 0 , 1 , · · · , τ − 1 do 4: v i,k +
= v i,k − γ v (cid:101) ∇ v F i (cid:0) u i,k , v i,k (cid:1) 5: u i,
k
+1 =
u i,k − γ u (cid:101) ∇ u F i (cid:0) u i,k , v i,k (cid:1) 6: update v + i = v i,τ 7: update u + i
=
u
i,τ 8: return u i , v i initialize i
i k v − 1 do 3: v i,k +
= v i,k − γ v (cid:101) ∇ v F i (cid:0) u, v i,k (cid:1) v i
= v i τ v initialize u i, 0 = u 5: for k =
1 , · · · , τ u − 1 do 6: u i,k +
=
u i,k − γ u (cid:101) ∇ u F i (cid:0) u i,k v + i (cid:1) 7: update u + i =
u i,τ u 8: return (cid:0) u
+ i , v
+ i (cid:1)
In order to analyze the convergence of the two algorithms, we make the following assumptions.
Assumption 1 (Smoothness) .
The function F i is continuously differentiable for each i = 1 , . . . , n , and there exist constants L u , L v , L uv and L vu such that for each i = 1 , . . .
, n , it holds that
• ∇ u F i ( u, v i ) is L u -Lipschitz with respect to u and L uv -Lipschitz with respect to v i ; • ∇ v F i ( u, v i ) is L v -Lipschitz with respect to v i and L vu -Lipschitz with respect to u .
Due to the deﬁnition of F ( u, V ) in (4), it is easy to verify that ∇ u F ( u, V ) has Lipschitz constant L u with respect to u , L uv / √ n with respect to V , and L uv /n with respect to any v
i .
We also deﬁne
[Equation]
which measures the relative cross-sensitivity of ∇ u F i with respect to v i and ∇ v F i with respect to u .
Assumption 2 (Bounded Variance) .
The stochastic gradients in Algorithm 2 and Algorithm 3 are unbiased and have bounded variance.
That is, for all u and v i ,
[Equation]
Furthermore, there exist constants σ u and σ v such that
[Equation]
We can view ∇ u F i ( u, v i ) , when i is randomly sampled from { 1 , . . . , n } , as a stochastic partial gradient of F ( u, V ) with respect to u .
The following assumption imposes a variance bound.
Assumption 3 (Partial Gradient Diversity) .
There exist δ ≥ 0 and ρ ≥ 0 such that for all u and V ,
[Equation]
With ρ = 0 , this assumption is similar to a constant variance bound on the stochastic gradient ∇ u F i ( u, v i ) ; 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 (cid:63) and denote ∆ F 0 = F (cid:0) u (0) , V (0) (cid:1) − F (cid:63) .
Further, we use shorthand V ( t ) =
( v ( t ) 1 , . . .
, v ( t ) n ) and ∆ ( t ) u
= (cid:13)(cid:13) ∇ u F (cid:0) u ( t ) , V ( t ) (cid:1)(cid:13)(cid:13) 2 , and ∆ ( t ) v = 1 n (cid:80) ni =1 (cid:13)(cid:13) ∇ v F i (cid:0) u ( t ) , v ( t ) i (cid:1)(cid:13)(cid:13) 2 .
For smooth and nonconvex loss functions F i , we obtain convergence in expectation to a stationary point of F if the expected values of these two sequences converge to zero.
We ﬁrst present our main result for FedSim (Algorithm 1 with LocalSim), proved in Appendix A.2.
Theorem 1 ( Convergence of FedSim ) .
Suppose Assumptions 1, 2 and 3 hold and the learning rates in FedSim are chosen as γ u = η/ ( L u τ ) and γ v = η/ ( L v τ ) with η ≤ min (cid:110) 1 12(1+ χ 2 )(1+ ρ 2 ) , (cid:113) m/n 196(1 − τ − 1 )(1+ χ 2 )(1+ ρ 2 ) (cid:111) .
[Table]
Then, ignoring absolute constants, we have
[Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Equation]
The left-hand side of (6) is the average over time of a weighted sum of E (cid:2) ∆ ( t ) u (cid:3) and E (cid:2) ∆ ( t ) v (cid:3) .
The right-hand side contains three terms of order O
(1 / ( ηT )) , O ( η ) and O ( η 2 ) respectively.
We can minimize the right-hand side by optimizing over η .
By considering special cases such as σ 2 u = σ 2 v =and m = n , some terms on the right-hand side disappear and we can obtain improved rates.
Tableshows the results in several different regimes along with the optimal choices of η .
Challenge in Analyzing FedAlt.
We now turn to FedAlt.
Note that the personal parameters are updated only for the m selected devices in S ( t ) in each round t .
Speciﬁcally,
[Equation]
Consequently, the vector V ( t +1) of personal parameters depends on the random variable S ( t ) .
This makes it challenging to analyze the u -update steps of FedAlt because they are performed after V ( t +1) is generated (as opposed to simultaneously in FedSim).
When we take expectations with respect to the sampling of S ( t ) in analyzing the u -updates, V ( t +1) becomes a dependent random variable, which prevents standard proof techniques from going through (see details in Appendix A.3).
We develop a novel technique called virtual full participation to overcome this challenge.
Speciﬁcally, we deﬁne a virtual vector (cid:101) V ( t +1) , which is the result if every device were to perform local v -updates.
It is independent of the sampling of S ( t ) and we can derive a convergence rate for related quantities.
We carefully translate this rate from the virtual (cid:101) V ( t +1) to the actual V ( t ) to get the following result.
Theorem 2 ( Convergence of FedAlt ) .
Suppose Assumptions 1, 2 and 3 hold and the learning rates in FedAlt are chosen as γ u =
η/ ( L u τ u ) and γ v = η/ ( L v τ v ) , with η ≤ min (cid:110)
[Equation]
Then, ignoring absolute constants, we have
[Equation]
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 1.
Comparing FedSim and FedAlt.
Table 1 shows that both FedSim and FedAlt exhibit the standard
O (1 / √ T ) rate in the general case.
Comparing the constants in their rates, we identify two regimes in terms of problem parameters.
The regime where FedAlt dominates FedSim is characterized by
[Equation]
[Table]
A practically relevant scenario where this is true is σ 2 v ≈ 0 and σ 2 u ≈ 0 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 + χ 2 ) , indicating that the rate of FedAlt is less affected by the coupling between the personal and shared parameters.
Our experiments in §4 corroborate the practical relevance of this regime.
The rates from Table 1 also apply for full personalization schemes without convergence guarantees in the nonconvex case (Agarwal et al., 2020; Mansour et al., 2020;
Li et al., 2021).
Our rates are better than those of (Dinh et al., 2020) for their pFedMe objective.
E XPERIMENTS
In this section, we 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 StackOverﬂow dataset, where each device corresponds to the questions and answers of one user on stackoverflow.com .
This task is 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 et al., 2020), a large-scale dataset with real images of global landmarks.
Each device corresponds to a Wikipedia contributor who uploaded images.
This task resembles a scenario where smartphone users capture images of landmarks while traveling.
We use a 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 a 28 × 28 grayscale image of a handwritten character and the output is its label (0-9, a-z, A-Z).
Each device corresponds to a writer of the character.
We use a ResNet-18 model, with input and output layers modiﬁed 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 classiﬁcation.
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 1a).
For the transformer, we use the ﬁrst 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 speciﬁc layers of a shared model (Figure 2a).
We use the transformer adapters of Houlsby et al.
(2019) and for ResNet-18, the residual adapters of Rebufﬁ et al.
Algorithms and Experimental Pipeline.
For full model personalization, we consider three baselines: i Finetune , where
each device
ﬁnetunes (using
SGD
locally)
its
personal full model starting
from
a learned common model, (ii) Ditto (Li et al.,
, which is
ﬁnetuning
with
(cid:96) 2 regularization, and, (iii) pFedMe (Dinh et al., 2020) which minimizes the objective (2).
All methods, including FedSim,
FedAlt and the baselines are initialized with a global model trained with FedAvg. [Table] [Figure]
4. E XPERIMENTAL R ESULTS
Partial personalization nearly matches full personalization and can sometimes outperform it.
Table 3 shows the average test accuracy across all devices of different FL algorithms.
We see that on the StackOverﬂow dataset, output layer personalization ( 25 . 05% ) makes up nearly 90% of the gap between the non-personalized baseline ( 23 .
and full personalization (25 . 21%) .
On EMNIST, adapter personalization exactly matches full personalization.
Most surprisingly, on GLDv2, adapter personalization outperforms full personalization by 3 .
5 pp (percentage points).
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 a 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 it harder to get these unseen classes right.
Such catastrophic forgetting is worse when ﬁnetuning on a 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 (Rebufﬁ et al., 2017).
Partial personalization only requires a fraction of the parameters to be personalized.
Figureshows that the number of personalized parameters required to compete with full model personalization is rather small.
On StackOverﬂow, personalizing 1 .
2% of the parameters with adapters captures 72% of the accuracy boost from personalizing all 5 .
7 M parameters; this can be improved to 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% of parameters.
The best personalized architecture is model and task dependent.
Table 3 shows that personalizing the ﬁnal transformer layer (denoted as “Output Layer”) achieves the best performance for StackOverﬂow, while the residual adapter achieves the best performance for GLDv2 and EMNIST. This shows [Table]
that the approach of personalizing a ﬁxed model part, as 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 ﬁnetuning matches the performance of pFedMe and Ditto on StackOverﬂow and EMNIST.
On GLDv2, however, pFedMe outperforms ﬁnetuning by 0 .
07 pp, but is still 3 .
5 pp worse than adapter personalization.
FedAlt outperforms FedSim for partial personalization.
If the optimization problem (3) were 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 is almost always better than
FedSim by a small margin, e.g., 0 .
08 pp for StackOverﬂow/Adapter and 0 .
3 pp for GLDv2/Input Layer.
FedSim in turn yields a higher accuracy than simply ﬁnetuning the personalized part of the model, by a large margin, e.g., 0 .
12 pp for StackOverﬂow/Output Layer and 2 .
55 pp for GLDv2/Adapter.
4. E FFECTS OF PERSONALIZATION ON PER - DEVICE GENERALIZATION
Personalization hurts the test accuracy on some devices.
Figure 4 shows the change in training and test accuracy of each device, compared with a non-personalized model trained by FedAvg.
We see that personalization leads to an improvement in training accuracy across all devices, but a reduction in test accuracy on some of the devices over the non-personalized baseline.
In particular, devices whose testing performance is hurt by personalization are mostly on 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 test accuracy also appear on the left side, signaling the beneﬁt of personalization.
Therefore, there is a large variation of results for devices with few samples.
Additional experiments (see Appendix C) show that using (cid:96) 2 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 to a worse average accuracy that is close to using a common model.
Other simple strategies such as dropout also do not ﬁx this issue.
An ideal personalized method would boost performance on most of the devices without causing a reduction in (test) accuracy on any device.
Realizing this goal calls for a 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.
[Figure]
D ISCUSSION
In addition to a much smaller memory footprint than full model personalization and being less susceptible to 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 signiﬁcant (especially when the personal parameters are only a small fraction of the full model), communicating only the shared parameters may have signiﬁcant implications for privacy.
Intuitively, it 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 at the devices.
For example, we speculate that less noise needs to be added to the communicated parameters in order to satisfy differential privacy requirements (Abadi et al., 2016).
This is a very promising direction for future research.
R EPRODUCIBILITY S TATEMENT
For theoretical results, we state and discuss the assumptions in Appendix A.
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 ﬁve 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 8 in Appendix B.
Finally, we will publicly release the code to reproduce the our experimental results.
E THICS S TATEMENT
The proposed framework for partial model personalization is immediately applicable for a range of practical federated learning applications in edge devices such as text prediction and speech recognition.
One of key considerations of federated learning is privacy.
Partial model personalization maintains the all the privacy beneﬁts of current non-personalized federated learning systems.
Indeed, our 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 ﬂip side, we also observed in experiments that personalization (both full or partial) leads to a 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 deﬁnitions of fairness in federated learning
Appendix Table of Contents
A Convergence Analysis: Full Proofs
A.1 Review of Setup and Assumptions . . . . . . . . . . . . . . . . . . . . . . . .
A.2 Convergence Analysis of FedSim . . . . . . . . . . . . . . . . . . . . . . . . .
A.3 Convergence Analysis of FedAlt . . . . . . . . . . . . . . . . . . . . . . . . .
A.4 Technical Lemmas . .
B Experiments: Detailed Setup and Hyperparameters
B.1 Datasets, Tasks and Models . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.2 Experimental Pipeline and Baselines . . . . . . . . . . . . . . . . . . . . . . .
B.3 Hyperparameters and Evaluation Details . .
C Experiments: Additional Results
C.1 Ablation: Final Finetuning for FedAlt and FedSim . . . . . . . . . . . . . . . .
C.2 Effect of Personalization on Per-Device Generalization . . . . . . . . . . . . . .
C.3 Partial Personalization for Stateless Devices . . . . . . . . . . . . . . . . . . .
A C ONVERGENCE A NALYSIS : F ULL P ROOFS
We give the full convergence proofs here.
The outline of this section is:
• §A.1: Review of setup and assumptions; • §A.2: Convergence analysis of FedSim and the full proof of Theorem 1; • §A.3: Convergence analysis of FedAlt and the full proof of Theorem 2; • §A.4: Technical lemmas used in the analysis.
A. R EVIEW OF S ETUP AND A SSUMPTIONS
We consider a federated learning system with n devices.
Let the loss function on device i be F i ( u, v i ) , where u ∈ R d 0 denotes the shared parameters across all devices and v i ∈
R
d i denotes the personal parameters at device i .
We aim to minimize the function
[Equation]
where V = ( v 1 , · · · , v n ) is a concatenation of all the personalized parameters.
This is a special case of (3) with the equal per-device weights, i.e., α i =
1 /n .
Recall that we assume that F is bounded from below by F (cid:63) .
For convenience, we reiterate Assumptions 1, 2 and 3 from the main-paper as Assumptions 1 (cid:48) , 2 (cid:48) and 3 (cid:48) below respectively, with some additional comments and discussion.
Assumption 1 (cid:48)
(Smoothness) .
For each device i = 1 , . . . , n , the objective F i is smooth, i.e., it is continuously differentiable and,
(a) u (cid:55)→ ∇ u F i ( u, v i ) is L u -Lipschitz for all v i , (b) v i (cid:55)→ ∇ v F i ( u, v i ) is L v -Lipschitz for all u , (c) v i (cid:55)→ ∇ u F i ( u, v i ) is L uv -Lipschitz for all u , and, (
d) u (cid:55)→ ∇ v F i ( u, v i ) is L vu -Lipschitz for all v i .
Further, we assume for some χ > 0 that
[Equation]
The smoothness assumption is a standard one.
We can assume without loss of generality that the cross-Lipschitz coefﬁcients L uv , L vu are equal.
Indeed, if F i is twice continuously differentiable, we can show that L uv , L vu are both equal to the operator norm (cid:107)∇ 2 uv F i ( u, v i ) (cid:107) op of the mixed second derivative matrix.
Further, χ denotes the extent to which u impacts the gradient of v i 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 (cid:48)
(Bounded Variance) .
Let D i denote a probability distribution over the data space Z on device i .
There exist functions G i,u and G i,v which are unbiased estimates of ∇ u F i and ∇ v F i respectively.
That is, for all u, v i :
[Equation]
Furthermore, the variance of these estimators is at most σ 2 u and σ 2 v respectively.
That is,
[Equation]
In practice, one usually has G i,u ( u, v i , z )
= ∇ u f i (( u, v i ) , z ) , which is the gradient of the loss on datapoint z ∼ D i under the model ( u, v i ) , and similarly for G i,v .
Finally, we make a gradient diversity assumption.
Assumption 3 (cid:48)
(Partial Gradient Diversity) .
There exist δ ≥ 0 and ρ ≥ 0 such that for all u and V ,
[Equation]
Sample data z t i k
i 7: i
1 = t i
− γ v G i v i k v t )
i k
z t ) i k ) 8: u t i k
1 = u t i k − γ u G
i u
i k
v i k
Update v
1) i = v t
i τ u
=
Update u
= i ∈
S t α u
i i aggregation 11: return u
T v v n
This assumption is analogous to the bounded variance assumption (Assumption 2 (cid:48) ), 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 et al., 2020).
Finally, it sufﬁces for the partial gradient diversity assumption to only hold at the iterates ( u ( t ) , V ( t ) ) generated by either FedSim or FedAlt.
A. C ONVERGENCE A NALYSIS OF F ED S IM
We give the full form of FedSim in Algorithm 4 for the general case of unequal α
i ’s but focus on α i = 1 /n for the analysis.
In order to simplify presentation, we denote V ( t ) =
( v ( t ) 1 , . . .
, v ( t ) n ) and deﬁne the following shorthand for gradient terms
[Equation]
For convenience
[Equation]
we restate Theorem 1 from the main paper.
Theorem 1 ( Convergence of FedSim ) .
Suppose Assumptions 1, 2 and 3 hold and the learning rates in FedSim are chosen as γ u = η/ ( L u τ ) and γ v = η/ ( L v τ ) with η
[Equation]
Then, ignoring absolute constants
[Equation]
we have
[Equation]
Before proving the theorem
[Equation]
we give the following corollary with optimized learning rates.
Corollary 3.
Consider the setting of Theorem 1 and let ε > 0 be given.
Suppose we set the learning rates γ u = η/ ( τL u ) and γ v = η/ ( τL v ) , where (ignoring absolute constants),
[Equation]
We have,
[Equation]
after T communication rounds, where, ignoring absolute constants,
[Equation]
Proof.
The choice of the constant η ensures that each of the constant terms in the bound of Theoremis O ( ε ) .
The ﬁnal rate is now O (cid:0) ∆ F 0 / ( ηε ) (cid:1) ; plugging in the value of η completes the proof.
We now prove Theorem
[Equation]
Proof of Theorem 1.
The proof mainly applies the smoothness upper bound to write out a descent condition with suitably small noise terms.
We start with some notation.
Notation.
Let F ( t ) denote the σ -algebra generated by (cid:0) u ( t ) , V ( t ) (cid:1) and denote E t [ · ] =
E [ ·|F ( t ) ] .
For all devices, including those not selected in each round, we deﬁne virtual sequences ˜ u ( t ) i,k , ˜ v ( t ) i,k as the
SGD updates in Algorithm 4 for all devices regardless of whether they are selected.
For the selected devices k ∈ S ( t ) , we have (cid:16) u ( t ) i,k , v ( t ) i,k (cid:17) = (cid:16) ˜ u
( t ) i,k , ˜ v ( t ) i,k (cid:17) .
Note now that the random variables ˜ u ( t ) i,k , ˜ v ( t ) i,k are independent of the device selection S ( t ) .
The updates for the devices i ∈ S ( t ) are given by
[Equation]
and the server update is given by
[Equation]
Proof Outline.
We use the smoothness of F i , more precisely Lemma 16, to obtain
F
[Equation] [Equation]
Our goal will be to bound each of these terms to get a descent condition from each step of the form
E t
[Equation]
where the O ( γ 2 u + γ 2 v ) terms are controlled using the bounded variance and gradient diversity assumptions.
Telescoping this descent condition gives the ﬁnal bound.
Main Proof.
Towards this end
[Equation]
we prove non-asymptotic bounds on each of the terms T 1 ,v , T 1 ,u ,
T 2 ,v and T 2 ,u , in Claims 4 to 7 respectively.
We then invoke them to get the bound
[Equation]
Note that we simpliﬁed some constants appearing on the gradient norm terms using
[Equation]
and
[Equation]
v
[Equation]
Our next step is to bound the last two lines of (11) with Lemma 8 and invoke the gradient diversity assumption (Assumption 3 (cid:48) ) as
[Equation]
This gives, after plugging in the learning rates and further simplifying the constants,
[Equation]
Taking full expectation
[Equation]
telescoping the series over t = 0 , · · · , T − 1 and rearranging the resulting terms give the desired bound in Theorem 1.
Claim 4 (Bounding T 1 ,v ) .
Let T 1 ,v be deﬁned as in (10) .
We have,
[Equation]
τm 2 nn (cid:88)
i =
(cid:13)(cid:13)(cid:13) ∇ v F i (cid:16) u ( t ) , v ( t ) i (cid:17)(cid:13)(cid:13)(cid:13)+ γ v m n n (cid:88) i =
τ − 1 (cid:88) k
=
E t (cid:20) χ
[Equation]
L u L v (cid:13)(cid:13) (cid:13)
˜ u
( t )
i,k − u ( t ) (cid:13)(cid:13) (cid:13)
+ L 2 v (cid:13)(cid:13) (cid:13)
˜ v ( t ) i,k − v ( t ) i (cid:13)(cid:13) (cid:13) 2 (cid:21) .
Proof.
Deﬁne T 1 ,v,i to be contribution of the i th term to T 1 ,v .
For i / ∈ S t , we have that T 1 ,v,i = 0 , since v ( t +1)
i = v ( t ) i .
On the other hand, for i ∈ S ( t ) , we use the unbiasedness of the gradient estimator
[Table]
[Table] [Table]
[Table] [Table]
[Equation]
(cid:17)(cid:13)(cid:13)(cid:13)
[Equation]
For the second term, we add and subtract ∇ v F i (cid:16) u ( t ) , ˜ v ( t ) i,k (cid:17) and use smoothness to get
[Equation]
∇
[Equation]
F i
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
(cid:13)(cid:13)(cid:13) ˜ v ( t )
i
[Equation]
v ( t ) i (cid:13)(cid:13)(cid:13) 2 .
[Equation]
[Equation]
Since the right hand side of this bound is independent of S t , we get,
[Equation]
and plugging in (12) and (13) completes the proof.
Claim 5 (Bounding T 1 ,u ) .
Consider T 1 ,u deﬁned in (10) .
We have the bound,
[Equation]
n n (cid:88) i =
τ − 1 (cid:88) k
=
[Equation]
(cid:20)
L 2 u (cid:13) (cid:13)(cid:13) ˜ u ( t )
i
[Equation]
k −
[Equation]
( t ) (cid:13) (cid:13)(cid:13)
[Equation]
[Equation]
u L v (cid:13) (cid:13)(cid:13)
˜ v ( t ) i,k − v ( t ) i (cid:13) (cid:13)(cid:13) 2 (cid:21) .
Proof.
Due to the independence of S ( t ) from ˜ u ( t ) i,k , ˜ v ( t ) i,k , we have,
[Equation]
where the last equality took an expectation over S ( t )
[Equation]
which is independent of ˜ u ( t ) i,k , ˜ v ( t ) i,k .
Now, using the same sequence of arguments as Claim
[Equation]
we have,
[Equation]
where the inequality ( ∗ ) follows from Jensen’s inequality as
[Equation]
Claim 6 (Bounding T 2 ,v ) .
Consider T 2 ,v as deﬁned in (10) .
We have the bound,
E t [ T 2 ,v ] ≤ 3 L v (1 + χ 2 ) γ 2 v τ 2 m 2 nn (cid:88) i = (cid:13)(cid:13)(cid:13) ∇ v F i (cid:16) u ( t ) , v ( t ) i (cid:17)(cid:13)(cid:13)(cid:13)
+ L v (1 + χ 2 ) γ 2 v τ 2 mσ 2 v 2 n +
[Equation]
γ
[Equation]
v τm 2 n
[Equation]
(cid:88) i =
τ − 1 (cid:88) k
=
E t (cid:20)
L 2 v (cid:13)(cid:13)(cid:13) ˜ v ( t ) i,k − v ( t ) i (cid:13)(cid:13)(cid:13) 2 + χ 2 L u L v (cid:13)(cid:13)(cid:13) ˜ u
( t ) i,k − u ( t ) (cid:13)(cid:13)(cid:13) 2 (cid:21) .
Proof.
We start with
E
[Equation]
=
[Equation]
(cid:17)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) ≤ γ
2 v τ τ − 1 (cid:88) k
=
E t
[Table]
i,
[Table]
u
[Table]
,k ,
[Table]
[Table] [Table]
,k
[Table]
≤ γ 2 v τ 2 σ 2 v + γ 2 v τ τ − 1 (cid:88) k
=
E t (cid:13)(cid:13) (cid:13) ∇ v F i (cid:16) ˜ u ( t ) i,k , ˜ v ( t ) i,k (cid:17)(cid:13)(cid:13) (cid:13)≤ γ 2 v τ 2 σ 2 v + 3 γ 2 v τ 2 (cid:13)(cid:13)(cid:13) ∇ v F i (cid:16) u
[Table]
) , v
[Table]
t )
i
[Table]
+ 3 γ 2 v τ τ − 1 (cid:88) k
=
E t (cid:20)
L 2 v (cid:13)(cid:13)(cid:13)
[Table]
v (
[Table]
i,k − v ( t ) i (cid:13)(cid:13)(cid:13)
+ χ 2 L u L v (cid:13)(cid:13)(cid:13) ˜
u ( t ) i,k − u ( t ) (cid:13)(cid:13)(cid:13) 2 (cid:21) .
Using
E t [ T 2 ,v ] =
L v (1 + χ 2 ) m 2 n
E t   1 m (cid:88) i ∈ S ( t ) (cid:13)(cid:13) (cid:13)
˜ v ( t ) i,τ − v ( t ) i (cid:13)(cid:13) (cid:13)  ≤ L v (1 + χ 2 ) m 2 nn (cid:88)
i =
E t (cid:13)(cid:13)(cid:13) ˜ v ( t ) i,τ − v ( t ) i (cid:13)(cid:13)(cid:13)
Plugging in the bound E t (cid:13)(cid:13)(cid:13)
˜ v ( t ) i,τ − v ( t ) (cid:13)(cid:13)(cid:13) 2 completes the proof.
Claim 7 (Bounding T 2 ,u ) .
Consider T 2 ,u as deﬁned in (10) .
We have,
[Equation]
Proof.
We proceed with the ﬁrst two inequalities as in the proof of Claim 6 to get
[Equation]
For T 3 ,j , (a) we add and subtract ∇ u F ( u ( t ) , V ( t ) ) and ∇ u F i ( u ( t ) , ˜ v ( t ) i,k ) , (b) invoke the squared triangle inequality, and, (c) use smoothness to get
[Equation]
For the ﬁrst term, we use the fact that S ( t ) is obtained by sampling without replacement to apply
Lemma 17 together with the gradient diversity assumption to get
[Equation]
1 m (cid:88) i ∈ S ( t ) ∇ u F i (cid:16) u ( t ) , v ( t ) i (cid:17) − ∇ u F (cid:16) u ( t ) , V ( t ) (cid:17)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) ≤m (cid:18) n − m n −(cid:19)n n (cid:88) i = (cid:13)(cid:13)(cid:13) ∇ u F i (cid:16) u ( t ) , v ( t ) i (cid:17) − ∇ u F (cid:16) u ( t ) , V ( t ) (cid:17)(cid:13)(cid:13)(cid:13)≤m (cid:18) n − m n −(cid:19) (cid:18)
δ 2 + ρ 2 (cid:13)(cid:13)(cid:13) ∇ u F (cid:16) u ( t ) , V ( t ) (cid:17)(cid:13)(cid:13)(cid:13) 2 (cid:19) .
Therefore,
T 3 ,j = 12 δ 2 m (cid:16) 1 − m n (cid:17) + 6(1 + ρ 2 ) (cid:13)(cid:13)(cid:13) ∇ u F (cid:16) u ( t ) , V ( t ) (cid:17)(cid:13)(cid:13)(cid:13)+n n (cid:88) i =
E t (cid:20)
L 2 u (cid:13)(cid:13)(cid:13) ˜ u ( t ) i,k − u ( t ) (cid:13)(cid:13)(cid:13) 2 + χ 2 L u L v (cid:13)(cid:13)(cid:13)
˜ v ( t ) i,k
[Equation]
v (
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
(cid:21) , where we also used the independence between S ( t ) and (˜ u ( t ) i,k , ˜ v ( t ) i,k ) .
Plugging this into the expression for E t (cid:107) u ( t +1) − u ( t ) (cid:107) 2 completes the proof.
Lemma 8.
Let F i satisfy Assumptions 1 (cid:48) -3 (cid:48) , and consider the iterates
[Equation]
and ,
[Equation]
for k = 0 , · · · , τ − 1 , where z k ∼ D i .
Suppose the learning rates satisfy γ u = c u / ( τL u ) and γ v =
c v / ( τL v ) with c u , c v ≤ 1 / √ 6 max { 1 , χ − 2 } .
Further, deﬁne,
[Equation]
and ,
[Equation]
where f ∈ (0 , 1] is given.
Then, we have the bound,
[Equation]
Proof.
If τ =
[Equation]
there is nothing to prove
[Equation]
so we assume τ >
[Equation]
Let ∆ k :
= A (cid:107) u k − u 0 (cid:107) 2 + B (cid:107) v k − v 0 (cid:107) 2 and denote by F k the sigma-algebra generated by ( w k , v k ) .
Further
[Equation]
let E k [ · ] = E [ ·|F k ]
[Equation]
We use the inequality 2 αβ ≤ α 2 /δ 2 + δ 2 β 2 for reals α, β, δ to get,
[Equation]
where the last inequality followed from the squared triangle inequality (from adding and subtracting ∇ u F i ( u 0 , v k ) and ∇ u F i ( u 0 , v 0 ) ) followed by smoothness.
Together with the analogous inequality for the v -update, we get,
[Equation]
where we have
[Equation]
and ,
[Equation]
Next, we apply Lemma 20 to get that A (cid:48)
≤ A/τ and B (cid:48) ≤ B/τ under the assumed conditions on the learning rates; this allows us to write the right hand side completely in terms of ∆ k and unroll the recurrence.
The intuition behind Lemma 20 is as follows.
Ignoring the dependence on τ, L u , L v , χ for a moment, if γ u and γ v are both O ( η ) , then A (cid:48) , B (cid:48) are both O ( η 3 ) , while A and B are O ( η ) .
Thus, making η small enough should sufﬁce to get A (cid:48) ≤ O ( A ) and B (cid:48) ≤ O ( B ) .
Concretely, Lemma 20 gives
[Equation]
and unrolling this recurrence gives for k ≤ τ −
[Equation]
where we used (1 + 1 /α )
α ≤ e for all α > 0 .
Summing over k and using the numerical bound e 2 < 8 completes the proof.
Remark 9.
We only invoked the partial gradient diversity assumption (Assumption 3) at iterates ( u ( t ) , V ( t ) ) ; therefore, it sufﬁces if the assumption only holds at iterates ( u ( t ) , V ( t ) ) generated by
FedSim, rather than at all ( u, V ) .
Input: Initial iterates u (0) , V (0) , Number of communication rounds T , Number of devices per round m , Number of local updates τ u , τ v , Local step sizes γ u , γ v , 1: for t = 0 , 1 , · · · , T − 1 do 2:
Sample m devices from [ n ] without replacement in S ( t ) 3: for each selected device i ∈ S ( t ) in parallel do 4:
Initialize v ( t ) i, 0 = v ( t ) i 5: for k = 0 , · · · , τ v − 1 do (cid:46) Update personalized parameters 6:
Sample data z ( t ) i,k ∼ D i 7: v ( t ) i,k +1 = v ( t ) i,k − γ v G i,v ( u ( t ) , v ( t ) i,k , z ( t ) i,k ) 8:
Update v ( t +1) i
= v ( t ) k,τ v 9:
Initialize u ( t ) i, 0 = u ( t ) 10: for k = 0 , · · · , τ u − 1 do (cid:46) Update shared parameters 11: u ( t ) i,
k +1 =
u ( t ) i,k − γ u G i,u ( u ( t ) i,k , v ( t +1)
i , z ( t ) i,k ) 12:
Update u ( t +1) i =
u ( t ) i,τ u 13:
Update u ( t +1) =
(cid:80) i ∈ S ( t )
α i u ( t +1)
i (cid:80) i ∈ S ( t )
α
i at the server with secure aggregation 14: return u ( T ) , v ( T ) 1 , · · · , v ( T ) n
A. C ONVERGENCE A NALYSIS OF F ED A LT
We give the full form of FedAlt in Algorithms 5 for the general case of unequal α
i ’s but focus on α i = 1 /n for the analysis.
For convenience, we reiterate Theorem 2 below.
Recall the deﬁnitions
[Equation]
and ,
[Equation]
Theorem 2 ( Convergence of FedAlt ) .
Suppose Assumptions 1, 2 and 3 hold and the learning rates in FedAlt are chosen as γ u =
η/ ( L u τ u ) and γ v = η/ ( L v τ v ) , with
[Equation]
Then, ignoring absolute constants
[Equation]
we have
[Equation]
Before proving the theorem
[Equation]
we have the corollary with optimized learning rates.
Corollary 10.
Consider the setting of Theorem 2 and ﬁx some ε > 0 .
Suppose we set γ u = η/ ( τL u ) and γ v = η/ ( τL v ) such that, ignoring absolute constants,
[Equation]
Then, we have,
[Equation]
after T communication rounds, where, ignoring absolute constants,
[Equation]
Proof.
We get the bound by balancing terms from the bound of Theorem 2.
The choice of η ensures that all the O ( η ) and O ( η 2 ) terms are at most O ( ε ) .
Finally, the smallest number of communication rounds to make the left hand side of the bound of Theorem 2 smaller than ε is ∆ F 0 / ( ηε ) .
We are now ready to prove Theorem 2.
Proof of Theorem 2.
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 (cid:101) ∆ ( t ) u as the analogue of ∆ ( t ) u with the virtual variable (cid:101) V ( t +1) : (cid:101) ∆
[Equation]
t )
[Equation]
F (cid:16) u (
[Equation]
[Equation]
) (cid:17)(cid:13)(cid:13)(cid:13) 2 .
Notation.
Let F ( t ) denote the σ -algebra generated by (cid:0) u ( t ) , V ( t ) (cid:1) and denote E t [ · ] =
E [ ·|F ( t ) ] .
For all devices, including those not selected in each round, we deﬁne virtual sequences ˜ u ( t ) i,k , ˜ v ( t ) i,k as the
SGD updates in Algorithm 5 for all devices regardless of whether they are selected.
For the selected devices i ∈ S ( t ) , we have v ( t ) i,k
=
˜
v ( t ) i,k and u ( t ) i,k = ˜
u ( t ) i,k .
Note now that the random variables ˜ u ( t ) i,k , ˜ v ( t ) i,k are independent of the device selection S ( t ) .
Finally, we have that the updates for the selected devices i ∈ S ( t ) are given by
[Equation]
and the server update is given by
[Equation]
Proof Outline and the Challenge of Dependent Random Variables.
We start with
[Equation] [Equation]
The ﬁrst line corresponds to the effect of the v -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 (cid:0) u ( t ) , · (cid:1) .
The latter is more challenging.
In particular, the smoothness bound for the u -step gives us
[Equation]
The standard proofs of convergence of stochastic gradient methods rely on the fact that we can take an expectation w.r.t.
the sampling S ( t ) of devices for the ﬁrst order term.
However, both V ( t +1) and u ( t +1) depend on the sampling S ( t ) of devices.
Therefore, we cannot directly take an expectation with respect to the sampling of devices in S ( t ) .
Virtual Full Participation to Circumvent Dependent Random Variables.
The crux of the proof lies in replacing V ( t +1) in the analysis of the u -step with the virtual iterate (cid:101)
V ( t +1)
so as to move all the dependence of the u -step on S ( t ) to the u ( t +1) term.
This allows us to take an expectation; it remains to carefully bound the resulting error terms.
Finally, we will arrive at a bound of the form
[Equation]
Next, we translate this bound from gradient E [ (cid:101) ∆ ( t ) u ] of the virtual (cid:101) V ( t +1) to E [∆ ( t ) u ] , which is the gradient computed at the actual iterate V ( t ) .
A careful analysis shows that we only incur a lower order term of O ( γ u γ 2 v ) in this translation.
Choosing γ u and γ v small enough will give us the ﬁnal result.
Analysis of the u -Step with Virtual Full Participation.
We introduce the virtual iterates (cid:101) V ( t +1) into the analysis of the u -step as follows:
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
(cid:124) (cid:123)(cid:122) (cid:125)
[Equation]
u
[Equation]
(cid:88) i = (cid:13)(cid:13)(cid:13)
˜ v ( t +1)
i − v ( t +1) i (cid:13)(cid:13)(cid:13)(cid:124) (cid:123)(cid:122) (cid:125) T 3 ,u .
The last two inequalities follow from Young’s inequality and Lipschitzness of V (cid:55)→ ∇ u F ( u, V ) respectively.
We have now successfully eliminated the dependence of the ﬁrst-order term T 1 ,u on V ( t +1) .
The virtual iterates (cid:101) V ( t +1) are now independent of S ( t ) .
This allows us to take an expectation w.r.t.
the sampling S ( t ) of the devices.
We bound each of these terms in Claims 11 to 13 below to get
E t
[Equation]
Note that we used the fact that 24 L u γ u τ u (1 + ρ 2 ) ≤ 1 to simply the coefﬁcients of some of the terms above.
The second term has also been referred to as client drift in the literature; we bound it with
Lemma 18 and invoke the assumption on gradient diversity (Assumption 3 (cid:48) ) to get
[Equation]
Plugging this back in, we get,
[Equation]
Note that we used 128 γ 2 u L 2 u τ u ( τ u − 1) ρ 2 ≤ 1 , which is implied by 24 L
[Equation]
τ
[Equation] [Equation]
Bound with the Virual Iterates.
We plug this analysis of the u -step and Claim 14 for the v -step into (14) next.
We also simplify some coefﬁcients using 128 γ v τ v L v χ 2 ( n/m − 1) ≤ 1 .
This gives us
E t
[Equation]
Taking an unconditional expectation
[Equation]
summing it over t = 0 to T − 1 and rearranging this gives
[Equation]
This is a bound in terms of the virtual iterates (cid:101) V ( t +1) .
However, we wish to show a bound in terms of the actual iterate V ( t ) .
Obtaining the Final Bound.
It remains now to relate (cid:101) ∆ ( t ) u with ∆ ( t ) u .
Using the Cauchy-Schwartz inequality and smoothness
[Equation]
we have,
[Equation]
, where the last inequality followed from Lemma 19.
Using
[Equation]
we get,
[Equation]
Therefore
[Equation]
we get,
[Equation]
where we used 32 η 2 χ 2 m n ≤ 1 / 2 , which is one of the conditions we assume on η .
Summing this up and plugging in (15) gives
[Equation]
Plugging in γ u = η/ ( L u τ u ) and γ v =
η/ ( L v τ v ) completes the proof.
The analysis of each of the terms in the u -step is given in the following claims.
Claim 11 (Bounding T 1 ,u ) .
We have,
[Equation]
Proof.
For i ∈ S ( t ) , we have that ˜ u
( t )
i,k = u ( t ) i,k .
Therefore, we have,
E
[Equation] [Equation]
[Equation]
γ
[Equation]
E t
[Equation]
i
[Equation] [Equation]
k = ∇
[Equation] [Equation]
( t
[Equation] [Equation]
(cid:17)(cid:43) .
Using that ˜ u ( t ) i,k is independent of S ( t ) , we get,
E
[Equation]
) , (cid:101) V ( t +1) (cid:17) , 1 n n (cid:88) i =
[Equation]
u
[Equation]
(cid:88)
[Equation] [Equation]
[Equation]
(cid:17)(cid:43)
=
− γ u τ u E t (cid:13)(cid:13)(cid:13)
[Equation]
u F (cid:16)
[Equation]
( t ) , (cid:101)
[Equation]
( t +1) (cid:17)(cid:13)(cid:13)(cid:13)− γ u
[Equation]
u
[Equation]
(cid:88) k
[Equation]
E
[Equation]
Invoking (cid:104) x, y (cid:105) ≤ (cid:107) x (cid:107) 2 /
+ (cid:107) y (cid:107) 2 / 2 for vectors x, y followed by smoothness completes the proof.
Claim 12 (Bounding T 2 ,u ) .
We have,
E t [ T 2 ,u ] ≤ 3 L u γ 2 u τ 2 u (cid:18) 1 + 2 ρm (1 − m/n ) (cid:19)
E t (cid:13)(cid:13)(cid:13) ∇ u F (cid:16) u ( t ) , (cid:101) V ( t +1)
(cid:17)(cid:13)(cid:13)(cid:13)+ 3 L 2 u γ 2 u τ u n n (cid:88) i =
τ u − 1 (cid:88) k
=
E t (cid:13)(cid:13)(cid:13)
[Equation]
t
[Equation] [Equation]
[Equation] [Equation]
Proof.
We use E (cid:107) z (cid:107)
= (cid:107) E [ z ] (cid:107)
+ E (cid:107)
z − E [ z ] (cid:107) 2 for a random vector z to get
E t [ T 2 ,u ] ≤ L u γ 2 u τ 2 u σ 2 u m + L u γ 2 u τ u τ u − 1 (cid:88) k
=
E t (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 1 m (cid:88) i ∈ S ( t ) ∇ u F i (cid:16) ˜ u ( t ) i,k , ˜ v ( t +1) i (cid:17)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) (cid:124) (cid:123)(cid:122) (cid:125) =: T (cid:48) k .
We break the term T (cid:48) k as
T (cid:48)
[Equation]
3 (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
i
[Equation]
( t ) (cid:16) ∇ u F
[Equation]
(cid:16)
[Equation]
u ( t
[Equation]
v ( t
[Equation]
1) i (cid:17) − ∇ u F
[Equation]
(cid:16) u ( t ) , ˜ v ( t +1) i (cid:17)(cid:17)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) + 3 (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 1 m (cid:88) i ∈ S ( t ) ∇ u F i (cid:16) u ( t ) , ˜ v ( t +1)
i (cid:17) −
[Equation]
For the ﬁrst term, we use Jensen’s inequality to take the squared norm inside the sum, then use smoothness and take an expectation over the sampling of devices to get
[Equation]
For the second term, we use the fact that S ( t ) was sampled without replacement (cf.
Lemma 17) and invoke the gradient diversity assumption (Assumption 3 (cid:48) ) to get,
[Equation]
[Equation]
(cid:101) V ( t +1) (cid:17) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
n
[Equation]
mn n
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
( t
[Equation]
(cid:17)
[Equation]
(cid:16)
[Equation]
(cid:101)
[Equation]
( t +1) (cid:17)(cid:13)(cid:13)(cid:13)
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
(cid:17)(cid:13)(cid:13)(cid:13)
[Equation]
To complete the proof, we plug these terms back into the deﬁnition of T (cid:48)
k and E t [ T 2 ,u ] to complete the proof.
Claim 13 (Bounding T 3 ,u ) .
We have,
[Equation]
m n (cid:17) .
Proof.
Since v ( t +1) i = ˜ v ( t +1) i for i ∈ S ( t ) , we have that
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
+
[Equation] [Equation]
(cid:13)(cid:13)(cid:13) 2 .
Since (cid:13)(cid:13)(cid:13) ˜ v ( t +1)
i − v ( t ) i (cid:13)(cid:13)(cid:13) 2 is independent of S ( t ) , we can take an expectation to get
[Equation]
i
[Equation]
t
[Equation]
t +1)
i
[Equation]
(cid:13)(cid:13)(cid:13) 2 .
Plugging in Lemma 19 completes the proof.
The analysis of the v -step is given in the next result.
Claim 14.
Consider the setting of Theorem 2 and assume that γ v τ v L v ≤ 1 / 8 .
We have,
[Equation]
Proof.
From smoothness, we get,
[Equation]
We bound the ﬁrst term as
[Equation]
[Table]
[Equation]
[Table]
Next, we observe that
[Equation]
i
[Equation]
u
[Equation]
v i )
[Equation]
We invoke this inequality to handle the second term as
[Equation]
( t ) , v
[Equation]
t )
i (cid:17)(cid:13)(cid:13)(cid:13)
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
(cid:13)(cid:13)(cid:13)
[Equation]
Plugging these bounds for T 1 ,v and T 2 ,v into the initial smoothness bound and using γ v L v τ v ≤ 1 /gives
[Equation]
We invoke Lemma 18 to bound the (cid:80) k
E t (cid:107) ˜ v ( t ) i,k − v ( t ) i (cid:107) 2 term, which is also known as client drift.
We simplify some coefﬁcients using 8 γ v τ v L v ≤ 1 to get
[Equation]
(cid:16) u ( t ) , v ( t )
i (cid:17)(cid:13)(cid:13)(cid:13) 2 + γ 2 v L v τ 2 v σ 2 v+ 4 γ 3 v L v τ 2 v ( τ v − 1) σ 2 v .
It remains to invoke that S ( t ) is a uniformly random sample of m devices from { 1 , · · · , n } and that ˜ v ( t +1)
i is independent of S ( t ) .
To this end, note that
[Equation]
Plugging in the previous bound completes the proof.
Remark 15.
We only invoked the partial gradient diversity assumption (Assumption 3) at (virtual) iterates ( u ( t ) , (cid:101) V ( t +1) ) ; therefore, it sufﬁces if the assumption only holds at iterates ( u ( t ) , (cid:101)
V ( t +1) ) generated by FedAlt, rather than at all ( u, V ) .
A. T ECHNICAL L EMMAS
The ﬁrst lemma involves smoothness of two blocks of variables; we use this in the proof of FedSim.
Lemma 16 (Block Smoothness) .
Suppose F i : R d × R d
i satisfy Assumption 1 (cid:48) .
Then, it holds that
[Equation]
Proof.
Using the L w -smoothness of F ( · , v (cid:48) i ) and the L v -smoothness of F ( w, · ) , we have
[Equation]
Summing the above two inequalities together gives
[Equation]
We can bound the ﬁrst inner product term on the right-hand side of the above inequality as
[Equation]
where the ﬁrst inequality is due to Cauchy-Schwarz, the second inequality is due to L wv -Lipschitz property of ∇ w F i ( w, · ) , the third inequality is due to the deﬁnition of χ in (5), and the last inequality is due to Young’s inequality.
Substituting the above inequality into (16) yields the desired result.
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 a 1 , · · · , a n ∈ R d be given.
Let S be a uniformly random sample of size m from this collection, where the sampling is without replacement.
Denoting the mean ¯ a = (cid:80) ni
=1 a i /n , we have,
[Equation]
Proof.
The statement is trivially true if m = 1 or m = n .
Therefore, we assume now that 2 ≤ m ≤ n − 1 .
Further, without loss of generality, we assume that ¯ a = 0 .
Finally, let S denote the set of all subsets of [ n ] of size m .
Note that |S| = (cid:0) nm (cid:1) .
We now have,
[Equation]
For the ﬁrst term, we have, (cid:88)
[Equation]
Likewise, for the second term, we use (cid:80) j (cid:54) =
i a j = − a i to get, (cid:88) i,
[Equation]
=
j
[Equation]
a
[Equation] [Equation]
= n (cid:88)
[Equation]
=
[Equation]
(cid:88)
[Equation]
(cid:88)
[Equation]
(cid:104)
[Equation]
Therefore, we get,
[Equation]
The next two lemmas are about the effect of the local updates in the local SGD literature.
The ﬁrst lemma has also appeared in (Karimireddy et al., 2020); we give the proof for completeness.
Lemma 18.
Consider f : R d → R which is L -smooth and ﬁx a w (0) ∈ R d .
Deﬁne the sequence ( w ( t ) ) of iterates produced by stochastic gradient descent with a ﬁxed learning rate γ starting from w (0) :
[Equation]
where g ( t ) is an unbiased (and independent of w ) estimator of ∇ f ( w ) with bounded variance σ 2 .
Fix a number τ of steps.
If γ ≤ ( √ 2 τL ) − 1 , we have the bound
[Equation]
Proof.
If τ = 1 , we have nothing to prove.
Assume now that τ ≥ 2 .
Let F ( t ) be the sigma-algebra generated by w ( t ) and denote E t [ · ]
= E [ · |F ( t ) ] .
We will use the inequality
[Equation]
We now successively deduce,
[Equation] [Equation]
[Table]
[Equation]
[Table]
[Equation]
Above, we used (a) the inequality 2 αβ ≤ α 2 /δ 2 + δ 2 β 2 for reals α, β, δ , (b) Eq.
(17), (c) Lsmoothness of f , and, (d) the condition on the learning rate.
Let C = 2 γ 2 τ (cid:107)∇ f ( w (0) )
(cid:107) 2 + γ 2 τσ 2 .
Unrolling the inequality and summing up the series gives for all t ≤ τ −
[Equation]
where we used the bound (1 + 1 /α )
α ≤ e for all α > 0 .
Summing over t and using the numerical bound e 2 < 8 completes the proof.
Lemma 19.
Consider the setting of Lemma 18.
If γ ≤ (2 τL ) − 1 , we have the bound
[Equation]
Proof.
Proceeding similar to the last proof (expect using δ = τ ) gives us
[Equation]
Unrolling and summing up the sequence completes the proof, similar to that of Lemma 18.
Finally, the last lemma is about constants.
Lemma 20.
Let γ u , γ v , L w , L v , χ, f ∈ R + and a natural number τ be given.
Denote
[Equation]
and ,
[Equation]
Suppose γ u = c u / ( τL u ) and γ v
=
c v / ( τL v ) with c u , c v > 0 satisfying
[Equation]
Then, we have that
[Equation]
and ,
[Equation]
Proof.
Note that it sufﬁces to show
[Equation]
and ,
[Equation]
Plugging in γ u , γ v , these are equivalent to
[Equation]
and ,
[Equation]
The assumption on c v implies that 6 χ 2 fc 3 v ≤ χ 2 fc v and 6 χ 4 c 2 v c u ≤ c u .
Therefore, the ﬁrst condition holds.
Similarly, the second condition holds too.
B E XPERIMENTS : D ETAILED S ETUP AND H YPERPARAMETERS
We conduct our experiments on three datasets from two modalities, namely images and text.
The datasets contain a natural, non-iid split of data which is reﬂective of data heterogeneity encountered in 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:
• §B.1 describes the tasks and their associated datasets and metrics.
• §B.2 describes the experimental pipeline as well as the baselines we compare to.
• §B.3 presents the hyperparameters of all the algorithms.
As discussed in §1, we take the weight α k to be proportional to the number of datapoints available on the device.
[Figure]
B. D ATASETS , T ASKS AND M ODELS
We consider three tasks motivated by real-world applications of federated learning.
The tasks are summarized in Table 2 of the main paper and the distribution of data across the clients is 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 ﬁrst 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 ﬁrst transformer layer in place of the word embedding layer owing to its large size.
(b) Output layer personalization : Motivated by Collins et al.
, this architecture learns a shared global representation but personalizes the prediction layer.
For the transformer model, we use the last transformer layer in place of the last prediction layer owing to its large size.
(c) Adapter personalization : We also consider a novel partial personalization architecture, where the full model is shared among all 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 Houlsby et al.
and residual adapters proposed by Rebufﬁ et al.
B.1. S TACK O VERFLOW FOR N EXT W ORD P REDICTION
Dataset.
The StackOverﬂow dataset comprises of questions and answers from the programming question-answer website stackoverﬂow.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 a good opensource benchmark for next word predictions in mobile keyboards.
We use the StackOverﬂow dataset provided by TensorFlow Federated.
Client Distributions.
Each client corresponds to one user on Stack Overﬂow; the data on the client corresponds to the questions and answers posted by this user.
We only consider clients with at least 100 training sequences and 10 testing sequences, where a sequence refers to either a question or an answer.
We use a ﬁxed subsample of 1000 of them.
Following Reddi et al.
(2021), we restrict the vocabulary to the top 10000 most frequently occurring words in the dataset.
We pad and truncate each sequence of each client to length 20 and consider at most 1000 training sequences on each client.
Model.
We use a transformer model (Vaswani et al., 2017) commensurate in size with BERT
Mini (Turc et 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 a fully-connected hidden dimension of 1024 .
The output layer is a causal language modeling head, i.e., a fully connected layer which assigns
[Table]
[Table]
a 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 summarized in Table 5.
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 a multiclass classiﬁcation problem to minimize the multinomial logistic loss, also known as cross entropy loss.
For evaluation, we 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. GLD V 2 FOR V ISUAL L ANDMARK R ECOGNITION
Dataset.
GLDv2 stands for Google Landmarks Dataset v2 (Weyand et al., 2020), which is a largescale 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 a scenario where smartphone users take photos of natural and architectural landmarks while traveling.
We use the federated version of the GLDv2 dataset introduced by Hsu et al.
(2020) 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 on each client as a testing set.
Model.
We use a ResNet-18 (He et al., 2016) model pretrained on ImageNet (Deng et al., 2009), with group normalization instead of batch normalization (Hsieh et al., 2020).
We resize all images to 224 × 224 .
We use two data augmentations for training: a random crop from 256 × 256 and a random horizontal ﬂip.
The model has 12 million parameters, which require around 49 megabytes of storage.
Partial Personalization Architecture.
The partial personalization architectures used are summarized in Table 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 classiﬁcation accuracy.
[Table]
B.1. EMNIST FOR C HARACTER R ECOGNITION
Dataset.
EMNIST (Cohen et al., 2017) is a 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 handwrote the digit/letter during the data collection process.
We only use those clients with at leasttraining points and 25 testing points: there are 1114 of such clients.
Model.
We use a ResNet-18 (He et al., 2016) model with group normalization instead of batch normalization (Hsieh et al., 2020).
We make two modiﬁcations to handle the smaller image size ( 28 × 28 × 1 as opposed to the 224 × 224 × 3 which the original ResNet was designed to accept): (a) we use a convolutional kernel of size 3 × 3 rather than the original 7 × 7 in the ﬁrst convolution layer, and, (b) we drop the ﬁrst 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 GLDv due to the architectural modiﬁcations we make for smaller images as well as the smaller number of classes.
Partial Personalization Architecture.
The partial personalization architectures used are summarized in Table 7.
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 classiﬁcation accuracy.
B. E XPERIMENTAL P IPELINE AND B ASELINES
There are three components in the training pipeline for all experiments:
(a) Non-personalized federated training: The ﬁrst step involves training a global model w g using the one-model-ﬁts-all approach of (1) with FedAvg variants.
(b) Personalized federated training: This optional second step involves training the shared parameters w together with the personalized parameters v k using a personalized federated learning approach.
We warm-start w, v k from the non-personalized model w g from the previous step.
(c) Final ﬁnetuning: The last step involves only ﬁnetuning the personalized parameters v k while the shared parameters w remain unchanged.
For step (b), we initialize v k for each k to be the appropriate part of w g for input/output layer personalization.
On the other hand, for adapters, we initialize v k to be equal to the same set of randomly initialized weights for each device k .
We consider the following baselines:
• Non-personalized : This denotes the performance of step (a) of the pipeline above,
i
e., nonpersonalized federated training with FedAvg variants.
• Full model personalization : We consider three baselines of personalization of the full model: (i) Finetune : The non-personalized model from step (a) of the pipeline above is ﬁnetuned locally on each client (step (c) of the pipeline).
Step (b) is skipped for this baseline.
[Table]
(ii) Ditto (Li et al., 2021):
The non-personalized model from step (a) of the pipeline above is ﬁnetuned locally on each client (step (c) of the pipeline) with (cid:96) 2 regularization (cid:107) v − w g (cid:107) 2 .
Step (b) is skipped for this baseline.
(iii) pFedMe
(Dinh et al., 2020):
The non-personalized baseline model from step (a) is trained further in step (b) to optimize (2) using the pFedMe algorithm of Dinh et al.
Finally the resulting model w is ﬁnetuned locally in step (c).
• Partial Model Personalization : We consider partial model personalization with three different
architectures, as deﬁned in §B.1.
For each personalization approach, we start with the nonpersonalized model in step (a), continue personalization in step (b) using either FedAlt or FedSim as the algorithm, and ﬁnally run step (c) for the local ﬁnetuning.
B. H YPERPARAMETERS AND E VALUATION D ETAILS
All the tuning of hyperparameters was performed on validation data, formed by holding out 20% of the training data on 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 of 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.1 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 §1.
The same metric is used for hyperparameter tuning and is reported in all the tables and plots, unless explicitly noted otherwise.
The ﬁnal hyperparameters we use are given in Table 8.
Rounds.
We start with the number of communication rounds (i.e., the number of calls to secure aggregation routine for the shared parameters), which is used to measure the progress of each algorithm.
For the non-personalized training, we use 1000 rounds for StackOverﬂow, 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 StackOverﬂow 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 et al., 2020), we select 50 devices per round for
[Table]
StackOverﬂow/GLDv2 and 10 per round for
, for the non-personalized as well as the
Local Updates and Minibatch Size.
Each selected device locally runs 1 epoch of mini-batch stochastic gradient descent locally for non-personalized as well as personalized federated training.
The ﬁnal ﬁnetuning at the end of personalized training is performed for 5 epochs.
We use a minibatch size of 64 for StackOverﬂow/GLDv2 and 32 for EMNIST for all settings.
Server and Client Optimizer Details.
We use FedAvg for EMNIST and FedAdam (Reddi et al., 2021) for StackOverﬂow and GLDv2.
We also use a global scheduler, which applies a schedule on the client learning rates across rounds, while the client learning rate within each round is held constant.
We use either a 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 0 .
An exponential scheduler halves the client learning rate once every ﬁxed number of rounds.
Both the client and server learning rates are tuned using the validation set.
Regularization Coefﬁcient for pFedMe and Ditto.
We tune the regularization coefﬁcient λ k = λ for pFedMe and Ditto using the validation data from the set { 10 − 4 , 10 − 3 , · · · , 10 0 } of possible values.
The tuned values are:
• StackOverﬂow: 10 − 3 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 E XPERIMENTS : A DDITIONAL R ESULTS
We now present the detailed experimental results.
We start with Table 9, which shows the standard deviations of Table 4.
We see that the numbers are fairly consistent across runs so the standard deviation is quite small ( 0 . 01 to 0 .
05 percentage points).
C. A BLATION : F INAL F INETUNING FOR F ED A LT AND F ED S IM
We now study the effect of the ﬁnal ﬁnetuning (step (c) of the experimental pipeline; cf.
§B.2) for
FedAlt and FedSim.
The ﬁnal ﬁnetuning has a minimal impact on partial personalization.
We see from Table 10 that the effect of the ﬁnal ﬁnetuning is much smaller than the improvements from personalization.
For instance, the improvements from ﬁnetuning are close to 0 for FedAlt on the StackOverﬂow dataset.
For GLDv2, the ﬁnetuning accounts for < 0 .
5 pp of improvement, whereas personalization overall accounts for 5 to 15 pp.
The ﬁnal ﬁnetuning is more important to FedSim than FedAlt.
Table 10 also shows that the ﬁnal ﬁnetuning helps FedSim more than FedAlt.
However, FedAlt still outperforms FedSim, as we saw in
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
Table 9.
Overall, this shows that FedAlt is a better algorithm than FedSim.
The ﬁnal ﬁnetuning helps
FedSim make up some percentage points in accuracy, but not enough to make up its gap with FedAlt.
C. E FFECT OF P ERSONALIZATION ON P ER -D EVICE G ENERALIZATION
Summary of all scatter plots.
All the scatter plots shown in the main paper are summarized in the violin plot of Figure 6.
We see from the leftmost ﬁgure that the training accuracies on all devices improve with personalization.
From the second ﬁgure, we see that the test accuracy of some of the devices reduces with personalization; this is true for both partial and full personalization.
From the third plot of Figure 6, we see that regularization does not mitigate this overﬁtting.
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 so at the expense of a 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.
Train Accuracy plots for devices.
From Figure 7, we see that personalization leads to a reduction in test accuracy on some of the devices beyond the initial non-personalized model.
The corresponding train accuracy plot is given in Figure 7.
We observe that the personalization always leads to an improvement in the training accuracy but not in the test accuracy.
The analogous plots for GLDv are in Figure 8, where the trends are similar.
Whether personalization helps a device or not depends on the random seed.
We see in Figurethat the shaded region for some of the devices intersects the dotted line at 0 .
In other words, personalization sometimes helps this device and sometimes hurts it, depending on the random seed.
This indicates that the best ﬁx in practice is to use A/B testing on the deployed model to choose whether to use the personalized model or the non-personalized one.
[Figure] [Figure]
Regularization and dropout do not mitigate this issue.
From the ﬁrst 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 a worse mean accuracy and does not mitigate the issue of personalization hurting individual devices.
From the second row of Figure 9, we see that the best dropout ( 0 . 3 in this case) leads to slight increase in average accuracy ( 0 .
18 pp).
It also reduces the number of devices hurt by personalization fromout of 1000 to 193 , but it does not ﬁx this issue.
Increasing dropout further only leads to a degradation of per-device statistics.
C. P ARTIAL P ERSONALIZATION FOR S TATELESS D EVICES
The algorithms we considered in this paper, namely FedAlt and FedSim, require the devices to maintain the personalized parameters v i ’s as state across rounds.
In cross-device federated learning settings, it is also interesting to consider stateless devices, which are not allowed to maintain state between training rounds.
[Figure]
[Figure]
We give preliminary experiments in this setting.
We modify the FedAlt and FedSim algorithms from the main paper so that the personalized parameters v i are reinitialized each time device i is chosen for participation.
We warm-start v i from the appropriate part of the non-personalized model trained in step (a) of the pipeline.
For adapters, we ﬁx a random initialization once, and reuse it.
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 ﬁnetuning for GLDv2, e.g., 0 .
24 pp with adapters.
The ﬁnal ﬁnetuning is crucial for stateless devices.
We see from Table 12 that the ﬁnal ﬁnetuning accounts for most of improvements in the stateless case.
For instance, for GLDv2, the ﬁnal ﬁnetuning accounts for 11 .
68 and 10 .
42 pp out of a total of 12 .
67 and 11 .
76 pp for FedAlt and FedSim respectively.
However, the personalized federated training (step (b) of the pipeline; cf. §B.2) still leads to an increase in accuracy of 1 to 1 .
34 pp.
[Table] [Table]