Partially Personalized Federated Learning: Breaking the Curse of Data Heterogeneity
Authors: Konstantin Mishchenko
Rustem Islamov
Eduard Gorbunov
Samuel Horváth
Venue: Accepted by TMLR ()

Abstract: We consider a partially personalized formulation of Federated Learning (FL) that strikes a balance between the flexibility of personalization and cooperativeness of global training. In our framework, we split the variables into global parameters, which are shared across all clients, and individual local parameters, which are kept private. We prove that under the right split of parameters, it is possible to find global parameters that allow each client to fit their data perfectly, and refer to the obtained problem as overpersonalized. For instance, the shared global parameters can be used to learn good data representations, whereas the personalized layers are fine-tuned for a specific client. Moreover, we present a simple algorithm for the partially personalized formulation that offers significant benefits to all clients. In particular, it breaks the curse of data heterogeneity in several settings, such as training with local steps, asynchronous training, and Byzantine-robust training.

Published in Transactions on Machine Learning Research (01/2025)
Partially Personalized Federated Learning: Breaking the Curse of Data Heterogeneity
Konstantin Mishchenko Meta, Prance"
konsto. misgeaili com
Rustem Islamov
University of Basel, Switzerlandd'
Eduard Gorbunov
Mohamed hin Zoyed University of Artificial Intelligence, United Arab Emirates
Samuel Horváth
Mohamed bin Zuyed University of Artificial Intellignnce, United Arab Emirates
Reviewed on OpenReview: https: // openreview. net/ forun? id- 4SMMCJ44YY
Abstract
We consider partially personalized formulation of Federated Learning (FL) that strikes balance between the flexibility of personalization and cooperativeness of global training. In our framework, split the variables into global parameters, which are shared across all clients, and individual local parameters, which are kept private. We prove that under the right split of parameters, is possible to find global parameters that allow cach client to fit their data perfectly, and refer to the obtained problem as overpersonalizel. For instance, the shared global parameters can be used to learn good data representations, whereas the personalized layers are fine-tuned for specific client. Moreover, we present simple algo- rithm for the partially personalized formulation that offers significant benefits to all clients. In particular, it breaks the curse of data heterogeneity in several settings, such as training with local steps, asynchronous training, and Byzantine-robust training.
1 Introduction
Federated Learning has emerged promising approach to address data privacy concerns and enable dis- tributed learning in various applications, such healthcare, finance, and mobile devices (Konečný et al., 2016). In this approach, multiple parties collaboratively train a machine learning model without sharing their raw data with each other, instead only sharing the model updates with a central server. Despite its potential, Federated Learning faces significant challenges in optimizing the models due to the hetero- geneity of the client's data, the heterogeneity of the devices involved training, and the communication constraints (McMahan et al., 2017).
Collaborative Federated Learning aims at uniting a number of private data holders in order to find together joint set of parameters x² that minimizes their private loss functions:
find x² s.t. fm(x") arg min (x) for all
(Non-personalizeed FL)
This formulation, however, is often too restrictive as there might exist common model x" that fits all participating clients. FL, therefore, often relics private fine-tuning personalization to find individual
"This work done when Konstantin Mishchenko at Samsung Al Center. This work was done when Rustem Islamov was Master's student Institut Polytechnique Paris.

---

Published in Transactions on Machine Learning Research (01/2025)
in for every client m. This leads to fully personalized formulation of FL:
find si, a..... s.t. fm (m)) arg (m()) for all m.
(Personalized FL)
This formulation offers lot of flexibility at the cost of needing to do more work for every client. Since there is no longer a global model æ trained to perform well on all clients, it is now necessary to find a personalized model for every new client. New clients, therefore, might experience bad performance of the initial model until they obtain enough data to personalize for their needs.
Thus, having no personalization might be too restrictive, while full personalization can be excessively ex- pensive and inefficient. Following recent works on personalization by Arivazhagan al. (2019) and Singhal et al. (2021), strike balance by splitting the parameters into global and private parts and formulate our problem as
find 0", wf, w s.1. Jm (0", wm)) x arg min fm (0, w)
0,00
for all m.
(Partially personalized FL)
Above, H* is the part of the parameters that shared by all clients, while parameters in w are specific client m.
1.1 Motivation and formulation
The main interest of our work is to show that the partially personalized formulation gives significant benefit to all clients as long as there exists a universal l" As a motivating example, consider the phenomenon in deep learning where we can train a big network on one task or dataset and fine-tune only its last few layers when working with new source of data. The layers that are not fine-tuned produce so-called representations. The overall objective is then to minimize (()(0, Xm), Wms Bm), where ((....) is a loss function, (00,-) is the features mapping parameterized by 0, Xm, Ym are the data of the m-th client, and wm fino-tuning parameters. Since the representation should fil all non-adversarial clients and send to clients only 0, we expect the m-th client to find approximately w (0) = argmin (0, w). Therefore, we want to find such that changing only wm is sufficient, which can be written as V Im (0, wx,(() = O, where Vi is the gradient of fm with respect to f that does not differentiate through w (0). Thus, our overall objective is to solve
find Or E Rd s.t. (0*)0") 0 for all m, where (0) 1Ifm(0 (0, w (0))
(1)
The values of private Wmn' 's are not available to the server to preserve the privacy of the clients. Moreover, to enable stateless FL applications where a client cannot maintain a state, we assume that a client might initialize its Wm randomly. To the best of our knowledge, this exact problem has not been studied in optimization literature. can also be seen as a first-order stationarity condition for the overparameterizze objective
min Ef (0, wm)],
6, w. Fm
which connects our work to the previous literate on overparameterization in stochastic optimization (Ma et al., 2018; Vaswani et al., 2019).
1.2 Our objective
Motivated by the scenario of representation learning, we ask about the potential benefits of partial person- alization. In summary, the key aspects of our work are follows:
2

---

Published in Transactions on Machine Learning Research (01/2025)
1. The problem is partially personalized with global parameters 0 (e.g., representations) and lo- cal/private parameters w (e.g., last layers).
2. The data is heterogeneous, but there are sufficiently many parameters in w to make the solution set of problem (1) non-empty. Similar to overparamcterization in deep learning (Allen-Zhu et al., 2019), we call this setting overpersonalized.
3. The clients are stateless by default, i.c., they might not be able to store their private parameters w, as is oflen the case in cross-device FL (Kairouz et al., 2021).
4. The clients should be able rum arbitrary local optimizers to fine-tune the private part of the model.
5. Our particular interest is in identifying the benefit of local training as well as the positive impact of client cooperation.
It natural ask if this objective reasonable at all. As the next proposition shows, it is always feasible if we allocate sufficiently many parameters in w.
Proposition 1. There always exists a split of parameters into global and local parameters 9 w1, W2,... such that there exists 0* satisfying Fm (0*) = 0 for all m.
Proof. The result is trivial: if we set wm. to be all parameters and 0 to be any variable such that fm does not depend on it, then naturally fm (0, w (0) = ming, Jm J6m(', ; w) and (0) = 0 for all
As indicated by Proposition 1, splitting the variables to remove the data heterogeneity is always possible. It non-trivial, on the other hand, to find the right split, as moving all parameters into w will eliminate any cooperation between the clients. The interesting case is, therefore, when includes sufficiently many parameters that affect fm- We leave the question of finding the right split for future work and focus instead the consequences of having it. In particular, we prove, under some mild assumptions, that the partially personalized objective can be cast as nonlinear equation:
find 0" s.t. F(0") 0,
(2)
where F' is a nonlinear operator. Moreover, we design algorithms capable of leveraging this property that demonstrates, for the first time, the provable benefits of partial personalization.
1.3 Motivating challenges
Device heterogeneity. So far we have been discussing how the data heterogeneity is a challenge for feder- ated learning. At the same time, heterogeneity in the devices participating training leads to unbalanced computation time of different clients. This issue also arose in classical distributed optimization, where it was proposed (Niu et al., 2011) to run asynchronous methods that do not ask the participating devices to synchronize their updates.
While asynchronous methods have been popular in practice in many applications, including Federated Learn- ing (Nguyen et al., 2022), their thcoretical utility is very limited in the setting of heterogencous data. This is an expected drawback as when all clients have different functions, clients that participate less will not give us sufficient information, leading to a biased solution. Nevertheless, in this paper, we show that in the context of learning representations, even with heterogeneous data, Asynchronous SGD would converge under arbitrary delays.
Byzantine attacks. The next challenge that we will consider is that of the potential presence of Byzantine clients, which is problem faced by many distributed systems (Lamport et al., 1982; Su & Vaidya, 2016; Lyu et al., 2020). Standard Federated Learning algorithms are vulnerable to Byzantine attacks even in the case of homogencous data. For general heterogencous problems, Byzantine robustness cannot be achieved in general. In this work, we develop a new method for the proposed problem formulation and show its provable robustness to Byzantine attacks, even in the case of heterogeneous data clients.
3

---

Published in Transactions on Machine Learning Research (01/2025)
Table 1: brief overview related work and conceptual differences to our approach. We say that local steps are procably helpful if the method achieves better communication complexity when increasing the number of local steps. method handles heterogeneous data its complexity does not depend on any heterogeneity constant. In "Local Training with GD", clients do not communicate at all and run Gradient Descent using local data (non-distributed method).
Provably helpful
Handles heterogeneous data
Algorithm
Personalized
Stateless
local steps
Local Training with GD
FodAvg (McMahan et al., 2017)
X
FedProx(Li et al., 2020a)
X
X
Scaffnew (Mishchenko et al., 2022b)
X
X
X
FedAlt/FedSim (Pillutla et al., 2022b)
V
X
FFGG (This work)
/
V
Personalization of representations. Beyond learning representations and personalizing the last layers, common framework is to personalize represenlations while the weights in the last layers can be shared. For instance, in medical applications, different scanners or sensors may have different intensities and contrast, resulting in a feature shift (Li et al., 2020b). The labels, however, are shared across hospitals, so we can still train personalized models with shared parameters. In that case, wm would be responsible for aligning the features, for instance, using batch normalization parameters as done by Li et al. (2020b), whereas 0 would represent the rest of the network.
Similarly, in pross-device Federated Learning, the features might change due to use of different phones, giving another motivation for personalizing the representations. As every phone model is equipped with its own version of a photo camera, the layers used to produce the embeddings may need to be readjusted, while the layers deeper in the network can remain the same. Likewise, a robot can be deployed in places with different lighting, so its visual signal may depend on the environment, suggesting that the first layers need to be adjusted. By and large, hardware and environment hetcrogencity will demand changes the part of the network giving us represenlations, whereas the decision layers can remain intact in such scenarios.
On final note, personalizing the first layers, which interact with the data, might boost privacy even than the personalization of the last layers. These layers can be used by clients to obfuscate the data source and distribution while relying on the server to get good universal predictors for generic obfuscated data. As our framework does not assume the exact meaning of 0 and w, all of our results are immediately applicable in the setting of personalizing the representation.
1.4 Related work
Methods for non-personalized FL. Most methods for Federated Learning stem from the Federated Averaging algorithm (FedAvg) of McMahan et al. (2017). FedAvg itself can be seen as variant of Local SGD, which received a lot of attention in the optimization literature (Mangasarian, 1995; Stich, 2019). The speed of convergence of Local SGD heavily depends on the amount of data hetcrogeneity, with much slower rates under arbitrary data dissimilarity (Khaled al., 2020). Other methods, such as FedProx (Li et al., 2020a) and Scaffold (Karimireddy et al., 2020b), were proposed to alleviate the issucs stemming from data heterogeneity, with provable acceleration if clients can maintain local states (Mishchenko et al., 2022b; Grudzien et al., 2022). To the best of our knowledge, however, the case of FL with stateless clients and heterogencous data remains challenging.
Personalization. FedAvg and other classical algorithms Irain a single model for all clients. Jiang el al. (2019) showed that FedAvg good at personalization, but this is an implicit property. Varions works have proposed alternative methods that run algorithms with explicit personalization, for example using model- agnostic meta-learning (Chen et al., 2018; Fallah et al., 2020) or using penalty- based objective (Hanzely & Richtárik, 2020). Unlike us, these papers consider full-model personalization and update all parameters.
4

---

Published in Transactions on Machine Learning Research (01/2025)
Partial personalization. More related to ours, some works split variables into personalized and non- personalized ones. In the context of Meta-Tearning with deep networks, Raghu et al. (2020) established that only the last few layers needed to be adapted for new tasks. The work of Pillutla et al. (2022b) is perhaps the closest to ours as they studied optimization properties of partially personalized FL. In particular, they established convergence of algorithms with alternating and simultaneous updates, but unfortunately their theory does not guarantec any benefit from using multiple gradient steps beyond reducing noise due to extra sampling.
Asynchronous optimization. Asynchronous SGD has been extensively studied for homogeneous data (Tsitsiklis et al., 1986; Agarwal & Duchi, 2011). Recently, it was shown that Asynchronous SGD converges regardless of the delays in the setting of homogeneous data (Mishchenko et al., 2022a; Koloskova et al., 2022). Moreover, for heterogeneous data, Asynchronous SGD was to converge to a neighborhood under gradient similarity (Mishchenko et al., 2022a) and to the exact solution under stochastic sampling of clients (Koloskova al., 2022; Islamov et al., 2024). However, to the best of our knowlodge, Asynchronous SGD is not guaranteed to converge for helerogeneous data without assuming structure in the delays.
Byzantine robustness. The first approaches to Byzantine-robust distributed learning were concentrated around the replacing of standard averaging in Parallel SGD by some other aggregation rule, which is more robust to the outliers (Blanchard et al., 2017; Yin et al., 2018; Damaskinos et al., 2019; El-Mhandi et al., 2018; Pillutla et al., 2022a). However, they turned out to be vulnerable to special Byzantine attacks (Baruch et al., 2019; Xie et al., 2020) that can be partially explained by the fact that the term "robust aggregation" " was formally introduced later by Karimireddy et al. (2021). In the case of homogeneous data, this discovery led to the development of provably Byzantine-robust distributed methods (Karimireddy et al., 2021; Gorbunov et al., 2022; 2023). In parallel, alternative approaches based on clients' elimination were developed in (Alistarh et al., 2018; Allen-Zhr et al., 2021).
In the arbitrary heterogeneous data case, it is impossible to distinguish regular clients from Byzantine ones. Moreover, even when the heterogeneity (in terms of the gradients' dissimilarity) is bounded by a constant, there is a lower bounded on the optimization error that can be achieved for a given fraction of Byzantine clients (Karimireddy et al., 2022), matching/neate (Karimireddy et al., 2022; Wu et al., 2020; Zhu & Ling, 2021; Gorbunov et al., 2023). Nevertheless, in the overparameterized regime, some methods can still converge to the exact solution (asymplotically) - this was shown in (Karimireddy et al., 2022; Gorbunov et al., 2023) for standard Federated Learning formulation. In our work, we continue this line of work by showing that for the proposed problem formulation, one can also achieve provable Byzantine robustness even in the heterogeneous case under similar assumptions about overparameterization.
matching upper-bounds were derived in
Bilevel optimization. The partially-personalizee FL formulation can be seen as a special case of bilevel optimization. Recently, bilevel optimization was also studied in the context of FL (Tarzanagh et al., 2022; Li et al., 2022). Tts main disadvantage is the requirement to differentiate through the inner-level objective, which is typically done by estimating inverse-Jacobian-veet -ectoor vector products. Morcover, algorithms for bilevel optimiza- tion usually require maintaining the private variable in memory and updating over multiple communication rounds, which not possible in stateless Federated Learning (Kairouz et al., 2021).
2 Technical Preliminaries
We say that function is L-smooth if its gradient is L-Lipschitz. We use V.fm(ty)) to denote the gradient with respect to the i-th argument of fm- Definition 1. For any client and parameters 0, we denote (i,(0) as any minimizer of the m-th client loss, i.e.,
Notation.
w", (8) E arg min fm. (8,
(3)
Assumption 1. We assume that the operator (0) = Vi w (0)) is + -cccoeereive in 0, i.e., for ony 0,,02
(F, - Fm. (92). 01 - -82) 2 (B) (.2))1
(4)
5

---

Published in Transactions on Machine Learning Research (01/2025)
The best way to view Assumption 1 is as generalization of convexity and smoothness to nonlinear operators. For instance, if Jm does not depend on w, then it is enough for it to be convex and L-smooth (Nesterov, 2013). However, our main interest is when parameters w are personalized, so below, we give a few examples with nontrivial dependence on w where provably holds.
Another assumption that we make regarding the overparameterization aspect of the problem. For example, if good data representations exist for all clients, then we can fine-tune w on each client to perfectly fit the training data. This is formalized below.
There erisls al least one 0" such that for any client m, the loss 50", ,w) achieves ils minimum with some optimal w (0*), i.e. fm(0", w (0)) ming, fm(0, w).
Assumption 2.
Notice that Assumption 2 requires that 0 does not represent personalization to be possible.
Now we give few examples for which Assumption holds. Note that we defer all proofs to the appendix.
Example 1. Let (0, w) be given as fm (0, w) = om (0) + & Al, 0 + Bm w ymll? where 6m is any convex Lg-smooth function. Then, Assumption 1 is satisfied with L = 2 mmxxLL, || (I- (I-BBIBB AAmll)
Example 1 is regularized version of the linear regression problem, which has served as litmus test for verifying if an assumption makes sense. Our assumption is thus validated on this simple example. At the same time, Assumption is satisfied for a wider range of functions that include the following example.
Example 2. Let fm be wifce-iifeerettibble, j-strongly convex and L-smooth in 0. Moreover, assume the cross-term the Hessian and the Jacobian of wx to be bounded as || 1/ (6, w)]l || S C1, and [/Vow (0)ll C2 with C₁O2 S 5. Then, Fm is cocoercive and -tstonngyy monotone.
Example 2 provides general recipe for an objective to satisfy Assumption 1. It is somewhat restrictive it requires smoothness of both f' and wi,, but these assumptions are, in fact, natural. The smoothness of Im is commonly assumed to make minimization of f' possible (Nesterov, 2013). The assumption on the smoothness of w/, (f) is directly related to the assumptions employed in the theory of min-max optimization; see, for example, (Nouiehed et al., 2019). Moreover, it is possible to ensure we, is Lipschitz by adding 전 [|]l| to fm, see Lemma B.1 of Nouiehed et al. (2019).
Let us also add small generalization of Example to the case of the composition of smooth nonlincar functions and linear transformations.
Example Let fm,(,u) w) be given as fm(0, w) = @m + e' (Amm + B, Vm), where t'm and on are any convex and -smoohh functions. Then, Assumption 1 is satisfied.
This example is particularly interesting from the perspective of paramotor-cfficient fine-tuning. Methods such as Low-Rank Adaptation (LoRA) (Hu et al., 2022), fine-tune linear layers by adding low-rank perturbations. Specifically, the linear layer is split into two parts, similar to how we compute A + Bm w in Example 3. Note that B, can have a rank much smaller than Am-
3 New Method: Fine-tuning Followed by Global Gradient (FFGG)
Having defined a formulation of partial personalization, it is very easy to derive an algorithm that will solve it. Under our assumptions, it is enough to update the global parameters 0 using the standard gradient- descent update. However, in practice, we need compute wx (0) using local gradient updates. This gives a double-loop method called Fine-tuning Followed by Global Gradient (FFGG), which we present in Algorithm 1. At round of FFGG, we sample set of clients C", and the clients in the set approximately minimize their loss functions w.r.t. the personalized parameters w and fixed global parameters 0" to get wh (Line 4). For example, clients can perform 7 local steps with stepsize "/v and find an approximate solution wh, == - of their local objective using the gradients 2 (: (see Algorithun 2). Then, the clients use wi to compute the gradient Vifm(,)) and send their update A to the server (Line 5). After that, the server averages the received vectors A and makes a gradient-descent type update of 0" with stepsize No (Line 7). The proposed method is applicable situations when clients are stateless since neither global parameters Or nor local parameters wf, are required to be stored client once the round is finished.
6

---

Published in Transactions on Machine Learning Research (01/2025)
Algorithm Fine-tuning Followed by Global Gradi- ent (FFGG)
Algorithm 2 Local GD fine-tuner
(to find wh R argmin Jm (0" .w)
Input: initialization d E IRd, stepsize " > 0 for ro 1,....... do
Input: stepsize To 0, number of local steps T; 0' E g
Sample batch of clients C
1: Initialize w randomly
3: for client m E C" do Solve winh s argmin fm w)
for O,...,T do
w um The V25m (0", w
4:
(E.g., using Algorithm 2)
end for
5:
A m end for
= (0'm(0" w
or+1 0r
10 पল Emeor Am
end for
Comparison with (Pillutla al., 2022b). We also notice that our method noticeably different from the ones proposed in (Pillutla et al., 2022b). From the methodological perspective, Pillutla et al. (2022b) studies how gradient-based methods popular in FL solve the considered problem. Our approach, in contrast, is based on the observation that the objective is similar to bilevel optimization with the restriction that cannot store auxiliary variables due to the stateless nature of the clients. Therefore, the method that we proposed trics to first find appropriate private paramcters (m) and only then update 9, whereas the FedSim method from (Pillutla et al., 2022b) updates 0 on each worker without first trying to see if a client's local objective can be minimized by updating only Wm-
There is more similarity between our FFGG and the FodAlt method from (Pillutla et al., 2022b) with T; (number of local steps w.r.t. 0). In particular, if we use gradient descent the local solver when optimizing Wm in FFGG, the methods become almost identical. However, the code of Pillutla et al. (2022b) only supports running the same number of updated with respect to 0 and Wm: and it does not seem that they have tested any alternative. Moreover, our method can use an arbitrary algorithm to solve the problem from Line 4, e.g,, Adam or its variations. In addition to being agnostic to the used sub-solver, our approach also supports stateless clients by design, whereas the stateless version of FedSim and FedAlt were tested empirically by Pillutla al. (2022b) as heuristic modifications.
4 Theory
We present the convergence theory for Algorithm and its variants in several cases. First of all, let us undersland how its idealistic version, which computes wit (0) exactly, would converge under our assumptions.
Theorem 1. Let Assumptions 1-2 hold and assume that we use exact updates, w = arg min, 60,, w). If we choose "o S +, then
L 1
min E [1)F( r<R
RoR
where 0* is anyy vector such that F(0*) 0.
The theorem establishes convergence of operator norms to 0, which implies that the global parameters 0" eventually approach 0-. we compare the result to the standard bounds on convergence of gradient descent, we can see that there no slowdown. In this sense, personalization completely removes the issues of data heterogeneity, which usually causes FL method to converge slower (Khaled al., 2019). In more practical scenarios, Algorithm 1 will compute w (f) inexactly, which will affect the convergence. The next theorem outlines what happens in that case.
Theorem (Informal). Let Assumplions 1-2 hold and fm be smooth and strongly conver in w. Tf T is large enough, then
[1F()) 0 11
min E r<R
R

---

Published in Transactions on Machine Learning Research (01/2025)
Algorithm 3 Asynchronous FFGG
Algorithm 4 Local FFGG
Input: initialization 0 R2, stepsizes Yw: 79 number of local steps N
Input: initialization 8 E &R,, stepsize Ywu Yo > 1: for r= 0. 1, do
for r 2, do
2: Sample batch of clients C" for client m E C do 0r.0 = 0" and wi R argmin (0", w) for k =, 1....... - do
3:
Receive update from client j. with delay dj
3:
0r+1 = 0" n r-
4:
Sample new client my and initialize wh ran-
5:
6:
wh w
m 7nm 22 (frk, wiik)
domly
grk+1 0r. 70 I (0rck W w
Solve Wn, argmin Jm, (0--1,,)
7:
8:
end for A o" gr.
E.g., using Algorithm 2 ++1 ViSm, (fr+1, (
9:
7:
10: end for
8: end for
11: 0"+1 0" Te Emeo AT
end for
We formulate Theorem 2 more rigorously in the supplementary material; see Theorem 6. The main mes- sage, however, remains the same. Moreover, we can see that the local updates are important, which is in stark contrast to the theory in (Pillutla et al., 2022b), where having more local updates does not improve convergence.
4.1 Breaking the curse in other settings
Asynchronous communications. We consider Algorithm 3, asynchronous variant of Algorithm 1. For this method, whenever one of the clients finishes local computations, the server immediately uses the client's update to perform step, and then that client begins local work starting from newly updated global param- eters. The following result holds in this setting.
Theorem 3 (Informal). Let Assumptions 1-2 hold and assume that we use exact updates. Assume the delays and the number of active clients are bounded. Then
20 0-1
min E (IF(0') reR
Ro
where 8 is defined in (15).
We highlight the fact that the proposed problem formulation allows for Algorithm 1 to converge an arbitrarily heterogeneous regime. The detailed formulation and proof of this result can be found in the supplementary material, see Theorem 7.
Byzantine-robust learning. To achieve Byzantine robustness, our method needs small but very impor- tant modification: in line 7, instead of averaging, apply agnostic (6, c)-robust aggregation (sce Definition 2 in the appendix): 0r+1_ = 0" - 7h ARAgg Ai 1/2 (see the detailed setup in Appendix D). For this version of the method, we derive the following result.
where {1,..., M) is the set of all clients, B < 5M
of them are Byzantines and
Theorem (Informal). Let operators Fm be lm cocoorriee,, F strongly monotone, and good clients com- pute stochastic estimates satisfying [[llgm(9)][? I Piull Fm. (0)11) If 5 is small enough, then for Byzantine- robust version of Algorithm we have
E - < (1 78/4 2
1100 - 101
We emphasize that for arbitrary heterogeneous problems, it is impossible to tolerate even one Byzantine client. In contrast, Theorem immlies that for the considered problem formulation, the proposed algorithm converges lincarly to the exact solution (asymptotically) even when data on clients heterogeneous and small amount of Byzantines takes part in the training. The complete formulation of the above result is deferred to the supplementary material, see Theorem 8.
8

---

Published in Transactions on Machine Learning Research (01/2025)
10
301
10"
PROS
10F
10"
10
10
10
10
10-1
10
10
15
500
1006
1600
2000
10"6
200
J00
600
800
10
2001
400
500
10-
200
400
600
800
# of communication rounds
of communication rounds
communication rounds
# of communication rounds
(b): 100
(c): 200
(d): 500
Figure 1: (a): convergence of FFGG varying the number of local steps of conjugate gradients solver т; (b-d): comparison of FFGG against Scaffold, Local GD, and L2GD varying the number of local gradient steps each method (denoted subsolver to make the comparison fair.
for all methods). Note that unlike in (a), FFGG uses gradient descent a
4.2 Benefit of cooperation
As discussed in the introduction, we could have chosen to fit the local data on each client individually. However, the drawback of such approach is that the clients will not benefit from cooperative training, Here, we show that having partial personalization allows to find paramcters that are useful for all clients. This implies that in expectation, a new client m that receives l as initialization will have better value of post-training risk fm(0", w (0)). We have the following result.
Theorem 5. Let Fm be monotone for all and Assumptions 1-2 hold, and assume we found 0 such that [F, 0. Then, 8" minimizes the expected risk Em wm (0))..
The key observation of Theorem 5 is that finding the solution to E Fm (0) = 0 is equivalent to minimizing E [f (0)). The obtained 8* turns out to be optimal for solving a more challenging bilevel optimization problem.
When we combine the result of Theorem 5 with Assumption 2, we get
(0", w (0)) argming,w fm (0, w) for all
This mcans that the combined partially personalized model 0* w (8*) fits the data of client perfectly. When the model is deployed using the found 0", all new clients benefit. from having part of the model pretrained optimally for them. This reduces the computation load, especially if 0 is big part of the model.
4.3
How to split the parameters
The derived theory gives us the following insights:
To break the curse of data heterogeneity, the personalized part of the model must be sufficiently large. Then, Assumption 2 is satisfied, and the training will be easier.
2. If the previous condition is satisfied, Theorem 5 suggests that every client will be able to fit the data perfectly by fine-tuning only w. Thus, the smaller the personalized part w is, the casier will be its fine-tuning.
Therefore, the balance between 0 and when splitting the parameters is crucial. In some cases, it is known that fine-tuning just the batch normalization layers can be sufficient (Li et al., 2020b). Our theory suggests that the key quantity to measure the quality of the split is (.)]l, which we can compute in practice. If it is too close to 0, will not be updated, so one should consider decreasing the number of personalized parameters. It seems less obvious how to detect that we should personalize more parameters. Several layers for personalization have been developped in practical works, see Section 5 for an empirical study.
5 Experiments
The detailed description of all experimental setups deferred the Appendix E.
9

---

Published in Transactions on Machine Learning Research (01/2025)
Table 2:
Test accuracy across different model variants and datasets. For FedAlt and FedSim, we report the numbers for the best performing option in each experiment.
Variant/ /esst acc. (%)
FEMNIST GLDv2
StackOverflow
FedAvg
93.18
51.43
23.82
Local Training
60.2
18.73+0.28
10.19+1.12
FFGG (Input Layer)
93.60zo.uz
51 21.25
24.11:0.02
FFGG (Oulput Layer)
93.5820.00
55.2020.04
24.9220.01
FFGG (Adapter)
94.26+o.m +o.03
64.93+0.04
24.80+0.01
FodAlt (Statcless, Best)
93.97 20.03
64.102o.14
24.94
FedSim (Stateless, Best)
93.89 =o.02
63.19+0.04
24.94=0.01
5.1 The more local work, the better the convergence
In our first experiment, we study the convergence A Algorithm with inexact gradient computation. We test it on the problem from Example 1, namely:
1
fm w) = um(()) - Am -B - Bmol?
2
1
2. (m) - 1 H bm 13,
where Hm; Am Ruxde Bm 1nxd. bm bm Um R" with 10000, de 100, d. 50. The number of clients is 32. All matrices generated from uniform distribution [0, 1], and then divided by the second dimension (i.e., de for Hm, Am and dw for Bm).
We use SciPy's (Virtanen et al., 2020) implementation of the Conjugate Gradient (CG) method to solve local subproblem in w and vary the number of inner steps t of CG (see Figure 1, (a)) The convergence with a small number local steps r = 10 is already sufficient to achieve an error as small as 10-4, Moreover, Algorithm 1 converges to the exact solution for 7 E (30, 40), i.c., without finding the precise solution of the local subproblem. This experiment shows that the convergence of FFGG indeed improves with increasing amount of local work, and the method is overall practical.
5.2 Comparison against other methods
Next, compare FFGG combined with Algorithm 2 as fine-tuner against non-personalized methods such Scaffold (Karimireddy et al., 2020b) and Local GD, and fully personalized method L2GD (Hanzely & Richtárik, 2020). For Scaffold, set outer and inner stepsizes to be equal to 0.5 and Lyr correspondingly, where Ly is a smoothness constant of fm- For Local GD the stepsize is equal to byz Finally, for L2GD we choose À = 0.1 and stepsize to be equal to (2M) max { L(1 - p)- " Ap-"), where == r-1 (we make such choice for p make the number of local steps be close т in expectation).
We test the convergence of the methods changing the number of local steps {100, 200, 500). Figure 1 (b-d) shows that FFGG outperforms other basclines in all cases. We also highlight that FFGG's convergence improves when we increase the number of local steps as it is predicted by theory.
5.3 Comparison on real-world datasets
Finally, we evaluate method real-world federated datasets: FEMNIST (character recognition), GLDv2 (Visual Landmark Recognition). and StackOverflow (next word prediction). We demonstrate that FFGG leads to a significant performance improvement compared to non-personalized FedAvg and exhibits better adaptability to heterogeneity compared to methods with the same personalized objective, such as FedAlt (Pil- lutla et al., 2022b).
As previously discussed, we do not address the optimal parameter split in this work. To achieve partial personalization, we follow the setup from (Pillutla et al., 2022b). However, we show that when overyersonal-
10

---

Published in Transactions on Machine Learning Research (01/2025)
isation occurs, i..., we can achieve zero training loss, FFGG outperforms all other methods by substantial margin. All experimental details and hyperparameter selections are provided in the appendix. For partial personalization, we consider three partitioning schemes:
Input-layer personalization: This architectural design customizes the imput layer to learn personal- ized representations, whereas the remaining part of the model is common to all clients. For predicting the next word, the initial transformer layer is personalized instead of the embedding layer.
Oulpul-layer personalization: This design learns a common representation but customizes the pre- diction layer. In a transformer model, we adapt the final transformer layer instead of the output layer for personalization.
Adapter personalization: Every client uses a personalized low-rank adapter to fine-tune the global model.
We also introduce an algorithmic extension to Algorithm 1 to incorporate local steps with respect to global parameters into the training loop. After receiving global model from the server, clients randomly initialize personalized parameters and perform one local epoch with respect to these paramcters to approximate (0). Following this step, we alternate between stochastic gradient steps with respect to global and local parameters. We only initialize Wm at the beginning of local training. This approach allows us to approximate wi, (0) with a single gradient step after initial fine-tuning. The pseudocode for this algorithm is provided in Algorithm 4. Our results are presented in Table The displayed values represent averages over three independent seeds. It is worth noting that for both datasets, FFGG leads to improvement of at least one percent in final test accuracy over non-personalized FedAvg. We also observe that Adapter is particularly useful technique for partially personalizing local models. The largest improvement, exceeding 13%, is observed for FFFG (Adapter) on the GLDv2 dataset. In this particular case, the final train accuracy for all clients is 100%, which aligns well with our theory as it implies that Fm (6*) = O. Additionally, since we employ the same experimental setup for personalization as Pillutla et al. (2022b), we also compare our results with their stateless version of FedAlt (notably, the stateless version is not analyzed in their work). We report the performance of the personalization technique that yielded the best results for FedAlt. It is noteworthy that FFGG consistently performs as well as, if not better than, FedAlt, outperforming FedAIt in two out of three datasets.
6 Limitations
In this section, we discuss some limitations of the proposed algorithms and analysis.
Since the proposed algorithms require the clients to approximately minimize their objective functions w.r.t. their personalized parameters (Line 4, Algorithm 1), it creates a computation and latency overhead in comparison to standard Parallel GS) However, if we use Local-GI (Algorithm 2) as a subsolver for the mentioned problem, then FFGG has no overheads FedAvg. Moreover, FFGG has even cheaper iterations than FedAvg since FFGG with Local-GD subsolver computes multiple gradients with respect to personalized parameters and one gradient with respect to global parameters, while FedAvg requires the clients to compute multiple gradients with respect to all parameters.
comparison to standard
The proposed methods require knowing the split of parameters, which might be non-trivial to obtain.
The analysis relies Assumption 1, which is closely related to convexity and smoothness. These assumptions are commonly used in the literature but are not common in practice so the main utility of the analysis is to gel a general understanding of the methods' behavior.
7 Conclusion
We proposed new analysis of partial personalization that shows its provable benefits in Federated Learning. The main takeaways can be summarized follows.
11

---

Published in Transactions on Machine Learning Research (01/2025)
We proved that the problem can always be made overpersonalized and the data hetcrogeneity slow- down can be completely eradicated.
We also illustrated this by showing that, in contrast to standard FL, asynchronous training with par- tial personalization converges precisely, and partial personalization can be made Byzantine-robust
Our theory also suggests algorithmic changes to how the training should be performed and allows for generic local solvers. Compared to the work of Pillutla et al. (2022b), our methods are stateless, and our theory does not require making stepsizes smaller than O (?/r), where is the number of local steps. Finally, our assumptions are satisfied for several natural classes of functions, highlighting that our theory is quite general.
There are several open questions that can be of interest to make personalization more practical. First of all, a direction that seems important to us how we can find optimal splits between and w to achieve both the spoed-up of removed data heterogeneity and make sure that clients benefit from cooperation. Secondly, parameter-effcient fine-tuning might bring even more speed-ups. Lastly, while the statistical effect of cooperation was left out of consideration in our work, it can bring new insights and deserves some attention.
References
Alekh Agarwal and John C. Duchi. Distributed delayed stochastic optimization. Advances in Neural Infor- mation Processing Systems, 24, 2011. (Cited page 5)
Dan Alistarh, Zeyuan Allen-Zhu, and Jerry Li. Byzantine stochastic gradient descent. Advances in Neural Information Processing Systems, 31, 2018. (Cited page 5)
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. In International Conference on Machine Learning, pp. 242-252 PMLR, 2019. (Cited page
Zeyuan Allen-Zhu, Faeze Ebrahimianghaaani, Jerry Li, and Dan Alistarh. Byzantine-resiliet non-convex stochastic gradient descent. In International Conference on Learning Representations, 2021. (Cited
page
Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary. Federated learning with personalization layers. ar preprint wr.Xi::912.000101 2019. (Cited on page 2)
Gilad Baruch, Moran Baruch, and Yoav Goldberg. A little is enough: Circumventing defenses for distributed learning. Advances in Neural Information Processing Systems, 32, 2019. (Cited on pages and 27)
Peva Blanchard, El Mahdi El Mhamdi, Rachid Guerraoui, and Julien Stainer. Machine learning with adversaries: Byzantine tolerant gradient descent. Advances in Neural Information Processing Systems, 30, 2017. (Cited on pages and 27)
Fei Chen, Mi Luo, Zhenhua Dong, Zhenguo Li, and Xiuqiang He. Federated meta-learning with fast conver- gence and efficient communication. arXio preprinl arXi::1802.07877 2018. (Cited on page 4)
Yudong Chen, Lili Su, and Jiaming Xu. Distributed statistical machine learning in adversarial settings: Byzantine gradient descent. Proceedings of the ACM on Measurement and Analysis of Computing Systems, 1(2):1 25, 2017. (Cited on pages 27 and 31)
Gregory Cohen, Saeed Afshar, Jonathan Tapson, and Andre Van Schaik. EMNIST: Extending MNIST handwritten letters. In 2017 international joint conference on neural networks (IJCNN) pp. 2921-2926. IEEE, 2017. (Cited on page 33)
Georgios Damaskinos, El-Mahdi El-Mhamdi, Rachid Guerraoui, Arsany Guirguis, and Sébastion Rouault. Aggregathor: Byzantine machine learning via robust gradient aggregation. Proceedings of Machine Learn- ing and Systems, 1:81-106, 2019. (Cited on page 5)
12

---

Published in Transactions on Machine Learning Research (01/2025)
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pallern recognilion, pp. Teee, 2009. (Cited on page 33)
El-Mahdi El-Mhamdi, Rachid Guerraoui, and Sébastien Rouault. The hidden vulnerability of distributed lcarning in byzantium. In International Conference on Machine Learning, pp. 3521-3530. PMLR, 2018.
(Cited on page 5)
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with theoretical
guarantees: A model-agnostic meta-learning approach. Advances in Neural Information Processing Sys- tems, 33:3557-3566 2020. (Cited on page 4)
Eduard Gorbunov, Alexander Borzunov, Michael Diskin, and Max Ryabinin. Secure distributed training PMLR, 2022. (Cited on page 5)
scale. In International Conference on Machine Learning, pp.
Eduard Gorbunov, Samuel Horváth, Peter Richtárik, and Gauthier Gidel. Variance reduction is an antidote to Byzantines: Better rates, weaker assumptions and communication compression International Conference on Learning Representations, 2023. (Cited on pages 5, 26, 27, and 29)
cherry on the top.
Michal Grudzieá, Grigory Malinovsky, and Peter Richtárík. Can 5th generation local training methods support client sampling? yes! Xiv preprint or Xix:2222 14370, 2022. (Cited on page 4)
Filip Hanzely and Peter Richtárik. Federated learning of mixture of global and local models. orXiv preprint arXiv:2002.05516 2020. (Cited on pages 4 and 10)
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognation, pp. 770-778, 2016. (Cited on pages 33 and 34) Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Federated visual classification with real-world data distribution. In Computer Vision ECdon-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part X 16, pp. 76-92. Springer, 2020. (Cited on page 33)
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Represenlations, 2022. URL https:
httpps:/opeereviieee http:///pppppeeeeeeeeeeeerataaaa net//formmiicmn//eeeeeee..
(Cited on page 6)
Rustem Islamov, Mher Safaryan, and Dan Alistarh. AsGrad: A sharp unified analysis of asynchronous-Se SGD algorithms. In Proceedings of The 27th International Conference on Artificial Intelligence and Statistics, volume 238, pp.
PMLR, 2024. (Cited on pages and 31)
Yihan Jiang Jakub Konečný, Keith Rush, and Sreeram Kannan. Improving federated learning personaliva- tion via model agnostic meta learning. arXiv preprint orXiv:1909.1244 2019. (Cited on page 4)
Peter Kairouz, H. Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Keith Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open prob- lems in federated learning. Foundations and Trends& in Machine Learning, 14(1), 2021. (Cited on pages 3 and
Sai Praneeth Karimireddy, Lie He, and Martin Jaggi. Byzantine-robust learning on heterogeneous datasets via bucketing. arXie preprint orXiv:2006. 09365, 2020a. (Cited on page 31)
Sai Pranceth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian U. Stich, and Ananda Theertha Suresh. SCAFFOLD: stochastic controlled averaging for federated learning. In Proceed- ings of the 37th International Conference on Machine Learning, pp. 5132 5143. PMLR, 2020b. (Cited on pages 4 and 10)
Sai Pranecth Karimireddy, Lie He, and Martin Jaggi. Learning from history for Byzantine robust optimiza- tion. In International Conference on Machine Learning, pp. 5311-5319. PMLR, 2021. (Cited on pages 5, 26, and 27)
13

---

Published in Transactions on Machine Learning Research (01/2025)
Sai Pranceth Karimireddy, Lie He, and Martin Jaggi. Byzantine-robust learning on heterogencous datasets via bucketing. In International Conference on Learning Representations, 2022. (Cited pages 26, and 29)
Ahmed Khaled, Konstantin Mishchenko, and Peter Richtárik. First analysis of local GD on heterogeneous data. arXiv preprint ar 1909.00499 2019. (Cited on page 7)
Ahmed Khaled, Konstantin Mishchenko, and Peter Richtárik. Tighter theory for local SGD on identical and heterogeneous data. In Proceedings of the 23rd International Conference on Artificial Intelligence and Statistics, pp. 11119-529.. PMLR, 2020, (Cited on page 4)
Anastasia Koloskova, Sebastian U. Stich, and Martin Jaggi. Sharper convergence guarantees for asynchronous SGD for distributed and federated learning, In Advances in Neural Information Processing Systems, 2022. (Cited on pages 5, 26, and 31)
Jakub Konečný, H. Brendan McMahan, Felix Yu, Peter Richtárik, Ananda Theertha Suresh, and Dave Bacon. Federated learning: Strategies for improving communication efficiency. In NIPS Private Multi-
Party Machine Learning Workshop, 2016. (Cited on page 1)
Loslic Lamport, Robert Shostak, and Marshall Pease. The byzantine generals problem. ACM Transactions on Programming Languages and Systems, 4(3):882 401, 1982. (Cited on page 3)
Junyi Li, Feihu Huang, and Heng Huang. Local stochastic bilevel optimization with momentum based variance reduction. arXiv preprint 2rXiv:2205.01 00688, 2022. (Cited on page 5)
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Systems, 2:429-450, 2020a. (Cited on page 4)
Xiaoxiao Li, Meirui Jiang, Xiaofei Zhang, Michael Kamp, and Qi Dou. FedBN: Federated learning on non-iid features via local batch normalization. In International Conference on Learning Representations, 2020b. (Cited on pages and 9)
Nicolas Loizou, Hugo Berard, Gauthier Gidel, Ioannis Mitliagkas, and Simon Lacoste-Julien. Stochastic gra- dient descent-ascent and consensus optimization for smooth games: Convergence analysis under expected co-coercivity. Advances in Neural Information Processing Systems, 34:19095-19108 2021. (Cited page 27)
Lingjuan Lyu, Han Yu, Xingjun Ma, Lichao Sun, Jun Zhao, Qiang Yang, and Philip S Yu. Privacy and robustness in federated learning: Attacks and defenses. arXiv preprint arXiv:2012.065: 2020. (Cited on page 3)
Siyuan Ma, Raef Bassily, and Mikhail Belkin. The power of interpolation: Understanding the effectiveness of SGD in modern paraparametrrieed learning. In International Conference Machine Learning, pp- 3325 3334. PMLR, 2018. (Cited page
Olvi L. Mangasarian. Parallel gradient distribution in unconstrained optimization. SIAM Journal on Control and Oplimization, 13(6): 1925, 1995. (Cited on page 4)
Horia Mania, Xinghao Pan, Dimitris Papailiopoulos, Benjamin Recht, Kannan Ramchandran, and Michael T. Jordan. Perturbed iterate analysis for asynchronous stochastic optimization. SIAM Journal on Optimiza- tion, (722:22002-22) 2229, 2017. doi: 10.1137/16M105700 (Cited page 24)
H. Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera y Arcas. caammiccation-eficiinnn learning of deep networks from decentralized data. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, pp. 1273-1282. PMLR, 2017. (Cited on pages and 4)
Panayotis Mertikopoulos and Zhengyuan Zhou. Learning in games with continuous action sets and unknown payoff functions. Mathematical Programming, 173:465-500 507, 2019. (Cited on page 27)
14

---

Published in Transactions on Machine Learning Research (01/2025)
Konstantin Mishchenko, Francis Bach, Mathieu Even, and Blake Woodworth. Asynchronous SGD beats minibatch SGD under arbitrary delays. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Newral Information Processing Systems, 2022a. URL https://openrevies net/ p
(Cited on pages and 30)
Konstantin Mishchenko, Grigory Malinovsky, Scbastian Stich, and Peter Richtárik. ProxSkip: Yes! local gradient steps provably lead to communication acceleration! finally! In International Conference Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 17-23 Jul 2022b. URL https: /ptp://///cceccs....... .mlr. www..prres//vv1//// https///1..2mmmshhheho222222222 hhm.mml.....mmm.....mm.mmmmmm... (Cited on page 4)
PMLR,
Philipp Moritz, Robert Nishihara, Stephanic Wang, Alexey Tumanov, Richard Liaw, Eric Liang, Melih Elibol, Zongheng Yang, William Paul, Michael I. Jordan, and lon Stoica. Ray: A distributed framework for emerging [AI} applications. In 13th USENIX symposium operating systems design and implementation (OSDI 18), pp- 561-577, 2018. (Cited on page 30)
Yurii Nesterov. Introductory lectures on conver optimization: A basic course, volume 87. Springer Science & Business Media, 2013. (Cited on pages and 18)
John Nguyen, Kshitiz Malik, Hongyuan Zhan, Ashkan Yousefpour, Mike Rabbat, Mani Malek, and Damitry Huba. Federated learning with buffered asynchronous aggregation. In International Conference on Arti- Sicial Intelligence and Statislics, pp. 3581-3607 -6007 PMLR, 2022. (Cited on page 3)
Feng Niu, Benjamin Recht, Christopher Re, and Stephen Wright. HOGWILD!: A lock-free approach to parallelizing stochastic gradient. descent. Advances in neural informalion processing systems, 24, 2011. (Cited on page 3)
Maher Nouiehed, Maziar Sanjabi, Tianjian Huang, Jason D. Lee, and Meisam Razaviyayn. Solving a class of non-convex min-max games using iterative first order mothods. Advances in Newral Information Processing Systems, 32, 2019. (Cited on page 6)
Krishna Pillutla, Sham M. Kakade, and Zaid Harchaoui. Robust aggregation for federated learning. 1EEE Transactions on Signal Processing, 10:1142-1154, 2022a. (Cited on pages and 27)
Krishna Pillutla, Kshitiz Malik, Abdel- Rahem- Mohamed, Mike Rabbat, Maziar Sanjabi, and Lin Xiao. Federated learning with partial model personalization. In International Conference Machine Learning, pp. 17716 17758. PMLR, 2022b. (Cited on pages 4, 5, 7, 8, 10, 12, and 31)
Aniruddh Raghu, Maithra Raghu, Samy Bengio, and Oriol Vinyals. Rapid learning or feature reuse? Towards understanding the effectiveness MAML. In International Conference on Learning Representations, 2020.
URL https: //tprr///e/eeeeeeeeeeeeee wwt.:/ettturiiiee ftnts/forumidddrke drr.gg//EE.................. (Cited on page 5)
Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and H Brendan McMahan. Adaptive federated optimization. arXiv preprint ar arXiw:2003.00200 2020. (Cited on pages 32 and 35)
Karan Singhal, Hakim Sidahmed, Zachary Garrett, Shanshan Wu, John Rush, and Sushant Prakash. Fed- crated reconstruction: Partially local federated lcarning. Advances in Neural Information Processing Systems, 34:11220- 11222, 2021. (Cited on page 2)
Chaobing Song, Zhengyuan Zhou, Yichao Zhou, Yong Jiang, and Yi Ma. Oplimistic dual extrapolation for coherent mon-mono variational inequalities. Advances in Neural Information Processing Systems, 33: 14303-14314 2020. (Cited on page 27)
Sebastian U. Stich. Local SGD converges fast and communicates little. In International Conference Learning Representations, 2019. (Cited page 4)
Lili Su and Nitin H. Vaidya. Fault-tolerant multi-agent optimization: optimal itcrative distributed algo- rithms. In Proceedings of the 2016 ACM synposium on principles of distributed computing, pp. 425-434, 2016. (Cited page 3)
15

---

Published in Transactions on Machine Learning Research (01/2025)
Davoud Ataee Tarzanagh, Mingchen Li, Christos Thrampoulidis, and Samet Oymak. FedNest: Federated bilevel, minimax, and compositional optimization. In Proceedings of the 39th International Conference Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 21146 21179. PMLR, 17-23 Jul 2022. URL https: https://proceeedigg. mlr. (rres//112//tazzaaaa2222 html. (Cited on page 5)
John Tsitsiklis, Dimitri Bertsckas, and Michael Athans. Distributed asynchronous deterministic and stochas- tic gradient oplimization algorithms. IEEE transactions automalic comtrol, (1(98:803 812, 1986. (Cited page 5)
Tulia Ture, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Well-read students learn better: On the importance of pre-training compact models. arXiv preprint arXic:1908. 08962, 2019. (Cited on page 33)
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances neural information processing systems, 30, 2017. (Cited on page 33)
Sharan Vaswani, Francis Bach, and Mark Schmidt. Fast and faster convergence of SGD for over- parameterized models and an accelerated perceptron. In The 22nd international conference on artificial intelligence and statistics, pp. 1195- PMLR, 2019. (Cited on page 2)
Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Ev- geni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, C Carey, îhan Polat, Yu Feng, Fric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy 1.0 Contributors. SciPy 1.0: Fun- damental Algorithms for Scientific Computing in Python, 2020. (Cited on page 10)
Tobias Weyand, Andre Araujo, Bingyi Cao, and Jack Sim. Google landmarks dataset. v2-a large-scale benchmark for instance-level recognition and retrieval. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pp.
2020. (Cited page 33)
Zhaoxian Wi, Qing Ling, Tianyi Chen, and Georgios B. Giannakis. Federated variance-reduced stochastic gradient descent with robustness to Byzantine attacks. IEEE Transactions on Signal Processing, 68:
2020. (Cited on page 5)
Cong Xie, Oluwasammi Koyejo, and Indranil Gupta. Fall of empires: Breaking byzantine-tolerant SGD by inner product manipulation. In Uncertainly in Artificial Intelligence, pp. 261-270. PMLR, 2020. (Cited on pages 5 and 27)
Dong Yin, Yudong Chen, Ramchandran Kannan, and Peter Bartlett. Byzantine-robust distributed learning: Towards optimal statistical rates. In International Conference on Machine Learning, pp. 5650-5659. PMLR, 2018. (Cited on page 5)
Heng Zhu and Qing Ling. BROADCAST: Reducing both stochastic and compression noise to robustify
communication-ellicennd federated learning. arXiv preprint ar 06685, 2021. (Cited page
16

---

Published in Transactions on Machine Learning Research (01/2025)
A Deferred proofs
A.1 Proof of cocoercivity for Example 1
Below we show that the function fm(0, w) = m0()) + Al 0 + Bm ym.l? satisfies Assumption
Proof. Notice that Om (0) does not depend on w; so
Vi (m Vem (0) (Am0 + B, (m (0) ymm)
Then, by convexity and La smoothhssss Φm, we have
(m (01, w (01)) Vifm (02, w (02)) 01 02)
(V)(0)) - (o((22, 81 02) (A (Am01 + (m Am 02 Bn w (02)), 01 02) 1 [o(1)) vo(02)]1 + (A), (Am 01 + Bm u (01) Am Bm m (02)), 01 02) Li
Let us find the value w (0). Differentiating fm(6. w) with respect to w and setting the gradient to 0, we get
B (Am) + Bn w (0) ym)
whence
B m B w (0) B (y/ Am and w (0) = B (3m Am.).
Substituting this into the previous lower bound, we get
(Vifm (0,, w (01) V Vifm w (02)), 62) (A (1 B,, B 0m Am02 + B, Ammm)),0 - 82) +
Vov())) vo(02)1
Lo
= -I- - B. B 01 - 02), 02)
[v)(0)) - vo(02)1
I
Let Bm UΣVT be the SVD decomposition Bm; then I- B B m symmetric positive semi-definite matrix. Therefore, AT -I B, as well. Thus, we obtain that the linear term in Fm is convex. Since it is the gradient of quadratic, is s-smooth with i = A (I- B, B Am Therefore,
UVE U(--E' UT a
Am is symmetric positive semi-definite
(Vifm (0,, um (0,)) V (62, um um (m2))), 01 02)
1 [Vo(1)) 2 + A B B Am (01 f2)ll
Ls
1
[ib(1)) - vo((22) + AT B B Am (01 02)11
max(Log .)
V 1ifmm( (0 wm (01)) V vifm(022 Wmm (02))
2 max f.)
which is exactly L) cocoercivity of Fm.
A.2 Proof of cocoercivity for Example 2
Now, study general function fm that has bounded derivatives.
Proof. Let us lower bound the Jacobian of Fm:
(0) Vi (0 w (0) (0) w (0))
> / + Vow (2 22 fm (0, w (0))
1 Vow (0) /
/ (0. w (0))
" I.
2
17

---

Published in Transactions on Machine Learning Research (01/2025)
This lower bound implies s-strong monotonicity. We also have similar upper bound:
VoFm (0) = (0, w/ + Vow (0)V w( ক (C
Since L 2 u4 we have C₁O₂ < 5タ - L, and
Voln x 2L1,
which implies F.. is (2I))Tipschitz. Combining this with strong monotonicity, get that it is also _ 3 cocoercive.
A.3 Proof of cocoercivity for Example 3
Example 3 is the hardest study. First, we state and prove the following standard result.
Proposition 2. Let 91 92 be L-smooth convex functions. Then yp = (0 + 422 is (2L)-cocoercive.
Proof. As can be found in standard textbooks, such as (Nesterov, 2013), convexity and smoothhess imply that both fp and 222 are T-cocoercive. Moreover, for any 01, 02 it holds
(Vp((1)) Vi((2)) 01 02) (V))(0)) V)r1 02) + (V))((1) Vip2(02), 02) (01) - Vpr( - (1)((())) - VV2(00)11
2 L
1
2 2T
+ (22(01) V
1
VP((1) - v=((2)11
2L
which is exactly what need to prove.
Now we procoed to prove cocoercivity of F. from Example 3.
Proof. Notice that @m (0) does not depend w, so
vm (0, w (0) V (0) + A (Am Bm ymm)
Let us find the value of w (0). Differentiating fm w) with respect to w and setting the gradient to 0, we get
B V (Am Bm " (0) Bmm = 0.
Let U1 Am + Bn w (01) Vm and U2 Amo B wt, (02) Vm Then,
(A Vnim (1 + Bm t (0:) - 3mm) A Vel (Am02 + Bm w (0z) - -ym) mmm), 01 02)
(A Velm - A V ( (U))) 01 -82)
eee (W1) - Velm Am (01 02))
(Vem (wi) Vilm (u2), (Am vm - (Am02 - vmm)
(V) Vmlm (u2), ) (Am Bm (01) 4m (Am02 Bm (02) vmm))
-VVm (u1) Velm (u2), Bm ( (01) Bn w (42))
Moreover, since B (un) = and BJ m V (u2) = O, we have
(1n - (u)m( Bn Bm (6:) B w (02))
B Veim - B, Volm (uz), wt, (01) w (02))
0.
18

---

Published in Transactions on Machine Learning Research (01/2025)
Plugging this back, get
(A ( A Vil (u2), 01 -82) (V)m(u)) V Vi 11 u2)
1 Vem((u) - Viim(u2)1?.
Therefore, Vifm(0, w (0)) is equal to the sum of two cocoercive operators. well.
Thus, Fm is cocoercive
A.4 Proof of Theorem 1
Proof. Since we assurne exact computation of w² (0") for all and m C", holds
A = Vifm (0", w (0") = (0")"".
Therefore, we have the following recursion:
100++1 1011 = 110" - -0"112
270
Σ meCr
(Fm (Fm( 00) +
Σ Fm((0)
|Cr
|O'|
MEC
(1) 110 1-
2710
Fm (0') +
70
Σ MEC
Fm (0r)
I. C
|C7
mEC
< 100 0-1
2710 L|C"
Σ |I
Σ (m
|C
mEE
MEC
7 1
100
1-1
Σ meC
FI. (07)
LIC
Taking expectation, get
< L (110 0-1 110 -0)
E [IF S
Σ
[mmm
|C"
70
mEo
Summing this bound over r = 0, .....,.e R 1, we get
R-1
0 8r 112
||0 1-1
L| 111
Ein [IF((')
Σ
E [IF(()))
<
F<R
R
Ro R
which completes the proof.
A.5 Proof of Theorem 5
Proof. Our goal is to show that all solutions to F(0) = 0 are also minimizers of Ef (0, w (0))). Since we assume that the functions are convex, it is sufficient to show that the gradient of E (0, w (0)) is equal to O. We have
w Vofm (6, w. (0) (i, WT, (0)) + Vi2f(0, w ff (,) (8)).
Note that by definition of wi,h (0), it holds 22fm (0, wth (0)) 0 since wi,, (0) optimal when the first argument of fm is fixed. Therefore,
Vofm " (0) fm (", (8)).
Let 8 be such that F(0*) = 0. Then,
F F(0) = FVV.fm (0, w (0)) VoR (0, w (0)) = 0.
19

---

Published in Transactions on Machine Learning Research (01/2025)
B Algorithm 1 with inexact gradient computation
We consider Algorithm 1 where A + (0") i.e., with inexact gradient computation. The analysis of the inexact version of Algorithm 1 requires additional assumptions on the problem which are listed below.
Assumption 3. There exist constants Lu and paw such that for any client m, the loss fm is Lu Lipschitz continuous and Pw stronnyyyy conver in w for any fixed B, i.e.,
|1 fm, w1) ViSm w2)l S Lu L w| w2ll (0, w1) 2 Jm (0, w2) (V2fm w2), wn w2) Mww 1 w2ll
(5)
(6)
If Assumption holds, then the standard result for Gradient Descent in w takes place
| (0')1 < (1 Mho Mw)ll w (0)) S (1 Nw 7wm Iw w (0")
(7)
where the inner stepsize Yo S L.
Assumption There exist constants A, C such that for any client m, the solution w (0) satisfies
I (0)1 S A| -
(8)
This assumption holds if the norm of Vew,, (0) is bounded by A, because then w is A-Lipschitz continuous, and consequently wt satisfies Assumption 4:
(0)] & ||,; (0) w,,, (0")) + Im (0*) S A| - 00 + I|w, (0*)]l I|-
Remark 1. For simplicity of explanation, let wht be initialized as zero and the cardinality of C" is fixed.
Remark 2. We provide the proof of Algorithm 1 where fine-tuning is performed using Local GD (Algo-
rithm 2). In fact, all clients may utilize any other method to solve a subproblem to approximate w (0") The only difference in the analysis is that we need to require ||w w (0') < (1 - p)]] w (0 |[2 and assume that p not too small or sufficiently large in order derive convergence. For example,
in the case of Example 1, we may use the Conjugate Gradient method, which more suitable for quadratic problems in w.
Theorem 6. Let Assumptions 1, and hold, set the stepsizes go 1, Nw I.,, and assume that the number of local iterations is lower bounded
Le max Ph
bR
R²R
T
2 2oo aR, 2 log
ro
log
2
where and b are defined
2beA to
b = 2L Then
4L | 00
8-
min E [1
<
r<R
R
20

---

Published in Transactions on Machine Learning Research (01/2025)
Proof. Due to inexactness of the update, 4n Vifm (0", (wi"" * Fm (0"). Thus, will be updated with biased estimate of Fm (f"). We start with unrolling |0++1 0²1
100+11 - 0-112 =110- - 0011
270
Σ
(Am 0" -0") +
70
Σ mEC
A
[Crl
m
meCr
2740
= =110- - 102
Σ (6))))))0 - -8")
(A - Fm(("), - 8")
|C*
|C"
mEC-
mEC
+
70
Σ
A FF + Fm
|C
meCr
110
0 112
2
Σ (I
2718
Σ
(A (m 8*)
|C7
|0"
meC
mEC
2
27
+
Σ
1 - Fm (0') +
Σ (IFm(()))
1C'
107
mE(
meC"
795
10" 0-112
70
Σ
Fm (0") 12
2
Σ
(A (m) -0") -0")
L|C''
|Cr
mEC
mEC
27
+ 1C"
Σ
A (0)
(9)
meCr
where the first inequality, we also use Young's inequality two times. Now we handle the third term (9)
taking expectation w.r.t to all probability events happened before iteration r
2719
27 C
Σ
( Fm (0"), 0" 0-) S
10 0-11 11 - (0))))
Cr'
mEC
mEC
20
Σ
110 - 0- 11 V fm (0r, win") V Vm (0r, wit (,,(("))) (01)
Cr
meCr
(5)
270 Lew
110 - 0r VIlwh w (1r)
10r
mEC
270 L
(1 Yonlher 1/2
Σ
110 - -0"1 || (0')(()) (0')
[rl
mEC
(1)
$ 2g - )wHHbw) 7/2 (All - 11 + +cll0" - -0"1))
(10)
where we use the assumption w = 0. Now we work on the last term in (9)
2nd [Crl
270 [rrl
Σ
D
20
Σ
I whih
(0", w Vi
1|2
mecr
meCr
(5) 2T 2
Σ
w/ (0") - wo:'" 112
|Crl
mECT
(7) 212 n
(1 )ro
I 112
[Crl
meC
(4)
2 2 )wolho) (2272110 0-112 + 2C
(11)
21

---

Published in Transactions on Machine Learning Research (01/2025)
Plugging (10) and (11) in (9), get
E, 1-112 & 110- -8 -0-12-
Σ
E, F (m))]]
L|C
meCr
2- Lu Mo (A 0011 C110" 11)
21. 70(1 Yod (2A2 - 1|2 C²).
Thus, laking full expectation, have
OSE E
Fm (0')
< R
(1100 8º11 - Ello 00 (1))
Cr
r=0 MEC
27g Lig (1 - Nw 7/2 )/2
n-1
+
(AT - 1-2 + CV Cv/ll 1
R
n
+
(1 Pmn)
Σ
(2A"F|10"
R
1=0
1
$
(1100 0-112 1112- 11)
R
R-1
+
(27 Tww Tem A(1 - Yollhw )r² 12 + 4rg 1.2 A2 - Nwokhw))
E110" 0"11
R
27 Log +
n
Follhw) T/22
E
Or 0r 1
4-0 02 nookho
(12)
R
p=
This implies that
ElloR 11 - |100 -11122
R-1
+ 2n A( Yor + 14 - -mortte)
B110r - 0-
0=
R 1
+ L C(1 - Poelhe)) T/2
Σ
1r - 0112 + 4R2/0 L2 C pwbkww
0=0
R-1 Bll0r 10'1
S 1100 - 001122+ + (ac
Ywul + a² YarHur
r=0
R-1
be Yott 7/2
Σ
Tw R. /Ellrr 0-113 7
(13)
where (we plug in stepsize values from the statement and use inequality (1 (1-z)a (e-a) aa)
a 2L b
C
Now we will show by induction that El|0r 1-1 < 4r where rā 2 1100 0-112, for any r. The base of
induction is trivial since ||e - 112 < T8 < 4r Assume that El 112 s r for all {O...., R 1}, then
it also holds for ERjR 0² 112. Indeed, the restriction
2n, log aR 20 log bR ra.
L log B²
7 2 max
Pw
Paw
22

---

Published in Transactions on Machine Learning Research (01/2025)
implies that by the basis of induction, the following:
"To P T//
El100 0-12 s aexp (2)
(- 2Tw
Pw
2L log a.
aexp
Pro
pw
aRr
ex
exxpp-logall (ogaR)Rt
0
aR
Similarly,
R-
a Yo fas7
1l|00 - 1- v a² exp( Yalkw R
r=0
Po /2))² R
$ a exp (- Hw
2w log at
Rr
2L
a² Rro
exppplo - aR)))
R
S r
R²R²
And we have the same bounds for the remaining two terms in (13):
be Yo Yu Hw
E
1 s bexp( T T 1 Rro
2 log
bit
S bexp
Pii
Rro
2Lw
bR
bRro
bexp
logg
Rro
Tū:
ro
bR/ro
and
B2 Lu log
be Nu T R b2 exp
rd
R
Lu:
BRR R
BR
= 62 exp
r
BRR/T
Hence, all four terms in (13) are smaller than vē- statement of induction. Note that restrictions on r logarithmically depend on R only, hence it is not a strong assumption.
Thus, with such choice of stepsize, we prove the
Now we establish the statement of the theorem. Using the statement of induction (12), we get
min [IF((()
1P
<
E
[((
R
1
E [I.. 11
R
mEC
4L 60
11
<
70 R
23

---

Published in Transactions on Machine Learning Research (01/2025)
C Asynchronous method
We formulate and prove the convergence of Algorithm 3 (which is the asynchronous version Algorithm 1) with exact computations only, ie., in line 7 the subproblem is solved exactly. However, the convergence of the incxact version can be derived in a similar way as for Algorithm 1 Appendix B. As defined Algorithm 3, we denote the client that finishes the computation iteration as jr and the newly sampled client as my,
To prove the convergence, we define (m, = max{j < r my = =m - the last itcration before iteration r when the update from client m was applied. Our analysis is based on the virtual iterates, also known as perturbed iterates, that were introduced by Mania et al. (2017). In particular, we consider the sequence defined recursively as
gr+1 = -yaFm = (8").
(14)
We also use initialization
00 = 8
Σ MECO
yo (80))
(15)
since all are initially sampled, clients will compute their gradients using 80 The set of active clients C" initialized with C is updated according to C++1 {m,) U (C" {i-)). We assume that O is always bounded, which is satisfied, for instance, when the total number of clients is finite. We assume that delays are bounded.
Assumption There erists a constant max such that for any client m at iteration the following inequality holds: |prev(m, - Tmax, i.e., all the delays are bounded by Tmax.
Theorem 7. Assume Assumptions 1 and 5 hold. Let the number of active clients is upper bounded by M. Assume the stepsize is such that S Mmaa// Then
min E (1E((")) < 20 11? reR
Proof. Let us first show the link between 0r and 6r by induction:
er 0--
Σ
Fm (fDrev(m.r)).
(16)
mEC
It is truc for the base r 0. Let us assume that it holds for r - and prove for r. We have
gr 8r = (gr-1
7/8 Fo,F,., ( (0--1- 700F
Fm (0pr) + 70 (0(Fmn ( Fj._1 (fprev(m..-1))).
We also remind the reader that prev (m--1,7) r- and C" (m,-1) (ir-1) 1)) Moreover, for the rest of active workers (those gradients still have not been applied) have prev(m, 1) prev(m, Thus, the above can be rewritten as
0"
Σ Yo Fm ((.rev(m.r))). mEC
Note that |C"| < M by the assumption of the theorem. This lemmn says that the difference between Or and Or is always equal to the of gradients that are being computed at iteration r. Having this link, we
24

---

Published in Transactions on Machine Learning Research (01/2025)
continue as follows
-00) E, | - 0-11]] = -10" - 112 -2ypEr (Fm, (0"), (m,,(0"), -00) 0-] + +1 [|| Fm, [0')
110" 0 12 - - [lFm, (0'm (0)),' 0- 0²)] +g [|| Fm., (0 [1)
+ 270 (F(0"), "0',,,"0 0 0r)
110 8-12
378 E, [[ m. (8)), H - 00) 70 E,
E, |( (8)), 0-1
2
2
+ E. [1 (0") }| + 27 (00)), 0 -0r)
(
110 0- 3 E, [ m. (0"),0' 00)
700 Er [1] [1))
2L
10
ET [I (0" [13] + 27o 2% ),0), - br)
2L
110" 102 3 (F(O"), 0- - - + +2 (F()), ),00 - Br)
2
110 0 1| (F(FF() 1r0- 0-) 70 (F(")) ) ) 00)
2
+ 27 (F(0", 0r 0r)
110- - --
10 |F(8)]12-
7o (F(("),0' 0 -0) + 2L
|| [[())
L
2
+ 21 110
00
= 110 - or 1|2
10
[00)
10 (F()"), 0 0²) + 2 10 0-
2L
2
where in (i) we use Assumption 1 and stepsize restriction 70 S 21. Shwartz inequality. Rearranging the terms, we have
in (ti) we use Assumption and Cauchy-
70
P(0))
< E, |10 - 11]
110- 1-1
78 (F()") 0 00) + 21-11110000---0 - 1011.
2L
2
After averaging over iterations from r = 0 to R 1 we get
aa 1 R-1
Σ
EIF(O)I
< 1000 - 1-
2L Σ E 110 0-12]
7
B BKF(O',, -0")
2L R
R
2R
Now we need to upper bound the third term. Using (16), we have
R- E 10 0-11
Σ
Σ mm
Fm (gprev(m.r))
P=0
2 < M
10
E
F m (gprev(m.r))11 [1P]]
r=0 meCr
R-1
As.1
S 20 ML
E
(Fm (fDr gprev(mer)- - -8-)
r=0 meCr
=
E
(F((0re((m)))) gerev(m.r) - -80]]
0= MEC
The term E [[F ((prev(m.r)) (Irev(m.r) -0*)] appears in the right hand side Tmax times at most. Indeed, it appears for all iterations between prev(m. and r, which is upper bounded by Tmax- Thus, we have
R-
1 110 - 001]] < ~ LTmax
R
EKF(O), -0
P=0
25

---

Published in Transactions on Machine Learning Research (01/2025)
If 70 S 2L./2M Tiaxx
1
then we have
R-1
R-1
Σ
110
= SL
Σ
E E( -8")
Thus, we derive
70 1 R-1
1100 10
2T
R-1
Σ
E
[1)01)))
Σ
B BKF(O'), -0-)
2L R
R
R
8L
r=0
R-1
70
Σ
E B(E(0'), 0- -
2R ==
1100 11
Finally, we get
< 20 1-
min E [1F(0')
ReR
Ro
Remark 3. We highlight the fact that if we use the stepsize 7a 2L.227
then the convergence is
min [1F((')) 4L 222 Faax
100
0-1
r<R
R
We observe square root dependency on Tmax- The result has been recently derived in homogeneous setting (Koloskova et al., 2022) for vanilla Asynchronous SGD.
D Byzantine-robust version
Preliminaries. We assume that among M clients participating in the training, there is a subset of clients B called Byzantines, i.e., clients that can (intentionally or not) deviate from the prescribed algorithm and are omniscient (i.c., they know the updates of other clients and aggregation rule applied by the server). More precisely, assumne that [M] = G U B, where G denotes the set of regular clients, |G] == G, |B] = B < 6M, where 8 < 1/2. The goal is to solve an instance of (2) with
F(0) = G mEG
(m).
(17)
where operators {fm {fm)mme Jme are defined as (0) V Im (0, w (0)).
Since the standard averaging is vulnerable to Byzantine attacks, we use robust aggregation rules in the sense of the following definition.
Definition 2 ((ō, ro-bou aggregator (Karimireddy et al., 2021; 2022; Gorbunov et al., 2023)). Let random
vectors ( # are such that there exists a subset g C [M] such that |g] = G 2 (1 8)n where < 1/2 and for some o 2 0 the following inequality holds: G(G-Tī Emnegg Em pm - x, [ S o² where the expectation is taken w.r.t. the randomness of Jmeg. Then, vector T is called (6, c)-uuu Aggregator ((6, c) RAgg) for some >> 0 and denoted as f RAgg(E1...... . if the following condition holds:
E ||
[| c502
(18)
where 7- = E Emeg Z4. If, in addition, the computation of independent of o², & is called (8, A)-Agnsti Robust Aggregator ((5, c) ARAgg) and denoted as T ARAgg (11 ... FM).
26

---

Published in Transactions on Machine Learning Research (01/2025)
Algorithm 5 Byzantine- Robuss Fine-tuning Followed by Global Gradient (BR-FFGG)
1: Input: initialization 0 RR, stepsizes Tw. 79 > 0, (6, -RAg
2: for r 0,1, 1...... do
3: for client m E g do
Compute and send gm( an unbiased estimate of (m) end for
for client B do
7: Send Im (0")
Byzantine clients can send anything to the server
8: end for
9: gr+1 = 0 - (0
10: end for
This definition is tight in a certain sense (see the details in (Karimireddy et al., 2021)) and is not satisfied for such defenses as Krum (Blanchard et al., 2017), geometric median (Pillutla et al., 2022a), and coordinate- wise median (Chen et al., 2017) that are known to be insufficient to ensure Byzantine-robusi ness (Baruch et al., 2019; Xie et al., 2020). For the examples of aggregators satisfying Definition 2 we refer to Gorbunov et al. (2023).
We also make additional assumption related to the overparameeterization.. Assumption 6. We assume that wny regular client E can compute an unbiased estimate gm g0) of Fm (0), e... [() (0), and for any 0 E TR.
E [Ilgmm (0)1 Piall (m))
(19)
addition, we assume that for aney E Rdd we have
1
Σ - MEG
||F(0)11 I Estm (F(0), 0*)
(20)
G
For example, inequality (19) is satisfied with Pmn when 9m (0) = (0), i.e., when regular clients compute Wm (0") and Fm, (f") exactly at each stop, and can be satisfied when the clients have overparameterized data (locally). Inequality (20) is satisfied with fxim S whenever Assumption 1 holds. However, saim can be much smaller than L if local operators {Fm hmgg are similar.
Finally, we made an extra assumption on structured non-monotonicity of operators Fmmt Assumption 7. We assume that for all m E g operators Fm are u-quassi strongly monotone, i.e., for all E Rd and H* such that Fm()") we have
(Fm (0),0 0*) 2 1110 1-112
(21)
Standard strong monotonicity, ie., (0,) - Fm (02), 01 02) 2 ul10: 02112 for any 01, 0g E Rd, implies condition from (21) (Mertikopoulos & Zhou, 2019; Song et al., 2020; Loizou et al., 2021) but the opposite implication is not always true. Moreover, as it is shown in (Loizou et al., 2021), an operator can be non- monotone but quasi-strongly monotone.
Theorem 8. Let Assumptions 1, 2, 6, 7 hold. Assume that
8 _
(22)
2c ( Pin + vtT)
Estm + (pmn 1)L)
70 S
(23)
4 (pm (pm-I)((+(emm)
+ 2c0 (p ctT)
Estm + (pin - 1)L))
Then, for any R 2 0 the iterates produced by Algorithm 5 satisfy
E [||0 - -0-12 = (1 70/1 R
1100 - 1º11
(24)
2
27

---

Published in Transactions on Machine Learning Research (01/2025)
Proof. To simplify the derivation, we introduce new vectors: g" 1
and g"
G Emeg Im (0)). Then, 0++1 = 0r Yoy" and
||0+ 8-1 = 110" 112 - 270 0" - -8)) + 10' 1|
0 I |10" 1-2 -270 0" -0") 27 0" -0" -8)
+ 27 I0'1 1 2 + 2 19" g'1l
$ (1 Y014
110" 1- 2 (g" A - 0r) + 2
+ 700
270
10" I'1?3
241
where in the second step we use ||a+b|] < |||a|2 |2|62 and, in the last step, we apply (a, S I|aa| 2a |b|l
that hold for any a, b E IRd and a > 0. Taking the expectation [,[[] w.r.t. the randomness coming from the r-th step and using E, [3'] = F(A"), ) we derive
E, [110+ - 101 -
< (1 + YoR
110" - 00-1 117-2210 (T0(F(0"), 0-) + 2, [|T'
2
+ 7o
2700 +
E, [lIsT --.
(25)
241
Next, use independence of f9m (0") (0m(0))mmm)) Immg:
E |1 1|] Er, Illg" +(00)12++111 (10(2)))12
Er [ll9m (0)) F'm (0 [[3] + 11
mEG
(19)
S
Pin
|| (0 112 [F
G2
mEG
Pmn
Σ |I 1 + (1+ Pin -1) |1F(0))12
1F(0)))P)) +1) (F(0)), -0).
G
G
G
mEG
(4)))2) S
(Pin [) (sim)
(26)
G
To upper-bound E, |10- 0" [11, we need to estimate PV, aid-1) E E, [l')m((") In (tr))1
PV, = G(G 1)
Σ mineo mar
E, (f + (0
Σ
(Fm)
G(G - 1)
m.neg minn
2 Σ E, [llgm (0')
2
1 Σ
F'm Fn)(
Fn FO
T
mEG
meg
HEG m.
Σ meg
E, |l 111) +
2
Σ
|| (1
Σ ((F
G(G - 1)
G-
mEg
mEg
(19) 2((G - 1)pmn + 1)
2G
(0'))
G --1 - 11
GGG-1)
meg
2((G 1)pmn +1)
Pm (0 - + +2(pm 1))1F(1))
G-1
meg
(20),(4)
2 2
((pu
1
+ G G-1
bsim + (pin 1)1,) (F(O"),0r -0))
(27)
28

---

Published in Transactions on Machine Learning Research (01/2025)
In view of Definition 2, this upper bound gives
Ex [11"" [ S coPV, (27)
20d ((mn +
--]
1
laim (pin 1)L) (F(0"), -0y. -0% (28)
G
Plugging (26) and (28) in (25), we get
Er ||0 0 11 < + 2014
110" 10
(Pm 1an--11 -1)(L falm) + +1)) (F(0),0-- 8*)
270
- 70
(29)
G
1)
((mm
Exim + (Pin - 1)1.) (F(0"), -0y
20070
2-10
+
2/1/
G G-1
(22)
(1 +
Y814
110" - 11
2
(Pin
1)([ + lalm) 1:)) (F(O)), -8)
270
70
G
((
1
+ 4cō70
Pin +
laim + (pin 111,) (F(0"), 0r 0*)
G-1
(23)
- [ + 70H)
110" 10"--011" - (F 0*)
(21)
< (1 Y014 110" 0-11 1|?
2
Taking the full expectation from the above inequality and unrolling the recurrence, we obtain the result.
The derived result establishes linear convergence to the exact solution (asymptotically, in expectation) with the possible presence of Byzantine clients. For simplicity, let us consider the case when Pin = 1. As mentioned earlier, this case corresponds to the exact computation of (0*) for all m G. Then, conditions (22) and (23) reduce to
8 5 2c (1
1
7o S
(1+
G-1 tsim
4 {(+
2cd (1 200
o-r) faim
In the worst case, faim = L that can be much larger than a and, thus, implies that 5 should be very small for the derived result. In the context of minimization, similar pathological behavior is observed in (Karimireddy et al., 2022; Gorbunov et al., 2023). In particular, the existing SOTA theoretical results under the assumption 2 >mmo fm (0, w)ll < f w)ll (Karimireddy et al., 2022; Gorbunov et al., 2023) require 8 S 1/en². However, when all functions {mm )me are I-smooth, have shared minimum, and is p-strongly convex, then, in the worst case, 13² = 1/u. Up to numerical constants and the differences between definitions of L and a in our work and in (Karimireddy et al., 2022; Gorbunov al., 2023), we get. the same worst-case upper- upper-bound bound for 8.
However, when faim & L, our condition on 8 can be very mild. For example, when the data on workers is similar to a certain extent, then fsim can be of the order of or even smaller. In this case, our condition on can become void, and the upper bound for will be determined by the type of aggregation rule (see the examples in (Karimireddy et al., 2022)).
E
Experiments
E.1 How does data heterogeneity affect the convergence?
In this section', provide the experiments Example 1 to demonstrate the robustness of the proposed FFGG algorithm varying the data heterogeneily of the optimization problem. To do so, we follow Section 5.1
Quu implementation this section available https: http::///gthubb.ccc/utttmeem wwt.huttt.ooooeeeeeeeeeeee http://eeen.ntaaaeeeeeeeeeeaaatr
29

---

Published in Transactions on Machine Learning Research (01/2025)
109
10"
10a
10
F1 8/ F16
10
10-1
10-2
10
10-2
-10
10-3
10-3
10
10-4
Local GD Scaffold
10
Local GD
Local GD
Scaffold
Scaffold
10
10
10
30
60
90
120
150
30
90
120
150
30
90
120
150
# Communication rounds
# Communication rounds
# Communication rounds
(a): <
(b): (= 40
(c): Ç 80
Figure 2: Comparison of FFGG against Scalfold, Local GD, and L2GD varying the data heterogeneily
constant Ç. FFGG uses gradient descent as subsolver with the same number of local steps as the other methods to make the comparison fair.
and construct the matrices Hm: Am, Bn and vectors Vm, bm as follows. First, sample matrices H. A, E nxxd and B Raxd from standard uniform distribution and rescale by the second dimension, i.e.
1 H.
1 [Unif(O, 1))mxd..
Unif 1)] m)),
B ~
dg
dw
Second, choose number Ç <> 0 and for each client m [M] define the matrices
H
Am = A A'
B'a
Hm = H
Bn = B BB
H l1
B 12
where
H A ~ 1)) B ~ [W ()] m))
Here the constant Ç controls the hetcrogencity of the problem. The higher the value of S, the more hetero- geneous the problem becomes. Finally, the vectors ym and bm are constructed by
Vm Am + Bm32 + 188, where y1 ~ [W(0, U2 ~ [W(0, 1)] .y (V(
bm Hm + tb2 where by ~ [W 1) (N(,,
and == 10-3, We set de 100, du 50, n 1000, M 32, and vary {20. 40, 80}. We consider the same set of algorithms as in Section 5.1 and set the theoretical values of the stepsize for all of them. All algorithms use local gradient updates with r - 20 local steps.
The results are reported in Figure 2. We observe that in all the cases FFGG algorithm achioves the gradient norm of order 10 6 after about 130 150 communication rounds, ie. the change in data heterogeneity has negligible impact on the convergence of FFGG. In contrast, the other three algorithms are affected more that they achieve gets worse with Ç. These results demonstrate the stability of
since the gradient FFGG in relation to the variations in the data heterogeneity of the problem and support the theoretical findings of our work.
E.2 The convergence of asynchronous FFGG
Next, demonstrate the convergence of the asynchronous version of FFGG summarized in Algorithm 3. We follow the experimental setup of (Mishchenko et al., 2022a) and test both synchronous and asynchronous versions of FFGG on the same problem as in the previous section. We set M 40, de = 400, dw = 50, n 1000, and Ç = 10. Both versions use - = {5, 10, 50, 200} local steps of GD. We use the Ray package (Moritz et al., 2018) to parallelize the execution and follow the official documentation for the implementation" of
paw.tttterraaaere.ee server. html Hasynchronous wwwprennnnupppratttereeee paraneter server-training
30

---

Published in Transactions on Machine Learning Research (01/2025)
10a
10a
10a
10F
10"
Angndtromuss F690
Syndrenssa FFGG
10'
Synstronoua MGG
10
asyndhoosss PRG
Syechrasssa
10"
Matchoowss FR00
Synchronzus FTG2
10'
10'
10'
10'
10"
10P
10"
10F
10
10-3
10-1
10-3
10
10"1
1000
2000
3000
4000
5000
10-1
1000
2000
3000
0000
5000
10
1000
2000
3000
4000
5000
10
1000
4000 # Computed updates
2000
3000
5000
Computed updates
# Computed updates
# Computed updates
Frequency ney
10
Frequenc ncy
10a
10a
Frequency
10
107
Frequere
10
10a
10'
10
10
10a
10
10'
10
250
300 Delay
750
1000
0
250
200 Delay
750
1000
0
250
500
730
1000
0
250
500
760
1000
Delay
Delay
(a):
7
(b): 7
10
(c): 50
(c): 200
Figure 3: First line: Comparison of asynchronous and synchronous variants of FFGG in terms of the number of computed update directions based on the local GD. Sccond line: Delays of updates using asynchronous FFGG.
asynchronous training. All delays appearing in the runs come from the execution on CPUs in contrast to previous works where the delays are simulated according to some distribution (Koloskova et al., 2022; Islamov et al., 2024). We tune both the inner and outer stepsize of FFGG by setting the stepsizes to theoretical values multiplied by a factor {0.25, 0.5, 1, 2,4..
We demonstrate the results in Figure 3 (first line) varying the number of steps of Local GD. We observe that the synchronous variant is faster than the asynchronous one. Moreover, the asynchronous FFGG improves by increasing the local steps similar to the synchronous FFGG. However, the improvement is slower than that of the synchronous variant. We report the delays throughout the optimization process in Figure 3 (second line). They suggest that the maximum delay is typically much larger than the number of workers in the system.
E.3 The convergence of FFGG with Byzantine workers
In our next experiment, we demonstrate the performance of the Byzantine-robust version of FFGG sum- marized in Algorithm& 5. To do so, consider the same data generation mechanism as in Section 5.1 with = 10000, de 100, dw 50, and M 42. To demonstrate the performance, we use the exact computa- tions on the worker side, i.c., each worker m computes locally argmin,,, fm. (fr-1, (00--1,))). w). We use coordinato-wise median with bucketing as an aggregation mechanism on the server (Chen et al., 2017: Karimireddy el al., 2020a). The 10 of all workers are Byzantine (which corresponds to ze 25% of all workers), and they flip the sign of the gradient. For each algorithm, we use the thcorctical values for the stepsizes. We report the convergence results in Figure 4. According to the plot, even in the presence of Byzantine workers FFGG algorithm still convergences linearly fast to the optimum. The convergence speed slightly slows down when there are Byzantine workers in training in comparison to non-Byzantine training as expected.
E.4 Expertimental details for real-world datasets
As mentioned in the main part, we adhere to the experimental arrangement outlined in (Pillutla et al., 2022b) for consistency. To provide comprehensive information, we present. the setup below.
Our experiments are conducted using two datascts encompassing two modalitics, specifically images, and text. These datasets feature a natural division of data that is non-iidd., mirroring the heterogeneity of data encountered in real-world Federated Learning scenarios. We provide detailed account of the experimental
31

---

Published in Transactions on Machine Learning Research (01/2025)
Ryrentte
Sycantan
10
10
200
400
600
600
1000
Commmnnizatinn rounds
Figure 4: The performance of FFGG algorithm with 25% of workers being Byzantine and without Byzantine workers.
setup and hyperparameters employed. We base our implementation on the publicly available code provided by Pillutla et al. (2022) 3
E.5 Datasets, tasks and models
We explore tasks inspired by real-world applications of Federated Learning: StackOverflow for next-word prediction and EMNIST for character recognition.
As part of our approach to partial personalization, we consider three partitioning schemes:
Input-layer personalization: This architectural design focuses on customizing the input layer to learn personalized representations while the remaining parts of the model are shared among all clients. Specifically, in the case of predicting the next word, we personalize the initial transformer layer instead of the embedding layer.
Oulpul-layer personalization: With this design, we aim to learn a shared representation while cus- tomizing the prediction layer. In the case of a transformer model, we adapt the final transformer layer instead of the output layer to achieve personalization.
Adapter personalization: In this scheme, each client utilizes a personalized low-rank adapter to fine-tune the global model.
These partitioning schemes serve strategies for incorporating partial personalization into Federated Learn- ing, allowing for different levels of customization while leveraging shared model across clients.
E.5.1 StackOverflow for next word prediction
Dataset. The dataset used for our task is derived from Stack Overflow, popular programming question- answer website. It consists of questions and corresponding answers. In the next word prediction task, the objective is to forecast the subsequent word based partial sequence of words within a question answer. This particular task serves as a valuable open-source benchmark for evaluating next-word prediction capabilities in mobile keyboards. For our experiments, we utilize the StackOverflow dataset made available by TensorFlow Federated
Client distributions. Each client our study corresponds an individual user on Stack Overflow, and the data available to each client comprises the questions and answers posted by that specific user. To ensure adequate amount of data for analysis, only include clients with minimum of 100 training sequences and 10 lesting sequences. Here, sequence refers to either question further narrow down the dataset by utilizing fixed subsample of 1000 clients.
answer posted by the user. We
Consistent with the approach doscribed in (Reddi et al., 2020), we limit the vocabulary to the top 10000 most frequently occurring words in the dataset. Additionally, we apply padding and truncation techniques
htttp://gthhu.ccmmaacbbrrrrrsrse hhttps://www. http:://wwww.tttooooooooo///eeee
32

---

Published in Transactions on Machine Learning Research (01/2025)
to standardize the length of each sequence within each client, setting it to 20, Furthermore, consider maximum of 1000 training sequences per client during our analysis.
Model. For our implementation, we employ a transformer model (Vaswani et al., 2017) that similar in size to BERT Mini (Ture et al., 2019). The model consists of translormer blocks, and each self-attention layer is equipped with attention heads. The transformer hidden dimension is set to 256, while the fully connected hidden dimension is 1024.
The model incorporates causal language modeling head, which refers to fully connected layer responsible for assigning scores to all possible vocabulary items, including special tokens.
Loss function and evaluation metric. During the training phase, we utilize the causal language mod- eling objective. This means that for each partial sequence, we treat the task of predicting the next word as multiclass classification problem and aim to minimize the cross-entropy loss.
For evaluation purposes, employ the top-1 accuracy metric, which measures the accuracy of predicting the correct word from the proper 10000-word vocabulary. This evaluation metric disregards special tokens such as padding, viitbbl-occablary terms, and beginning/end of sequence markers.
E.5.2 GLDv2 for visual landmark recognition
Dataset. Google Landmarks Dataset v2 (GLDv2) (Weyand et al., 2020) is large-scale image dataset. This dataset comprises pictures of well-known landmarks across the globe, all of which were captured and uploaded by contributors to Wikipedia. Although the image dimensions vary, the most prevalent size is 800 by 600 pixels.
The primary objective of the visual landmark recognition assignment is to pinpoint the landmark depicted in the image. This task mirrors real-life situations where individuals use their smartphones take pictures of natural architectural landmarks during their travels. We utilize the federated version of the GLDv2 dataset by (Hsu et al., 2020), which includes 2028 landmarks and is provided by TensorFlow Federated.
Client distributions. Every client is associated with specific Wikipedia user and includes all images contributed by that user. We only incorporate the 823 clients that have a minimum of 50 datapoints. We don't utilize the original lest set from GLDv2 for evaluation, as originates from distinct clients. Rather, we allocate 50% of the data from each client to be used a test set.
Model. We employ ResNet-18 (He et al., 2016) model, which has been pretrained the ImageNet dataset (Deng et al., 2009). We use group normalization in place of batch normalization. All images are resized to dimensions of 224 by 224 pixels. Our training incorporates two data augmentations: a random cropping to a 256 by 256 size and a random horizontal flip.
Loss function and evaluation metric. We use the cross-entropy loss. The model's effectiveness is evaluated based on its classification accuracy.
E.5.3 EMNIST for character recognition
Dataset. The FMNIST dataset (Cohen et al., 2017) serves as character recognition dataset. The ob- jective identify images containing handwritten digits or letters, with a total of 62 possible options encompassing lowercase and uppercase letters (a-z, A-Z) as well as digits (0-9).
The images within the dataset are grayscale and have dimensions of 28 28, resulting in total of 784 pixels. For our experiments, we utilize the EMNIST datasct made available by TensorFlow Federated.
Client distributions. In our study, each client represents an individual 'writer," " referring to the human subject who contributed by hand-writing the digit or letter during the data collection phase. We specifically consider clients that have minimum of 100 training points and 25 testing points, resulting in total of 1114 eligible clients for analysis.
33

---

Published in Transactions on Machine Learning Research (01/2025)
Table 3: Hyperparameters for each dataset/tane /task.
Hyperparameter
StackOverflow
GLDv2
EMNIST
Batch size
64
64
32
Devices round
50
50
10
Local epochs
Server optimizer
FedAdam SGD
FedAdam SGD
FedAvg
Common
Client optimizer
SGD
Global scheduler
Linear
Linear 10% of rounds
Exponential N/A
Warm-up
10% of rounds
L.R decay rounds
N/A
N/A
Max. grad. norm.
0.1
N/A
N/A
# Rounds
1000
2500
2000
Non-personalized training (step 1. of the pipeline)
Server learning rate
10
2 10
1.0
Client learning rate
1
10
0.5
# Rounds
500
600
500
Personalized training (step 2. of the pipeline)
Server learning rate
5 10-5
2 x 10-5
1.0
Client learning rate
10-1
10-3
10
# Epochs Optimizer
Local finetuning (step 3. the pipeline)
SGD 10-1
SGD 10-3
SGD 10-2
Client learning rate
Model. To address the smaller image size (28 x 28 1) in our datasct, which differs from the 224 224 size that the original ResNet was designed for, we utilize ResNet-18 model (He et al., 2016). However, we make two modifications to accommodate this smaller size.
Firstly, we adjust the convolutional kernel size in the first convolution layer from the original 7 x 7 to 3 3. This modification allows the model to process the input images appropriately.
Secondly, we omit the first pooling layer that is present in the original ResNot architecture. By removing this layer, we ensure compatibility with our image size and maintain the effectiveness of the model for our specific task.
Loss function and evaluation metric. We use the cross-entropy loss. We evaluate the performance of the model using its classification accuracy.
E.6 Experimental pipeline and baselines
There are three components in the training pipeline for all experiments:
Non-personalized federated pre-training: The first step involves training a global model with- out any personalization using FedAvg that we use to initialize (1).
2. Partially personalized federated training: This is the main training step that describe in detail in Section 5.3.
3. Final finetuning: client.
The last step involves only finetuning the personalized parameter w (0) for each
E.7 Hyperparameters and evaluation details
The hyperparameters we given in Table 3.
34

---

Published in Transactions on Machine Learning Research (01/2025)
Evaluation metric. Our main evaluation metric for both next-word prediction and image classification tasks is the weighted average of test accuracy across all clients. This weighted average takes into account the number of test examples from each client, allowing for a comprehensive assessment of model performance. This evaluation metric is equivalent to an unweighted accuracy achieved by aggregating all the data centrally.
Rounds. We utilize the concept of communication rounds, which refers to the number iterations during which the shared parameters are securely aggregated, to track the progress of each algorithm. In the case of non-personalized training, set the number of rounds to 1000 for StackOverflow, 2000 for EMNIST, and 2500 rounds for GLDv2.
For personalized training, we initialize the model with the parameters obtained from the non-personalized training and continue training for an additional 500 rounds for both StackOverflow and EMNIST datasets, and 600 rounds for GLDv2.
Devices per round. We assume that all devices are accessible and selections are made in a uniformly random manner. Consistent with the approach described (Reddi et al., 2020), we choose devices per round for /aackOverllow/GLDDv and 10 devices per round for EMNIST. This selection process applies to both non-personalized and personalized training for all datasets.
Local updates and minibatch size. In both non-personalized and personalized federated training, each selected device executes one epoch of mini-bate batch stochastic gradient descent locally. Following this, during the final fine-tuning stage of personalized training, we perform five epochs of training.
For the SeerlOverllow and GLDv2 datasets, we utilize a mini-batch size of across all settings. As for the EMNIST dataset, the mini-batch size is set 32 for all confiuraaions..
Server and client optimizer details. For the EMNIST dataset, we utilize the FedAvg algorithm, while for the StackOverflow dataset and GLDv2, employ FedAdam (Reddi et al., 2020). Additionally, we incorporate a global scheduler that applics schedule to the client learning rates across multiple rounds while maintaining a constant learning rate for each client within a round.
There are two types of schedulers we use: linear scheduler and an exponential scheduler (referred to as "stepLR" in PyTorch). The linear scheduler involves a linear warmup, applicable, until reaching the maximum learning rate, followed by linear decay 0. On the other hand, the exponential scheduler reduces the client learning rate by half after fixed number of rounds.
35


Reviews:
[{"review_id": "WMpUaYSrfo", "reviewer": "Reviewer_attF", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:  \n\n-\tThis paper proposes a novel approach to FL by introducing partially personalized parameters, which includes both globally shared and client-specific local parameters, to address the challenge of data heterogeneity across different clients.  \n-\tIt provides theoretical proof that under the right parameter split, it is possible to find global parameters allowing each client to fit their data perfectly, an idea the authors refer to as \"overpersonalized,\" which is a meaningful contribution to the field of personalized FL.  \n-\tIt analyzes the algorithm's performance under asynchronous training conditions and in the presence of Byzantine attacks, asserting that the method can converge under arbitrary delays and remain robust even with heterogeneous data, which is a significant advancement for the security and reliability of FL systems.  \n\nWeakness:  \n\n-\tThe manuscript urgently needs a dedicated methodology section for clarity and readability. This section should explicitly describe the algorithms in the Partially Personalized Federated Learning framework with step-by-step explanations, pseudocode, or flowcharts. Any modifications to existing algorithms must be detailed, emphasizing their differences and rationale. Visual aids (e.g., figures) can be incorporated to clarify complex ideas, and key points should be reiterated for emphasis. Theoretical discussions should follow only after establishing this practical foundation.  \n-\tWhile the paper presents a theoretically sound approach, its experimental validation is insufficient. The comparative analysis lacks breadth, falling short of evaluating the proposed algorithm against a wide range of state-of-the-art FL methods. To strengthen its claims, the paper should include a comprehensive set of benchmarks and demonstrate superiority in practical scenarios.  \n-\tThe concept of parameter splitting between global and local components is introduced but lacks a detailed explanation. The choice of parameter split is critical for performance, affecting model complexity, client-side storage, and personalization effectiveness, yet the paper offers no detailed methodology for determining the optimal split. Addressing this gap with a more nuanced optimization approach would significantly enhance the framework's practical utility.  \n-\tElaborate on how the proposed Partially Personalized Federated Learning approach differs from and improves/surpasses existing methods in a comparative manner, highlighting the unique contributions and improvements over the state-of-the-art.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces a Partially Personalized Federated Learning framework to address data heterogeneity in FL by splitting model parameters into global and local components. It claims that an optimal split allows for global parameters enabling perfect data fitting at each client, termed as \"overpersonalized.\" The proposed algorithm shows benefits in various FL settings, including local steps, asynchronous training, and Byzantine-robust training. The paper also discusses the practical applications of partial personalization in representation learning and heterogeneous environments and analyzes the method's convergence and robustness against Byzantine attacks under data heterogeneity.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "-\tIntroduce a dedicated methodology section that clearly outlines the algorithms and their step-by-step operations within the Partially Personalized Federated Learning framework. Include pseudocode or flowcharts for better understanding and use visual aids to elucidate complex concepts. Ensure that any algorithmic modifications are thoroughly explained, with emphasis on their distinctions from existing methods and the reasoning behind them.  \n-\tBroaden the experimental section to include a more comprehensive comparative analysis against a wider array of state-of-the-art Federated Learning methods. This expansion will provide a stronger empirical foundation for the paper's claims and demonstrate the practical superiority of the proposed algorithm.    \n-\tAddress the critical gap in the optimization of the parameter split between global and local components. Provide a detailed analysis or propose a methodology for determining the optimal split, which is essential for balancing model complexity, client-side storage requirements, and the effectiveness of personalization in the framework.   \n-\tThe manuscript should include a clear and concise comparative analysis that explicitly differentiates the proposed method from existing work in the field. This section should articulate the novel aspects of the approach and demonstrate its advantages, providing specific examples of how it improves upon or addresses limitations of previous methods.", "broader_impact_concerns": "Not applicable.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "lYRwTYIPEZ", "reviewer": "Reviewer_XKnK", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The paper addresses an important problem of partially personalized FL to address data heterogeneity.\n2. The paper proposes various modification to the proposed algorithm to adapt it to asynchronous learning and byzantine robustness.\n3. The clients are stateless by default.\n4. The paper presents a comprehensive theoretical analysis.\n\nWeaknesses/Questions:\n1. The paper is hard to follow. Please add a separate section for methodology and clearly describe the algorithms and proposed modifications/variations of algorithms before moving into the theory.\n2. The proposed algorithm has computation and latency overheads. For example, in algorithm 1, step 4 requires multiple iterations before 1 iteration of FedAvg which increases the compute and latency. Similarly Algorithm 4 requires multiple backward steps. The paper doesn't discuss any limitations/over-heads of the proposed method.\n3. The paper has limited experimental results. Not all the claims are backed up by the experimental results. \n4. In table 1, why is FedProx not listed under \"Handles data heterogeneity\"?\n5. What is the type of and level of data heterogeneity (if any) used in the experiments (Table 2)?\n6. Does the results in table 2 use Algorithm 1 or 4?\n7. In figure 1, $\\tau$ is related to the local steps to tune local parameters and this is a hyper-parameter that is specific to FFGG. So, why do the curves wrt to other algorithms (scaffold, local GD etc) are different across graphs b, c, and d? Please provide some clarity on how to interpret the figure.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces a partially personalized formulation of Federated Learning. The paper claims that the data heterogeneity slowdown can be completely eradicated via over-personalization. The paper proves that, in contrast to standard FL, asynchronous training with partial personalization converges precisely, and partial personalization can be made Byzantine robust. \n\nThe proposed algorithm divides the model parameters into global and local (personalized) set. The training includes three main steps:\n1. Non-personalized federated pre-training: Here the global set of parameters are trained using FedAvg without any personalization (the local set of parameters are frozen in this stage).\n2. Partially personalized federated training: In this stage, both the global and local parameters of the model are trained using Algorithm 1. Here, the following two steps are repeated until convergence -  a) The local parameters are randomly initialized and fine-tuned based on the local objective for $\\tau$ iterations, and b) perform one FedAvg step on the global parameters with number of local steps set to 1.\n3. Final fine-tuning: This step involves only fine-tuning the local (personalized) parameters for each client.\n\nAn extension to Algorithm 1 to incorporate local steps with respect to global and local parameters into the training loop is presented as Algorithm 4.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Please add a separate section for methodology.\n2. Add a separate limitation section discussing the overheads of the proposed algorithm.\n3. Include some experimental results to support the claims on data heterogeneity, asynchronous training and byzantine robustness. \n4. Since FedAvg (no personalization) and Local GD (only personalization) are the two extremes of the setup, it is useful to add Local GD results in Table 2.\n\nMinor changes:\n1. The reference in point 2 of section E.3 is missing (indicated as ??).", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "ijhSAMcQkq", "reviewer": "Reviewer_8Vd9", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. The theoretical analysis in this paper seems solid and comprehensive. However, one suggestion I have is that the author could consider summarizing their theoretical contributions into takeaway messages so that the audience can directly understand how it contributes to the community compared with existing works.\n\n2. This paper identifies an interesting problem - over-personalization, which can be seen as finding the right amount of parameter split between the shared parameters and private parameters. This question seems intuitive, and this paper proposes theoretical insights to guide this process.\n\nWeaknesses:\n\n1. Authors should highlight how stateless property is achieved using their proposed method, but this point is not clear in the current version of the paper.\n\n2. From a high level, the difference between this work and [1] is identified in Table 1; however, in detail (e.g., methodology level), how does the proposed method differ from the one proposed in [1]?\n\n3. Is the studied overpersonalization regime client-dependent? Intuitively, based on different amounts of local data, as well as the complexity of the distribution, the number of private parameters should be client-dependent. Is this point reflected in your analysis or algorithm?\n\n[1] Federated learning with partial model personalization, ICML 2022.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies the problem of Partially Personalized Federated Learning, which refers to the Federated Learning scheme where part of the model parameters are tuned in a client-dependent manner. This can be seen as an analog of pre-training and fine-tuning and is a trending paradigm for Federated Learning in recent years [1]. Specifically, this paper improves upon [1] by (1) proposing a stateless optimization scheme where the clients are assumed to be unable to maintain the state of the model and (2) relaxing the assumption on the step size. In addition, this paper also identifies the problem of finding the appropriate split to control the size of shared parameters and private parameters. \n\n[1] Federated learning with partial model personalization, ICML 2022.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Summarize and highlight the contributions of the theoretical analysis, differentiating them between those \n\n2. Highlight how the stateless is achieved in your proposed method.\n\n3. Authors might consider using a separate section to introduce the methodology part.", "broader_impact_concerns": "Not applicable.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "crF4xDDhpY", "reviewer": "Reviewer_cUMT", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is well written and generally easy to follow. The idea is interesting and the work shows that there is some benefit to having partial personalization. \nHowever, in some places, further explanation is required:\n- When comparing to the work of Pillutla et al. in the experiments, the comparison seems to be only with the FedAlt. Despite that in the paper by Pillutla et al., FedAlt almost always outperforms FedSim, however, since here the work is using the stateless clients and possibly changing a little in the problem, it would be good to see if FedSim would do better. Moreover, the difference between the FDGG used in this work and FedAlt seems to be very small, so it would be good to understand better why one would prefer to use this method.\n- In the experiments, the FDGG method that performs best uses one partitioning scheme or another depending on the dataset. It is important to try to explain why one is preferable to another (or some intuition), what are the possible factors, and how would one decide which to use when using this method.\n- In the proof of proposition 1, it is unclear to me how one would pick such a $\\theta$ that all $f_m$ do not depend on.\n- In section 3.2, on what basis do you pick the values of the parameters in these experiments?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies the problem of partial personalization for federated learning. The effects of having partial personalization, where the variables are split into local private parameters and global parameters that are shared among the clients, while the clients are stateless, on the accuracy of the model are studied.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "It is important to respond to the questions in the previous section. Additionally, here are some minor comments:\n- Subsections 1.3-1.5 seem a little out of place. If those are related works, put them under that section, while if they are contributions, then put them under contributions. Having them as just subsections in the introduction is a little confusing.\n- There are some typos in the text, please proofread.", "broader_impact_concerns": "I do not see any ethical implications of the work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper analyzes an approach for stateless personalized federated learning by considering modeling via a combination of global and local-only parameters. The authors analyze their approach in several scenarios, including communication-efficient training regimes such as local-updating & asynchronous training scenarios, as well as in the presence of malicious clients. The reviewers provided significant suggestions to improve the presentation of the work and better position the work with respect to prior art in personalized FL, which the authors have incorporated in their revision. Although some uncertainty remains regarding the empirical improvements of this approach relative to existing methods, reviewers were nonetheless impressed with the theoretical insights in the work and uniformly recommended acceptance of the paper.
Decision: Accept (poster)