Straggler-Resilient Personalized Federated Learning
Authors: Isidoros Tziotis
Zebang Shen
Ramtin Pedarsani
Hamed Hassani
Aryan Mokhtari
Venue: Accepted by TMLR ()

Abstract: Federated Learning is an emerging learning paradigm that allows training models from samples distributed across a large network of clients while respecting privacy and communication restrictions. Despite its success, federated learning faces several challenges related to its decentralized nature. In this work, we develop a novel algorithmic procedure with theoretical speedup guarantees that simultaneously handles two of these hurdles, namely (i) data heterogeneity, i.e., data distributions can vary substantially across clients, and (ii) system heterogeneity, i.e., the computational power of the clients could differ significantly. By leveraging previous works in the realm of representation learning (Collins et al., 2021; Liang et al., 2020), our method constructs a global common representation utilizing the data from all clients. Additionally, it learns a user-specific set of parameters resulting in a personalized solution for each individual client. Furthermore, it mitigates the effects of stragglers by adaptively selecting clients based on their computational characteristics, thus achieving for the first time near optimal sample complexity and provable logarithmic speedup. Experimental results support our theoretical findings showing the superiority of our method over alternative personalized federated schemes in system and data heterogeneous environments.

Published in Transactions on Machine Learning Research (10/2023)
Straggler-Resilien Personalized Federated Learning
Isidoros Tziotis
isidoros 13 130uterss.ed edu
The University of Texas at Austin
Zebang Shen
zelang.shensüinffellae
ETH Zurich
Ramtin Pedarsani
rambin BeleiiiBeceuccbbb..oo edu
The University of California, Santa Barbara
Hamed Hassani
hassoni@seas. wpenn. edu
The University of Pennsylvannia
Aryan Mokhtari
makkhlari@@ouutinnnnkeee eduu
The University of Texas at Austin
Reviewed on OpenReview: https: /ttpsr///ppereevie.t net/ forum? /at.g//phh.ttddd//eeeeee/e///aaa
Abstract
Federaled Learning is emerging learning paradigm that allows training models from samples distributed across a large network of clients while respecting privacy and communication restrictions. Despite its success, federated lcarning faces several challenges related to its decentralized nature. In this work, we develop a novel algorithmic procedure with theoretical speedup guarantees that simultancously handles two of these hurdles, namely (i) data heterogeneityy i.e., data distributions can vary substantially across clients, and (ii) system heterogeneity, i.c., the computational power of the clients could differ significantly. By leveraging previous works in the realm of representation learning (Collins et al. 2021 Liang
et al. 2020), our method constructs global common representation utilizing the data from all clients. Additionally, it learns a user-specific set of parameters, resulting in a personalized solution for each individual client. Furthermore, mitigates the effects of stragglers by adaptively selecting clients based on their computational characteristics, thus achieving for the first time near-optimal sample complexity and provable logarithmic spcodup. Experimental results support our theoretical findings showing the superiority of our method over alternative personalized federated schemes in system and data heterogeneous environments.
1 Introduction
Due to growing concerns on data privacy and communication cost, Federated Learning (FL) has become an emerging learning paradigm as it allows for training machine learning models without collecting local data from the clients. Due its decentralized nature, a major challenge in designing efficient solvers for FL is hetcrogeneity of local devices which can be categorized into two different types: (i) data heterogeneity where the underlying data distributions of clients vary substantially, and (ii) system helerogeneily where the computational and storage capabilities of devices differ significantly. In fact, it has been observed that the seminal Federated Averaging (FedAvg) method suffers from slow convergence to high quality solution when facing highly heterogeneous datasets (McMahan et al. 2017 as well as heterogeneous systems (Li et al. 2020
Kairouz et al. 2021
In this paper, we aim address these two challenges simultancously by introducing generic framework that includes algorithms with robust performance in the presence of those forms of clients' heterogeneity. Inspired by prior works in the literature of FI. (Collins et al. 2021 Liang et al. 2020 that utilized representation

---

Published in Transactions on Machine Learning Research (10/2023)
learning theory to tackle data heterogeneity, we propose meta-algorithm that produces personalized solutions and handles data heterogeneity by leveraging a global representation shared among all clients. Further, our method circumvvntt the delays introduced due to the presence of stragglers by carefully selecting participating nodes based on their computational speeds. In carly stages, only a few of the fastest nodes are chosen to participate and sequentially slower devices are included in the training process until the target accuracy is achieved. Although the disproportional selection of nodes raises fairness and accuracy concerns, we highlight that our method achieves speedup without compromising the resulting solution. The most significant contribution of our work is achieving near-optimal sample complexity in regimes with data and system heterogeneity, alongside a provable logarithmic speedup guarantee in terms of running time. Next we summarize our contributions:
1. SRPFL Algorithm. We propose Stsslieer-Resilientt Resilient Personalized Federated Learning (SRPFL), adaptive node participation meta-algorithm that builds upon subroutines that fall into the repre- sentation learning framework (Collins et al.) 2021 Liang et al. 2020) enhancing their resilience to stragglers and performance.
2. Logarithmic Specdup. Assuming that clients' speeds are drawn from the exponential distribution, we prove that SRPFL guarantees logarithmic speedup in the linear representation setting, outperforming established, aaraggle--prone prone benchmarks while maintaining the state of the art sample complexity per client m = O(d/N + log(N))), where d and N denote the feature vector size and number of active nodes. Our results hold for non-convex loss functions, heterogencous data and dynamically changing client's 's speeds.
Numerical Results. Experiments on various datasets (CIFAR1O, CTFAR100, EMNIST, FEMNIST, 14nt support our thcoretical results showing that: (i) SRPFL significantly boosts the performance of different subroutines designed for personalized FL both in full and partial partieipation settings and (ii) SRPFL exhibits superior performance in system and data heterogeneous settings compared to state-of-the-art basclines.
1.1 Related Work
Data heterogeneity. In data heterogeneous settings, if one aims minimizing the aggregate loss in the network using the classic FedAvg method or more advanced algorithms, which utilize control-variato techniques, such as SCAFFOLD (Karimireddy et al. 2019), FEDGATE (Haddadpour et al., 2021), FedDyn (Acar lt al. 2021) or FEDSHUFFLE (Horváth et al.) 2022) the resulting solution could perform poorly for some of the clients. This is an unavoidable hurdle due to the fact that no single model works well for all clients when their underlying data distributions are diverse. A common technique that addresses this issue is fine-tuning the derived global model to each local task by following a few steps of SGD updates (Wang et al. 2019 Yu et al. 2020). Based on this observation, Fallah et al. (2000b) showed that one might need to train models that work well after fine-tuning and showed its connections to Model-Agnostic Meta-Learning (MAML). In (Cho et al. 2022 Balakrishnan et al. 2021 novel client-sampling schemes were explored achieving increased efficiency in regimes with data heterogeneity. Another line of work for personalized FL is learning additive mixtures of local and global models (Deng et al. 2020: Mansour et al.] 2020 Hanzely and Richtárik 2020). These
methods learn local models for clients that are close to each other in some norm, an idea closely related to multi-task FI. (Smith et al., 2017 Hu et al. 2021). The works (Chen al.] 2022: Lee et al. 2022 studied the interplay of local and global models utilizing Bayesian hierarchical models and partial participation, respectively. An alternative approach was presented by Collins et al. (2021), where instead of enforcing local models to be close, the authors assumed that models across clients share common representation. Using this perspective, they presented FedRep method that provably lcarns this underlying structure in the lincar representation setting. Building upon the idea of a common representation Zhu el al. (2021) and [Jiag and Lin (2022) proposed federated methods that can handle data heterogeneity while exhibiting robustness to distribution shifts. Recently, a novel framework was proposed allowing the comparison of personalized FL methods under various metrics (Wu et al., (022).. In all of the aforementioned methods however, a subset of clients participate regardless of their computational capabilities. Thus, in the presence of stragglers, the
2

---

Published in Transactions on Machine Learning Research (10/2023)
speed of the training process significantly goes down as the server waits, at every communication round, for the slowest participating node to complete its local updates.
System heterogeneity. Several works have attempted to address the issue of system heterogeneity.
Specifically, asynchronous methods, which rely on bounded staleness of slow clients, have demonstrated
significant gains in distributed data centers (Xie et al. 2019 Stich 2019 So et al. 2021). In FT. frameworks, however, stragglers could be arbitrarily slow casting these methods inefficient. In an attempt to manually control staleness, deadline-based computation has been proposed (Reisizadeh al. 2019) as well as aggregation of fixed number of models per round (Nguyen et al., 2022). However, in the worst case scenario the performance of these methods is still determined by the slowest client in the network. Active sampling is another approach where the server aggregates as many local updates as possible within a predefined time span (Nsshio and Yonetani, 2019 In a different line of work the effects of stragglers are mitigated by utilizing computation/grant /radient coding schemes (Tandon et al. 2017: Wang et al., 2018 2020b). In (Cho et al.) 2021 clients use hetcrogencous model-architoctures to transfer knowlodge to nodes with similar data distributions while Yang et al. (2022) proposed a new framework where clients are free to choose their participation scheme. More recently, normalized averaging methods were proposed in (Wang et al. 2020a Horváth et al.) 2022) that rectify the objective inconsistency created by the mismatch in clients' updates. FedLin (Mitra et al. 2021j instead utilizes gradient correction and freor- mechanisms to circumvent the speed- accuracy conflict. A novel approach to mitigate the effects of stragglers was proposed by Reisizadeh et al. (2022). ) where clients are selected to take part in different stages of the training according to their computational characteristics. Alas, all of the above methods yield improvement only in data- data- homogeneous settings and they are not applicable in regimes with data heterogeneity.
2 Problem Formulation
In this section, we introduce our setting and define the data and system heterogeneity model that we study. Consider the FL framework where M clients interact with central server. We focus supervisod, data-heelerogeneous setting where client i draws samples from distribut D;, polentially with D; DD Further, consider the learning model of the i-th client as q; Rd + 3 which maps inputs x; E RR to predicted labels (x,) E y. The objective function of client i is defined as f((g:)= = B. (x.), (xx.),))), y:)],, where the loss f: y x y - R penalizes the gap between the predicted label g: (x.. and true label y:. In the most general setting clients aim to solve
M fl(e),
1
min ( JE
(1)
M
with Q the space of feasible tuples of mappings (g1, qaa) Traditionally in FL., methods focus learning single model --- qM that performs well on average across clients [Li et al. 2020 McMahan et al.
2017). Although such solution may be satisfactory in data-homogeneous settings, it leads to undesirable local models when the data distributions are diverse. Indeed, in the presence of data heterogeneity the loss functions fi have different forms and their minimizers could be far from each other. This justifics the formulation in [ and necessitates the search for personalized solutions that can be learned in federated fashion.
Low Dimensional Common Representation. There have been numerous examples in image classification and word prediction where tasks with heterogeneous data share common, low dimensional representation, despite having different. labels (Bengio et al., 2013 LeCun et al. 2015 Pillutla et al. 2022). Based on that, a reasonable choice for Q is a set in which all d share a common map, coupled with a personalized map that fits their local data. To formalize this, suppose the ground-truth map can be written for each client as g:=h,o o o where of is a shared global representation which maps d-dimensional data points to a lower dimensional space of size i and h, RR* -,, which maps from the lower dimensional subspace to the space of labels. Typically kedd and thus given any fixed representation o, the client specific heads h2 R are easy to optimize locally. With this common structure into consideration, ① can be reformulated minge 37 heads. This formulation leads to good local solutions, if the underlying data generation models for the clients share a low dimensional common representation, i.e., y/ h o o² (xx) + 24, where 2; is some additive noise.
1 .M min fo fh,,, oΦ), where & is the class of feasible reprosentation and H is the class of feasible
3

---

Published in Transactions on Machine Learning Research (10/2023)
Upshte be reprncttzics
Uptar the sprosemttion
Updine teal on of
Updare lsaal of
sy
M-3
M
M& M-3
Selncing Saturaade
Sclaslag
actire codes
Nadior
N5-5 M-3
mmpky
X
Stage
Stage
Figure 1: Classic FL schemes for solving P
Figure 2: SRPFL for solving
2
The server and clients collaborate learn the common representation é, while locally each client learns their unique head h4- Since clients often do not. have access to their true data distributions, instead of minimizing their expected loss, they settle for minimizing the empirical loss associated with their local samples. Specifically, we assume client has access to St samples xf,,>? ..., x,, J, and its local empirical loss is f(l, Φ) 5, E. l(1 o ((x;), Hence, the global problem becomes
min GE M
Σ min f((h) NEH
Σ lhe ((x)),n)) S.
1
System Heterogeneity Model. In most FL settings, thousands of clients participate in the training process, each with different computational capabilities. Thus, fixed computational tasks such as gradient computations or local model updates require different processing time, for different clients. Formally, for each client i E [M], we denote by T. the time required to compute a local model update. When a subset of clients participate in the learning process, the computational time of each round is determined by the slowest, participating node. Naturally, as the number of nodes in the network grows, we expect the number of stragglers to increase. This phenomenon calls for the development of rrsggler-resilient methods in system- heterogeneous settings.
3 Straggler-Resilien Personalized FL
In the shared representation setting, we face the challenge of finding an algorithm that coordinates server and clients in order to learn common reprosentation and set of personalized parameters in federated and straggler-resilient fashion. To this end, we propose method that tackles problem (2) with limited sample access and provably superior performance over naive, straggler-prone methods. Specifically, we propose the Straggler-Resilien Personalized Federated Learning (SRPFL) meta-algorithm, designed to mitigate the effects of system heterogeneity in environments with non-convex loss functions and heterogencous data while accommodating a variely of methods as subroutines. In a nutshell, SRPFL iteratively solves problem (2) while adaptively increasing the set of participating nodes based on their computational capabilities. As a result the process of learning the common representation across clients' tasks is accelerated, without compromising the resulting accuracy.
To simplify the exposition, henceforth we denote by A some federated algorithm of choice, designed to solve (2). As depicted in Figure [ in standard FL frameworks, out of all M clients in the network, the server often selects uniformly at random a subset of size N. Subsequently, a few iterations of algorithm A are performed to approximately solve a version of (2) which corresponds to those N selected clients i.c., min N Σ minn,eH fe(h, o Φ). In every following stage, a new subset of N nodes is sampled and the same process is repeated. Although such a procedure eventually learns the global representation across all tasks, it is susceptible to delays cansed by stragglers, as the server has to wait for the slowest client (among the N selected ones) to complete its updates. Hence, when N is large the training procedure could become prohibitively slow.
4

---

Published in Transactions on Machine Learning Research (10/2023)
SRPFL takes different approach in order to mitigate the effects of straggling clients. An overview of the selecting scheme is provided in Figure 2 At each stage, N clients are randomly selected, but only a small, fast. subset of them is used to solve their corresponding learning problem. More precisely, suppose that at stage r, each client i in the sampled subset of size N, is associated with a computational time T.". For simplicity we assume that the nodes are re-indexed al every slage so that they maintain a decreasing ordering w.r.t. their times, i.c. ' S T S ... S TTy Initially, only the no fastest clients, {1,2,.., no), are included in the learning procedure, with no much smaller than N. At every communication round, the set of no nodes perform the model updates indicated by algorithm A, to solve their version of (2). i.e. minges no E ming, fi o)) We note that during this stage of the algorithm, the server waits only for the slowest client among the participating ones, i.e, client Ro which is potentially much faster than node N.
Remank 3.1. In practice, the knowledge of clients' computational power is not required to figure out the no fastest nodes. Instead, the server sends the global representation model to all N sampled clients and updates the common represenlation once the first no new models are received. Indeed, these representations belong to the fastest no nodes.
Once the current stage terminates, a new batch of N clients is sampled and the 2n1o fastest nodes are chosen to participate in the learning process. Since speeds vary across stages, consecutive sets of participating nodes could have small or no overlap. However, the representations learned in previous stages still operate as good starting points for subsequent stages which is possible since nodes are homogeneous w.r.t. their representations (see Section 2 Thus, utilizing the representation model learned from the previous stage, nodes (1, continue the learning process deriving a model of improved accuracy, The procedure of geometrically increasing the number of participating nodes continues until the target accuracy is achieved. Hence, SRPFL uses the data of stragglers only at the latest stages of the algorithm when an accurate approximation is required. Remark 3.2. For simplicity of exposition we assume that the set of N sampled nodes maintains connectivity to the server throughout each stage. However, our analysis remains unaffected even if new set of nodes is sampled at every round.
We proceed to characterize the class of federated algorithms able to employ SRPFL to enhance their performance in system heterogeneous environments. Any iterative method that solves an instance of ? can be combined with our adaptive node participation scheme, however in this paper, we focus on a broad class of alternating gradient- based update methods, presented in (Collins et al.) 2021). As the name suggests, in each round, clients update their heads and representation models in alternative fashion. After certain number of gradient updates is completed, all clients send their derived representations to the server where the models are averaged and broadcasted back to the clients. Next, we rigorously illustrate this procedure.
Alternating Update Scheme. At round t, the server communicates a common representation & to the clients and a subset of them It, are selected lo participate by performing the following updates.
Client Head Update. Each client i E I' performs Th local gradient-based updates optimizing their head ,..., Ta client updates their head parameter, given the received representation d'. Concretely, for model as follows
1 GRD (f: (M.--) (00)
;
1 .1)
(3)
Client Representation Update. Once the updated local heads nTTT are obtained, each client i exccutes To local updates on their representation parameters. That is for s s= l,
= GRD (Fi (M (00=,0810-)) (0000) .06-1,100--1111 08.01117)
(4)
In the above expressions, GRD(f, h, n) captures the generic notion of an update of variable h using the gradient of function ff with respect to h and step size n, This notation allows the inclusion of large class of algorithms such as Gradient Descent with momentum, SGD, etc.
Server Updale. Each client i sends their derived representation models of to the server, where they are averaged producing the next representation model +t+1
Coupling SRPFL with a generic subroutine that falls into the Alternating Update Scheme, gives rise A Algorithm Every stage is characterized by participating set of size 2" no, denoted by I'. At the
5

---

Published in Transactions on Machine Learning Research (10/2023)
Algorithm SRPFL
Input: Initial number of nodes no; step size ny; number of local updates for head Thi number of local updates for representation Tg. ,..,,N 1.
Initialization: <- no, @o. R1 T
for =,,1,2, 0.,,,....,II((NN Iog(N/no) do
Φ 0r
5: for 1, do
6: Server sends representation 6²-1 to N clients sampled from [M].
7:
for E do
Client i initializes h,-0 f- hi 111% and runs Tu updates hę = h 17 fi ( @"
8:
Client i initializes o, + 0--1 and runs To updates of =d= - of 17
10: Client i sends Te to the server.
11:
end for
12:
for each client i i I" do
h4tTh + h 1,7. 14: end for
13:
15: Server computes of t Ei= LTo
16: end for
17: Server sets + ( 2n) and Ør+1 + 6².
18: end for
beginning of cach round the server provides representation model to the participating clients (line 6). The clients update their models (lines 8 and 9) and sent their representations back to the server where they are averaged producing a new global model. The numbers of local updates The To depend on the subroutine method of choice and the number of rounds per stage is denoted by Ty. At the end of every stage the set of participating nodes is doubling in size until a set of size N reached (line 17).
Remark 333 summarizes the technical novelties our work and highlights crucial benefits enjoyed by SRPFL. Remark 3.3. Reisizadeh et al. 2002)) proposed similar participation scheme, however their approach differs from ours and their results apply to significantly more restrictive regimes. Specifically, in (Reisizadeh et al. 2022) the analysis heavily relics on deriving conncction between the ERM solutions of consecutive stages. In order lo control the statist ical accuracy of the corresponding ERM problems (i) data homogeneity across all clients is necessary and further (ii) clients who participate in early stages are required to remain active and connected to the server in subsequent stages, maintaining fixed computational speeds throughout the whole training process. (iit) The results of their analysis hold only for strongly convex loss functions and (iv) their stage termination criterion requires the knowledge of the strong convexity parameter. The above restrictions are detrimental in the FL regime and severely undermine the applicability of the resulting algorithm.
In our work we follow a different approach controlling -in terms of principal angle distance a quantity analogous to statistical accuracy, therefore directly connecting the common representation (and overall solution) at every stage to the ground truth representation. This novel approach allows our algorithm to accommodate (i) data heterogeneity, (ii) clients with dynamically changing speeds or equivalently clients that are replaced by new ones at every round, and (iii) non-convex loss functions. Additionally, a major part of our technical contribution focuses (iv) analytically deriving the optimal number of rounds per stage, thus producing simple and efficient doubling scheme.
4 SRPFL in the Linear Representation Case
Our theoretical analysis focuses on specific instance of ( where clients strive to solve a linear represent a ation learning problem. Concretely we assume that f. is the quadratic loss, & is a projection onto a k-dimensional subspace of Rdd given by matrix E Rdxk and the i-th client's head h;, is a vector w, E R We model local data of client i such that yi w T + 2,, for some ground truth representation B Rddk local heads W E k and 2 o2) capturing the noise in the measurements. Hence, all clients' optimal solutions lie

---

Published in Transactions on Machine Learning Research (10/2023)
Algorithm PedRep--RPFF (Linear Representation)
1: Input: Step size y) Batch size m; Initial nodes fo;
2: Initialization: Client [N] sends to server:
P;: +. (x
n - ni-
Server finds UDUT + rank-k SVD( E P,).
for 1, log( (nn)) do
Server initializes representation B f- U.
6: for 1, do
7:
Server sends B to N clients sampled from [M].
8:
for i (1,.,n do
Client i samples fresh batch of m samples.
Client i computes w arg min ft w*
10:
Client i computes BT.+1 + B VBB Brr) and sends it back to the server.
11:
12:
end for
Server computes Bra+1 1 n Σiez B 14: Server computes Br.t1 RR Rr + QR(B..+1
13:
15: end for
16: Server sets U <- Br++1 and n min [N, 2m]
17:
end for
in the same k-dimensional subspace. Under these assumptions the global expected risk is
M
Ex..B [(v. - w,, 1=1
min B,W
(5)
2M
where W = [w], wil E x is the concatenation of the client-specific heads. Since the true distributions Dj's are unknown, algorithms strive to minimize the empirical risk instead. The global empirical risk over all clients is
Σ fw, B 2Mm
1
m
v w/ [Tx](
(6)
where is the number of samples at each client. The global loss in (6) nonconvex and has many global minima, including all pairs of (W*Q, Q (V=Q=I,BQQ)) where Q E R is some invertible matrix. Thus, the server aims to retrieve the column space of B", instead of the ground truth factors (W", B*). To measure closeness between column spaces, we adopt. the metric of principal angle distance (Jain et al. 2013).
Definition Let matrices B, B2 E daxk and B orthonormal matrices s.t. span(Bi,.1 = span( (Ba) The principle angle distance between the column spaces of B, and B is defined to be dist [B, B2) = |B B2lla-
and span( (Ban
Federated Representation Learning (FedRep) an alternating minimization-desceenn algorithm, recently proposed in (Collins et al. 2021) for the Linear Shared Representation framework. SRPFL coupled with FedRep gives rise to Algorithm 2 Below we highlight the main points of interest.
In the initialization phase (lines 2 and 3) model of bounded distance from the optimal representation is obtained, via the Method of Moments (Tripuraneni et al. 2021). Subsequently, at every round t, client samples a fresh batch of samples {x , ". 47mm from its local distribution (line 9) and thus the corresponding
loss function becomes f((w, o B)== 2m E w.B.. 22 Utilizing the global representation provided by the server, client i computes the optimal head w; (line 10). Fixing the newly computed head, client proceeds to update its global representation model with one step of gradient descent (line 11) and transmits it back to the server. As depicted in lines 13 and 14 the parameter server averages the models received and orthogonalizes the resulting matrix before broadcasting it to the clients, a component of crucial importance required in our analysis.
Mapping this method back Algorithm T we note that the number of representation model updates Ta is set to 1, whereas the number head updates 7h is sufficiently large derive (approximately) the optimal solutions. This imbalance is designed to take advantage of the inherent. structure of our problem where the
7

---

Published in Transactions on Machine Learning Research (10/2023)
size of wy's is significantly smaller than d. Finally, we point out that the mumber of the communication rounds per stage is a small and priori known to the algorithm constant, specified by our analysis.
Remark 4.2. Although our proposed method utilizes FedRep as backbone, our analysis and framework differ substantially from the ones in (Collins et al. 2021). Specifically, Collins et. al. (2021) assume access to (i) infinite samples, (ii) without the presence of noise. Additionally, (iii) the number of participating nodes remains fixed throughout the training and (iv) the focus lics solely on handling heterogencous data with system heterogeneity being an orthogonal direction to their work. In contrast, our analysis requires only (i) finite and (ii) noisy samples, and our theoretical results (Theorems 4.7 and 4.9 revolve around (iii) participating subsets of different sizes and (iv) regimes where both data and system heterogeneity is prevalent.
4.1 Theoretical Results
In this subsection, provide rigorous analysis of FedRep-SPPFL in the linear representation setting. First, we present the notion of Wall Clock Time (WCT) which is the measure of the performance for our algorithm. Subsequently, we illustrate the contraction incquality that determines the rate at which the distance to the optimal representation diminishes. We conclude showing that FedRep-SRPFL out performss its straggler-prone variant by a factor of (llo N), under the standard assumption that clients' computational times follow the exponential distribution. Before we proceed, we introduce the necessary notation and the assumptions.
Eg = = 1 - dist (B".B").
(7)
1
"max, x I maax
max
Wi
(8)
nn
1
min,
min IE
Jmin
W
(9)
Assumption 4.3. (Sub-gaussian design). The local samples x; Rdd are i.i.d. with mean 0, covariance La and are La-sub-gausiaan, ie. Ele²"x Assumption 4.4. (Client diversity). Let min,+ defined in 99. be the minimum singular value of any matrix that can be obtained by taking rows of 1 W* Then min, 0. V
x1 "xe]xx1] /v1l/22 for all v Rd
Specifically, our theoretical analysis requires Assumption 1.4 to be satisfied for every 2 no- Assumption 4.4 implics that the optimal heads, of the participating clients, span R.a. This is true in many FL regimes as the number of clients is usually much larger than the dimension of the shared representation.
Remark 4.5. In this work we consider client speeds being independent of the local data available to them. This is a natural assumption since the computational power of each client crucially depends on their device characteristics (battery, CPU, etc.), whereas any connection to their local data is unclear. However, in the presence of strong correlation between data and system heterogeneity, Assumption 4.4 may not hold, which can be seen as a potential limitation of our work and an interesting future dircction to explore.
Assumption 4.6. (Client normalization). The ground-truth client specific parameters satisfy || vk for all i E [x] and B" has orthonormal columns.
Assumption 4.6 ensures the ground-truth matrix W is row-wise incoherend, Lee. its row norms have similar magnitudes. This is of vital importance since our measurement matrices are row-wise sparse and incoherence is common requirement in sensing problems with sparse measurements.
Wall Clock Time. To measure the speedup that our meta-algorithm enjoys we use the concept of real time or WCT as described below. FedRep-SRPFL, runs in communication rounds grouped into stages. Consider such a round l at stage r, with nodes 21, ) participating in the learning process. Here ny denotes the slowest participating node. The expected amount of time that the server has to wait for the updates to take place is E. [TT... Put simply, the expected computational time of the slowest node acts as the bottleneck for the round. Further, at the beginning and at the end of every round, models are exchanged between the server and the clients. This incurs an additional, fixed communication cost C. If 7, communication rounds take place al. every stage r, then the overall expected WCT for FedRep-SRPFL, E SPSP) > (log(N/mo)) T, (E [TT.]] [0]+++0 CC. Similarly, the total expected runtime for FedRep can be expressed in terms of the total number of rounds, TER, E. [T - [TK] Te).iaking the ratio of these quantities derives the desired speedup guarantee.
8

---

Published in Transactions on Machine Learning Research (10/2023)
CIFAR1O, M- 100, M==100,Shardd Shard=$
CIFAR100, N=100, Shard=20
EMNIST, No 100. Shardm5
FEMNIST, 100, Shard-3
Testing Accuracy
60-
80
Testing Accuracy
18 -6 88 3g
60
Testing Accuracy
80
90
70
2
60
70-
-50
10 MARe
660
40
240
30
30
20
38
20
20
10
0.00
0.95
1.90
2.85
3,30
0.00
0.95
1.90
2.85
3,80
0.00
0.72
1.44
2.16
2,88
0.00
0.72
1.44
2.16
2,88
Total Time (C.T.-0)
Total Time (C.T.-D)
Total Time (C.T.-0)
Total Time (C.T.-0)
CIFAR1D, N-1 00, Shard-
CIFAR100, N=100, Shard-20 20
EMNIST, N=100, Shardw5
FEMNIST. 100, Shard-3
Testing Accuracy
55
60
80
Accuracy
50
%
40
Testing Accuracy
80
90
60
70
60
.80
70-
30
50
40
20
40
30
20
10
20
0.00
1.75
3.50
5.25
7,00
0.00
1.65
3.30
4.95
6.60
0.00
1.42
2.84 4.26
5.68
0.00
1.42
2.84
4.26
5,68
Total Time (C.T.-10)
Total Time (C.T.-10)
Total Time (C.T.-10)
Total Time (C.T.=10)
CIFAR10, N-100, Shard-
CIFAR100, N-100, Shard-20
EMNIST N-100, ShardwS
FEMNIST, N-100, Shard-3
Testina Accurac
Ttstting Accuracy
Testing Accuracy
2 -80
2 70
60-
50
40-
60
70
60
16 30
50
20
40
30
16
300
20
20
0.000 0,812 1.624 3,248
0.795
2.305 3,180
0.772
2.316
0.000 0.772
1544
3,088
Total (C.T.=100)
Total (C.T.=100)
Total Time (C.T.=100)
Total Time (C.T.=100)
Figure 3: Numerical results on CIFAR10, CIFAR100, EMNIST, FEMNIST with full participation (M= = N) in the fixed computation speeds setting. 'Shard' denotes the number of classes per client. 'C.T.' denotes the communication cost per round.
Contraction Inequality. Theorem 4.7 captures the progress made between two consecutive rounds of FedRep-SRPFL. It follows that the rate of convergence to the optimal representation is exponentially fast, provided that the number of participating nodes and the batch size are sufficienlly large.
Theorem 4.7. Let Assumptions 4.6 hold. Further, let the following inequalities hold for the number of participating nodes and the batch size respectively, n Z no and m 2 Co
(1+02 max {log(NN), d/no]: for
Fg
some absolute constant co. Then FedRep- -SRPL with stepsize n S 85 satisfies the following contraction inequality:
dist "g+t,B") dist (B', v
(10)
w.p. least T' exp (-90 min (d, Log(N))) where a En min,, 5
Here T denotes the total number of communication rounds which is logarithmic w.r.t. the target error The initial representation computed by the Method of Moments satisfics dist (1, for some constant Cм. Since Eg is strictly greater than zero inequality (10) ensures contraction.
Remark 4.8. Theorem 4.7 suggests that the server can learn the ground-truth representation before some of the clients update their local heads or participate in the learning process. This might raise concerns about fairness or accuracy, however it is important to highlight that such concerns are unfounded. This is because, after obtaining it the server shares the ground-truth representation with all the clients in the system. Thus, even if a client& was not selected in the representation learning process, it can still optimize its low-dimensional head w; E & using its local samples through a few local updates (given that is is a small constant.). Consequently, the derivation of the ground-truth model benefits both the clients that already participated in the learning procedure as well as new clients who opt to join the federated system at later time.
Logarithmic Speedup. Algorithm 2 sets off with no participating clients and follows a doubling scheme so that at stage r only the fastest 2" ng nodes contribute to the learning process. Thus at stage r, incquality (10) can be written as:
dist
dist vI a
(11)
2r 1
9

---

Published in Transactions on Machine Learning Research (10/2023)
CIFAR1O, M=100, Shard=5
CIFAR100, M=100, Shard=20
EMNIST, M=100, Shard=5
FEMNIST, M=100, Shard=3
80
60
Testing Accuracy
70-
50
Testing Accuracy
-90
80
Testing Accuracy
80
70
60
40
50
40
FLRPP
20
30
50
60
30
30
Tecleg-
10
20:
10
0.00
1.09
2.18
3.27
4.36
0.00
1.09
2.18
3.27
4,36
0.00
1.09
2.18
3.27
4.36
0.00
1.09
2.18
3.27
4.36
Total Time (C.T.=0)
Total Time (C.T.=0)
Total Time (C.T.=0)
Total Time (C.T.=0)
CIFAR10, M-100, Shard-5
CIFARI0D, M-100, Shard-20
EMNIST, M-100. Shard-5
FEMNIST. M-100, Shard-3
Testing Accurasy
70-
Testing Accuracy
50
60
Testing Accuracy
80
Testing Accuracy
90
60
40
70
30
60
40
40
20
30
10
20
0.00
1.51
3.02
4.53
6,04
0.00
1.69
3.38
5.07
6.76
0.00
1.51
3.02
4.53
6.04
0.00
1.51
3.02
4.53
6.04
Total
10)
1M
(C.T.=10)
Total
(C.T.=10)
CIFAR10,
S
100,
Shard-20
M=100, Shard=5
FEMNIST,
80-
60
Testing Accuracy
70-
50
Accurae
290-
80
70
40
50
30
70
240
50
40
20
30
10-
20
0.000
0.534 1.068 1.602 2,136
0.000 0.709 1.418 2.127 2,836 Total Time (C.T.=100)
0.000 0.534 1.068 1.602 2.136
0.000
1.068 1.602 2,136 Total (C.T.=100)
0.534
Total Time (C.T.=100)
Total Time (C.T.=100)
Figure 4: Numerical results on CIFARIO, CIFAR100, EMNIST, FEMNIST with partial participation (N M/5) in the fixed computational speeds setting. 'Shard' denotes the number of classes per client. "C.T. denotes the communication cost per round.
We note that the second term on the r.h.s. is an artifact of the noisy measurements. Utilizing geometric series properties we can deduce that in the limit, contraction 111) converges to (1- This implies that the achievable error of our algorithm is lower bounded by / (1----)((---i000 since the total number of stages is at most log (N/no).
To illustrate the theoretical benefits of SRPFL we compare Algorithm 2 to FedRep. One can distill FedRep from Algorithm 2 by disregarding the doubling scheme and instead at cach round, randomly sampling N nodes to part part icipate. For fair comparison between the two methods we set the larget error small enough so that the contribution of all N nodes is necessary. Specifically, we express the error as
=
with V2 v2>E>1.
(12)
N (1 (1--a a) v1 no
Intuitively, one should expect FedRep-SRPFL to vastly outperform straggler- pronee FedRep as approaches v2 (large error), since in this case the biggest chunk of the workload is completed before FedRep-SRPFL utilizes the slower half of the clients. In contrast, FedRep experiences heavy delays throughout the whole training process due to the inclusion of stragglers at every round, On the contrary, as è approaches (small error), the amount of rounds spent by FedRep-SRPFL utilizing N clients increases. In this case one should expect the speedup achieved by FedRep-SRPFL in early stages, to eventually become obsolete. Theorem 4.9 provides a rigorous exposition of these insights.
Theorem 4.9. Suppose that al each slage the client's 's compulational times are i.i.d. random variables drown from the exponential distribution with parameter X. Further, suppose that the expected communication cost per round is C = A. for some constant c. Finally, consider the target error e given in 12 Then, we have
T/SPPeL] El Shetee
O loR(
+logg
10

---

Published in Transactions on Machine Learning Research (10/2023)
CIFAR10, M=100, Shard=5
CIFAR100, 1=100, Shard=20
EMNIST, N=100, Shard=5
FEMNIST, N=100, Shard=3
80
60
Testing Accuracy
70
Testing Accuracy
80
Testing Accuracy
60
Tazhaz
40
70
50
LOCEFEL
40
Testing
30
50
16
30
20
30
30
20
10
10
0.00
0.71
1.42
2.13
2,84
0.00
0.71
1.42
2.13
2,84
0.00
0.55
1.10
1.65
2,20
0.00
0.55
1.10
1.65
2,20
Total Time (C.T.-0)
Total
Time (C.T.-0)
Total Time (C.T.-0)
Total Time (C.T.-))
CIFAR10, N=100, Shard=s
CIFAR100, N=100, Shard=20
EMNIST, N=100, Shard=5
FEMNIST, N=100, Shard=3
80
90
70
Testing Accuracy
70
60
50
Testing Accuracy
90
60
Accurac
70
40
50
60
9400 440
30
50
Testing
40.
40
30
20
20
10
20
10
0.00
1.51
3.02 4.53
6,04
0.00
1.51
3,02
4.53
6,04
0.00
1.25
2.50
3.75
5,00
0.00
1.25
2.50
3.75
5,00
Total (C.T.=10)
Time (C.T.=10)
Time (C.T.=10)
Total
(C.T.=10)
CIFARIO, N=100, Shard=5
CIFARI0O, N=100, Shard=20
EMNIST, N=100, Shard=5
FEMNIST, N 100, Shard=3
Testing Accuracy 7700
80
60-
60
Ascuracy
50
70-
Accuras
70
50
& 60
30
50-
940
30
20
40
30
20
10
0.000 0.788 1.576 2.364
3,152
0.000
0.788 1.576 2.364 3.152
0.000 0.755
1.510 2.265 3,020
0.000
0.755 1.510 2.265 3,020
(C.T.-100)
Total Time (C.T.-100)
Total (C.T.-100)
Total (C.T.-100)
Figure 5: Numerical results on CIFAR1O, CIFAR100, EMNIST, FEMNIST with full participation (M = N) in the random (dynamic) computation speeds setting. 'Shard' denotes the number classes per client.
"C.T. denotes the communication cost per round.
Theorem 4.9 establishes (log N) speedup for our method compared to its straggler prone benchmark. This result holds when speeds are drawn once at the beginning of the process as well as when new speeds are drawn at every round thus rendering our method versatile in broad class of settings.
Remark 4.10. Our analysis is crucially intertwined with the representation learning framework where the presence of a shared, global representation serves as common ground across setaahheteoogeeouss clients allowing us to show that intermodiate solutions constitute good starting points and substantial progress is achieved between stages. Despite FedAvg being a general-purpose algorithm not designed for representation learning, it was recently shown to recover the ground-truthe representation in the case of multi-task linear regression (Collins et al. 2022). thus casting FedAvg potential candidate subroutine for our method.
Remark 4.11. The initialization phase of Algorithm 2 requires one-time exchange of information between the clients and the server. Although this process reveals only the sum of outer products of local samples, it can be
further fortified using differential privacy techniques, such as the ones in (Jain et al. 2021 Shen et al. 2022
5 Experiments
In our empirical study we consider the classification tasks for the CIFAR10, CIFAR100, EMNIST, FFMNIST and Scnt140 datascts. We conduct experiments under the full and partial participation scheme with different computation speed distributions comparing the performance of our proposed method against other state-of- the-art benchmarks, Due to space limitation, in this section we present the results for the image classification tasks in the full and partial participation regime with fixed and dynamic computation speeds. Additional results and extensive discussion can be found in Appendix
The first benchmarks under consideration are FedRep (Collins et al. 2021 and Local-Global AvdAvg (LG-FedAvg) (Liang et al.] 22020). These foderated methods utilize mixture global and local models to derive personalizod solutions with small global loss. Coupling these algorithms with our proposed doubling scheme gives rise to PedRep-SRPFL and LG-SRPFL, respectively, which are proposed algorithms
Baselines.
11

---

Published in Transactions on Machine Learning Research (10/2023)
CIFARIO, 1=000, Shard=$
CIFAR100, M=100, Shard=20
EMNIST, M=100, Shard=5
FEMNIST, M- 100, Shard=3
Testing Accuracy
80-
60-
Testing Accuracy et
900
90
70-
50-
80
70
60-
40
60
50-
30
60
40
40
20
50
30
Fuskegir
10
10
20
0.000
0.779
1.559 2.338 3.118
0.000 0.779
1.559 2,338
3.118
0.00
1.09
2.18
3.27
4.36
0,000
0.779
1.559 2.338 3.118
Total Time (C.T.=0)
Total Time (C.T==0)
Total
Time (C.T.=0)
Total Time (C.T.=0)
CIFARIO, M-100, Shard-5
CIFAR10D, M~100, Shard
EMNIST. M- 100, Shard-5
FEMNIST, M- M-100, Shard
90
Testing Accuracy
70
50-
Testing Accurag Cu
290
70
60
40
80
60
50-
30
240
40
20
60-
30
10
50
10
3.000 0.784 1.567 2.351
3,135
0.000
0,784
1.567 2.351
3,135
0.000
0.784
1.567 2.351
3,135
0.000
0.784
1.567 2.351 3,135
lotal Time (C.T.-10)
Total Time (C.T.-10)
Total Time (C.T.-10)
Total Time (C.T.-10
CIFAR10, M- 100, Shard-5
CIFAR100, M-100, Shard 20
EMNIST, M- 1000 Shard
FEMNIST, M- 1000 Shard-3
Testing Accuracy
80
70
Testing Accuracy
50
Testing Accuracy
90
Accuracy
80
90-
60
40
60
50
30
70-
50
40
20-
60-
30
30
3.000 0.822 1.644 3,288
10
50
10
0.000
0.822
1.644 2.466 3,288
0.000
0.822
1.644 2.466 3,268
0,000
1.644 2.466 3,288
Total Time (C.T.=100)
Total Time (C.T:=100)
Total Time (C.T.=100)
Total Time (C.T.=100)
Figure 6: Numerical results on CIFARIO, CIFARI00, EMNIST, FEMNIST with partial participation (M = N) in the random (dynamic) computation speeds setting. 'Shard' denotes the number of classes per client. 'C.T? denotes the communication cost per round.
for the numerical experiments. We further compare our methods with FLANP (Reisizadeh et al., 2022 and FedAvg (McMahan et al.) 2017) with and without fine-tuning. Note that the finc-tuning variants are also considered as they often lead to better performance in data-heterogeneons settings (Wang et al., 2019 Yu
et al. 2020). Our final baseline is the HF-MAML algorithm from (Fallah al. 2020a) which is a Model Agnostic
Meta tearning-based method producing high quality personalized solutions in lata-heterogeeeeuus settings.
Data allocation. To ensure that our data allocation is heterogeneous we randomly split the data points among the clients in a way that each client can only observe a specific subset of classes. For instance, in the CIFARIO dataset where there are in total 10 different classes of data points, each client is only assigned data from 5 different classes, which we refer to as Shards; see first column of Figure 3 We also make sure that the test set for each client is consistent with the samples they have access to at training time, e.g., if client only obscrves samples with labels 1, 4, during training, then at test time they are only asked to classify samples from the same classes. Further details and different allocation schemes are presented in Appendix C
Simulation of System Heterogeneity. We consider two different types of client speed configurations: Fized computation speeds. In this configuration we sample a value for each client from the exponential distribution with parameter X, once at the beginning of the training process. These personalized values capture the computational time of every client and remain fixed thronghout the training procedure. In addition to the computational time, our methods suffer a fixed communication cost at every round. Figures 3 to each row depicts the effects of different values of communication cost on the convergence of the algorithms under consideration (C.T. 0, 10, 10
Dynaicc computation speeds. In this configuration every client samples at every round their processing times from personalized exponential distribution that remains fixed throughout the process Specifically, at the beginning of the training process we sample for each client i parameter X; from the uniform distribution over [1/M,1] 1]. Subsequently, at every round we sample a new computational time for each client i from the exponential distribution with parameter A, Similarly to the former setting an additional, fixed communication
For reference the computational times are sampled with A=1.
12

---

Published in Transactions on Machine Learning Research (10/2023)
cost is incurred at every round contributing to the overall running time. As illustrate in Figure 5
(full
participation) and G (partial participation %2)). the experimental results under this configuration are qualitatively similar to the ones presented in Figure 3 (full participation) and Figure Sauureepar partial participation
%20) for fixed computation speeds, with SRPFL providing substantial speodup over straggler prone variants.
Results and Discussions. From the numerical results in Figures g to we distill the following takeaways: 1) Coupling SRPFL with FedRep exhibits consistently superior performance across different datasets and communication time regimes compared to all proposed basclines. 2) Applying SRPFL to personalized FL solvers (FedRep and LG-FedAvg) significantly enhances their efficiency. 3) The speedup achieved by meta-algorithm is more significant for smaller values of communication time. Concrotely, we observe that the gap between FedRep-SRPFL and FedRep as well as LG-SRPFL and LG-FedAvg diminishes as the communication cost increases (plots in the same column C.T.= 0, 10, 100 ). This is unsurprising since our method improves the computational cost of the training and thus when the communication cost dominates the overall running time, the benefits of SRPFL are less apparent. 4) FedRep-SRPFL vastly outperforms the fine-tuning variant of the previously proposed FLANP, especially in regimes with high data heterogeneity (FEMNIST).
6 Conclusion
In this paper, we proposed SRPFL, straggler resilient FL. meta-algorithm with near-optimal sample complexity and provable logarithmic speedup guarantees in regimes with data and system heterogeneity. Our method leverages ideas from representation learning theory to compute global representation model along with local client heads, thus deriving personalized solutions for all clients. In SRPFL the participating clients are selected in an adaptive manner. In early stages fast nodes are prioritized and progressively slower nodes are included in the training process, therefore mitigating the effects of stragglers without compromising the quality of the solutions. Our numerical results illustrated the benefits of SRPFL when coupled with different personalized FL. methods such as FedRep and LG-FedAvg. Furthermore, our experiments support our theoretical findings showcasing the superior performance of FedRep-SRPFL. compared to state-of-the-art FL methods.
Acknowledgments
The research of I. Txiotis and A. Mokhtari is supported in part by NSF Grants 2019844 and 2112471, ARO Grant W911NF2110226, the Machine Learning Lab (MLL) at UT Austin, and the Wireless Networking and Communications Group (WNCG) Industrial Affiliates Program. The research of Z. Shen and H. Hassani is supported by the NSF Institute for CORE Emerging Methods in Data Science (EaCORF) as well as The Institute for Learning-enabled Optimization at Scale (TILOO)). The rescarch of R. Pedarsaní is supported by NSF awards 2003035 and 2236483.
References
Acar, D. A. E., Zhao, Y., Matas, R., Mattina, M., Whatmough, P., and Saligrama, V. (2021). Federated learning based on dynamic regularization. In International Conference on Learning Representations.
Balakrishnan, R., Li, T., Zhou, T., Himayat, N., Smith, V., and Bilmes, (2021). Diverse client selection for federated learning via submodular maximization. In International Conference Learning Representations.
Bengio, Y., Courville, A., and Vincent, P. (2013). Representation learning: A review and new perspectives. IEEE Transactions on Pattern Analysis and Machine Intelligence, 35(:1)1798 -828.
Chen, H., Ding, J., Tramel, F. W., Wu, S., Sahu, A. K., Avestimehr, S., and Zhang, T. (2022). Actperfl: Active personalized federated learning. In ACL 2022 Workshop on Federated Learning for Natural Language Processing.
Cho, Y. J., Wang, J., Chiruvolu, T., and Joshi, G. (2021). Personalized federated learning for heterogeneous clients with clustered knowledge transfer. arXie preprinl urXiv:2109.08119
Cho, Y. J., Wang, J., and Joshi, G. (2022). Towards understanding biased client selection in foderated learning. In International Conference Artificial Intelligence and Statistics, pages 10351-10375. PMLR.
13

---

Published in Transactions on Machine Learning Research (10/2023)
Collins, L., Hassani, H., Mokhtari, A., and Shakkottai, S. (2021). Exploiting shared representations for personalized lederated learning. In Proceedings of the 38th International Conference on Machine Learning, Proceedings of Machine Learning Research. PMLR.
Collins, L., Hassani, H., Mokhtari, A., and Shakkottai, S. (2022). Fedavg with fine tuning: Local updates lcad to representation learning. Advances in Neural Information Processing Systems, 35:10572-10586,
Deng, Y., Kamani, M. M., and Mahdavi, M. (2020). Adaptive personalized federated learning. ar preprint arXiv:2009.13//
Fallah, A., Mokhtari, A., and Ozdaglar, A. (2020a). On the convergence theory of gradient-basee model- agnostic meta-learning algorithms. In International Conference on Artificial Intelligence and Stalistics. PMLR.
Fallah, A., Mokhtari, A., and Ozdaglar, A. (2020b). Personalized federated learning: meta-learning approach. arXiv preprint orXiv:2002.0799 0794488.......
Haddadpour, F., Kamani, M. M., Mokhtari, A., and Mahdavi, M. (2021). Federated learning with compression: Unified analysis and sharp guarantees. In International Conference on Artificial Intelligence and Statislies. PMLR.
Hanzely, F. and Richtárik, P. (2020). Federated learning of a mixture of global and local models. arXiv preprint trXii:2002.000 05516.
Horváth, S., Sanjabi, M., Xiao, L., Richtárik, P., and Rabbat, M. (2022). Fedshuffle: Recipes for better use local work federated learning. arXiv preprint 1rXie:2244.1812.
Hu, S., Wu, Z. S., and Smith, V. (2021). Private multi-task learning: Formulation and applications to federated learning. arXiv preprint ar arXiv:2108. 12978.
Jain, P., Netrapalli, P., and Sanghavi, S. (2013). Low-rank matrix completion using alternating minimization.
In Proceedings of the Porty-Fifth Annual ACM Symposium Theory of Compuling, STOC '13, page 685-674, New York, NY, USA. Association for Computing Machinery.
Jain, P., Rush, J., Smith, A., Song, S., and Guha Thakurta, A. (2021). Differentially private model personalization. Advances in Newral Information Processing Systems, 34:29723 29735.
Jiang, L. and Lin, T. (2022). Test-time robust personalization for federated learning. arXiv preprint 22.512000009
Kairouz, P., McMahan, H. B., Avent, B., Bellet, A., Bennis, M., Bhagoji, A. N., Bonawitz, K., Charles, Z. Cormode, G., Cummings, R., al. (2021). Advances and open problems in federated learning. Foundations and Trendsm in Machine Learning.
Karimireddy, S. P., Kale, S., Mohri, M., Reddi, S. J., Stich, S. U., and Suresh, A. T. (2019). Scaffold: Stochastic controlled averaging for on-device foderated lcarning. arXiv preprint arxiv:1910. 06378.
LeCun, Y., Bengio, Y., and Hinton, G. (2015). Deep learning. Nature, 521:436-44
Lee, S., Sahu, A. K., He, C., and Avestimehr, S. (2022). Partial model averaging in federated learning: Performance guarantees and benefits. arXiv preprint arXiv:2201.002 033789
Li, T., Sahu, A. K., Zaheer, M., Sanjabi, M., Talwalkar, A., and Smith, V. (2020). Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Systems, 2:429 450.
Liang, P. P., Liu, T., Ziyin, L., Allen, N. B., Auerbach, R. P., Brent, D., Salakhutdinov, R., and Morency, L.-P. (2020). Think locally, act globally: Federated learning with local and global representations. arXiv preprint ar xrXet:201.001522
Mansour, Y., Mohri, M., Ro, J., and Suresh, A. T. (2020). Three approaches for personalization with applications to federated learning. arXiv preprint ar xrxix2200. 10022.0066 10619.
14

---

Published in Transactions on Machine Learning Research (10/2023)
McMahan, B., Moore, E.,, Ramage, D., Hampson, S., and Arcas, B. A. y. (2017). Commmiccaton-Efffdiatt Learning of Deep Networks from Decentralized Data. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, Proceedings of Machine Learning Research. PMLR.
Mitra, A., Jaafar, R., Pappas, G. J., and Hassani, H. (2021). Lincar convergence in federated learning: Tackling client heterogeneily and sparse gradients. In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W., editors, Advances in Neurol Information Processing Systems, volume 34, pages 14606-14614 Curran Associates, Inc.
Nguyen, J., Malik, K., Zhan, H., Yousefpour, A., Rabbat, M., Malek, M., and Huba, D. (2022). Federated learning with buffered asynchronous aggregation. International Conference on Artificial Intelligence and Statistics, pages 3581-3607. PMLR.
Nishio, T. and Yonetani, R. (2019). Client selection for federated learning with heterogeneous resources in mobile edge. In ICC 2019-2019 IEEE International Conference on Communications (ICC), pages 1-7. TEBE.
Pillutla, K., Malik, K., Mohamed, A.-R., Rabbat, M., Sanjabi, M., and Xiao, L. (2022). Federated learning with partial model personalization. In International Conference on Machine Learning, pages 17716-17758. PMLR.
Reisizadeh, A., Taheri, H., Mokhtari, A., Hassani, H., and Pedarsani, R. (2019). Robust and communication- efficient collaborative learning. In Proceedings of the 33rd International Conference on Neural Information Processing Systems, pages 8388 8399.
Reisizadeh, A., Tziotis, I., Hassani, H., Mokhtari, A., and Pedarsani, (2022). Straggler- resilient federated learning: Leveraging the interplay between statistical accuracy and system heterogencity. IEEE Journal on Selected Areas in Information Theory.
Shen, Z., Ye, J., Kang, A., Hassani, H., and Shokri, R. (2022). Share your representation only: Guaranteed improvement of the privacy-utility tradeoff in federated learning. In The Eleventh International Conference Learning Representations.
Smith, V., Chiang, C.-K., Sanjabi, M., and Talwalkar, A. S. (2017). Federated multi-task learning. Advances in neural information processing systems.
So, J., Ali, R. E., Güler, B., and Avestimehr, A. S. (2021). Secure aggregation for buffered asynchronous federated learning. arXiv preprint 10.::210000217
Stich, S. U. (2019). Local sgd converges fast and communicates little. In ICLR 2019 ICLR 2019 International Conference on Learning Representations.
Tandon, R., Lei, Q., Dimakis, A. G., and Karampatziakis, N. (2017). Gradient coding: Avoiding stragglers in distributed learning. In Precup, D. and Tch, Y. W., editors, Proceedings of the 34th International Conference on Machine Learning, Proceedings Machine Learning Research. PMLR.
Tripuraneni, N., Jin, C., and Jordan, M. (2021). Provable meta-learning of linear representations. International Conference on Machine Learning, pages 10434-10443. PMLR.
Vershynin, R. (2018). High-DDimesiooal Probabilily: An Introduction with Applications in Data Science. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press.
Wang, J., Liu, Q., Liang, H., Joshi, G., and Poor, H. (2020a). Tackling the objective inconsistency problem in heterogeneous federated optimization. Advances in neural information processing systems.
Wang, K., Mathews, R., Kiddon, C., Eichner, H., Beaufays, F., and Ramage, D. (2019). Foderated evaluation of on-device personalization. Xiu preprint ar aiv: 1910 10252.
15

---

Published in Transactions on Machine Learning Research (10/2023)
Wang, S., Liu, J., and Shroff, N. (2018). Coded sparse matrix multiplication. In Dy, J. and Krause, A., editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 5152 5160. PMLR.
Wang, S., Liu, J., and Shroff, N. B. (2020b). Fundamental limits of approximate gradient coding. ACM SIGMETRICS Performance Evaluation Review, 48:21 - 22.
Wu, S., Li, T., Charles, Z., Xiao, Y., Liu, Z., Xu, Z., and Smith, V. (2022). Motley: Benchmarking heterogeneity and personalization in federated learning.
Xie, C., Koyejo, S., and Gupta, 1. (2019). Asynchronous federated optimization. arXiv preprint 1Xix:1903.03934
Yang, H., Zhang, X., Khanduri, P., and Liu, J. (2022). Anarchic federated learning. In International Conference on Machine Learning, pages
PMTR.
Yu, T., Bagdasaryan, E., and Shmatikov, V. (2020). Salvaging federated learning by local adaptation. arXiv preprint trXf::2002.044 04788.
Zhu, C., Xu, 7., Chen, M., Konečný, J., Hard, A., and Goldstein, T. (2021). Diurnal or nocturnal? lederated learning of multi-branche networks from periodically shifting distributions. In International Conference on Learning Representations.
16

---

Published in Transactions on Machine Learning Research (10/2023)
A
Appendix
Before we dive into the analysis we provide the following useful definition.
Definition A.1. For a random vector x Rd and a fixed matrix A Rd xd2, the vector A is called x is subgaussian with subgaussian norm O (||A l1yll)) for all y E R i.. || || Alll-sunn lugguugannsiin if y Ax
E exp (y A AT)]] I exp /111//-
We study the performance SSRERL with FedRep as the subroutine of choice. The first part our analysis focuses on a single round and extends the analysis in Collins et al. (2021). We assume that there are N clients in the network and round a subset It of them participate in the learning procedure with cardinality 2 no log(NN)-Ouxx.-
Without loss of generality we assume that the clients are indexed from fastest to slowest thus the clients that participate in the learning process are all i E [n]. Each client i, draws a batch of [1+=")A*p /x* Nog(N) fresh, ii.d. samples at every round. We denote by X E Rdxm and Y R" the Fg
2d
m 2 co
matrix of samples and the labels for client i, such that the rows of X are samples (x,,,.,x'*) By Z R we denote the noise in the mcasurements of client i, with 2., ~ (O, o²). Let B E Raxk and W E ER stand for the optimal representation and the concatenation of optimal heads respectively. The hat denotes that matrix is orthonormal i.e. its columns form orthonormal set. Similarly, B' E Raxk and W E N denote the global representation and the concatcnation of the heads at round t. wi's and wf's denote the optimal heads and the heads at round t which constitute the rows of W and W respectively. Furthermore we define min, == minze(N), a mmin 1 W and max,, = [Z] F v LW where Wz is formed by laking the rows of W indexed by I. That is maxx, and min, are the maximum and minimum singular values of any submatrix W that can be obtained throughout the course of our algorithm. Notice that by assumption 4.6 each row of W has norm Vk, so / acts 1 Will = Vk. Finally we define / Since we focus on a single round the time index can dropped for simplicity. Further, henceforth we drop the subscripts It on W
/
normalization factor such that
First derive the update scheme our algorithm follows. Notice that the empirical objective function given in can be expressed via matrices X; and Y,,
(N W) 2mn
Y; ..,,BB,))
(13)
Further, computing the gradients we derive
Σ VB (Y, - BB,
VB
1
x (,,Ww, Y.) w.., (14)
2mn
mn
1 2mn
V w.
Σ
(Y, (V,-X,Bw,)
x (x,Bw, N YY))
(15)
jml
and since us
1 m
L B,
is invertible with high probability by Lemma A.4 solving for the minimizer gives
w = (LBTxXX,B)
B
(16)
Thus our update scheme with stepsize 17 is
Vi [n] w =
(1
B x X,
Y
(17)
m
m
n
B B Σ X
(x, YY)) w.
(18)
mn
l=1
B+,
(19)
where QR denotes the QR decompasition and Y, X,B'w; + Z,
17

---

Published in Transactions on Machine Learning Research (10/2023)
Lemma A.2. For every client the update for w, can be expressed as follows
B B'w F, G.,
(20)
where F; and G, are defined in equations
24 and 25 respectively.
Proof. Further expanding
17 can write
W
B
1 B
+
X
1
B Z.
(21)
m
m
B
B
- B-) w/
m
m
B x X,B
x
(22)
B B' G.
(23)
where we define
F
1 - B X
X
XB
x
B
X,B
L x B, B-) wi:
(24)
m
G
L x
B
(25)
m
We further have the following immediate corollary. Corollary A.3. Let W F and G be the matrices with rows the concatenation w. F,, and G., respectively. Then
W = W
-
(26)
Our first goal is to control the norm of w In order to achieve that we provide lemmas that bound the norms of F; and G; extending the analysis in Collins et al. (021) and Jain et al. (2013).
Lemma A.4. Let 3/22 (lgg(N) for some absolule constand c, then with probability at least exp (-115k3 log(N))
Vm
LBTXTX,B X.B 21-8
Vi E [n],
Tmin m
(27)
1t. follows that with the same probability
E [n],
Tmax ('BTxTX,,)
X,
-1
(28)
1-8
Proof. First notice that we can rewrite
B
X,B
Σ
B
B
x
B
x
(29)
m
For all i E [n].j E [m] we define v B x such that each vž is an i.i.d. || 1mm Blla-subgaussian random variable (please see the definition of A a-subgaussian in Definition A.1 and thus by equation (Theo)) 4.6.1) in Vershynin (2018) we obtain the following bound for any m k
(V
Umin
xxxxx1B) m
B
21-
(30)
m
18

---

Published in Transactions on Machine Learning Research (10/2023)
with probability least 1 exp (-?²) and e some absolute constant. We set == 12k9/² log(N) vk and
61 = h2// 3/2 1009///oog((0 (og(N) and the above bound becomes
X XB 21 51,
mmin
(31)
m
with probability at least 1 exp (-h (12k)log(NN
log(NV) -1) )
Further notice that
(-144k² exp (-k (12)))og())-1) exp ( 1104444-1og((N)) log(N) + 24k (4klog(N) 1))
(32)
S exp (- (-120k log(N))
(33)
Thus taking Union Bound over i E [n] we have that for all E [n]]
B
-
Mmn
21 81
(34)
m
with probability at least
- exp (-20kk log(NN) 2 exp 115 log(N))
(35)
Choosing sufficiently large derives the statement of the lemmaa.
Lemma A.5. Let H. = ( (
AB - -1 B and == c 1312 Vlo((N)
X,
for an absolute con-
m
Vm
sland c. Then with probubilily at least
exp
(- 115k2 log(N)) we have
Σ
H. w 11 11 8² || 1| dist (B.B')
(36)
Proof. In order to argue about the quantity H,
B x,
1 X, (BBT - -La) B we define matrix v
U== L==Lx, X;
BB -I.) B such that its j-th row, Uj = mm
[BB T La) x is subgaussian with norm 1 Bx has B x such that its j-th V, vm
B
mm
T -BB--1a) Ia). Similarly define V 1 mm
at most
norm at most
1 B. We are now ready to use a concentration argument similar to Proposition (4.4.5) in ym
Let sh-1 denote the unit sphere in dimensions and N₂ the 1/4-net of cardinality 9*, From equation (4.13)
in Vershynin (208)) have
( (B X; [B -Is) B'B -|IuV/
EuvJ y
2 P.E.
(37)
= 2 max p.E.N
Σ u.) (vj.g)
(38)
j=1
B (BBT - -1)) L (B. B-)
By definition (p, uy) and (V are subgaussians with norms and 1 Vi
vm
B
B
vim respectively and thus for all [m]] the product (p, uy) (V,,y) is subexponential with
norm at most dist (B, (B. for some constant C¹. Note that
C'
E (p, u,) (v.y)) =p_ ( Ia BB
BBy y-
(39)
19

---

Published in Transactions on Machine Learning Research (10/2023)
and thus use Bernstein's 's inequality to bound the sum of mean subexponential random variables, for any fixed pair p.y E Na:
m2
s2m²
P
Σ
p. uy) (vy.8) 2 s exp
-2 min
(40)
dist? ?
(B, B-
dist B B-
8²
$ exp
-2m min
8
(41)
dist²
(B. B
dist
[B, B
for constant c2 Thus taking Union Bound over all the point the net we derive
s2
P Ep.yEN 2 > (p,.))(yy.y)) 2 2s
s 9 exp
Cym min
dist (B,
dist (B, B.)
(42)
Since m > Ch2 log(NN, by setting s = dist (B. / Cka log(N) 4m
and 38 we obtain
P
B X.T x; BB Ia) B B111 2 dist (B, B+) V
IT)BB
Ch2 log(N)
S 92h exp
exp
-02m dist?
g2
(B,
(43)
< 92h exp (-C cymn? log(NN)
(44)
exp (-120k² (-120)>log(V)))
(45)
for sufficiently large C. Using Union Bound again over all participating clients we get
P
ViE [n]
Hall S dist (B, (B.B*)
Ch2 log(N)
2 - nexp (-120k² (0120))oog(N))
(46)
2 1 - exp (-115k² (-115k2log(N)) log(N))
(47)
The above also implies
IP
H 11 < Cdist? BB.
2.2 log(N))
2 1 exp (-115k log(N))
(48)
1=1
k
1 log(N)
PP
W
H
C
w
dist?
[B B-)
exp (-11kk2 log(
(49)
2
Finally notice that
W 2
H. H.
<
H
1
| 2
Σ H
(50)
71
where used Assumption 4.6 and the fact that the rank of W is k. Combining this with and choosing sufficiently large we derive the result.
Building the previous lemmas bound the norm of F:.
20

---

Published in Transactions on Machine Learning Research (10/2023)
Lemma A.6. Let ó == 3,112 /log()) ym
for some absolute constant and for all E [n] let F; given by Further let matria F E Buuk such that its rows are the concatenation of F, 'ss. Then with probability at least - exp (og))))ogg(N) we have
24
8
Vi [n]
F
1 1- dist [B, B¹) ||w; 1
(51)
F I
dist
B. B. || ll
(52)
1
Proof.
-112
F
B
x X,
H [w;
(53)
m
8²
dist?
[B. B¹ ||wil2
(54)
5)
which holds for all i E [2] with probability at least 1 - exp (-110k² log(A)) by using Union Bound on the
failure probability of 28 and
(4 Similarly, we have
-1
F
F
2
L m
B x X,
(55)
1
H. w 11
(56)
(1 8)2
82
dist? (B, (B) /w112
(57)
(1 8)²
which holds with probability at least - cxp (-10k log(N)) taking Union Bound on the failure probability on 28 and 36).
We now lurn our atlention on deriving a bound for ||G.ll2-
Lemma A.7. Let 8 33/2 (lg(NN) for some absolute constant and for all i [n] let G, given by
25
Vm
Further let matrix G E 1Rn such that its rows are the concalenation of G 's. Then with probability at least - exp (-110k-Pog((M)) we have
8
Vi [n]
||G:|] ll2 -
(58)
8
8 |G|| v
n
(59)
Proof. Notice that we can write
G
1 BBBx,, m
B Z,
(1 B x XB m
(60)
m
m
B is an i.i.d. zero mcan subexponential with norm at most C | Ch for some constant C2. Once again denote by s the unit sphere in k
and since 2f V (0, (0) we can conclude that for all i.j.
dimensions and by N₂ the 1 1/4-net with cardinality 9a. Using Bernstein's is inequality and Union Bound over
all the points on the net we follow the derivations from Lemma A.5 to get
2s s gh+t exp (-ey min { 8 ))
PP
(61)
04 02
21

---

Published in Transactions on Machine Learning Research (10/2023)
Since C2k² (N), by setting = o²
/Cck
C2k2 log[N]
derive
P
1 Σ
X
2 0²²
C2h² (N)(N)
S ghtt1 exp (-eam
8²
(62)
2
ghtt exp (-C2 cak? log(N))
(63)
S exp (-115h² log(N))
(64)
for sufficiently large C2. Choosing large enough and taking Union Bound over all E [n] we can obtain
P vi E
[n]
- = x
<
s 1 n nexp (- 1-115k2 log((N))
(65)
S exp (-1131² log(N))
(66)
(67)
Finally taking Union Bound over the failure probabilities of
28 and 67
we get
(L X,
$ 5
vi E [n] |G:1l2
x
(68)
6
with probability at least. nexp (-110k2 (N)( It follows that with the same probability
5 G.ll < - 8 o |G||
(69)
For all i E [n] we define = Bwt - B w. The following lemma provides upper bounds on the norms of w. and qi
Lemma A.8. Let 8 :== kar2 (lg(NN) for some absolute constant e and 8 5/(1 - 8). Then with probability at least exp (-15 log(N)) have
ym
E [n]
W ll2 & 2 + 028
(70)
Further with probability at least 1 - exxp (- 105k2 log(N)) we have
Vi [n]
||q:l12 1 s 22 dist (B.B*) + 028
(71)
Proof.
w/ 11,1,22 - B
|| + ||Fill2 + ||
(72)
|w ll + w/lll dist (B, B 802
(73)
< 2 + 802
(74)
where the first incquality comes from (20 and the third from Assumption 4.6 For the sccond incquality take Union Bound the failure probability of (51 and (58) and thus the above result. holds with probability at least 1 exp (-107k² log(N)) Taking Union Bound for all i e [n]] we get that with probability
at least 1 exp (-105k log0)))og((N))
[ [w < 22 + o28
(75)
22

---

Published in Transactions on Machine Learning Research (10/2023)
and the first result of the lemma follows. For the second part we have
q 2l2 = Bw B w;
- B B'w; + BF, + BG, B"w;
(76)
2
[BB T - -Id B B'w; +
BF
+
BG
(77)
[ __
B'
|w; ll2 + |F:
||Gill2
(78)
dist
(B. B²) || 1l2 dist
(B. B*
|| 1l +
(79)
S dist
(B.B*)
2 + 0
(80)
where the first inequality comes from 20 For the forth inequality we take Union Bound over the failure
probability of (51) and (58) and thus the above result holds with probability at least - exp (-107k? log(N)).
Taking Union Bound for all [n] get that with probability at least 1 exp (-105k² log(A))
vi [n] |1q:112 $ 22 dist (B.B*) + o²8
(81)
k3/2 /log(N)
Lemma A.9. Let 8 = =c vmn
for some absolute constant a and 8 8/(1 6). Then with probability
at least exp (-105d) exp (- 105k2 log(N)) we have
0² (Vb++ )o²) Vd
x TT
c.
(82)
mm
vmn
Proof. Let S denote the unit spheres in d and l dimensions and Na. the 1/4-nets of cardinality gd and gh, respectively. By equation 4.13 in Vershynin (2018) we have
L X Z.W T mi
2
max yEN,
Σ iil
X
Z.W
y
(83)
mn
= max PENAYEN,
T
x w! +T
y
(84)
= max PENa.EE
x x (w:+ .))
(85)
mn
Notice that for any fixed y and Vi E [m],j [m] the random variables mm (x).P) (w y) are i...d. zero mean subexponentials with norm most event
[w! for some absolute constant C/. Conditioning on the
mn
E =
n {
/|
11, s 2V + 80²)
(86)
which holds with probability at least 1 exp (-g0)))log(N)) by Lemma A.8 we can invoke Bernstein's 's inequality to get
P
(x/. wt y) 2 20/81 < exp
-cammmin 0d (2V + 028) 8 o² (2Vk 028
x2
8
(87)
23

---

Published in Transactions on Machine Learning Research (10/2023)
dCa by setting o2(2vittōo") va-ca the above quantity simplifies as follows 8ymmn
Since
d.Ca ng
o² 2v + 802 vd.
s2
P
x/, (wT .y)
S exp
-amm
(88)
mn
mm
at (2vR 0
< exp
(-O3 c4 c..d) d)
(89)
I exp -110d)
(90)
for C3 large enough. Taking Union Bound all points p, the No, Ne and using
we derive
o² (Vk+ 80²) vdl
PP
Σx
x Z
w.
+T
2 C3
El I goth exp (-110d)
(91)
mn
v mm
exp (-105d)
(92)
and removing the conditioning & get
n
o²
(VK+ 12+0002 Vdl
mn X L.W.
2 VCs
S exp (-105d) + P (ET)
(98)
Vmn
S exp (-105d) + exp (-105k² log(N)) logN))
(94)
Choosing large enough and taking the complementary event derives the result.
1/32 NoR(N) for some absolute constant a and 8 5/(1 8). Then with probubility
Lemma A.10. Let ;= c at least exp (-100d) exp (-100k2 log(N))
Vmn
have
dist (B, & + Vhōo² + (802) 2
Vd
Bw,,- B'wi) - (Bw.- -B BBw:)) wxt sc.
X X;
m
/mn
(95)
Proof. Let us define the event
Eg = n (1 wt $2Vk 80² n l1qill2 s dist (B. BB 2 - 800
(96)
which happens with probability at least 1 exp ( (- 100k log(N)) by Union Bound and Lemma A.8 For the rest of this proof we work conditioning on event &g. Recall that qi = Bw B'w; and thus we can write
x x; (Bw,- B'w;)) (Bw, - BBww)) wtT m
X X.
Σ wwttt)
(97)
qi
x
w
+T
Σ q.w.T
m
(98)
24

---

Published in Transactions on Machine Learning Research (10/2023)
Let 8t-1 , denote the unit spheres in d and & dimensions and No. the 1/4-nets of cardinality 9d and 9r, respectively. By equation 4.13 in Vershynin 2018) we have
72
m
x q
x x.W.
1 " n
Σ
q:wp"
2 max PeNayyNN n
1 -p
n m
q.w.
n m
m
72
(99)
71
(x) (p.x)) (w. (p. qi) (WT, .y))
2 max pENayENN mm
(100)
Notice that for any fixed p,y the products (2 are iid. subgaussians with norm at most
č |/ and (p.x/) are ii.d. subgaussians with norm at most 82 lllll δ. Hence un-
der the event & the product (x xx1, (p.))pp.)))) x (wty)) are subexponentials with norm at most
ci dist [B, BB) k + vkdo² + (Bo2) for some constant C/ Also note that
E
x q:) p. (p.xt) (w. y) - (p. qi) (w. v) = 0
(101)
and thus applying Bernstein's inequality get
P
n m
x q:
(p p. x
(w,
1
Σ p.. q.) (w/,)) 2 Eg
17
n
x²
I exp
C5 min
dist (B. xo + vkddo2 + (802) 2'
dist (B, B²) k + Vi8do2 + (802) 2
(102)
4Cs > dCs by setting s Vord dist [B (b+B*)++ V + no
Since m
and taking Union Bound over all
2/mm
N E Ne we derive
VC . dist [B, k + vido² + (862)) 2)
x, X (Bw,- -Bww)) (Bw, - -B'w:) w.
P
Eg
Vmn
- -mns2
I ggd+ exp
dist (B, (B.B*) k + Viso²2 + (80z)
21
gd+h exp
(-C4 os .d)
(103)
gd+h exp (-10dd)
(104)
exp (-100)
(105)
25

---

Published in Transactions on Machine Learning Research (10/2023)
choosing a large enough constant C1. Recall that P (Ef) 1 exp (-100k² (=100)+log(NN). Hence by removing the
conditioning on &g we get that with probability at least 1 exp (-100d) - -exp (-100k² log(N))
vd dist (B, B²) k + vkdoo (802) 2
n2
x X; (Bw,- - B'w;)) - (Bw, - B B.::)) w,t
se.
m
Vmn
(106)
for sufficiently large c.
Having set all the building blocks we now proceed to the proof of Theorem 4.7
Theorem 4.7. Let Assumptions 4.3 4.6 hold. Further, let the following inequalities hold for the number of participating nodes and the batch size respectively, 2 no and
(1+02)k*' max (log(N), d//n)), for
2 Co
E
some absolute constant co. Then FedRep-SRPFF with stepsize n inequality:
satisfies the following contraction
S mms.
dist (B ,B*) dist dist(B', BB) VI- a+
(10)
(1 0)
w.p. at least T exp (-90 min (d, h2 log(N)))) where a = on min, S 7
Proof. First let us recall the definition of ó == c 3/22 Tog(N N)
for some absolute constant c and 8 = 6/(1 - 6). Further notice that for our choice of m and sufficiently large we have the following useful inequality
Vm
Eg
1
1 20
28 <
I
(107)
1
k0
1 0²
From the update scheme of our algorithm
18 have
B "7
x x, X,Bw w.t XxxxB x,B w.B'::W.T X Zxtt
(108)
mn
B 1)
Σ m
X.X,, x, (Bw:- - -B'w;) (Bw, - Bww:)) w.t
x
n
π n (w B'w:) w."' + 1 n "
x Z.w.T
(109)
n
n m
where we added and subtracted terms. Multiplying both sides by B we get
B B B-
T
Σ
1 Xx,x, x, (Bw, -B'w;)
(Bw, B (B:w:) wxtT)
n
m
Bi Bw - B B'w:))
w.t
+ 2
B
x Z. Z.
(110)
n
m
1=1
B B
I
!
Σ
,; w.!""
+ "
Σ m
X Z.
a i=1
! BB.
Σ
n
1
x, x;
(Bw,t-Be -B (Bw: B.w:)) w.t
(111)
n
m
i=1
26

---

Published in Transactions on Machine Learning Research (10/2023)
where the second equality holds since B = 0. Recall that from the QR decomposition of B we have B = B+R+. Hence multiplying by (R+)-1 and taking both sides the norm we derive
n w w.
|!! ((+t)
dist
B
B
I
1,_
2 B1 n
X
Z.W
(R+)-1)
n
2
m
+
T
Σ
12
[ x x; m
(Bw,- -Bw,-B'w, -Bww:) - (Bw: - B.wi)) w.,T)
RR
n
2
i=1
(112)
Let us define
A1 = dist BB.
I
17
w w
(113)
n
A2 =
D B1
x Z.
(114)
1=1
A3 =
B
1 x x; (Bw:- - B.w:) (Bw, B B.::) w. )
(115)
so that the following inequality holds
dist (B+,B*) S (A1 + A2 As) |
(116)
For the rest of the proof we will work conditioning on the intersoction of the events
& = n wt 112 s 2vk + 802 n | l22 S dist (B.B*) 2 + 800]
(117)
1=1
& =
|F| F - dist (B,B*)
||
n |G|| F S _
(118)
cc. o² VV + S0²) Vd
L,w :1 T
(119)
mi
Vmn
cvd dist (B,
B³ k. vkôo² + (802)
E5
X x
(Bw,- - B'wwi) - (Bw, B BBw:)) wx+ 5
n
m
Vmn
(120)
which happens with probability at least. - exp (-90d) - -exp (-90k2 log( by Union Bound on the failure probability of 12
59
82
155 and 96).
We will now provide bounds for each of the terms of interest in 116 starting from A1. Notice that by 26
have
Amax (W+++)=
w
W B F
(121)
2 W
? - 4
||
1 + ||
(122)
< 2 W 113 d
dist (B.BB) 482
W 1 + 0
(123)
| +m)
(124)
4n ( + 1)
(125)
27

---

Published in Transactions on Machine Learning Research (10/2023)
where in the last. inequality we use the fact that | vn- max,e- Since v ( (maxx + 1)= the matrix W W is positive definite. Thus we have
R 2
2
I
w
W
S 1
min
(W+TW+
(126)
2
n
7
Amin
(W"B'B+F+ (WBBB(BBBBBB
(127)
71
6 1
T
W'B'B) - - (F) - mmnn())) (G)
min
n
2n
maxx
F FW B)
mmax (F G) + mmax (W
(128)
n
1 1)
min
(W*)
o min
B B +2 B T B (mmax (F Ww) max (G (G+))
n
27
max (F) FF)mmxx( (G)
(129)
n
< 1 - 1- δ min,+ amin (B* (B B)B + 27 (1 + [[G|2) [ [ 21 [1 -|G]|2) (130)
where used that the norms of B and B are 1 since the matrices are orthonormal and min, m (W*). Recall that we operate under &3 and thus we can further write
Is 7 W n
W
S 77 o2 min. min B (
+ 2n n
dist (B.B:) 8 || + vndo²) |
2
2m dist (B, B²) 8 Vr [*
+
(131)
n
W 12
W 112
< 1 1) T2 min,+
min B
+ +2n
8
80
+
82
(132)
n
vn
/n
Eud min.. _²
< 1 1) o2 min,
min
BT
B
+ 3n
2062 maxx*
max,+
(133)
< 1 - min
min
B
B
+
6
1
n7
Ego
(134)
where we upper bound dist (B.B*) by 1, 8 S
E
and use Assumption 4.4
the third incquality.
x0
1t
Further by the definition of Eg = 1 - dist? 2 B B
I o oiin
(B .B) we have
I TW W+
< 1 - E min,+ + gn
(135)
and it follows immediately that
A S dist (B.B*) (1 min,
1 min,+
+
6 1
(136)
Further since we operate under E 119 and
B =1 we have
Vd Vmn
Aa noo² (vi++)
(137)
and since we operate under &s 120 we obtain
As S ne (dist B & Vio vk++) vd
(138)
mmn
28

---

Published in Transactions on Machine Learning Research (10/2023)
Combining 116
1360.
137
and 138 get
z dist (B, (1 (1-5, 5 1 nyck
Vd
dist (B+, B²)
||t
(mn)
vd
+ ne (v++1) (o² + +)
(Rt)-1)
(139)
mm
The last part of the proof focuses bounding (R+)
Let us define
S = Σ X Xi (Bw,+ - B'w,) w;
(140)
1=1
X Z.
(141)
m
and hence 108 takes the form
B = B 11 + " E
(142)
and also
(BTS+STB)
+ " "B"BE+ET
" STS
1/2 n
S E)
17²
E
E (143)
n
n
n
n
I ! n
(B
+ BBTE+EE
n
SES++S
r + n
E
E +
n
n
STS
(144)
By Weyl's inequality and since R B we derive
o 2 RR)21- Amax (B
"
Amax
B'E+E.B.
" n
Amax
S
(145)
Let us further define
R !
S
(146)
Amax
n
R2 Amax
17
I
(147)
n
R3
= Amax
"
B E + E
(148)
n
So that we can succinetly rewrite the above inequality as follows
(R))
2 1- Ri Ra R3
(149)
We work bound separately each of the three terms.
2n max P p:Ilp1l;=1
R1
T
Sp
(150)
x X, (Bw, - B'wi)) (Bw, B B'w:) w. )]
27 p T
max
p: p1 n
1=1
2n p TB
(Bw - B'w;)) w.
+ max
P
(151)
p:I|p||==1 n
29

---

Published in Transactions on Machine Learning Research (10/2023)
and since operate under Es 120 the above simplifies to
R S 2n
B dist
B, B' k + vk
+
vd
max 1pl2
21
P
B T
Bw - B B'ww)) w.
p (152)
mm
p:
3nc dist B, BB k + Vk
vd
+ max
2n
p
T
B
Bw - -B B:w:) wxt
(153)
mm
p:lpp1,=11
We focus on the second term and using
20 we get
2/
p
Bw - -B'w;) wt T
p
21
tr
Bw - -B'w,) w, pp 'pp TB
(154)
n
27
Bw - -Bww;)
F. )
T
pp T
tr
(155)
We bound each term separately and to this end we define
2m n
Ti
tr
Bw B'w;) w; BB Bpp Bpp B
(156)
n
T3
2m
tr
(w -B'w;) F. BT
(157)
n
T3
2m
n
(w -Bww;) G pp B
(158)
n
such that
155 can be expressed
2n
B
Bw - B'ww)) wit
Ti +
(159)
Further expanding Ti we have
T
2n
n Σ
B w/ + BF.WT + BG..; - -B'www;]]
B= BT Bpp TB
(160)
n
2m tr
B
(BBT -Is) Σ B (w:w;: B Bpp
T
2y +
Σ
(BF.,;
BTB B
tr
n
n
+ 71
2m tr
(G: B T B
(161)
211 tr
L (FT +G W
(162)
71
21
([[F]] + | p [W'|]
(163)
where in the first equality expand wt via 20 and in the third equality that B
BB
Ia]
and F F. GT Σ G. The incquality is obtained by noticing that the norms of the
1=1
orthonormal B, B one and also
Ipp 1l, s Conditioning on Es (118 we can further simplify as follows
2m dist (B.B*) 8 || 11 802 ([]l))
Ti 5
(164)
n
<
Eno min,+
(165)
10"
30

---

Published in Transactions on Machine Learning Research (10/2023)
We now turn our attention to T2
T2
21
tr
BB B B',; + BF, + BG; - B'w;) F. pp T
(166)
27 tr
BB - L1) Σ B';; F. pp T
2m tr +
2m tr
B T
B B
B
F.F
pp
+
B
B
B
G F
pp
(167)
2m ir
B (1B(FFFL F G F) ]
(168)
71
2n
[IFI G |F]]
(169)
where in the third equality we used that B
I2] = O and in the forth that the norms of the orthonormal matrices is 1 as well as the norm of pp Following the same calculations for T3 we get
BB
2n
T3
G 2
| (F
(170)
and thus summing the two terms get the following
27
T2 + S
[[ + |G
(171)
n
Again conditioning on &g 118 we derive
T2 + T3 S 20 (8 ([W))l l12 vna?))
(172)
s 2nj²² (d + ot)
(173)
< 17 Frod min.
(174)
10
Hence combining
153
165
and
174 we get a bound for R₁
Vd
R1 S 37e (k + vk)
+ 5 n
(175)
Vmn
Vd
1
6n c.
+ " Enoo
(176)
Vmn
We work in similar fashion to derive the bound on Ra
n² Amax (E S EE)
(177)
2ny²
max
p SSSEP
(178)
n²
pil|pli==1
2
x X; (Bw, - BB'w)) - (Bw+ - -B'w:)
w TT
Σ
n
-
x
zwwt))
D
max
n
m
2n² 2
max
Σ
Bw?t - B'w;)) w;
Σ
n
1
xJ Ziw? TT
p
(179)
n²
p:Ilp! =1
m
i=1
n
< 2n)
x. X, (Bw," - BB'w;)
(Bw, - B'wi))
w. T
Σ
[ Z, (T)
m
mn
2
+ 2n²
Σ
[ww -Bw::) w;
1 - x
Z.w.T)
(180)
2
(181)
31

---

Published in Transactions on Machine Learning Research (10/2023)
Since we work conditioning on the event E Es further derive
Vil
dist dist
(B,B*) k viô0² + (802)
a2
Vk (v
)02) )02)
vdl
R2 S 2n/²
c
c
Vmn
Vmn
o² VT S02)
vd
+ 2y72
Σ Bw -B" B'w;) wit
(182)
Vmn
Vd
(oVko?
Vd
2n²²
coTor Vd vmn
S 3r/²
ck
+
11 lll || 12
(183)
Vmn
Vmn
And since we also condition on &
117
we finally get
R2 < k
d
2n²²
Σ
12 + 80²)²
eeVko2 Vd Vmnn
(184)
mn i=1
d
ck+o² Vd Vmn
<
3n e 0
+m72
(185)
mn
The last term we need to bound is R3
R3 =
17 Amax
E B + B E
(186)
20y²
p B ''p
(187)
n² p: p:lppll =
27
B
Z W,
(188)
n
m
a² (Vk + 802) Vd
< 2nc
(189)
Vmn
S eneeeoo T0o2 Vd
(190)
mmn
Combining 149) with 176), 185
and 190) we derive
Vd
Ego min,. 37p21t02 d mn
(oktos
Vd
Vd
0 muin
R 2 1 67 c.
k
9n2
--nnoooo Too"
mn
mm
v /mn
(191)
2 1 14nck 3 a
Vd
3 2kio
d
1 5 Ego min,
(192)
Vmn
mn
vd Vmn
2 1 - 15nck # 0²
min,
(193)
where the last incquality holds since vmn 2 ovd. We can now combine contraction inequality
139 and 193
obtain the
dist 1 [ .B. < dist (B.B)) (1 -n nckk
vd /mn
|(
(1- 11 11-1mnch
Vd
- 11/2
15nck 0²
5 Eg
min,+
Vmn
-1/2
+ mje (Vi++) (o²+1) Vd
1 15nck 00 Vd
En
(194)
Vmm
/mn
32

---

Published in Transactions on Machine Learning Research (10/2023)
We divide and multiply by no and using our bounds on and n the previous inequality further simplifies,
-1/2
vd
vd
dist BttBB) I dist (B.B.) 1 1 n) "u min, + nek
1 15nck
Ego
V mno
n mng
5
min,
no
112
+ np (vk+1)
+ (o2
vd
1 15nck à 0²
vd
1 1j Ego
V mmg no
V mmno
min,+
mg
(195)
I dist (B.B*) (1 Ego min,+ )
(1 Eno minu,
-1/2
(V
no
(1-2
-1/2
+
n) Eno min,
2 Ego min,+
4n
2
(196)
(B. ((1
1 11 1 EnnoE min,+)
( En min.
< dist
(197)
R.
En min,+)
where in the second incquality we used that for our choices of m and the following incquality holds 15nck (2) the result.
Vd
< To Eio min. Taking Union Bound over the total number of iterations T we derive
ymmo
]
Corollary A.11. Recall that our algorithm starts at stage O with no participating clients and doubles the number of participating clients at every subsequent slage. Thus, by slightly abusing nolation, we can reformulate the contraction inequality of Theorem 4.7 at stage r follows
1
distt S ddsst/1
with a
(198)
22 -11--a)
B Appendix
In the second part of our analysis we compute the expected 'Wall Clock Time' of our proposed method and compare it to the corresponding 'Wall Clock Time' of straggler- FedRep. We prove that when the computational speeds are drawn from the exponential distribution with parameter and the communication cost given by C- c,, (for some constant c), then SRPFL guarantees logarithmic speedup. Recall that in Corollary A.11 we get. the following simplified version of the contraction inequality
dist + S dist
a
1 with S 4
(199)
22
For the rest of this section w.l.g. we assume that the clients are re-indexed at every stage so that the expected computation times maintain a decreasing ordering i.c. Vr [[T]] < [Ty| < ..., S E [TK]] For simplicity henceforth we drop the stage index r. Notice that the decreasing ordering of the computation times combination with (199) imply that SRPFL initially benefits by including only few fast nodes in the training procedure. However,
the distance diminishes the improvement guaranteed by the contraction inequality becomes negligible and thus our method benefits by including slower nodes, thus decreasing the second term of the r.h.s. of (199)
Let us denote by X; the maximum distance for which the optimal number of participating nodes (for SRPFL) is no 2t. This definition immediately implies that X, = too To compute each X, we turn our attention on mcasuring the progress per unit of time achieved by SRPFL, when 2" ng nodes are utilized. This ratio stage r can be expressed as
dist. dist á (2)(1-a)
(200)
E [ng22 + CC
33

---

Published in Transactions on Machine Learning Research (10/2023)
Notice that by (199) the nominator captures the progress per round while the algorithm incurs E [Tng2r] computation and C communication cost. Similarly the ration when 2 no nodes are used is given by
dist + dist /
(201)
11 +
Based on the above incqualities we can now compute the optimal doubling points (in terms of distance) and thus the values of s's. Subsequently, we compule the number of iterations SRPFL spends in every stage.
Lemma B.1. For all let X; denote the marimun distance for which the optimal rumber of nodes for SRPFL is no 2¹. Then the following holds
(E [Tg22=] ] +C)
a
v2.
Vi 0
X;
1+ E [
(202)
22 a)
E (Tng2
Xo = to0
+00
Further SRPFL spends at each stage r at most t" communication rounds such that
<V((E|T
T -B Tmoar
2 log
[[ -B
T Taoar
t
(203)
log( I-a
Proof. For each stage r let us compute the point where the transitioning between 2" 2 ng and 2 -g occurs. That is the distance at which SRPFL benefits by doubling the number of participation nodes to 2 Thus equating the two ratios in 200 and (201] we get
X,+1 X - (1)
X,41 X,+IVI (1-+)(1-o)
(204)
E [ +C
E [T022+1 +C +C
E. [Tno2r] + (1
1
X,+1 (B (Tn,ett1) - B [Tng2))
E [T [T
122(1- (1- a) v1 a)
/2
v2
(205)
(E [ +C)
X 125 - a)(1 vI -
1+
(206)
E [ [Tno2")
Let us now compute the number of rounds t" (henceforth denoted by t) required in stage r. That is the minimum number of iterations that SRPFL needs lo decrease the distance from X, to X,-1 using only the n²2" fastest participating nodes. Thus, starting off at X, and following (199 for rounds have
x x
(VI a)
(207)
1=0 2 -
As stated above we want to find the minimum number of rounds such that we reach the next doubling point L.e. want large enough such that
X 2 X,(VI-110
(VI a)'
(208)
2
1 /1
2 X,
a a)¹ +
a
(209)
2r a)
v1
34

---

Published in Transactions on Machine Learning Research (10/2023)
where in the last incquality use geometric series properties. We proceed to solve for t by rearranging and using (202) and the fact that X, >>>>> (=(1-a(1-----)
2 (2" a) a X
(vI-a)
(210)
22 a1 - v1 X, -
(E Tmoar +C] v
70 T T2g2++1
Tn
(211)
(E Ingar-1 +C)
2
E T[ -B Tagar
v2
E [ 2 T 1] [T E Tngg
(212)
(213)
taking the logarithm both sides derive the required amount of rounds
V2(
T
m02r-1
Tager
log
[mger
E
T apar
(214)
log( 1-a
(215)
The following lemmas compute the 'Wall Clock Time' that SRPFL and FedRep require order to achieve target accuracy e. As discussed in Section 4.1 for fair comparison we consider accuracy of the form
=
a N no
with
v2>2>1.
(216)
(1
a) (1 a)
When takes values close to 22 we expect SRPFL vastly outperform FedRep and as takes values close the performance gap diminishes.
Lemma B.2. Suppose at each stage the client's computational times are i. random variables drown from the exponential distribution with parameter X. Further, suppose that the expected communication cost per round is C= , for some constant c. Finally, consider target accuracy e given in Then the expected
Wall Clock Time' for SRPFL is upper bounded as follows
[(cc+ 1) + 4log(
E [TRRFII] & log N
(217)
log( 1-a 1))
7
Proof. First we upper bound the expected cost suffered by our method until the distance between the current representation and the optimal representation becomes smaller than Xiog(M/o)r i.e. the cost corresponding to the first log Zno
N
stages of SRPFL (denoted by E [TSRPFI.)
35

---

Published in Transactions on Machine Learning Research (10/2023)
log N
2mg
E [TAPFL]
t' [B [Tno2]] + C)
(218)
log( Zng
log
v2(E
T
m02
1
-E
Tnggi
T
ng
2 (E [ +C)
(219)
log
1-a
/2
-E
Tx
log
N logí
kmg
12.
T
-
2 (EE) [mo 21 +C)
*
log(
1
1-a
v2(E]TN)
T
log
E
-B
x
2(E [Tx/a] + C)
(220)
log(
214 Since the computational times of the clients come from the exponential distribution it is straightforward to derive the following bounds
where used
E [Tv]
(x//a]
(In(N/2) + +1 S log(N)
(221)
1
1
1
N
4
1
E Tx/a
E TT//
+
(222)
N/2
+1
N/2+2
3N/4/
A
3N
3A
N
2
E [T/2a] E [Tyr -
(223)
4
N
2X
1
1
N
4
1
E [x/4]] E [T//)]
+
(224)
3N/4+1
3N/4 +2
TNJ8
8
3N
6Λ
Making use the above bounds the expression in
220 further simplifies to
log
E [Tno C) log(3/2)
log(3 log(N))
2
+ + (E [T//2] +C)
(225)
log( r
log( -
i=1
log N
5
. 2m
1og(3 log(N))
E [T + log
C
+ 2
(E [Tx/a] + C)
(226)
log(
log( -
1-a
Further, notice that
1 N/2
[N//2]
Σ X²
(227)
N/2 +
and similarly [x]j]] < 1 :.
IE Tx/a S 4 #: EE Tx/sol S
and so on. Thus,
log
N 2ng
E [Too22:] [2ng] + E [mm]] +... N/4 < 7 i=1
(228)
Combining the bounds from 227 and 228 and substituting = ot in expression 226 we derive the following bound
2 2log logg log(N))
E [
SRPPL]
S
(elog(N/ (N/2n0) + +1 +
(229)
log
I-a
log
T-a (c 1)
6(c + 1)
1 <
1/oo)
(230)
log
36

---

Published in Transactions on Machine Learning Research (10/2023)
Having derived an upper bound on the cost suffered by SRPFL on the first log attention on bounding the cost incurred from logt ) until the target accuracy e is achieved (denoted by I [SSPri)) ). Recall that.
stages we now turn our
2no
a
(E TN/a + E [T
+C](1
vs
Xlog( N/mo)
1
(231)
(1
a)(1 vI a)
[T//2]
2mg
and further during the last stage of SRPFL, N clients are utilized deriving the following form in the contractions inequality from (199)
1 with a as 4
dist & dist /
(232)
N (1
We first compute the number rounds required this second phase of the algorithm. Starting with distance \oo)))/))) and following the contraction in 232 for rounds, we derive current distance at most
X log
2 (V1 a)¹ +
a
(V1 - a)
(233)
==
N
a)
ag
1 vI-
X
(V1
a) "
(234)
log
v mg
1 v1 VI-a
(1 a)
E [Tyaal + C)
v2)
+ (
a
V2
1+
(235)
E [Tw] E [Tx/a]
a)(1 vI
where in the first equality use geometric series properties and in the second we substitute according to 231 D. Using the fact that (E(E(TN// + C1 5/2 s [c+ 1)(2 - 1) and E [Tw] E [TN/2] s à log N
expression
235 is upper bounded by
a
vI ((22+
(c 1)(22 (c+1)(22-1) 1 (V1 (1-aa) (1-
(236)
a)(1
log N
na
V1
(237)
(1 [1- a1 - v1
(1 (1 - vi-a)
The above implics that the number of rounds the sccond phase is the smallest so that the target accuracy is achieved, i.e.,
2
a
a (1 a)(1 v1 -1
VI-a"
(238)
N
a)(1 v1 a)
N
ng
no
Further recall that from 216 the accuracy can be expressed in terms of
N (1 a) (1-v1-a" with v2>0>1.
(239)
Combining the above and solving for we derive the required number of rounds for the second phase
2 log
t t
(240)
log(
The expected cost during phase 2 be computed follows
37

---

Published in Transactions on Machine Learning Research (10/2023)
looí 2-1
E SRPFL.]
S( [Tx) +C)
+1
(241)
log( 1-a
1
2 2
<(I((M) + 1 11
+1
(242)
log I-a
log( T-
<4log(N)
(243)
log( I-
4
1 Ā
<log(NV
log
(244)
log
7--1
Summing the two quantities of interest we can derive the promised upper bound on the 'Wall Clock Time' of SRPFL.
E [TRRrr]] = E SAPFL]] + [TrrFF
(245)
6(c + 1)
1
4
1 1 ė-
I noo(
log(NN
log
(246)
log
log
T-a
< log N (0(e+)) + 4 logl
8-
1
(247)
log(
(og(
A
Having computed an upper bound on the expected 'Wall Clock Time' of SRPFL we proceed to compute an lower bound on the expected WWll Clock Time' of FedRep.
Lemma B.3. Suppose at each stage the client's 's computational times are i..dd random variables drawn from the exponential distribution with parameter X. Further, suppose that the expected commmuiication cost per round is C = x for some constant c. Finally, consider target accuracy e given in 12 Then the expected
Wall Clock Time' for FedRep lower bounded as follows
log 2lo
E Feddteep 2 loog
(248)
1 log
1-a
Proof. First we compute the number of rounds required by FedRep achieve the larget accuracy. Recall that FedRep utilizes N clients each round deriving the following form of the contractions inequality from 199
1
dist + < dist
with as 4
(249)
N (1 a)
Starting with distance equal and following the contraction in 249 for rounds, we derive current distance most
1 vI (1 a) a
(V1
a)*
a)' = (V1-
(250)
N no
1 a)
using the properties of geometric scries. Further recall that from
216 the accuracy can be expressed
with
v2> 6> 1.
(251)
N (1 a) (1 v1 - a)
n0
38

---

Published in Transactions on Machine Learning Research (10/2023)
The above imply that the number of rounds is going to be the smallest that guarantees that the target accuracy has been achieved that is
a
2 (vI a)
a
1 - vI-a a) - vI-a
(252)
(1 a) (1 vl
N ng
70
We use the fact that and solve for t. Thus, we derive
N (1 (1-10) (1-a)(1 v1 a) - a >0 for a S 1/4 and all reasonable values of to rearrange
no
2 loo
log N
+
(253)
log
log
I-a
Multiplying the number of rounds with a lower bound on the expected cost incurred per round, results in the desired lower bound on the expected 'Wall Clock Time' suffered by FedRep:
2 2
2-1
E [ Feddtep
2 E [x]
(254)
log
T-
log N + 2 2lo
2-1
2 log
(255)
log
1
À
T-a
Combining the results of Lemma B.2 and Lemma
B.3 obtain Theorem
4.9
Theorem 4.9. Suppose that at each stage the client's computational times are i. random variables drown from the erponential distribution with parameter X. Further, suppose that the expected communication cost. per round is C S.for some constant Finally, consider the target error given in 12 Then, we have [[TsapFL] E[TraRe
log( (2)1
0
log (og(N)+Hog( (1)
Proof.
[SnPFL]
6(c+
log(
( = 0
log
2-1
(256)
E FedRep
log N 2 log
log(N) + log
c-1
c--1
S 1 c, with probability at least 1 - ((mn) effectively without increasing the overall sample complexity, The formal statement and proof identical to Theorem 3 in (Collins et al.] 2021 and is onitted.
Remark B.4. The initialization scheme in Algorithm 2
guarantees that dist (is BB)
C More on Experiments
Hyperparameters and choice of models. We set the hyperparameters following the work of (Collins et al. ) Specifically, for the implementation of FedRep and FedRep-SRPFL we use SGD with momentum where the momentum parameter is set 0.5 and the local learning rate 0.1. Further, similarly to (Collins et al. 2021 set the local learning rate to 0.1 for all other methods under consideration, which obtains optimal performance. We fix the batch size to 10 for all our implementations. The number of local epochs is set to 1 for CIFAR10 with N = 100 and to 5 for the rest of the datascts. In terms of the choice of the neural network model, for CIFAR1O, we use LeNet-5 including two convolution layers with (64, 64) channels and three fully connected layers where the numbers hidden neurons are (120, 64). The same structure is used for CIFAR100, but the numbers of channels in the convolution layers are increased to (64, 128) and the numbers of hidden neurons are increased to (256, 128). Additionally, dropout layer with parameter 0.6 is added after the first two fully connected layers, which improves the testing accuracy. For EMNIST and
39

---

Published in Transactions on Machine Learning Research (10/2023)
CIFAR10, 1=100, Shard=10
EMNIST, M=100, Shard=10
SENT140, M--33 Shard-2
0,80
Testing Accuracy
70-
Testing Accuracy
60
Ferh
80
0.75
50
LGfccing
70-
0.70
FLAAT
40
TLAR
60
0.65
30
Feshiggt
0.00
1.29 2.58 3.87 5,16
0.00
1.29
2.58 3.87
5.16
0.60
0.00
2.07
4.14
6.21
8.28
Homogeneous (C.T. - 0)
Homogeneous (C.T. - 0)
Homegene ous (C.T. 0)
CIFAR10, M=100, Shard=10
EMNIST, M=100, Shard=10
SENT140, M-83, Shard-2
Testing Accuracy
Testing Accuracy
90
Testing Accuracy
0.80
60
80
0.75
50
-70
0.70
60
60-
50
0.65
0.00
1.71 3.42 5.13
6,84
48.00
1.71
3.42 5.13
6.84
0.60
0,000
0.332
0.664
0,996 1.328
Hamageneaus (C.T.
10)
Homogeneous (C.T. 10)"
Homageneous 10/"
CIFAR10, M=100, Shard-10
EMNIST, M-100, Shard-10
SENT140, M-83, Shard-2
Testing Accuracy
770
Testing Accuracy
.90
Testing Accuracy
0.80
60
80
0.75
50
70-
0.70
60
60
-30
50
0.000 0.554 1.108 1.662 2,216
0.000 0.554 1.108 1.662 2.216
0.000 1.457 2.914 4.371 5,828 0.60
Homogeneous (C.T. 100f'
Homegeneous (C.T. 1007"
Homogeneous (C.T. 100Y"
Figure 7: Numerical results on CIFARIO, EMNIST, Sent140 with full participation (M= N) in the fixed computation speeds setting. 'Shard' denotes the number of classes per client. 'C.T. denotes the communication cost per round.
FEMNIST, we use MLP with three hidden layers with (512, 256, 64) hidden neurons. For Sent140 we use two-layer bidirectional LSTM with dimension 256 and dropout rate 0.5 followed by a fully connected layer of dimensions 5 and a classification head. Further, we use the standard glove embedding of dimension 100 and vocabulary of size 10000.
For the needs of SRFRL, we split the neural network model into two parts, the customizod head h; and the common representation 6. In our experiments, we simply take the customized head to be the last hidden treated as the common representation. Note that LG-FedAvg and
layer and the rest of the parameters LG-FLANP have different read/represeenattion split scheme and the head is globally shared across all clients while local version of the representation is maintained on every client. For all included datasets, i.e. CIFAR1O, CIFAR100, EMNIST, FEMNIST and Sent the common head include the last two fully connected layers and the rest of the layers are treated as the representation part.
Datasets. We include five datasets in our empirical study: CIFAR10 which consists of 10 classes and a total number of 50,000 training data points, CIFAR100 which consists of 100 classes and the same amount of data points as CIFAR10, and EMNIST (balanced) which consists of 47 classes and 131,600 training data
points. Note that in Figures 3 7
we use the first 10 classes from EMNIST following (Collins et al.] 2021
For FEMNIST, we use the same selling as (Collins et al. 202T with the exception that in Figures 3 we allocate each client 150 data points and in Figure 6 we allocate client i, (100 + u,) samples, with U; uniformly distributed random variable in [0,50] The sentiment1.0 dataset contains 1, 600, 000 tweets annotated negative positive and they can be used to detect sentiment. For this dataset perform pre-processing splitting the samples across clients both in homogeneous well as a heterogeneous manner. The number of allocated samples to clients follow the log-normal distribution. During our training procedure, we perform the data augmentation operations of standard random crop and horizontal flip on the first two datasets and and perform no pre- pre-processing for the last one.
Homogeneous Setting. Apart from providing shared representation model to simultaneously address the hurdle of data heterogeneity. In the less challenging
resillence benefits our method takes advantage of the
40

---

Published in Transactions on Machine Learning Research (10/2023)
FEMNIST, M=100, Shard=3
SENT140, M=83, Shard=2
SENT140, M=83, Shard=2
CIFARIO, M=100, Shard=5
Testing Accuracy
80-
Testing Accuracy
0.80
0.75
Jeship
HLP
Teating Accurasy
0.75
0.80
Testing Accuracy
80
70
60
Fochip
S1P
0.70
0..7
50
-40
TUAP
0.65
0.65
30
Festip
20
20
Fushipif
0.60
0.60
0.00
1.29
2.58
3.87
5,16
0.00
0.58
1.16
1.74 2.32
0.00
2.07
4.14 6.21
8,28
0.00
1.29
2.58
3.87
5.16
Total Time (C.T. m 0)
Heterogeneous (C.T. - 0)
Stragglers % (C.T.==
0)
Total Time (C.T. - 0)
FEMNIST, M=100, 3hard=3
CIFAR10, M=100, Shard=5
SENT140, M=83, Shard-2
SENT140, M-83, Shard-2
Testing Accuracy
80
Testing Accuracy
0.80
Tosting Accuracy
0.80
Testing Accuracy
70
.00
60
0.75
60
550
0.70
0.70
40
80
0.65
0.65
20
20
0.60
0.00
1.71
3.42
5.13
6.84
0.00
1.83
3.66
5.49 7.32
0.000
0.332
0.664
0.996 1,328
0.00
1.71
3.42
5.13
6,84
Total Time (C.T.
10)
Meterogeneous (C.T.
10)"
Straggiers
10)
Total Time (C.T. = 10)
FEMNIST, M=100, 3
CIFAR10. M-100, Shard-S
SENT140, M-83, Shard-2
SENTI40,
M-83
Shard-2
Testing Accuracy
Testing Accuracy
0.80
Testing Accuracy
0.80
Testing Accuracy
80
0.75
70
60-
0.70
0.70
50
40
40
0.65
0.65
20
20
0.60
0.60
0.000
0.554
1.108 1.662 2,216
0.000
1.308
2.616 3.924 5,232
0,000
1.457
2:914 4.371 5.828
0.000
1.108 1.662 2,216 Total Time (C.T. 1001
0.554
Total Time (C.T. (C.).
Heterogeneous (C.T. 100y
Stragglers % (C.T. -000""
Figure 8: Numerical results on FEMNIST, Sent140, CIFARIO in the fixed computation speeds setting. From left to right: papalaaced/Imbalaneeee /Straagleess Saanglless//Corrrlattee classes per dlient.C.T. denotes the communication cost per round.
/Correlated Heterogeneity. 'Shard' denotes the number of
homogeneous setting one would expect FedRep-SRPFL. to lose some of its competitive advantage. However, our numerical results in Figure indicate that our method exhibits a stable performance while maintaining an edge over the other baselines. We note that LG-SRPFL outperforms FedRep-SRPFL in the Sent dataset where LG- (G-FedAvg appears to be better-suited than FedRep. Importantly, our doubling scheme continues enhance both subroutines providing clear stragglev-resilience benefits even in data homogencous environments.
Additional Regimes. In Figure we observe the extent at which the benefits of our meta-algorithm persist in various regimes. Imbalanced Dalasels. In the two left columns we explore how the performance of our doubling scheme degrades for different levels of imbalanced local datasets. In the FEMNIST dataset. we allocate to each client i, (100 + u.) local samples where u, random variable uniformly distributed in 50). In this setting the numbers of local samples are sufficiently close and as result the benefits of our scheme are prevalent for all different values of communication cost. As we allow the clients to have more diverse numbers of local samples however, the speedup provided by our doubling scheme diminishes. Indeed this is the case in the 140t140 dataset where we allocate data to clients in a more dispropotional manner following the log-none normal distribution. We point out that in settings with high diversity simply doubling the number of clients not sufficient. Instead, to achieve optimal performance one needs to select consecutive participating sets of clients with cumulative number of samples that double from one stage to the next. Different levels of stragglers. We further consider the setting where clients are split into two categories. The first category consists of typically fast clients whose computational values come from the exponential distribution with À = 0.1. The sccond category consists of clients who are typically straggling, i.c. sampling their computational cost from the exponential distribution with À = 10. On the third column of Figure S we compare FedRep, LG-FedAvg and their straggler-resilne resilient variants for different percentages of clients coming from the latter category. Unsurprisingly, the performance of the methods under considcration dogrades as the percentage of stragglers increases however our doubling scheme to some extent mitigates this effect. Correlated helerogeneity. In this setting we assign data to clients depending their speeds. More specifically,
41

---

Published in Transactions on Machine Learning Research (10/2023)
split clients into groups based on the computational times and respectively we group samples into 3 distinct groups based on their labels. Subsequently, we allocate mutually exclusive (with respect to their labels) groups of data to specific groups of clients. The rightmost column of Figure S indicates that the benefits our doubling scheme persist, although to a lesser degrec as the correlation between data and system heterogeneily grows.
42


Reviews:
[{"review_id": "ecLbbyXDoS", "reviewer": "Reviewer_rm5K", "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": "My overall impression of the paper is that it has a good core (simple, in a postiive way) idea coupled with strong theoretical analysis. However, the paper has the following weaknesses.\n\n1. Its claims in the abstract & introduction that overstate what the paper does.\n1. The theoretical assumptions that are arguably too strong (as best as I can tell, though I would love to discuss with the authors about this).\n1. An algorithmic choice that potentially goes in the face of the data minimization principle underlying federated learning.\n1. The empirical results have some potential issues, especially related to how datasets are partitioned across clients, and a lack of reproducibility.\n\nI discuss these in greater detail below.\n\n## Strengths\n\nThe fundamental idea, the doubling mechanism, is interesting theoretically and empirically. The algorithm has a lot of intuitive appeal (ie. wait for the fastest $n$ clients to finish, and proceed), and its core simplicity is a virtue. In particular, the doubling mechanism can ostensibly be applied to any federated learning setting in which there are (1) sufficiently many clients available and (2) sufficiently many can participate in a given round. In particular, it is not actually tied to the FedRep method (though some of the wording in the first 2 sections obscures that fact, more on that below). The method can be applied to other methods, including FedAvg and LG-FedAvg (the latter of which is actually implemented in the empirical results).\n\nOf course, simple mechanisms often still require nuanced theoretical analysis, and this is no difference. I think that the focus on the \"linear representation\" setting is actually a strength of the paper. It allows for much more fine-grained control over the optimization behavior, and therefore over the runtime speedups that the method can incur. Moreover, the theory involves a ton of great high-dimensional probability results that I enjoyed reading. I will caveat that i have not verified the theory 100%, though I believe that the theory holds **with the assumptions made in the main body** (more on that below).\n\nFor the empirical results, the authors do a good job of benchmarking a wide swath of methods. I do not believe that this is strictly necessary (as the paper's primary contributions are, in my opinion, theoretical) but this was nice to see nonetheless.\n\n## Weaknesses\n\n### Accomplishments versus Claims\n\nOne of the core themes of the abstract and introduction are that this work addresses both data heterogeneity and system heterogeneity simultaneously. From the introduction:\n```\nWe propose Straggler-Resilient Personalized Federated Learning (SRPFL), an adaptive node-participation method which accommodates gradient-based subroutines and leverages ideas from representation learning theory to find personalized models in a federated and straggler resilient fashion.\n```\nThis statement struck me as incongruous with what the paper is doing. The representation learning in Algorithms 1 and 2 are all from the FedRep algorithm from (Collins et al., 2021). This is freely discussed in later sections of the work, but section 1 completely removes this connection from the discussion.\n\nMore generally, I don't know that I agree that this paper really addresses data heterogeneity, as prior work has shown that FedRep does that. Rather, the contribution of this work is the doubling mechanism, which could be applied to any FL algorithm (whether or not it uses representation learning). This brings me to a fundamental question:\n\n**Question:** Why does this paper couple the study of the doubling mechanism to the FedRep algorithm? Are there theoretical obstacles to studying algorithms such as FedAvg when using this mechanism? Are there aspects of FedRep that specifically motivated the doubling mechanism?\n\nAs it stands, I would greatly prefer that the authors focus on what they do provide and study. The question of data heterogeneity is also related to the following section, as I believe that the theoretical assumptions help mitigate issues at the intersection of data & system heterogeneity.\n\n### Theoretical Assumptions \n\nThe authors make a number of assumptions, 4.1 - 4.4, that I think are worth considering. Many of these appear in (Collins et al., 2021). However, they take on new meaning under the doubling mechanism proposed by this paper.\n\nPrimarily, I am concerned with Assumption 4.3 (aka \"client diversity\"). This assumption states that if you take $n$ rows of a matrix $W^*$ (each row of which corresponds to a client), then you get a full rank matrix. As best as I can tell, this is the same $n$ from the doubling mechanism. This is the core issue: In work like (Collins et al., 2021), it is reasonable to assume that the $n$ rows are drawn randomly, and therefore this statement need only hold with high probability. In this work, $n$ is not random, it is simply the $n$ clients who finish fastest. The work seems to make no assumptions on this distribution. Thus, it is reasonable to imagine settings where the data heterogeneity is tied to the system heterogeneity, in which case this assumption is **much** stronger than it seems on first glance.\n\nFor example, if we consider a simple setting with 3 clients, and $n = 2$, and one of the clients is always the slowest, then this assumption seems strange in practice. In fact, you can construct counter-examples in this vein to Theorem 4.5, but these counter-examples fail Assumption 4.3 if you assume that it must hold for all possible subset of size $n$.\n\nAs discussed above, this relates back to data heterogeneity. As best as I can tell, this \"full rank\" assumption limits the amount of data heterogeneity possible. This is not necessarily bad, but it's certainly a limitation of the paper and probably deserves more discussion. I would love feedback from the authors about this.\n\n### Data Minimization and SVD Initialization\n\nOne of the core tenets of federated learning (and indeed, one of the reasons it exists as a computational paradigm) is data minimization - the server should have access to no more data than is completely necessary. However, the initialization phase of Algorithm 2 involves sending sums of outer products of data. This seems potentially (not definitely) problematic from a data minimization perspective. I am no expert on privacy mechanisms for distributed SVD computations, so it is possible that there is some way that this could be integrated with something like differential privacy. I do not think that this integration is necessary, but I do believe that the authors need to discuss why this initialization phase is not tantamount to sending data in the clear to the server. In short, why is this initialization operation acceptable in the context of federated learning? There is no single answer to this, but I think it warrants discussion.\n\n### Empirical Investigation and Datasets\n\nI will preface this with saying that the empirical investigations are not the core of this paper, at least in my mind. That being said, they exhibit some characteristics that unfortunately reduces how convincing they are. First and foremost, two of the datasets that the paper mentions, FEMNIST and SENT140, have a natural federated structure, as discussed in the Leaf benchmark (Caldas et al., 2017). However, a close read of this paper suggests that this was not used, and instead the data was partitioned by some random process. Can the authors comment on why this natural federated structure was not used?\n\nSecond, I will mention briefly that the choice of optimizer hyperparameter seems somewhat arbitrary to me. In particular, all methods use the same learning rate (whether or not they do partial model participation or global optimization), which I believe needs some form of justification. On a more minor note, this work says that it sets hyperparmeters following the work of (Collins et al., 2021), but this doesn't seem to be the case. The latter uses SGD with momentum and a learning rate of 0.5, while the former uses SGD with a learning rate of 0.01.\n\nThese issues are not huge, but they are things that stuck out to me upon reading closely.", "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 straggler nodes in the context of federated learning with partial model personalization. More specifically, it proposes a sampling mechanism based on doubling in which:\n\n* At each round, the server sends a model to $N$ clients.\n* The server only waits for the fastest $n \\leq N$ clients to return an update.\n* The parameter $n$ is doubled at a specified cadence.\n\nThe authors then study the application of this \"doubling mechanism\" to the FedRep algorithm (Collins et al., 2021), which does partial model personalization. The authors study this algorithm theoretically in a \"linear representation\" setting, and show that under various assumptions (including a kind of incoherence across clients, a \"full-rank\" assumption, and an exponential distribution runtime assumption), that this incurs a logarithmic speed-up in runtime.\n\nFinally, the paper does some comparisons on image datasets (ie. (F)EMNIST and CIFAR-10(0) datasets) and the SENT140 dataset. These experiments show that in many settings (varying amounts of heterogeneity, and varying communication costs), this FedRep + doubling algorithm converges faster than other methods, especially FedRep without the mechanism.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "## Critical Recommendations\n\n1. Please clarify the contributions of this paper, especially disambiguating the doubling mechanism proposed and the representation learning and personalization approach of FedRep. Additionally, if there is any facet of the doubling mechanism that is intertwined with or reliant upon representation learning mechanisms (and is therefore not amenable to FedAvg), please discuss. Finally, please discuss obstacles (theoretically or otherwise) to studying FedAvg with the doubling mechanism (note that this might be a duplicate request with the previous sentence).\n1. Please clarify the significance and limitations of Assumption 4.3. In particular, is this assumed to hold over all possible size $n$ sets simultaneously? If so, how does this impact the possible data heterogeneity possible, especially statistical correlation between data and system heterogeneity.\n1. Please discuss why the initialization phase of Algorithm 2 is compatible with the general principles of federated learning, especially data minimization.\n\n## Other Recommendations\n\n1. Please clarify why the experiments do not use the \"natural\" client partitions of the FEMNIST and SENT140 datasets.\n1. Please discuss how the learning rate for the optimizers were chosen, and why it is reasonable to use the same learning rate across all methods.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "TqxMbWPuso", "reviewer": "Reviewer_5MxQ", "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 paper studies important issues in federated learning: it considers data/system heterogeneity and aims the issues of stragglers. It is easy to follow and well-organized. \n2. The paper conducts a theoretical analysis of the proposed algorithm to show the speedup guarantee. \n\n\nWeakness:\n1. The idea of using shared representation and client-specific head parameters for personalized FL is not new, e.g., see (Jiang & Lin, 2023, URL: https://openreview.net/forum?id=3aBuJEza5sq), (Zhu et al., 2022, URL: https://openreview.net/pdf?id=E4EE_ohFGz). As far as I can tell, the main contribution seems to be the adaptive participating scheme, which I think also has some issues as detailed next.     \n\n2. To tackle the issue of stragglers, the proposed algorithm only selects a subset of the fastest clients to participate in each round. This may raise several issues: \n(a) In practical federated learning systems, the participation of clients is usually out of the central server's control. This significantly limits the adoption of the proposed algorithm in practice. \n(b) While the proposed algorithm geometrically increases the number of participating nodes, certain nodes such as stragglers may never participate. This can be problematic especially when the data distribution of stragglers differs from the fast nodes (as these clients never got the chance to update their local head). As also indicated by the authors, such a disproportional selection of nodes raises fairness and accuracy concerns.\n\n3. In the experiments, the computational time of all clients is determined uniformly (i.e., time follows an exponential distribution with parameter uniformly sampled from [1/M,1]). This setup is not comprehensive enough as it intentionally avoids the issue of 2(b) mentioned above. Moreover, only overall accuracy on all clients is reported. It is necessary to compare local accuracy for personalized FL. More experiments should be conducted (see next section).\n\n", "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 proposes an algorithm for personalized federated learning that tackles data/system heterogeneity and is resilient to stragglers. Specifically, personalization is attained by learning a global representation followed by user-specific parameters. Straggler-resilience is achieved by adaptively and disproportionally selecting participating clients over time. Theoretical results show that the proposed algorithm can attain near-optimal sample complexity and logarithmic speedup in certain settings. Experiments on several datasets show that the algorithm can attain better performance than baselines.  ", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. More experiments need to be conducted to validate the proposed algorithm, including:\n(a) settings when the computational times of clients are not uniformly distributed, e.g., when stragglers are always the same set of clients, when certain classes are not accessible to the fastest clients, etc.\n(b) in addition to the overall accuracy, the performance of local clients also needs to be reported. \n(c) comparison with more baselines, especially those tackling personalization and heterogeneity.\n\n2. Because the idea of using shared representation mapping and client-specific head parameters has been leveraged in many prior works, it is essential to highlight the novelty of the proposed algorithm. \n\n3. The theoretical results are only for a special case with linear representation, I wonder whether it can provide some insights for non-linear cases.  ", "broader_impact_concerns": "The proposed algorithm may raise unfairness issues for devices with large computational time. ", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "skze692xEg", "reviewer": "Reviewer_wdpT", "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": "Pros:\n1. The paper is well-written and easy to follow. The background introduction is detailed and the commons and differences of the proposed framework have been clearly illustrated.\n2. The theoretical analysis shows the proposed method could achieve O(log N) improvement over FedRep although it is based on a linear representation cases.\n3. The experiments on different level system heterogeneity show the potential of the proposed method to achieve a faster convergence over other baselines. \n\nCons:\n1. The proposed setting is not that practical and affects its novelty. Although the paper mention it is possible to get the computation power of every client, I still think it is a very big assumption. The network communication could affect the overall aggregation as well so that the slowest client might not be the slowest always. Also, I think it is not trivial to achieve the rank every easily.\n2. The theoretical analysis is done in the linear way that is not practical as well. \n3. The experimental results should be included for the data-homogeneous settings to show it could still achieve a stable performance. ", "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 proposes a new framework to handle the data heterogeneity problem as well as system heterogeneity at same time. Specifically, to deal with the data heterogeneity problem, it adopt a framework that include a personalized head for every client and a shared embedding among all the clients. To deal with the system heterogeneity problem, instead of random sampling several candidates uniformly, based on the computation power, it select the top n_0 clients in every round. The theoretical analysis shows that in the linear representation case, the proposed SRPFL method could achieve O(log N) speed up compared with FedRep. The experiments done in several datasets with different system heterogeneity show the proposed method could achieve a faster convergence than other baselines.  ", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Typos:\ndist(B^0, B*) in eq 7.\n\n1. Move the experiments of Figure 5 or 6 in the appendix to the main paper. \n2. Add experiments on the data-homogeneous settings. ", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: All reviewers were ultimately in favor of this work appearing in TMLR---they agreed that this was an interesting approach, and appreciated the simplicity of the procedure and corresponding theoretical analysis. The reviewers also suggested several useful improvements to the paper, including more clearly explaining the main contributions, discussing potential limitations of the assumptions made (particularly Assumption 4.3), and including experiments in additional regimes including in homogeneous settings as well as scenarios when there may be correlations between data/systems heterogeneity. The authors have already thoughtfully addressed many of these suggestions in their revision, though we would suggest that they take a final pass to ensure that the feedback has been thoroughly incorporated---particularly in terms of clearly explaining key assumptions and potential limitations of the approach in addition to its benefits (e.g., in terms of data minimization/privacy, correlated heterogeneity, etc).
Decision: Accept (poster)