Distributed Personalized Empirical Risk Minimization
Authors: Yuyang Deng
Mohammad Mahdi Kamani
Pouria Mahdavinia
Mehrdad Mahdavi
Venue: NeurIPS 2023 poster (2023)
Keywords: distributed learning
heterogeneous data
heterogeneous system
convergence analysis

Abstract: This paper advocates a new paradigm  Personalized Empirical Risk Minimization (PERM) to facilitate learning from heterogeneous  data sources without imposing stringent constraints on computational resources shared by participating devices. In PERM, we aim at  learning  a distinct model for each client by personalizing the aggregation of local empirical losses by effectively estimating the  statistical discrepancy among data distributions, which entails  optimal statistical accuracy for all local distributions and overcomes the data heterogeneity issue.  To learn personalized models at scale,  we propose a distributed algorithm that replaces the standard model averaging with model shuffling to simultaneously optimize 
PERM objectives for all devices. This also allows to learn distinct model architectures (e.g., neural networks with different number of parameters) for  different clients, thus confining to  underlying  memory and compute resources of individual clients. We rigorously analyze the convergence of proposed algorithm and  conduct experiments  that corroborates the effectiveness of proposed paradigm.

Distributed Personalized Empirical Risk Minimization
Yuyang Deng Pennsylvania State University yzd82@psu. edu
Mohammad Mahdi Kamani
Pouria Mahdavinia Pennsylvania State University
Wyze Labs mmkamani @alumni psu. edu
pxm54260psu. edu
Mehrdad Mahdavi
Pennsylvania State University m2m6160pst edu
Abstract
This paper advocates a new paradigm Personalized Empirical Risk Minimization (PERM) to facilitate learning from heterogeneous data sources without imposing stringent constraints on computational resources shared by participating devices. In PERM, we aim to learn distinct model for each client by learning who learn with and personalizing the aggregation of local empirical losses by effectively esti- mating the statistical discrepancy among data distributions, which entails optimal statistical accuracy for all local distributions and overcomes the data heterogeneity issue. To learn personalized models at scale, we propose distributed algorithm that replaces the standard model averaging with model shuffling to simultaneously optimize PERM objectives for all devices. This also allows us to learn distinct model architectures (e.g., neural networks with different numbers of parameters) for different clients, thus confining underlying memory and compute resources individual clients. We rigorously analyze the convergence of the proposed algo- rithm and conduct experiments that corroborate the effectiveness of the proposed paradigm.
Introduction
Recently federated learning (FL) has emerged as an alternative paradigm to centralized learning to encourage federated model sharing and create framework to support edge intelligence by shifting model training and inference from data centers to potentially scattered-and perhaps self-interested systems where data is generated [1]. While undoubtedly being better paradigm than centralized
learning, enabling the widespread adoption of FL necessitates foundational advances in the efficient use of statistical and computational resources encourage a large pool of individuals
corporations
to share their private data and resources. Specifically, due to heterogeneity of data and compute resources among participants. it necessary, if not imperative, to develop distributed algorithms that are i) cognizant of statistical heterogeneity (data-awa awareness) by designing algorithms that effectively deal with highly heterogeneous data distributions across devices; and ii) confined to learning models that meet available computational resources of participant devices (xystem-awwarensss))
To mitigate the negative effect of data heterogeneity (non-IIInesss, two common approaches clustering and personalization. The key idea behind the clustering-based methods [2, 3, 4, is is to partition the devices into clusters (coalitions) of similar data distributions and then learn single shared model for all clients within each cluster. While appealing, the partitioning methods are limited to heuristic ideas such clustering based on the geographical distribution of devices without taking the actual data distributions into account and lack theoretical guarantees postulate strong assumptions on initial models or data distributions [4, 5]. In personalization. based methods [6, 7, 2, 9, 10, 11]. the idea is to learn a distinct personalized model for each device alongside the global model, which
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

can be unified as minimizing a bi-level optimization problem [12]. Personalization aims to learn model that has the generalization capabilities of the global model but can also perform well on the specific data distribution of each participant suffers from a few key limitations. First, as the number of clients grows, while the number of training data increases, the number of parameters to be learned increases which limits to increase in the number of clients beyond a certain point to balance data and overall model complexity tradeoff- a phenomenon known as incidental parameters problem [13]. Moreover, since the knowledge transfer among data sources happens through a single global model, it might lead to suboptimal results. To see this, consider an extreme example, where half of the users have identical data distributions, say D, while the other half share data distribution that is substantially different, say -D (e.g. two distributions with same marginal distribution on features but opposite labeling functions). In this case, the global model obtained by naively aggregating local models (e.g., fixed mixture weights) converges to a solution that suffers from low test accuracy on all local distributions which makes it preferable to learn model for each client solely based on its local data or carefully chosen subset of data sources.
Focusing on system heterogeneity, most existing works require learning models of identical ar- chitecture to be deployed across the clients and server moodeel-homogenetty [14, 15], and mostly focus on reducing number [15] or size [16, 17, 18] of communications or sampling handling chaotic availability of clients [19, 20]. The requirement of the same model makes it infeasible to train large models due to system heterogeneity where client devices have drastically different computational resources. A few recent studies aim to overcome this issue either by leveraging knowledge dis- tillation methods [21, 22, 23, 24] or partial training (PT) strategies via model subsampling (either static [25, 26), random [27]. or rolling [28]). However, KD-based methods require having access to public representative dataset of all local datasets at server and ignore data heterogeneity in the distillation stage to a large extent. The focus of PT training methods is mostly on learning a single server model using heterogeneous resources of devices and does not aim at deploying model onto each client after the global server model is trained (which is left as a future direction in [28]). The aforementioned issues lead to a fundamental question: "What is the best strategy to learn from peeterogeneoue data sources to achieve optimal accuracy w.E.t. each data source, without imposing stringent constraints on computational resources shared by participating devices?".
We answer this question affirmatively, by proposing a new data&system-awaar paradigm dubbed Personalized Empirical Risk Minimization (PERM), to facilitate learning from massively fragmented private data under resource constraints. Motivated by generalization bounds in multiple source domain adaptation [29, 30, 31, 32], in PERM we aim to learn distinct model for each client by personalizing the aggregation of empirical losses of different data sources which enables each client to learn who to learn with using an effective method to empirically estimate the statistical discrepancy between their associated data distributions. We argue that PERM entails optimal statistical accuracy for all local distributions, thus overcoming the data heterogeneity issue. PERM can also be employed in other learning settings with multiple heterogeneous sources of data such as domain adaptation and multi- task learning to entail optimal statistical accuracy. While PERM overcomes the data heterogeneity issue, the number of optimization problems (i.e.. distinct personalized ERMs) to be solved scales linearly with the number of data sources. To simultaneously optimize all objectives in a scalable and computationally efficient manner, we propose novel idea that replaces the standard model averaging in distributed learning with model shuffling and establish its convergence rate. This also allows us to learn distinct model architectures (e.g., neural networks with different number of parameters) for different clients, thus confining to underlying memory and compute resources of individual clients, and overcoming the system heterogeneity issue. This addresses an open question in [28] where only single global model can be trained in a model- goodel-hetrrogeeouus setting, while PERM allows deploying distinct models for different clients. We empirically evaluate the performance of PERM, which corroborates the statistical benefits of PERM in comparison to existing methods.
2 Personalized Empirical Risk Minimization
In this section, we formally state the problem and introduce PERM as an ideal paradigm for learning from heterogeneous data sources. We assume there are N distributed devices where each holds a distinct data shard S, = {( } with n, training samples that are realized by a local
source distribution D, over instance space = = X x V. The data distributions across the devices not independently and identically distributed (non-IID or heterogeneous), i.c., Dj # D2 #
2

---

+ DN, and each distribution corresponds to a local generalization error or true risk L;(h) ((r.))),, (1)((z), 3)]. 1, 1.2.... ..,..... NN on unseen samples for any model E H. where H. is the hypothesis set (e.g., a linear model or a deep neural network) and : y x y - R + is given convex or non-convex loss function. We use (h) = (1/nz) ((w,)ss e (h(2)),) B)) to denote the local empirical risk or training loss at ith data shard S; with n, samples.
We seek to collaboratively learn a model or personalized models that entail a good generalization on all local distributions, i.e. minimizing true risk L.(.), i = =I..... for all data sources (all-for- all [33]). A simple non-personalized solution, particularly in FL, aims to minimize (weighted) empirical risk over all data shards in a communication-effiiiee. manner [34]:
arg min Ex (h)) with p AN.
(WERM)
where AN (p E R E: p(1) 1] denotes the simplex set.
has been shown that a single model learned by WERM, for example by using fixed mixing weights p(i) = n,/n, where n is total number of training samples, or even agnostic to mixture of distributions [35, 36], while yielding a good performance on the combined datasets of all devices, can suffer from poor generalization error on individual datasets by increasing the diversity among distributions [37, 38, 39, 40J. To overcome this issue, there has been a surge of interest in developing methods that personalize the global model to individual local distributions. These methods can be unified as the following bi-level problem (a similar unification has been made in [12]):
N h4 = arg minscH Σnm (() (BERM)
arg mink, ha EH Ohi@h)) subject to
where @ denotes the mixing operation to combine local and global models, and Fy is a modified local loss which is not necessarily same as local risk L. By carefully designing the local loss F; and mixing operation @ we can develop different penalization schemes for FL including existing methods such as linearly interpolating global and local models [11. 21. multi-task learning [10] and meta- learning [9] as special cases. For example, BERM reduces to zero- perr--prsonnaizattonn objective WERM when 1,@h. = h. and F; = 2. At the other end of the spectrum lies the zero-collaberation where the ith client trains its own model without any influence from other clients by setting h, h, =., The personalized model with interpolation of global and local models can be recovered by setting h, - ad,, + (1- hh. and F, = Σ,. While more effective than a single global model learned via WERM, personalization methods suffer from three key issues: i) the global model is still obtained by minimizing the average empirical loss which might limit the statistical benefits of collaboration, ii) overall model complexity increases linearly with number of clients, and iii) same model space is shared across servers and clients.
To motivate our proposal, let us consider the empirical loss ΣM (h) in WERM (or the inner level objective in BERM) with fixed mixing weights a E AN. and denote the optimal solution by ha The excess risk of the learned model ho on ith local distribution D; w.r.t. the optimal local model hn = arg minpee LD, (h) (i.e. all-for-one) can be bounded by (informal) [31]
N
a()) Σ
C. (ma) S C. (ht) + Σ (H) + 2 Σ a(j)discH D.) +
(GEN)
where Ry (H) is the empirical Rademacher complexity w.r.t. S,, and discH D₂) a pseudo- distance on the set of probability measures on E to assess the discrepancy between the distributions D; and D, with respect to the hypothesis class H. as defined below [29]:
Definition 1. For a model space H and D, D' two probability distributions on E X ),
discæ (D,D') sup [Ee-D(E(h. RCH
{)) EggwD (e(h,
Intuitively, the discrepancy between the two distributions large, if there exists a predictor that performs well on one of them and badly on the other. On the other hand, if all functions in the hypothesis class perform similarly on both then D and D' have low discrepancy. The above metric which is a special case of a popular family of distance measures in probability theory and mathematical
3

---

statistics known as integral probability metrics (IPMs) [41], can be estimated from finite data by replacing the expected losses with their empirical counterparts (i.e. C, with L.).
From GEN, can be observed that mismatch between pairs of distributions limits the benefits of ERM on all distributions. Indeed, the generalization risk w.r.t. D, will significantly increase when the distribution divergence terms discu (D,. J,D,,, D:) are large. It leads to an ideal sample complexity 1/ /n where 7 = + f22 +.......................... nn is the total number of samples, which could have been obtained in the IID setting with (x = 1/N when the divergence small as the pairwise discrepancies disappear. Also, we note that even if the global model achieves a small training error over the union of all data over over parametrized setting) and can entail good generalization error with respect to average distribution, the divergence term still remains which illustrates the poor performance of the global model on all local distributions D, i = ......... This implies that even personalization of the global model as in BERM can not entail a good generalization on all local distributions as there is no effective transfer of positive knowledge among data sources in the presence of high data heterogeneity among local distributions (similar impossibility results even under seemingly generous assumptions on how distributions relate have been made in multisource domain adaptation as well [42]).
Interestingly the bound suggests that seeking optimal accuracy on all local distributions requires choosing distinct mixing of local losses for each client i that minimizes the right-hand side of GEN. This indicates that in an ideal setting (i.e. all-for-all), we can achieve the best accuracy for each local distribution Dj by personalizing the WERM, i.c., (i) first estimating at, i = 1, 1.2..... .N for each client individually, then (ii) solving a variant of WERM for each client with obtained mixing parameters:
arg min NEH.
Σ 1= a,) for i i 1, N.
(PERM)
By doing this each device achieves the optimal local generalization error by learning who to learn with based on the number of samples at each source and the mismatch between its data distribution with other clients. We also note that compared to WERM and BERM, in PERM since we solve different aggregated empirical loss for each client, we can pick a different model space/model architecture H4 for each client to meet its available computational resources.
While this two-stage method is guaranteed to entail optimal test accuracy for all local distributions D;, however, making scalable requires overcoming two issues. First, estimating the statistical discrepancies between each pair of data sources (i.e., a;, = 1,.... N) is a computing burden as it requires solving O(2) difference of (non)-conve) functions enough samples form each source to entail good accuracy on estimating pairwise discrepancies [41]. Second, we need to solve N variants of the optimization problem in PERM, possibly each with a different model space, which is infeasible when the number of devices is huge (e.g., cross-device federated Icarning). In the next section, we propose a simple yet effective idea to overcome these issues in a computationally efficient manner.
distributed manner and requires
3 PERM at Scale via Model Shuffling
In this section, we propose method to efficiently estimate the empirical discrepancies among data sources followed by a model shufiling idea to simultancously solve N versions of PERM to learn a personalized model for each client. We first start by proposing a two-stage algorithm: estimating mixing parameters followed by model shuffling. Then, we propose a single loop unified algorithm that enjoys the same computation and communication overhead BERM (twice communication of FedAvg). For ease of exposition, we discuss the proposed algorithms by assuming all the clients share the same model architecture and later on discuss the generalization to heterogeneous model spaces. Specifically. we assume that the model space H. is a parameterized by a convex set W C Rd and use fi( : [.())= = (..)ess ((w; (x, v)) to denote the empirical loss at ith data shard.
3.1 Warmup: a two-stage algorithm
We start by proposing two-stage method for solving N variants of PERM in parallel. In the first stage, we propose an efficient method to learn the mixing parameters for all clients. Then, in stage two, we propose a model shuffling method solve all personalized empirical losses in parallel.
Stage I: Mixing parameters estimation, the first stage we aim to efficiently estimate the pairwise discrepancy among local distributions to construct mixing parameters a;, = 1,2,... N. From

---

generalization bound GEN and Definition I, direct solution to estimate at; is to solve the following convex- monex-noocconave minimax problem for each client:
0 = arg min OEAN
N Σ
alj) max (f(( - [5(2] + Σ j=1 att) /ng wew
(I)
where we estimate the true risks in pairwise discrepancy terms with their empirical counterparts and drop the complexity term it becomes identical for all sources by fixing the hypothesis space H and bounding it with a computable distribution-inn innuutionnnddeeddeent quantity such as VC dimension [43], or can be controlled by choice of H through data-dependent regularization. However, solving the above minimax problem itself is already challenging: the inner maximization loop is a nonconcave (or difference of convex) problem, most of the existing minimax algorithms will fail this problem. To our best knowledge, the only provable deterministic algorithm is [44], and it is hard to generalize it stochastic and distributed fashion. Moreover, since we have N clients, we need solve N variants of (I), which makes designing a scalable algorithm even harder.
To overcome aforementioned issues, we make two relaxations to estimate the per client mixing parameters. First, we optimize an upper bound of pairwise empirical discrepancies sup,, (f:(w) [w)] in terms of gradient dissimilarity between local objectives ([)) ffy quantifies how different the local empirical losses are and widely used in analysis of learning from heterogeneous losses as in FL [46]. Second, given that the discrepancy measure based on the supremum could be excessively pessimistic in real-world scenarios, and drawing inspiration from the concept of average drift at the optimal point as a right metric to measure the effect of data heterogeneity in federated learning [47], we propose to measure discrepancy at the optimal solution obtained by solving WERM, i.e., w" := arg (1 Σinf(((w) fi(w). By doing this, the problem reduces to a simple minimization for each client, given the optimal global solution. These two relaxations lead to solving the following tractable optimization problem to decide the per-client mixing parameters:
[45], which
arg min gi(w", ;0) aa == ;=1 a()IVV - (f )12 aEAN
D
j=1 a(j) /ng
(2)
where we added a regularization parameter À and used the squared of gradient dissimilarity for computational convenience. Thus, obtaining all N mixing parameters requires solving a single ERM to obtain optimal global solution and N variants of (2). To get the optimal solution in communication-reduce manner, we adapt the Local SGD algorithm [48] (or FedAvg [14]) and find the optimal solution in intermittent communication setting [49] where the clients work in parallel and are allowed to make K stochastic updates between two communication rounds for R consecutive rounds. The detailed steps are given in Algorithm A1 in Appendix B for completeness. After obtaining the global model " we optimize over a in gi ( .x) using Ta iterations of GD to get á; Actually, we will show that as long as " converge to w", o,, i I,..., N converges to solution of (2) very fast. Our proof idea is based on the following Lipschitzness observation:
||ag (w*
( -a (w0") 4L Σ
(21V (f 4L² 1| - w =='l²)) [|"" w01l
where og (w) == arg min ge( (w, a) and Kg := ma / (2X) is the condition number of g: (w.-) where T'max ma n,, The Lipschitz constant mainly depends on gradient dissimilarity at optimum. As "h tends to w", the o (-) becomes more Lipschitz continuous, i.e., the coefficient in front of | -="l w=1 getting smaller, thus leading to more accurate mixing parameters.
To establish the convergence, we make the following standard assumptions.
Assumption 1 (Smoothness and strong convexity). We assume fi(x)'s 's are L-smooth and p- strongly convex, i.e.,
Vo, y ([))) v.((g))l < L L| - -yl--
Va.y f(()) 2 f1(x) + (Vfi()), y æ) + jully x0|2
We denote the condition number by k = L/p. Assumption (Bounded variance). The variance of stochastic gradients computed at each local function is bounded, i.e., Vi E [NV, Veu W,ElIV fi(w; E) fi((w)) [ < 62.
5

---

Assumption 3 (Bounded domain). The domain W C R is a bounded convex set, with diameter D under l metric, i.e., vw, w' W, ||w w'll S D.
Definition 2 (Gradient dissimilarity). We define the following quantities to measure the gradient dissimilarity among local functions:
N
(t)(w) = -V(((()) (((w) = N 3=1
G)(ww)
Ç : sew [ [V - (1/N) Σ j=1 (.))"
The following theorem gives the convergence rate of estimated discrepancies to optimal counterparts.
Theorem 1. Under Assumptions 1-3, if we run Algorithm Al on F(w) := NN f with
y θ (RR(RK) for R rounds with synchronization gap K. for Rig = 1/(Mmmi))) holds that
pRK
Ta Kig
D2
RC²
8
E
50 c R
exp
+ 27 )1
RK
12 R
P2NKK
Vi E [N].
An immediate implication Theorem is that even we solve (2) at w² the algorithm will eventually converge to optimal solution of (2) at w". The core technique in the proof, as we mentioned, is to show that for parameter within small region centered at w*, the function og (w) becomes 'more Lipschitz". The rigorous characterization of this property is captured by Lemma 3 in appendix.
Stage 2: Scalable personalized optimization with model shuffting. After obtaining the per client mixing parameters, in the second stage we aim at solving N different personalized variants of PERM denoted by (()), v), ( v).... l(āN, v) to learn local models where
min D (d;,, NEW
1 N N
Σ ;=1 (.))
(3)
Here we devise an iterative algorithm based on distributed SGD with periodic averaging (a.k.a. Local SGD [48]) to solve these N optimization problems in parallel with no extra overhead. The idea is to replace the model averaging in vanilla distributed (Local) SGD with model shuffling. Specifically, shown in Algorithm the algorithm proceeds for R epochs where each epoch runs for N communication rounds. At the beginning of each epoch the server generates a random permutation a, over N clients. each communication round ý within the epoch, the server sends the model of client to client i; = (i j) mod N in the permutation a, along with cx(), ). After receiving model from the server, the client updates the received model for K local steps and returns it back to the server. As it can be seen, the updates of each loss 4 (q, v), i 1, ,,
.N during an epoch is equivalent sequentially processing individual losses in (3) which can be considered as permutation-bast By interleaving the permutations, we are able to simultaneously optimize all N objectives. We note that the computation and communication complexity of the proposed algorithm is the same as Local SGD with two differences: the model averaging replaced with model shuffling, and the algorithms run over permutation of devices. The convergence rate of Local SGD is well-established in literature [50, 51, 52, 53, 54], but here we establish the convergence of permutation -based variant which is interesting by its own.
SGD but with the different that each component now is updated for K steps.
Assumption (Bounded Gradient). The variance of stochastic gradients computed at each local function is bounded, i.e., Vi E [N], [[V]i < G.
We note that the Assumption 4 can be realized since we work with bounded domain W. Theorem Let Assumptions 1-4 hold. Assume of is the solution of (2). Then if we run Algorithm I
on the &; obtained from Algorithm AI, then Algorithm with 17 = 0 [og(NKK KR
will output the solution 0G Vi E [N), such that with probability at least 1 - p. the following statement holds:
aR
D2L
L82
L +
E e.) (aa v" (a)))) 0
+
+
LG log(1/p)
NKR
12
Ta
KC( +
82
)
+ Nop
exp
Kig
+ 2 (w="()")1'"
11222

---

Algorithm 1: Shuflling Local SGD
Input:
Clients 1, I,..., N, Number of Local Steps K Number of Epoch R, mixing parameter
Ō1,,..,, NNN
Epoch for ..., R 1 do
Server generates permutation o, [M] (M].
parallel for i = 1, ..... N do
Client sets initial model v.-0 U.
for j= NN do
Set indices i = or((i+)) mod N).
Server sends x. to Client iz.
v. vf = SGD- -UD-Upd vf , 2, K, a.).
Client does projection: w Pw( Output: 8, (w
( w/t)), Vi [N].
SCD-Update(v. [), 1/ K ox)
Initialize rd
for = O,... K do
v² = o nar( NVI
Output: ok
where K. = L. H: of Algorithm Al. That is, to guarantee (, 0.) (b S we choose R = L82 1 (()0)"")))))))))
Pmua and Kig YNG, and the expectation taken over randomness 2A
O maxx
{
}) and Ta = 0 (ig log
Lmn
pe
The above theorem shows that even though we run the optimization on ( v), our obtained model é, will still converge to the optimal solution of ((x;;, v). The convergence rate is contributed from two parts: convergence of &; (Algorithm AI) and convergence of personalized model é; (Algorithm 1). Notice that, for the convergence rate of ú; we roughly recover the optimal rate of shuffling SGD [55], which is (1/R²). However, suffer from a O R) term since each client runs vanilla SGD on their local data (the SCD- procedure in Algorithm 1). One medication for this variance term could be deploying variance reduction or shuffling data locally at each client before applying SGD. We notice that there is a recent work [56] also considering the client-level shuffling idea, but our work differs from it in aspects: 1) they work with local SGD type algorithm and the shuffling idea is employed for model averaging within a subset of clients, while in our algorithm, during each local update period, each client runs shuffling SGD directly on other's model 2) from theoretical perspective, we are mostly interested in investigating whether the algorithm can converge to the true optimal solution of ( v) if we only optimize on a surrogate function 4 ( v).
One drawback of Algorithm is that we have to wait for Algorithm A1 to finish and output á;, so that we can proceed with Algorithm 1. However, if we not satisfied with the precision of a; may not have a chance to go back to refine it. Hence in the next subsection, we propose to interleave Algorithm 1 and Algorithm A1, and introduce single-loop variant of PERM which will jointly optimize mixture weights and learn personalized models in an interleaving fashion.
3.2 A unified single loop algorithm
We now turn to introducing a single-stage algorithm that jointly optimizes a,s and v,s as depicted in Algorithm 2 by intertwining the two stages in Algorithm A1 and Algorithm in single unified method. The idea is to learn the global model, which is used to estimate mixing parameters, concurrent to personalized models. At each communication round, the clients compute gradients on the global model, on their data, after the server collects these gradients does a step mini-batch SGD update on the global model, and then updates the mixing parameters. Then we proceed to update the personalized models similar to Algorithm 1. We note that, unlike the two-stage method where the mixing parameters are computed at the final global model. here the mixing parameters are updated adaptively based on intermediate global models.
Theorem 3. Let Assumptions to to be satisfied. Assume OE, is the solution of (2). Then if we run Algorithm 2 with 7) = [og(NN R2 and 7 = θ [log(NKR KR3) it will output the solution ba ail
aR
7

---

Algorithm 2: Single Loop PERM
Input:
Clients 1, ,,,,, N, Number of Local Steps K Number of Epoch R, Initial mixing parameter 0 _, -,.., 2 a = [1/N, N/N]. Epoch for = 0, ...,
do Server generates permutation or [N] + [M]. parallel for Client 1, .... N do
Client sets initial model vi:0 u.
for j = ..... N do
Set indices i; or ((i+)) mod N).
Server sends v. to client ij.
of UpD-tedate r-) 7. 2j, K, a;).
// Personalized model update
Client i does projection: u. wº Pw(D,
un*+1 Pw(w" - 7 ,ेx Σi= M Σド= Vii(w" ((,))
// Global model update
Compute a: r+1 by running Ta steps GD g g,w aa)
a/ update
Output: 0; = Pw -1////(1///////// (1/L)V. of a; = a. Vie [N].
SGD SGD-Updateev. K,
Initialize 1/
for 0, ...., do
0' 001 NVI
Output: x
vi E [N], such that with probability at least
p. the following statement holds:
LD2
NL
+ L8²
E E[D(a); 0.) -b(aa,; w;)1
+
G'Nlog(1//p)
NKR3
R
12
μR
DG
Ta
Lk k2g 22g (wn") )82
+ 2
LÒ
+ R
exp
+
R
Kgg
μ2M
where i. L n1: 18g
Penax 2A
Rib VNG and the expectation is taken over the randomness of stochastic
samples in Algorithm 2. Thut is, guarantee E[(((x) 0.) D(at, vt)] S e, we choose M
O
Li 1222k2n2
R O maxx
L82 x. "a..""R 12 (D) }) and Ta = O Kgg log
LR²
1rcc
Juc
Compared to Theorem 2, we achieve slightly rate, since need large mini-batch when we update global model w. However, the advantages of the single-loop algorithm are two-fold. First, mentioned in the previous subsection, have the freedom optimize &e; to arbitrary accuracy, while in double loop algorithm (Algorithm A1 Algorithm 1), once we get áG, we do not have the chance further refine it. Second, in practice, single-loop algorithm is often easier to implement and can make better use of caches by operating on data sequentially, leading to improved performance, especially on modern processors with complex memory hierarchies.
3.3 Extension to heterogeneous model setting
In the homogeneous model setting, we assumed shared model space W for clients and the server. However, in real-world FL applications, devices have diverse resources and can only train models that match their capacities. We demonstrate that the PERM paradigm can be extended support learning in model- hetero-eetro settings. where different models with varying capacities are used by the server and clients. Focusing on learning the global optimal model to estimate pairwise statistical discrepancies, we note that by utilizing partial training methods [28], where at each communication round sub-model with a size proportional to resources of each client is sampled from the server's global model (extracted either random, static, or rolling) and is transmitted to be updated locally. Upon receiving updated sub-models, the server can simply aggregate (average) heterogeneous sub- model updates sent from the clients to update the global model. We can consider the complexity of

---

PERM
Locatzed FedAug
Perood Vwstat Noo.
82.6
Peniorl Valdalaa Lana.
Perfed/wg
80.0
8.00
(FodME
PERM
Loca Ized FedAvg
0.50
75.0
PerF adAvg
0.50
PFedME
0.46
70.0
67.5
040
036
Number of Communicaiina
Number of Conmaniiaiion
(a) Personalized Accuracy
(b) Personalized Loss
Figure 1: Comparative analysis of personalization methods, including our single-loop PERM algo- rithm, localized FedAvg, perFedAvg, and PFedME, with synthetic data. The disparity personalized accuracy and loss highlights PERM's capability to leverage relevant client correlations.
models used by clients when estimating mixing parameters by solving modified version of (2) as:
N
Σ a(j) VC(H))/m + Σ EOUIIVS(( o w²) - (mj(m) O w ""Il + > aa)i//// j=1
where we simply upper bounded the Rademacher complexity w.r.t each data source in (GEN) with VC dimension [57]. Here m, the masking operator extract a sub-model of the global model to compute local gradients at client i based on its available resources. By doing so, we can adjust mixing parameters based on the complexity of underlying models, as different sub-models of the global model (i.e., m, @ w" versus m C w²) are used to compute drift between pair of gradients at the optimal solution. With regards to training personalized models with heterogeneous local models, as we solve distinct aggregated empirical loss for each client by interleaving permutations and shuffling models, we can utilize different model spaces W;, i 1,.... N for different clients that meet their available resources with aforementioned partial training strategies.
4 Experimental Results
In this section we benchmark the effectiveness of PERM on synthetic data with 50 clients, where notably outshone other renowned methods evident in Figure 1. Our experiments concluded with the CIFARIO dataset, employing a 2-layer convolutional neural network, where PERM, despite a warm- up phase, demonstrated unmatched convergence performance (Figure 2). Additional experiments are reported in the appendix. Across all datasets, the PERM algorithm consistently showcased its robustness and unmatched efficiency in the realm of personalized federated learning.
Experiment on synthetic data. To demonstrate the superior effectiveness of our proposed single-loop PERM algorithm compared to other existing personalization methods, we conducted an experiment using synthetic data generated according to the following specifications. We consider a scenario with total of N clients, where we draw samples from the distribution N(1 Σ.) for half of the clients, denoted by i E [1, N, and from (1,, Σ,) for the remaining clients, denoted by i ( N). Following the approach outlined in [58], adopt a uniform variance for all samples, with ΣR, --1.2 Subsequently, we generate a labeling model using the distribution N (po E.).
Given data sample e Rdd, the labels generated follows: clients 1, .... NV assign labels based on y sign(w Tw), while clients 2 + 1, 1,.., N assign labels based on y sign w (-w æ). For this specific experiment, we set μز = 0.2, A2 0.2, and f 0.1. The data dimension is d = 60. and there are 2 classes in the output. We have a total of 50 clients, each generating 500 samples following the aforementioned guidelines. We train logistic regression model on each client's 's data.
To demonstrate the superiority of our PERM algorithm, conducted performance comparison against other prominent personalized approaches, including the fined-tuned model of FedAvg [14] (referred to as localized FedAvg). perFedAg [9], and pFedME [7]. The results in Figure 1 highlight PERM's efficient learning of personalized models for individual clients. In contrast, competing methods relying on globally trained models struggle to match PERM's effectiveness highly heterogeneous scenarios, as seen in personalized accuracy and loss. This showcases PERM'S exceptional ability leverage relevant client learning.

---

90.0
0.65
FEAM
6875
Localized Fedlug
0.60
PurFelley
85.0
pFcuME
0.55
82.5
0.50
Vel Valdae
27.5
0.40
75.0
0.35
72.5
0.030
0.25
Number of Communications
Number Communications
(a) Personalized Accuracy
(b) Personalized Loss
Figure 2: Comparative analysis of our single-loop PERM algorithm, localized FedAvg, PFedMe, and perFedAg, on CIFARIO dataset and a 2-layer CNN model. Each client access only classes of data. PERM rapidly catches up after 10 rounds of warmup without personalization involved.
Experiment CIFARIO dataset. We extend our experimentation to the CIFARIO dataset us- ing a 2-layer convolutional neural network. Dur- ing this test, 50 clients participate, each limited to data from just 2 classes, resulting in pro- nounced heterogeneous data distribution. We benchmark algorithm against PerFedAvg. PFedMe, and the localized FedAvg. As illus- trated in Figure PERM demonstrates supe- rior convergence performance compared to other personalized strategies. It's noteworthy that PERM's initial personalized validation is signif- icantly lower than that of approaches like PerFe- dAvg and PFedMe. This discrepancy stems from our choice to implement 10 communica- tion rounds as warm-up phase before initiating personalization, whereas other models embark on personalization right from the outset.
phode
Figure 3: Runtime of different algorithms in limited environment. We compare PERM (single loop), PerFedAvg. FedAvg, and pFedMe. PERM has a minimal overhead over FedAvg and is com- parable to other personalization methods.
Computational overhead. demonstrating the computational efficiency of the proposed PERM algorithm, we present a comparison of wall-clock time completing one round of communication of PERM and other methods. Each method undertakes 20 local steps along with their distinct computations for personalization. As depicted in Figure 3, the PERM (single loop) algorithm's runtime is compared against personalization methods such as PerFedAvg, FedAvg, and pFedMe. Remarkably, PERM maintains notably minimal computational overhead. The run-time is slightly worse due overhead of estimating mixing parameters.
5 Discussion & Conclusion
This paper introduces new data&syste--awan paradigm for learning from multiple heterogeneous data sources to achieve optimal statistical accuracy across all data distributions without imposing stringent constraints on computational resources shared by participating devices. The proposed PERM schema, though simple, provides an efficient solution to enable each client to learn a personalized model by learning who to learn with via personalizing the aggregation of data sources through an efficient empirical statistical discrepancy estimation module. To efficiently solve all aggregated personalized losses, propose model shuffling idea to optimize all losses in parallel. PERM can also be employed in other learning settings with multiple sources of data such as domain adaptation and multi-task learning entail optimal statistical accuracy.
We would like to embark on the scalability of PERM. The compute burden on clients and servers is roughly the same as existing methods thanks to shuffling (except for extra overhead due to estimating mixing parameters which the same as running FedAvg in a two-star stage approach and an extra communication in interleaved approach). The only hurdle would be the required memory server to maintain mixing parameters, which scales proportionally the square of the number of clients, which can be alleviated by clustering devices which we leave as future work.
10

---

Acknowledgement
This work was partially supported by NSF CAREER Award #2239374 NSF CNS Award #1956276.
References
[1] Peter Kairouz, H. Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Ar- jun Nitin Bhagoji, Keith Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, Rafael G. L. D D'Oliveira, Salim EI Rouayheb, David Evans, Josh Gardner, Zachary Garrett, Adrià Gascón, Badih Ghazi, Phillip B. Gibbons, Marco Gruteser, Zaid Harchaoui, Chaoyang He, Lie He, Zhouyuan Huo, Ben Hutchinson, Justin Hsu, Martin Jaggi, Tara Javidi, Gauri Joshi, Mikhail Khodak, Jakub Konečný, Aleksandra Korolova, Farinaz Koushanfar, Sanmi Koyejo, Tancrède Lepoint, Yang Liu, Prateek Mittal, Mehryar Mohri, Richard Nock, Ayfer Özgür, Rasmus Pagh, Mariana Raykova, Hang Qi, Daniel Ramage, Ramesh Raskar, Dawn Song, Weikang Song, Sebastian U. Stich, Ziteng Sun, Ananda Theertha Suresh, Florian Tramèr, Pranceth Vepakomma, Jianyu Wang. Li Xiong, Zheng Xu, Qiang Yang, Felix X. Yu, Han Yu, and Sen Zhao. Advances and open problems in federated learning. Foundations and Trends@) in Machine Learning, 2021.
[2] Yishay Mansour, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh. Three approaches for personalization with applications to federated learning. arXiv preprint arXiv:2002. 10619, 2020.
[3] Chengxi Li, Gang Li, and Pramod K Varshney. Federated learning with soft clustering. IEEE Internet of Things Journal, 9(0):7773-7782, 2021.
[4] Avishek Ghosh, Jichan Chung, Dong Yin, and Kannan Ramchandran. An efficient framework for clustered federated learning. Advances in Neural Information Processing Systems, 33:19586- 19597, 2020.
[5] Jie Ma, Guodong Long, Tianyi Zhou, Jing Jiang, and Chengqi Zhang. On the convergence of clustered federated learning. arXiv preprint arXiv:2202.06188 2022.
[6] Hubert Fichner, Tomer Koren, Brendan Memahan, Nathan Srebro, and Kunal Talwar. Semi-
cyclic stochastic gradient descent. In Proceedings of the 36th International Conference Machine Learning, PMLR, volume 97, 2019.
[7] Canh T Dinh, Nguyen Tran, and Tuan Dung Nguyen. Personalized federated learning with moreau envelopes. Advances in Neural Information Processing Systems, 33, 2020.
[8] Yutao Huang, Lingyang Chu, Zirui Zhou, Lanjun Wang, Jiangchuan Liu, Jian Pei, and Yong Zhang. Personalized federated learning: An attentive collaboration approach. arXiv preprint XrXiv:2007.03797, 2020.
[9] Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning: meta-lcarning approach. Advances in Neural Information Processing Systems, 2020.
[10] Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet S Talawalarr Federated multi-task learning. In Advances in Neural Information Processing Systems, pages 4424-4434, 2017.
[11] Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Adaptive personalized federated learning. arXiv preprint arXiv:2003. N.XV::2003./3461. 2020.
[12] Filip Hanzely, Boxin Zhao, and Mladen Kolar. Personalized federated learning: A unified framework and universal optimization techniques. arXiv preprint arXiv:2102.0977 2021.
[13] Tony Lancaster. The incidental parameter problem since 1948. Journal of econometrics, 95(2):391-413, 2000.
[14] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas Communication-effnn
learning of deep networks from decentralized data. In Artificial intelligence and statistics, pages 1273-1282. PMLR, 2017.
[15] Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International Conference Machine Learning, pages 5132-5143. PMLR, 2020.
[16] Jenny Hamer, Mehryar Mohri, and Ananda Theertha Suresh. Fedboost: A communication- efficient algorithm for federated learning. In International Conference on Machine Learning, pages 3973-3983. PMLR, 2020.

---

[17] Farzin Haddadpour, Mohammad Mahdi Kamani, Aryan Mokhtari, and Mehrdad Mahdavi. Federated learning with compression: Unified analysis and sharp guarantees. In International Conference Artificial Intelligence and Statistics, pages 2350-2355 PMIR, 2021.
[18] Yan Sun, Li Shen, Tiansheng Huang, Liang Ding, and Dacheng Tao. Fedspeed: Larger local interval, less communication round, and higher generalization accuracy. In The Eleventh International Conference on Learning Representations, 2023.
[19] Haibo Yang. Xin Zhang, Prashant Khanduri, and Jia Liu. Anarchic federated learning. In International Conference on Machine Learning, pages
PMLR, 2022.
[20] Shiqiang Wang and Mingyue Ji. A unified analysis of federated learning with arbitrary client participation. Advances in Neural Information Processing Systems, 35:19124-19137, 2022.
[21] Chuhan Wu, Fangzhao Wu, Lingjuan Lyu, Yongfeng Huang, and Xing Xie. Communication- efficient federated learning via knowledge distillation. Nature communications, 13(1):1-8. 2022.
[22] Chaoyang He, Murali Annavaram, and Salman Avestimehr. Group knowledge transfer: Feder- ated learning of large cnns at the edge. Advances in Neural Information Processing Systems, 15:14068-14085 2020.
[23] Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. Advances in Neural Information Processing Systems, 33:2351-2363, 2020.
[24] Sohei Itahara, Takayuki Nishio, Yusuke Koda, Masahiro Morikura, and Koji Yamamoto. Distillation-bass semi-supervised federated learning for communication-efficcae collaborative training with non-iid private data. IEEE Transactions on Mobile Computing, 22(1):191-205, 2021.
[25] Enmao Diao, Jie Ding, and Vahid Tarokh. Heterofl: Computation and communication efficient federated learning for heterogeneous clients. arXiv preprint arXiv:2010.0126. 2020.
[26] Samuel Horvath, Stefanos Laskaridis, Mario Almeida, Ilias Leontiadis. Stylianos Venieris, and Nicholas Lane. Fjord: Fair and accurate federated learning under heterogeneous targets with ordered dropout. Advances in Neural Information Processing Systems, 34: 34:12876-12889 2021.
[27] Sebastian Caldas, Jakub Konečny, H Brendan McMahan, and Ameet Talwalkar. Expanding the reach of federated learning by reducing client resource requirements. arXiv preprini arXiv: 10, 2018
[28] Samiul Alam, Luyang Liu, Ming Yan, and Mi Zhang. Fedrolex: Model- hodel-hheerogeeeeuus federated learning with rolling sub-model extraction. Advances in Neural Information Processing Systems, 35:29677-29690. 2022.
[29] Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jen- nifer Wortman Vaughan. A theory of learning from different domains. Machine learning, 79:151-175, 2010.
[30] Yishay Mansour and Mariano Schain. Robust domain adaptation. Annals of Mathematics and Artificial Intelligence, 71(4):365 -380, 2014.
[31] Nikola Konstantinov and Christoph Lampert. Robust learning from untrusted sources. In International conference machine learning, pages 3488-3498 PMLR, 2019.
[32] Koby Crammer, Michael Kearns, and Jennifer Wortman. Learning from multiple sources. Journal of Machine Learning Research, 9(8), 2008.
[33] Mathicu Even, Laurent Massoulié, and Kevin Scaman. On sample optimality in personalized collaborative and federated learning. NeurIPS 2022-36th Conference on Neural Information Processing System, 2022.
[34] Jakub Konečný, H Brendan McMahan, Felix X Yu, Peter Richtárik, Ananda Theertha Suresh, and Dave Bacon. Federated learning: Strategies for improving communication efficiency. arXiv preprint arXiv: 1610.00610 2016.
[35] Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh. Agnostic federated learning. In International Conference Machine Learning, pages 4615-4625, PMLR, 2019.
12

---

[36] Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Distributionally robust federated averaging. Advances in neural information processing systems, 33:15111 2020.
[37] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smithy. Feddane: A federated newton-type method. In 2019 53rd Asilomar Conference on Signals, Systems, and Computers, pages 1227-1231. IEEE, 2019.
[38] Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank J Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for on-device federated learning. International Conference on Machine Learning, 119:5132-5143, 2020.
[39] Farzin Haddadpour and Mehrdad Mahdavi. On the convergence of local descent methods in federated learning. arXiv preprint (0tiv:190.,444 2019.
[40] Tao Yu, Eugene Bagdasaryan, and Vitaly Shmatikov. Salvaging federated learning by local adaptation. arXiv preprini arXiv:2002.047558 2020.
[41] Bharath K Sriperumbudur, Kenji Fukumizu, Arthur Gretton, Bernhard Schölkopf, and Gert RG Lanckriet. On integral probability metrics,), phi-divergeenes and binary classification. arXiv preprint arXiv:0901. 2009.
[42] Steve Hanneke and Samory Kpotufe. no-free-lunch theorem for multitask learning. The Annals of Statistics, 5066):311-33143 2022.
[43] Shai Shalev-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algorithms. Cambridge university press, 2014.
[44] Zi Xu, Huiling Zhang, Yang Xu, and Guanghui Lan. A unified single-loop alternating gradi- ent projection algorithm for canconvex-concan and convex-nonconcave minimax problems. Mathematical Programming, pages 1-72, 2023.
[45] Yatin Dandi, Luis Barba, and Martin Jaggi. Implicit gradient alignment in distributed and federated learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 1454-6462, 2022.
[46] Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data, arXiv preprint arXiv:1806.0053: 2018
[47] Jianyu Wang, Rudrajit Das, Gauri Joshi, Satyen Kale, Zheng Xu, and Tong Zhang. On the unreasonable effectiveness of federated averaging with heterogeneous data. arXiv preprint arXiv:2206.0472. 047223 2022.
[48] Sebastian U Stich. Local sgd converges fast and communicates little. In International Conference on Learning Representations, 2018.
[49] Blake E Woodworth. Jialei Wang, Adam Smith, Brendan McMahan, and Nati Srebro. Graph oracle models, lower bounds, and gaps for parallel stochastic optimization. In Advances in neural information processing systems, pages 8496-8506, 2018.
[50] Blake Woodworth, Kumar Kshitij Patel, and Nati Srebro. Minibatch vs local sgd for heteroge- neous distributed learning. Advances in Neural Information Processing Systems, 33:6281-6299 2020.
[51] Konstantin Mishchenko, Grigory Malinovsky, Sebastian Stich, and Peter Richtárik. Proxskip: Yes! local gradient steps provably lead to communication acceleration! finally! In International Conference on Machine Learning, pages 15750-15769. PMIR, 2022.
[52] Honglin Yuan and Tengyu Ma. Federated accelerated stochastic gradient descent. Advances in Neural Information Processing Systems, 33:5332-5344, 2020.
[53] Farzin Haddadpour, Mohammad Mahdi Kamani, Mehrdad Mahdavi, and Viveck Cadambe. Local sgd with periodic averaging: Tighter analysis and adaptive synchronization. In Advances in Neural Information Processing Systems, pages 11080-1 1092, 2019.
[54] Eduard Gorbunov, Filip Hanzely, and Peter Richtárik. Local sgd: Unified theory and new efficient methods. In International Conference on Artificial Intelligence and Statistics, pages 3556-3564. PMLR, 2021.
[55] Kwangjun Ahn, Chulhee Yun, and Suvrit Sra. Sgd with shuffling: optimal rates without component convexity and large epoch requirements. Advances Neural Information Processing Systems, 15551666577555 2020.
13

---

[56] Yae Jee Cho, Pranay Sharma, Gauri Joshi, Zheng Xu, Satyen Kale, and Tong Zhang. On the convergence of federated averaging with cyclic client participation. arXiv preprint arXiv:2302.03 109, 2023.
[57] Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar. Foundations of machine learning. MIT press, 2018.
[58] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. arXiv preprint arXiv:1812.06177, 2018.
[59] Sebastian Caldas, Peter Wu, Tian Li, Jakub Konečnj, H Brendan McMahan, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for federated settings. arXiv preprint arXiv: 18/2.01097, 2018.
[60] Tianyi Lin, Chi Jin, and Michael I Jordan. On gradient descent ascent for nonconvex-concave minimax problems. arXiv preprint arXiv:1906.00331, 2019.
[61] Markus Schneider. Probability inequalities for kemel embeddings in sampling without replace- ment. In Artificial Intelligence and Statistics, pages 66-74. PMLR, 2016.
14

---

A Additional Experiments
In addition to experiments synthetic and CIFAR 10 datasets reported before. we have also conducted experiments on the EMNIST dataset, highlighting PERM's capability to derive superior personalized models by tapping into inter-client data similarities. Additionally, further insights emerged from tests on the MNIST dataset, revealing how PERM's learned mixture weights adeptly respond to both homogeneous and highly heterogeneous data scenarios.
Experiment on EMNIST dataset In addition to the synthetic and CIFARIO datasets discussed in the main body, we run experiments on the EMNIST dataset [59], which is naturally distributed in a federated setting. In this case, we chose 50 clients and use a 2-layer MI.P model, each with 200 neurons. We compare the PERM algorithm with the localized model in FedAvg and perFedAvg [9]. As it can be seen in Figure 4, PERM can learn a better personalized model by attending to each client's 's data according to the similarity of the data distribution between clients. The learned values of a, in Figure 5, show that the clients are learning from each others' data, and not focused on their own data only. This signifies that the distribution of data among clients in this dataset is not highly heterogeneous. Note that, since we are using a subset of clients in the EMNIST dataset for the training (only 50 clients for 100 rounds of communication), the results would be sub-optimal. Nonetheless, the experiments are designed to show the effectiveness of different a algorithms. As it can be concluded, in terms of performance, PERM consistently excels beyond its peers, demonstrating exemplary results on various benchmark datasets.
3.0
PERM
Vemoon Wi Wilthoo Metamng
2.8
Lacklized FadAAg
ParFadRug
60
Penoona Welsatoon
2.4
65
2.2
60
2.0
PERM
65
Localized FedAvg
PerFedAvs
40
0
23
80
100
120
0
20
40
60
80
100
120
Number Communications
Number of Communccaions
(a) Personalized Accuracy
(b) Personalized Loss
Figure 4: Comparative Analysis of Personalization methods, including our single-loop PERM algorithm, localized FedAvg, and perFedAg, with EMNIST dataset. The disparity in personalized accuracy and loss highlights PERM's capability in leveraging relevant client correlations.
Figure 5: The heat map of the learned a values for the PERM algorithms on the EMNIST dataset with a 2-layer MLP model. The weights signify that clients mutually benefiting from one another's s data, which also highlight that the distri- bution of data is not significantly hetero- geneous in this dataset.
0.200
0.100
0122
2.05
0.000
0.05
24 27 Cleris
42 45
The effectiveness of learned mixture weights To show the effectiveness of the two-stage PERM algorithm, as well as the effects of heterogeneity on the distribution of data among clients on the learned weights a in the algorithm, we run this algorithm on MNIST dataset. We use 50 clients, and the model is an MLP, similar to the EMNIST experiment. In this case, we consider two cases: distributing the data randomly across clients (homogeneous) and only allocating class per client (highly heterogeneous). As it can be seen from Figure 6, when the data distribution is homogeneous the learned values of a as diffused across clients. However, when the data is highly heterogeneous, the learned values will be highly sparse, indicating that each client is mostly learning from its own data and some other clients with partial distribution similarity. Notably, the matrix predominantly exhibits sparsity, indicating that each client selectively leverages information solely from subset of other clients. This discernible pattern reinforces the inherent confidence that each client effectively learning from limited but strategically chosen group of clients.
15

---

0.16
0.14
0.7
35
0.12
0.8
15
0.10
18
0.5
21
0.08
24
27
04
0.08
35
30
08 86
0.94
38
36
0.02
42
45
12161122
0.00
Clients
(a) Homogeneous distribution
(b) Highly heterogeneous distribution
Figure Comparing the performance of two-stage PERM algorithm in learning a values on heterogeneous and homogeneous data distributions. We use MNIST dataset across 50 clients with homogeneous and heterogeneous distributions.
B Proof of Two Stages Algorithm
In this section we provide the proof of convergence of two-stage implementation of PERM (computing mixing parameters followed by learning personalized models via model shuffling using permutation- based variant of distributed SGD with periodic communication).
B.1 Technical Lemmas
Lemma 1. Define v² (a:) == arg min ((a, v). and assume (( .) is p-strongly convex and V. v) is L Lipschitz in a. Then, v' (-) is k- k-Lisschitz where = L/p.
Proof. The proof is similar Lin al's result on minimax objective [60]. First, according to optimality conditions we have:
v ee (a), V,b(a, v' (a))) 2 0,
(v - v- (a'), V2 v' v'(a')) 2 0
Substituting v with va()) and (a) in the above first and second inequalities respectively yields:
(v" (a), V20(a, v2 2 0,
(a) - o (a'), V2 v' 20
Adding up the above two inequalities yields:
(0 (a') v² (a), V2 v' V, (a'))' 20, Since ((x, -) is μ strongly convex, have:
(4)
(v" (a') v (a), V, v²
(a)) 2 (al - o'))]1
(5)
Adding (4) and (5) yields:
(o" (a') 8" (a:), V2 80 (a')) - V2 80 (a')) 2 mlle"(aa) - 2a)22
Finally, using L smoothness of & will conclude the proof:
(ll) (')aall a'l - va)l2
- xl|a a'l 2 [l'((a')+''(a))
- v
Lemma 2 (Optimality Gap). Let ((a, v) be defined in (3). Ler = Pw(ü b(a ú)). we assume each fs is L-smooth, p-stroogly convex and with gradient bounded by G. then the following statement holds true:
4NG²) L
( 0) - +(a", v²) < 2T v² (2)))1 + 2 I+
| 1
where Kg.
VNG v" argmineww "(a" v).
16

---

Proof.
First we show that v. is VNG Lipschitz in a. To see this:
N
N
Vv,t(a, =
Σ (())
E d(V))(
S VVG ||a aill-
Hence due to Lemma 1, we know v" (a) is ki = NNG Lipschitz. According to property of projection, we have:
0 (U 0 L( 6) + V (. 8))
= u ④.
L(u
+ (w 0, V. 8) V, 8)) 8) + V a 2)
Ti
22
For Tú, we notice:
1
(o 0, L(b - +)+ + V,0(0,, 8)) = [p 0, 0) +
(ö 0,
8) + (v - è, V,0(aa", 6))
1
= [L (b - ü, 0 8)
L
|
6||
+ (v ú1 v,t(a",
1
s <(1l0 l||? + l
|| أ|| (u ج v.b(a" 8))
where at last step we used Young's inequality. To bound convexity of ((a, -):
we apply the L smoothness and μ strongly
(o - V. 0;)) (v ü, v.D(a", 6)) + (ū o, v. ,))
L || 011 H | v1l + (<x",, 7) -0 00)
< 1baa D(a", ,)
2
2
$ <(aa", v) 0(a" ,0
P || w11 +
L 0 b1
2
2
Putting above bound back yields:
v é, "
8) D(a",, v)) ((a",, ()
1
10 v1
(3b L)
21
2
Now we switch to bounding T2. Applying Cauchh-Schward yields:
L
L
4
(o 0, -. 0:)) < |/ 6,11 + 4
|
0,11
T, V( 0)) vo
4
L
L
4NG²
||u 1.11
||
0,113
||á: a;l?
4
4
L
where at last step we apply VNG smoothness of ( v). Putting pieces together yields:
0 5 o (a;; v) (a 0,) + |6; - 0 + ||u - 8/11 + 4NG²
l|á; -
2
Re-arranging terms and setting v = v² (a*) = arg min ((a", v)) yields:
4NG2
( t) "(a", v²) < L || v"
|| a²1
L
At last, due to the Ra Lippshitzzeess property of of 4r (-) as shown in Lemma 1, it follows that:
L v" (23) < L 22))2 + Lilee - v² (a("a)
22 5 - va))l + 2x a²1
as desired.
17

---

Algorithm A1: Discrepancy Estimation at Optimum
Input: Number of clients N, number of local steps K number of communications rounds R
for r = 0,...,P R-1 1-1 do
parallel for client i = 0, ..., N 1 do
Client initializes model w w for t = 0. K - do
w (+ un r.d Vii :f where 5 is mini-batch sampled from S..
Client sends w. to Server.
Server computes w Pw N EE w
Server broadcasts w all clients.
Server computes Oie i 1,2, .,, N by running Ta steps of GD on gi (w o²).
Output: .N..
B.2 Proof of Convergence of Theorem
In this section we are going prove the result Theorem 1. To this end, we need show that mixing parameters we compute by first learning the global model and then solving the optimization problem in objective (2) (as depicted in Algorithm Al) converges to optimal values. Notice that in Algorithm A1 we do not solve gi (w" .)) directly, but optimize gi (w a)) on a for Ta iterations of GD. Hence, firstly we need to show that optimizing the surrogate function will also guarantee the convergence of output of algorithm ā to a" by deriving a property of the objective in (2). Formally the property is captured by the following lemma.
Lemma 3. Let g(w. a) = Σj= Oj [lV]: (f + EΣ DEj1aā//mg and oz(ww)
arg min AN g(w,a, c)). Let w be the output of Algorithm Al. Then the following statement holds:
(w (13""\) S K Σ 22]V vf(("))1" 4L2 + (2=l==1)))) AL 1 w0ll
w²1²)
where Kig =
Pmnax 2A
Proof. Define function
N =x Σ, az2y + + E as/n²
W(z,a) = a(z,a)=
(6)
Apparently, W(z, a)) is linear in z and 2 strongly convex in z. Next we show that V (22 a)
Lipsehitz in w. To see this,
-V) a)ll [|[ ... 2N] 12 2x1Il < || z'|l-
Then, according to Proposition 1, (z) == arg millo W(z, a) is Kg lipschitz in z where Amar i.e., (z) anyy (z')ll Ki || 2111. Now, let consider the objective (2):
Ki
2A
g(w,ar a) Σ,
1= aj [i(() [f)(w)) + +xEx j=1
ag/ns
We define og (w) arg millasd (lw,a)
We set
R R) Vf(w 112 ..... [V - VfN p. 112
x" [IV -//---/(w" ..... [Vi(w")- VSN(w=")11
18

---

Then we know that
|a (w)( (2"(2")] =1a
(z - (>)) S kig -=11-
(7)
N
2 gig
Σ V -1/((=') 12 - VV
(8)
+ Σ (Vi.
V.f((w)) v.f((w)))
(Vfi(wB
V/f(20) V---v-f'">> + (fz(w)))
2
Vf(ww Vf( Vf(("") Vff(w") || "LL2 w w"ll
Nig
Since
(w (2 Vfi(w=)) Vif(")) )1 + 2L w w"1 we conclude
that
( S x Σ
22l vl""")|1" 4L2 1| //oR-=w'l' 4L w w01
==1
With above lemma, show the convergence of & a², we do the following decomposition
|| 2 s 2 211a () + +2
- a''
N - V[r(w")|l [|
< 2 )ma) K + 2mg Σ -
+ 4L2 1 w" 11²) 4L 10 w0'l
Now it remains to show the convergence of Local SGD last iterate w to optimal solution w". By convention, we use w' * Σm w to denote the virtual average iterates.
Lemma 4 (One iteration analysis of Local SGD). Under the condition of Theorem 1, the following statement holds true for any & [T]]
E
w r.1+
w
1| < (1 - WE 1r w" 1|
(2y- - 2) (F)())((w)) (F(w") F(w":t))
w wi + 222 82
+ (7L+2
N
Proof. According to updating rule in Algorithm A1, we have the following identity:
N
E 11001 - w0112 Bl -1
w01 - 27E
Vi(wf"i wi w
+ +
fo (wolit)
N
i=1
E 1/o0'
w -27
Vi(w))) wiit
N
T
t 02
+ E
Σ
V
of 125 222))
(10)
N
T2
19

---

For T1, since each fy is L smooth and μ strongly convex, we have:
N
-27
Vi((w("),
w² w0
-27
Vf((w)) w' -wet w..d w=
N
N
S -27
Vii(un)), with -. + unn 207
N
[iww') - Ic((0"it) - 14
2 11 - w0 112 + 2 11 -0000=000:10
< 27
N
==1
Due to Jensen's inequality we know:
N Ei= Ehr | - w021 S - 1| w01l" 1|2 Hence we
know:
27
N
Vii(w(('), whi -----w"
w² 22 F(w0)) =(w=" 1 w -w= 201 + N
1
10 200 1p))
For T2, have:
N
N
1
E
V
Vi((o
= -2E
Vi((w - VF(Wo))))
+ 2E V /
N
N
2L
E w -0rit1² + 4L (F(w"tt) - F(w")).
N
Now, plugging T1 and T2 back to (10) yields:
E w w" w011 s (1 - 1 1r0it
w² w012 -022 207-477- L) (F(w") -F("*)))
N
/woit - -00 + +12 202 82
+ +
N
N
2m
Lemma 5. [50, Lemma 8] For the iterates {w]''} generated in Algorithm A1, the following statement holds true:
1 N
w o'''/ = +6K2222
Lemma 6 (Last iterate convergence of Local SGD), Under the conditions of Theorem 1, the following statement holds true for the iterutes in Algorithm Al:
fl w01 (1 111-m) py) w w w²1|²
+
1
(3K'28² 6K32252) + 752
My
P.N
Proof. We first unroll the recursion in Lemma 4 from = K to 0, within one communication round:
K 1
E [||W".K - -w= w2 = (1 -1)-)))B111 1000 w=
(1
- -F(w*)))
1=0
K
Σ
(1-
(1-
(z 2
1
Σ
N
|
rit
w
K 1
Σ
(1 -1- wy) K -t 82 (1-
N
N
t=0
i=1
1=2
Since we choose 7 < Plugging in Lemma yields:
1
we know Σk- (1 - 117) k (27 - 42 I.). (F(w") - - 2 0.
2T,
E. I/ww w0012 (1 17) R 1100
w011² +
1
1
(7 2, 2 + 752
WN
20

---

Algorithm A2: Shufiling Local SGD (One Client)
Input: Clients O, b.... N 1, Number of Local Steps K Number of Epoch R. Mixing parameter
Epoch for .... R do
Server generates permutation or [M] [N].
Client initial model vr: =
for j = 0, ...,. . - do
Server sends vt to Client (,(
yr SGD SGD-Update(o'44 1), (,), K, o)).
Client does projection: wh+1 Pw(
Output: = R
SGD- GGD-Update(* 1);
Initialize rd = V
for = 0, ,., K-1 do
='= na(j) (
Output o.
Plugging in y log(R (K)) RRK
gives the convergence rate:
E
w B||w
w
11
11=
S
[l w= w=1
E.
+
8 1222
+
C2 12
+
82 bRNRK
RK
which concludes the proof.
Equipped with above results, we are now ready to provide the convergence of main theorem.
Proof of Theorem I. The proof simply follows from Lemma 3: l|á 12 & 2 2l - ai (w)
01² + )|²
[1 - (10"
s 2(1 pma) To g Σ 21[V - (1)(2")11 4I2 1 -=00-=0(1())))) 1/ wn
$ 21 )ma) To + 8TaY (2)(20)) + 4NI2 /w -====")))) w1'l())
1|
-w'l"
Plugging in the convergence of |w w
from Lemma and the stepsize na
by for yields:
Ta + "n(((("))
D2
5²
S
PNNK ))
Ella, R =
[exp(-
Kg
RK
+
(22PK
p²
B.3 Proof of Convergence of Shuffling Local SGD
In this section, we are going to prove the convergence of proposed shuffled variant of Local SGD (Theorem 2). The whole proof framework follows the analysis of vanilla shuffling SGD, but notice that there are two differences. First, in vanilla shuffling SGD, in each epoch, algorithm only updates on each component function f once, while here we have to take K steps of SGD update on each component function. Second, we are considering weighted sum objective in contrary to averaged objective in [56], which means we need to rescale the objective when we apply weihhout- concentration inequality. Even though our algorithm solves models for N clients, for the sake of simplicity, throughout the proof we only show the convergence of one client's model. The algorithm from one client point of view is described in Algorithm A2, where we drop the client index for notational convenience.
Ji=1 is obtained by Proposition 1. Assume sequence {wo' LK
w' = ===-11 VnaNVf(w" ;: t = I,..,,,
21

---

then we have
1 T+1 w w0t1_ =w"
++1
(I aIIa-aNNtte yoNNS((w") a naN'' so K
0=00==
0-==t
where 8 V((w";g") v((w"), and by convention, we define IT A iifa < b.
Proof. According to updating rule, we have:
w ++ w w² w" noNV (w":
w w noNV (ww/)/w')))))))) NON8 w w naNVV (ww)((w)) nON "aan(WV((WW)))) v)(w")) noNNo'
Since f is L. smooth, and according to Mean Value Theorem, there is matrix H, satisfying
μI 3 H. 3 LI, such that V)(w")) V H.w w²)). Hence we have:
with w = (I - NH,HH) (w' w - no No.NVS(w") - No..60
Unrolling the recursion from f to will conclude the proof.
The following lemma establishes the updating rule of models between epochs r and r + 1. For notational convenience, whenever there is confusion, drop the superscript r in o".
Lemma (One epoch updating rule). Let v and v + be two iterates generated by Shuffling Local SGD (Algorithm A2), then the following updating rule holds:
j+1
w++1= = v' Σ I QI J
r-+1
8,).
where
K
Qj
Σ
II (I - m((ot)))H.)) ná(o())),,
r= t'=K-1
K-1r+1
6y:= = Σ II (I a(o(i))NNHH) (
T=U P't
by convention, we define I1 A₁ lifa
Proof. According to Proposition 1, we have
K a- ((0))(())))) ((co())
wn++1 n
K-1rt1
ΣI a- ((NTaao))))) 0 F==
Plugging our definition of Q, and 8; yields:
0 J-ー t" v. o" Q, V
Following the same reasoning in the proof of Proposition will conclude the proof.
Lemma 8 (Summation by parts). Let A, and B, be complex valued matrices. Then the following fact holds:
N
N-1 An) Σ B,.
A,B
AN Σ
B₁
Σ
(A
j=1
n=1
22

---

Proposition 2 (Spectral bound of polynomial expansion). Given collection of matrices {At} and [B.]. such that A, 3 LI and B; 3 LI, the following bound hold:
e(h - 1)
II aA.)) l=1
(aT.)m,
m=1
h-1
ch - 1)
1)
m
h-1
<e(h <ehh-1 1) m
aA)
II (B,)
(aL) m
(bL)
m
m
Proof. We start with proving the first statement. Expanding the product yields:
h-1
II aAi) I+ Σ
(-1)ma" m
E
II
Am
1S DI h) mi
Hence have:
h-1
h-1
I aA) -1 = Σ
a
Σ
II
Am
S
(aL)
m
m=1
|S|=m.
ISCC (I.....hb m'ES
m==
According to the upper bound for binomial coefficients:
(M--)) S [((k=1)
m
we have:
m
h
e(h - 1)
(aL)
(aL)
m=1
m=1
Then we switch to the second one. Using the same expanding product yields:
h
h
[I
aA,)
II [I (B,)
h-1
m
Σ
II
Am
Σ
(-1)mmm
Σ
II Byn
18 =m,ss E(l....)) m
18 SIS (I.....) m'
1
h-1
h
(aL)
(bL)
m
m
m=1
e(h -))
h-1
eehh (1hh-1)
1)
(aL)
+
(bL)
m=1
The following concentration result is the key to bound variance during shuffling updating. The original result holds for the average of gradients, and we will later on generalize it to an arbitrary weighted sum of gradients.
Lemma 9 (161, Theorem 2)). Suppose 2 2. Let 91 ... 9n. E Rdd satisfy ||, < G for all j. Letg- - 1 Σm= 9j. Let E S, be uniform random permutation of n elements. Then, for i I n, with probability at least we have
8(1
1-1 log
Σ 9olt)
g
S
S
Lemma 10 (Concentration of partial sum of gradients). Given uniformly randomly generated
permutation o, and simplex vector ar, if we assume each supoew [[V]] & G, then the following statement holds true:
Σ (()())) 1=
EG/8nlog(Lyp)
n (1/p) vd(à, w*)l
N
23

---

Proof. The proof works by re- re-writing weighted sum of vectors to average of the these vectors:
Σ 1=0
((o))))V)))(("")
1 = N
a(oti))V..cc
(())))))
+n ve((,,
N
S / 8m log((11)) + N v
Proposition 3 (Spectral norm bound of Q). Let Q, be defined in (11). Then the following bound for the spectral norm of Qj holds true for all j E [N]:
1Q S ()((()))NK( + N/)
Proof. The proof can be completed by writing down the definitin of Qj and applying Cauchy- Schwartz inequality:
10
II ( na(O())N
Σ
II (I-
(I-
t'=K-1
K-1
S {n(((())
Σ
I (H
T=0
P=K-1
nn(o(i))
Σ
II (L+N(())))
P=K-1
(((
The last step due to we choose v such that 7 S K..
The following lemma establishes the bound regarding cumulative update between two epochs, namely, -1+ o". particular, Lemma below shows that: (a) shuffling Local SGD, our update from v" to o r approximates performing N K times of gradient descent with á(j) foo) (v"), namely, the bias is controlled, and (b) the update itself bounded, and can related to the norm full gradient.
Lemma 11. During the dynamic of Algorithm A2, the following statements hold true with probability at least p:
(a)
N
N Q, (a()((()) NNNK Σ
a j=1
S 10y2
)
e
v")
4R 2R-c)
)
2
23
G2 log(1/p).
4R--
(b) for any N' such that 0 N' <N
N'
N
Q (o(
33nNNK (IVo(á, + +Gv8Nloog1/1
where
Q= Σ
II
( n((o()))HH) nā(o())...
(11)
P==K1
24

---

Q Proof. We start with proving statement (a). Let A, a(o()))
and B, V(oli))Voao)))
applying the identity of summation by parts yields:
QN-1 Q (oo))(o')= a(o(N- - 1)) Σ a(c())Vfoc(()))
N-1
Σ
Q n+1
Qn
Σ
a
((oo(i))foo)(o'))
ā(o(( + 1))
a(o(n))
2, VoVoo))) 9NK Σ a())
2
N a(o(N 1)) NNNKI
QN-1
+2
N-1
Σ
Qu+1 alo(n 1))
Qu a(o(n))
Σ
Valo())V
5=1
According to Proposition 2, we have:
++1
K-2-T
<(K -2- nā(O6))NL)
II
(I - [ -1
S
Σ
m
N/ K-1
m
Since we choose 7) S ANKRE. we have:
r+1
2-2-T Σ
II
(I- (
S
e
S
(12)
4Rm
4R
M==--1
m=
where we use the fact that -Kーー-
m
K
2-2-T
m
Hence we know:
4fim
47t 1- c/4R
Ti s (
a(o(N - NKKI
QN-1
Σ a(d(i))Vaa)("""
N
T+1
Σ
II
(I- nāláa(N ())NH)H))))))) nN MNK1
Σ
Voo()))
#=
K-1
5TNN2K
Σ 7=0
II (I - nāna(o ( - Σ ((a(i))Vooo)(e
#=
N
S N2²
Σ a(o(i))
4R e
Thus have:
2 [V
Ti S N2
4R
For T2, we first examine the bound of
Qut1 a(o(++1)
Q a(o(n))
Quti
Σ
II aI nálo(n 1))NH,))
n/-
Σ II aI a- ===K1
a +
a(a(n))
P=K-1
K
II (I (I- n-laa(+ 1))((x+1))NHH))) 1))NHH)
II (I - (((a(())))HHH))
P=K-1
K-1 r+1
K r+1
-N
Σ
II
(I (1)
II (I- ()-
7-0 M==-1
P=k-1
<nN>
R-1
Σ
K-2-1
Σ
c(K 7)
ná(o()))NL)
K-2- Σ
[(K 2 2- (ck--2-)) nā(a(n (nn(((n+n)))) ((((o(++1))N)) ) (7)
m
m
D=D
m=1
m
25

---

where we evoke Proposition 2 at last step. Given that S ANKRL we have:
K-1
2-2-T
K-2-T
Qut1 1l)(n+1))
Qn ((o(n))
nN
((())))
e
m
Σ
a +1)1)
4Rm
4Rm
a(a(n))e
a(o(n+1))
S nNK
4R
4R
where we use the reasoining in (12). Hence for VT:
N-1
((((n))e
á(o(n + (1))e
V
Σ a(o())V(as(""))
4R e
4R-6
N-1
S NNNK
Σ
(á(o(n)) a(o(n 1))) (Gv 8n log(1/p) + N v
AR
n=l
2e
SNNK (GV8N log(1/p) [V)(av"))
4R
where at last step
Lemma 10. (GGSN) + ([V v"))(( I*)
we
evoke
So we can conclude T2 S
2
N²?
2e 4R-c
Putting the bounds of Ti and T/
together will conclude the proof for (a).
Now switch proving (b). Once again by the summation of parts identity we have:
Qn
E (QV(o())(') = a(o(N)))
(o))
Σ (((s)))V
a(o(n + 1))
a(a(n))
Taking the norm of both side yields:
N²
N' (o)((()))))(
QN ((T(N")
Q VoVn(()))"))
B
Qn a(o(n))
Σ facy())
((((++1)) 1))
Plugging our developed bound for ||Qv. and
a(o(at))) a(o(n)) yields:
N²
QN ((O(NN))
<
Σ
N' + N (V(()v")))
S MNKK1 WNL)*
(G
where at last step evoke Lemma 10. And for C, the similar reasoning:
-
Qp+1
Qn a(o(n))
Σ
Σ ((ori)\(o)(+"))
a(p(n+1))
N'
< Σ
N e
(á(a(n + 1)) + a(o(m)))
+ n [V v")
4R -
N
==1
< N K
(GV8Nlog (og( [Vo(a,,''')'')
4R - e
26

---

Putting these pieces together yields:
N
S (IVo(á, v")) e²Ill GVSM 8N log(1/p)
Lemma 12. During the dynamic of Algorithm A2, the following statements hold true with probability at least p
N-1 1
Σ
I (I-
Qu+iHn+) H Σ 22 QoVoo(i(
S 186 18
+
Proof. We first apply Cauchy- Schwartz inequality:
2+2
a
a1-
Qu+1H Hn+1 Σ
Qj (o())
J'=N
==1
H, )
n+2
II
H
Σ
Q ((oo)(e0)
J=
N-1 L
a
5(1+ NNK N=) 2N
Σ
á(o(n 1))
Σ Q focs)oo)
NeLL Σ
á(o(n 1))
Σ Q Q
We proceed by applying the bound from Lemma 11 (b):
VoVoo()
33nnNK (IVe(), v")l (ogg((/)))
Therefore, follows that:
n42
Il [I 11) Qn+1Hn+1 Σ Q, (QV P=N
N-1 S SeNNKL2 Σ á(o(n +1) 1))-)))nNne (IV@() (GI
s 5383 (IVb(á, w²)l G(/Nll/g
Lemma 13 (Noise bound). During the dynamic of Algorithm A2, the the following statement for gradient noises holds true with probability at least p.
N j+1
Σ
-- H3,
SnNKe
i-IT-N
where
K-17H1
6y:= = ΣI (l 1=
27

---

Proof. According to triangle and Cauchy-Schwande inequalities we have:
N
i+1
Q j=
Σ
II I1- (
j=lj=EN
(1 + (INKOO))NKK(+NN))5))) ||8y1|
(1871
N
(1+ (((0())))(((++)))))))))) NA(O()) (1 " L 8
S N..KE88
B.4
Proof of Theorem 2
Proof. For notational convenience, let us define
n j+1
g" = Σ Il
N j+l
8" = Σ
H,
-Iy=NN
Then we recall the updating rule of (Lemma 7):
0 0+ Pw (v" g" - 8*)
Hence we have:
E Bl - (a)]l = E PP (0" - 9 - -88 - v''(a))
< E 1|00 - g" 8" =')(1)?
< 1100 - v" (a=)l -2 v² 00 (d)) + EIgg"1 + +R 1
s E 1|00 - 8" (a))l 2B(nNKVO( v²), v" - va())) - 22 N - (a((
E ||"""|1 + 118"12
Now, applying strongly convexity of ( :) and Cauchy- Cauchy-Schwaree Schwartz inequality yields:
| 80 (a)ll s (1 - pm |0- -- (a)]l - NKK (a))I=-NNKE(((( v²) (((),
+ WK Ellgg - "NKV (oa
+ E ||"" + E 111'
$ (1 2 N E 1100 (/kk)\l~nVKK/(aaaaa - v( và)))
1
+
K Ellg" "KV
+ 2E 1g1" - NNKVO o' 0')1 + 2B NV + ++E187 1
Since ((á,) :) is smooth, we have: E [[ < 22LE 200 (a))). Therefore,
we have:
E El - (a)]l < (1- (1-)]<<<\]n]]]]] [ - 22)61
(NNK - 1 (()),"((a) va)))
(13)
- +2 +2)REI1 Elgg" -NVVV +B101111
2
(14)
28

---

Now, we examine the term 11g" VVVV vz || First according to summation by part
(Lemma 8) by letting A₁ IT (I Q₂, and B, Qj (o))) have:
N j+1
g" = Σ -T- (1=-00))0
1=1 P=N
n+
Σ A, Σ
22 (o - Σ
II [I Q₂H,) (I-))
II (I- Q,H,,)
Σ Q VoV
j=1
V'=N
j'=N
N-1 Σ
Σ Q (o(s)())
II (I- HyH,,)
Qnt1Hn-1 Σ Qj (ob))))
n=1
P'=N
Hence we have:
11
nNKVv(,,""))
N Q, (,ooo ((o(ii)V/
N-1
n+2
N
-
Σ
II (I Hy
Q Ho E9, (oo
Q
|=1
P=N
j=l
)1
N-1
n+2
2
nNK a(o(i)V\.'((')> Σ 20 (.
+
Σ
II (1 Q₂H,y)
Q Σ Q, (oon)
1=1
1=1
n=1
P=N
==1
20n 20 k²
+ 36c 366"
[V((,,"))"" +256/ k²
)
2
0 log([1/p)
4R
1R 1R--)
2116 1 No K I N
(8)
20=
+ 36e 360 of' "'""K*"
2L 2L vaa
4R eR-e
4R - )
211 AM 2 N x²
0² log(/)))
where in (1) we apply Jensen's inequality, in (2) we plug Lemma (a), and Lemma 12, and in (3)
we use the I--smoothness of 4. Plugging above bound back in (19) yields:
+1 (a)ll
(11 N 1100 - (a))l + 22 22
NNK N2 2NNNK +2)
2202² N²?2
- 36e 8 NN
B(0(á, v²) $(6, v' 06)))
1R-e)
( +2) (24te"n
N* (244e"n'"""*' 25"
))
G Goolog(1/p).
4R
4 Ingl (NKK) and large enough epoch number: P.N KR
Since we choose 7)
40 (1++)) 16 16 (oooooVKR)) 64k 64% 64klog(VKK))) VNKR)
R 2 max
know that Ti < 0. We thus have:
E || 1 - - (a=)
$ (1 -
+ (2 +2)
1
2246"
N
AR )
C2 ( (2Nlog(1/))
29

---

Unrolling the recursion from r = R to 0:
E t 2" (a:
I (1 2 umNK) R - - 2
1
))
2 +2) (4886 No k 512
C
G log(1/p).
4R AR-
Plugging in our choice of 7) will conclude the proof:
E 0" (a)
E
2
52
NN (PNlog(1/P)
E - (a))) <
NKR
R²R
Finally, according to Lemma 2 we can complete the proof:
4NG2
((a 0.) : s 2L 1, 00 (as))l
2
Iā; aa
LD2
L82
LG log(L/))
+
+
NKR
4=R
H²R
Ta
D2
RG?
LŌ
+ K5 (w L2
cxp
Kig
RK
1²R² 1ºNRK)
where we plug in the convergence result from Theorem at last step.
C Proof of Convergence of Single Loop Algorithm
In this section, we turn to presenting the proof of single loop PERM algorithm (Algorithm 2) where the learning of mixing parameters and personalized models are coupled. Compared to Algorithm A2, here during the optimization of model, the mixing parameters are also being updated. As a result, we need to decouple the two updates which makes the analysis more involved. We begin with some technical lemmas that support the proof of main result.
C.1 Technical Lemmas
Proposition 4 (Basic Properties of SGD on Smooth Strongly Convex Function). Let wt be the tth iterate of minibatch SGD on smooth and strongly convex function F, with minibatch size M and learning rate 7/. Also assume the variance bounded by 8. Then the following statements hold true after T iterations of SGD:
2yk82
VF(W"" 112 = 2L (1 - pY)" (F(w²) - F(w")) +
(15)
M
w- 2 (1- - (F(w") =F(w")) + 273kE2
7282
E m T+1
(16)
M
M
2 (1 uY) (F(w")) F(w=)) 2
82
E Elw"" - w0 1|
(17)
12
122
Lemma 14 (Bounded iterates difference of a)). Let {o]] be iterates generated by Algorithm 2, then under conditions of Theorem the following statement holds:
||a - 1p 1
1--1 Kig
To
+0 On (K3
(, - wN)" (F(w)) - F(w")) + 3xx22
7282)
M +
M
Proof. Define
z" = [[VVii(w'') -Vi(w=)) ,..., (V/(w") - VSN(w=))1
30

---

According to updating rule of a in Algorithm 2 and Lemma we have:
| - aj 11 - 3 3 (w )|² + 3 a (w"-1) agz (w (w")ll + 3 a; 112
s 6(1 Mg TM + ox 31 (w"-1) og (w0")|
s 6(1 Mgllaa Ta 3m ||=-- 201
N
S 6(1 - )37laa* TT + 3
V[i((n)) - Vfz(w") + Vii(w="))) (1)(w" 1|]= 1|20" - w²111
where the third inequality follows from (8).
Since
[V V[z(2")] 5
|| li fur") - Vfz ([z(w") I] + 2L 1|w0" w=Il I|. we can conclude that
|a 1-1 12 S 6(1 Pa Ta
121.2 K ==1
8 (8IVJ((w")) - "V/''''') 81. 1|e0 -=1 81.2 w +8.? Σ
-ww=l w0(l²) 1100" -w-- w011
Ta
+0 10x2212 [12 || [[0"""("0'100000001111
<
Ta
n9x52 7202
<
1
+0 (
(1 py)" (F(w)") - -F(""))))))))))))))))))))>)))-
+
Kg.
M
M
where at last step we plug in Proposition 4 (16).
U
Lemma 15 (Convergence of a). Let { be the mixing parameters generated by Algorithm
Then under the conditions of Theorem 3, the following statement holds:
||á - a2 a < 2
1 - To
Re R5 ) 2 (1 - (1-17)² + 27
82
iE [N]
Kg
12M
Proof. We notice the following decomposition:
||0; - x x²1 = || ar 1
2 <2||a! aw 1)| + (w0)) - ag ag (w0)) )l|
< 2(1 - - + 0 (N² ((((w)) ) + N2 |w0t - -woll²) AL ( 1"=""="11" w11)
Kg
1 TT + 0
82
2
Ki (10 (1 - n²)²² + 27
Kg
H2M
where in the second inequality we apply Lemma 3, and in the third inequality we use Proposition (17).
C.2 Proof of Theorem
Proof. According to Lemma 2, we have:
2(a 6;) (aa 0/)) < 2L 0R v'(a))? 2222L - [d;1=
4NG²
We first examine the convergence of ||0 - váx)l Applying
inequality yields:
|| v² 0"aa"t1))1 < 1 1+
|| 0² (a'')) (1 4a 2) 10 (a (a"+1)
v" (a"))
4a 2
1 + 1
<
2 11 (1 4a 2) nāllarrt =a1l
4a
(18)
31

---

where pnNK and last step is due that v" (a) is Kg. Similar to the proof of Theorem first define
/NG Lipschitz, proven in Lemma
i+1
g Σ
II (I
J'=N-1
i+1
8" Σ
II (I- H.H))))))
J==-11
Then recall the updating nile of
o-+1 = Pw - -0)
Hence we have:
E | - v (a')l = F/p (v" g" - 8 - v'"(a")))
$ E 1| g" 8" (a')2
< E 110" v² (a (a)) 2g(g", o" - v² Elg + E1881
s E 1|0" v' (a"))l 2 v)), v' - (a"))
2E(g" VV'' v²), v- (a وو 1
Now, applying strongly convexity of o(a", -) and Cauchy- Schwartz inequality yields:
- - (a")l s (1 W
o' (a") 12 N)K w²) $(a", x" (a)))
VNKV + - o'
2 M
+ E ||"0||
1
1
1
2 (uNKK E
- v (a))]
N Elgg" NVKV 0²)1
+ 2E و1g - VKV o' v²)] NV ,')) E
where in the first inequality we applied Cauchy-Schwartz inequality and strongly convexity. Since ((a", :) is L smooth, we have: E [[""",a", "²'l| s 2LE(#(a", v²) $(a", v*("))) Therefore, it follows that:
E E 2" (a)]l² - 1 11-2uNNK E v"'''" -nNN - N
v²)
2umNKK +2
Ellg" - NNVKV 0'' E 1
(19)
Now, we examine the term l1g NNKV 22)l in the right hand side of abovee inequality. First according to summation by part (Lemma 8): let A₁
i+1 (I Q2 H,)) and
B, (ov)) then have:
N i+1
g"
Σ II
1=1 P=N
N
Σ AJB (Q)oo(
Σ
a- Hy)
1,, Σ Q, (o(o(()))))
Hy)
1=1
1=1
J'=N
F=N
N-1
Σ Q, (o)
Σ
Il
-I- Hy)
I
QutiHn-1 Σ Q Q.Vlfa
n=1
/'=N
32

---

Hence we have:
11
N
j+1
)N - Σ II (I- (I-Q,H;)Q,
v (o))
NNK Σ a(o())V(au((0))
N
21Viooo)0 Σ
n+2
II
H.. Qn H Σ Q, V (oon)("))
"
j=1
P'=N
j=0
2
NNK a(o()))(aa)()))))
Σ
Q2 Sogi(
3=1
3=1
N-1
n=
2
Σ
II (I Q2HH,)
Q Qu+1Hn+1 Σ Q, (,Vlocy)("))
P'NN
(2)
2==
[V
N2
)
2
G² log(1/p)
4R
4R e
244c 1/ A N L (oog(p/p)
2012N2
+ 36c N
21 2I((((,)")
$( va))))
4R -
4R AR-e) ) GoNlog(L//))
N K
N
where in (1) we apply Jensen's inequality, (2) we plug in Lemma 11 (a), and Lemma 12, and in (3)
we use the L-smoothness of &b Plugging above bound back in (19) yields:
1 00 (a")
< (1 2 Ju K)B E 1|0" v'(a')1 K
nNK - -ANK-A2N2K2L ( +2)
220m2
e
366 36E" (
4R
x EB 1--((aa -4(a", -4(a",0"(aa))
1 +2)
244c ' N= 256
1P-)) -
E
)²
G2 N log(1//)
4R
Since we choose
4log(v MongvNKKI R)
and
RNKR
3
64x² log / / K R)
R. 2 max
c,
8
33

---

hence we have:
E 1 r+ - 0-(20)1
(1 -
22 R, v") - D( v- (a)))
20
1R AR-e) )
2 +2)
(246G))\*< 226m
G2 (og(L/))
1 1 J B'10"= N
1
+2)
(244eS NO N K*
)
e
))
G2 log(1//)
2
4R
Putting above inequality back to (18) yields:
| -
(
11
110 -
+ 4a 2 "/ll"" aa11
4a
+
1
+2)
256
(1R-e)
G2 G2Nlog(11/))
2
1
1100
4a
+2
2 +2)
1
2f
2
(R AR==)) ) G2 G2Nlog(1/1)
(1--1)
(171(1- (17)) DG + 2n002
7282
+0
pNNK
Kg
M
M
where at second inequality we plug in Lemma 14. Unrolling the recursion from r = R to 0, and (R3)
plugging in 1) ANKR
yields:
| R
('(
R 1
1
11000 -''0)(1))1 N
2 +2)
N
+ 256 (
(R -R-e) (AR-)) )²) G2 G2Nlog(1/P)
8 11
+0
2
Σ
R
1
R-r
1
Ta
+ Ne =x)22(((0))
(L - (1-1))) DG + x38
+ x282)
M ))
W K
4a
Kgg
M
r=0
5
11000
)12
+0
+
k
G GoNlog(L/p) +
8²
NKR
Re2t
μR
+
R
log(N
R-r
((1-
1
To
+ kig
2 2 (ww")
(?² DG
n2m02
+
2282 M
R
Kig
M
1=0
34

---

Plugging in 7= log (R)
yields:
akt
1100 (7)))
((
n²
N
82
l
v0 =²²=0
+
Gogloo(1//))
NKR3
R2
1²R²
p.R)
log(NKRR) log
+0
Rip.
2 1 R
Σ
Kg
R
R
log(NKRR))
R-T
Ta
Rig
L2 5 (:w2"))))))))))))))))) )2202
Σ
+
R
Kig
M
r=
D2
16
N
8² +
00
+
Go
NKR
122
1²
PR
Kig
2
2 kig 1
w0))
DD
(1--))
k n² C, (w))
+
Rig
R
Kig
M
Since ú; R and o; R we have the convergence of || v*(a)) gence rate together with the convergence of || a²112 from Lemma 15:
Plugging this conver-
| 50
00
2(1
1 To
1
+0
( I1, (1
1-1n)
27
8²
lig
11
12M
<ō
Ta
0
R
kn )T. 1
8²
1 RM
Kig
1
together with applying Lemma 2 leads to:
4NG² L
((a é.) (aa 2t) < 2L 1|ó -(((())
2
| 1
0
LD
+Ō
NL
G'Nlog(1/p))
LO
NKR3
R
p²R
μR.
k2 I Ei
DG
1
To
L x2 x2 2
(w
+0
nh LR
1
+
R
Kig
MM
4NG2
Ta
"g5((
1:27
( )KL
82
+
2 L
+
L
Kg
R.
1²RM
0
LD
+Ō
((12)
(R--L
NL
G2Nlog(1/p)
L8²
NKR3
R²
R²
+0
Righ2 hig
37 (w")
DG
R T.R2
(1-
-11
1
Ta
L2x2² kg mig Zo (:
R
1ig
p M
thus completing the proof.
35


Reviews:
[{"review_id": "BK6mAIvfoU", "reviewer": "Reviewer_DEoF", "summary": "The paper proposes a novel learning framework for personalizing empirical risk minimization in a given federated learning problem where the data is heterogeneous (non-IID).  The core idea of the paper is to use optimization relaxations and data shuffling to achieve near-optimal solutions for the general NP-hard personalized federated learning formulation. The paper supports the framework's effectiveness by providing generalization bounds for strongly convex losses. Further, it evaluates the performance of the method on synthetic data where each client learns a logistic regression model.", "questions": "- What is the current work's main computational and theoretical advantage to [1]? The theoretical assumptions seem the same in both papers. Is there any improvement in the generalization bounds?\n\n\n[1] Fallah, A., Mokhtari, A., & Ozdaglar, A. (2020). Personalized federated learning: A meta-learning approach. arXiv preprint arXiv:2002.07948.", "limitations": "- The main limitation of the paper is the reliance on the strong concavity of the model in both theory and experiments. \n\n- There is no experiment on real datasets where each client can have its unique model or architecture. \n\n- There is no experiment on large-scale datasets where the underlying model is non-smooth and non-convex (neural networks with ReLU activation function).", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper reviews the most important limitations of the current federated learning literature on heterogeneous data, including clustering and personalization ideas.  Further, it demonstrates the main problems of the personalization frameworks (WERM and BERM). \n\n- The theoretical study of the problem when the loss of each client is strongly convex is rigorous, and it offers reliable convergence guarantees for large-scale problems (huge number of data points and limited number of devices)\n\n- The experiment results on the synthetic data are impressive, showing a significant improvement of the local losses compared to the baseline existing approaches.", "weaknesses": "- While the paper claims the method can be used for cases where different models have different architectures, the reliance on the strong concavity and smoothness of $f$ is a serious limitation (neural networks with ReLU activation function are not convex nor smooth).  \n\n- The experiments are only performed on the synthetic datasets where the underlying model in each client is logistic regression (convex). Therefore, it is not clear that the performance of the proposed algorithm is superior to other baselines if the data is real and the models are non-convex. \n\n- While one of the motivations for relaxing is scalability, there is no experiment on many data points and clients. One suggestion is to report the algorithm's runtime aside from the other approaches in the current experiments. \n\nBased on the above points, it is understandable that the theory might not be easily extendable to non-convex losses. However, it is crucial to add experiments on real data and with many data points on diverse neural network architectures for different clients.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable comments! We will try to address your concerns as follows.\n\n**The reliance on the strong concavity and smoothness is a serious limitation**\nStrongly convexity and smoothness are standard assumptions in the proof of shuffling-based optimization algorithm. Our objective is indeed a bi-level optimization: in inner-level we have to optimize $ \\alpha $ to get an approximate solution for $ \\alpha^*_i$, and then in outer-level we optimize model based on the approximate $ \\alpha$. If we remove convexity of $f_i$, then our inner problem will be nonconvex. \n\nOur empirical results (included in Appendix C of main submission), demonstrate the effectiveness of PERM on non-smooth and non-convex objectives as well and we believe extending our theory to establish the convergence without smoothness or strong convexity is an interesting future work. \n\n \n**More experiments (real datasets, model heterogenity, non-smooth, non-convex)**\nAs it was mentioned in Global comment, we do have empirical results on EMNIST, which is a real-world federated dataset, and the MLP model, which is a non-convex model. We also provide experiments on CIFAR10 in attached PDF, Figure.2. Due to space limit we only included part of results in main body and reported these results in Appendix C, and apologize for inconvenience.   \n\n**Scalibity and reporting the algorithm's runtime**\nThanks for the comment and suggestion. We discuss the scalability of PERM in comparison to existing methods in Global answer.  To summarize, the compute burden on clients and server is roughly same as existing methods thanks to shuffling (extra overhead due to estimating mixing parameters).  \nPer you suggestion, we reported the **wall-clock time** of PERM and other methods in Figure 1 of attached PDF, demonstrating that the run-time is slightly worse  as mentioned above.  \n\n**Advantage to meta-learning approach Fallah et al**\nAs discussed in Section 2 of main body, Fallah et al fall in the category of BERM, in which they try to learn a global model (meta model) on the average loss, then personalized the \n (meta)global model to adapt to local distribution. As we discussed in Section 2, when the clients' distribution is extremely heterogeneous, there is no good global model for personalization (please note that knowledge transfer among data sources still happens through the global model which is learned by optimizing vanilla empirical loss). For example, given a $x$, if the first half of clients are with labeling function $y=f(x)$ while the second half with labeling function $y=-f(x)$, the global model learned on the average of their empirical risks cannot have good performance. This is also verified by our experiments, where in high heterogeneous scenarios like synthetic datasets other personalization approaches like PerFedAvg fail. Even in lower heterogeneous cases such as EMNIST, our approach still outperforms PerFedAvg as we effectively estimate how much each data source contributes to training each model based on relatedness of their distributions. \n\nAs for generalization of PERM, from algorithm-independent generalization (Theorem 1), PERM can potentially entail optimal generalization (except the error due to  our relaxation to estimate the mixing parameters). We note that while [1] establishes the convergence rate of optimization algorithm of proposed meta-learning method, no generalization guarantees are provided.  We believe understanding the generalization of both methods through the lens of stability  would an interesting future work.  "}, {"review_id": "TnWPMHaiiq", "reviewer": "Reviewer_S5TK", "summary": "This paper proposes an optimization strategy for federated learning (FL). A key challenge of FL is data heterogeneity, where each client device has different data distributions. Unlike a naive FL, which minimizes the empirical risk simply weighted over clients (called WERM), the proposed PERM first minimizes the empirical risk for each client individually. To mitigate the scalability problem in PERM, this paper proposes a practical algorithm using model shuffling. Experiments show better accuracy and loss compared to other personalized FL methods.", "questions": "Training time compared to other personalized FL methods. This would be an important aspect and might be related to a limitation of the proposed method. Also, I'm curious about the discussion of scalability against the number of clients, compared with the existing methods.", "limitations": "As described in the Weakness section, I suspect the proposed method may lack scalability.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "+ PERM seems theoretically sound, achieving the optimal solution for each client. The performance seems to be better than representative personalized FL methods.\n\n+ In addition to data heterogeneity, the proposed method can be easily extended to heterogeneous devices and models since PERM minimizes the loss of each client individually.\n\n+ This paper explains the theoretical aspects very clearly. Especially Section 2, which describes the key theoretical part, is carefully written and quite easy to understand, even for readers a bit outside the field (like me). ", "weaknesses": "- The theoretical/practical merit compared to other personalized FL methods is not clearly mentioned (although experimental results show that the proposed method achieves better performance compared to representative methods of personalized FL.)\nEspecially, I am curious about the scalability of the proposed method. In the introductions, I understand that existing personalized methods lack scalability. However, if my understanding is correct, the proposed method also suffers from increasing the number of training data and devices. Rather, PERM theoretically has O(#devices^2) loss functions, as mentioned in the paper, implying that scalability is a significant limitation of PERM. \n\n- The overall description was very clear and easy to read, but I sometimes struggled because of missing definitions and confusing notation. For example:\n    * Weight parameter p vs. \\alpha might be better for consistency. The definition of \\alpha \\in \\Delta_N does not appear the first time it is used.\n    * The procedure of PERM (L133-134 and Eq.(PERM)) was a bit unclear. In general, it was unclear to me how to determine \\alpha_i before knowing the result of each client (i.e., h minimizing Eq.(PERM)) until I read the next section. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your compliments and valuable comments! We will try to address your concerns as follows.\n\n**The scalability of the proposed method**\n We discuss the scalability of PERM in comparison to existing methods in Global answer.  To summarize, the compute burden on clients and server is roughly same as existing methods thanks to shuffling (except extra overhead due to estimating mixing parameters which is same as running FedAvg in two-stage approach and an extra communication in interleaved approach).  To evaluate this empirically, we reported the **wall-clock time** of PERM and other methods in Figure.~1 of attached PDF, demonstrating that the run-time is slightly worse due to overhead of estimating mixing parameters.\n\nWe are also sorry for the typos and will fix them in the revised version. Thanks for pointing them out!\n \n "}, {"review_id": "OojUbJc433", "reviewer": "Reviewer_x8t7", "summary": "This paper studies a distributed training strategy to obtained a personalized model that shares a bases for all clients while an adaptor for each clients.", "questions": "See above", "limitations": "See above", "rating": 5, "confidence": 1, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The overall derivation of the algorithm is theory-driven and reasonable.\n\nGood experiment result in the simulation dataset.", "weaknesses": "This paper limits the mixing ops by a weighted summation, which great limits the scope and advantage of the proposed method.\nDoes your algorithm work for a more general class of mixing ops?\n\nThe experiment using the synthesis data looks great but as a paper aims at a very practical setting, could you show experiment on less-toyish benchmark？", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable comments! We will try to address your concerns as follows.\n\n**Limitation of weight sum based mixing method**\nIn multi-source learning, weighted summing source domains and learning on it is the most popular and natural method [1,2], due to its simple form, easy implementability and excellent performance. We agree that there are other methods to share the knowledges of source domains, but it is beyond the scope of this paper. We will consider it as a promising future work. Thank you for the suggestion!\n\n\n**Experiments on less-toyish practical cases**\nWe discuss this issue in Global comment. We have already reported results on real data sets in Appendix C of main paper. We apologize for the inconvenience. Moreover, we provide new experiments on CIFAR10 in the attached PDF.\n\n[1] Konstantinov, Nikola, and Christoph Lampert. \"Robust learning from untrusted sources.\" International conference on machine learning. PMLR, 2019.\n[2] Mansour, Yishay, et al. \"A theory of multiple-source adaptation with limited target labeled data.\" International Conference on Artificial Intelligence and Statistics. PMLR, 2021."}, {"review_id": "gFZEgKKUdL", "reviewer": "Reviewer_XAoP", "summary": "This paper proposes a new paradigm for learning from multiple heterogeneous data sources to achieve optimal statistical accuracy across all data distributions without imposing stringent constraints on computational resources shared by participating devices. The proposed Personalized Empirical Risk Minimization (PERM) schema provides an efficient solution to enable each client to learn a personalized model by learning who to learn with via personalizing the aggregation of data sources through an efficient empirical statistical discrepancy estimation module. The paper rigorously analyzes the convergence of the proposed algorithm and conducts experiments that corroborate the effectiveness of the proposed paradigm.", "questions": "1: Could you please provide additional details regarding the implementation of the proposed algorithm, including the specific hyperparameters utilized in the experiments?\n\n2: How does the performance of the proposed paradigm compare to existing approaches in scenarios with highly heterogeneous data distributions?\n\n3: What is the scalability of the proposed paradigm with respect to the number of participating devices? Are there any potential computational or communication bottlenecks that could hinder its scalability?\n\n4: Why were real data experiments not conducted in this study?", "limitations": "See weaknesses.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The strengths of this paper include:\n1: The proposed PERM schema is a novel and practical methodology that addresses the data heterogeneity issue in distributed learning.\n\n2: The paper provides insightful theoretical analysis of the proposed algorithm and rigorously proves its convergence.\n\n3: The paper is well written and organized. ", "weaknesses": "1: Inadequate evaluation and lack of verification on real data: Although the paper includes experiments showcasing the effectiveness of the proposed paradigm, it would greatly benefit from more comprehensive evaluation, particularly using real-world data.\n\n2: The practical usefulness of Theorem 2 is limited due to its dependence on parameters related to the oracle models, which are typically unknown in real-world scenarios.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive comments! We will try to address your concerns as follows.\n\n**Inadequate evaluation and lack of verification on real data**\nWe discuss this issue in Global comment. We have already reported results on real data sets in Appendix C of main paper and reported new results  in attached PDF as part of responses. We apologize for the confusion. \n\n**Details of implementation**\nFor the synthetic dataset, we train using logistic regression with an $\\ell_2$ regularization of $5\\times10^{-4}$. The learning rate follows a convex decay schedule, initialized at $0.01$ and reduced by $10\\%$ each iteration. We use a batch size of $128$ and perform $10$ local gradient descent steps on each communication round before model aggregation. The synthetic dataset itself is generated as detailed in the paper. For the EMNIST and MNIST datasets, we follow the same hyperparameters, except, we use the initial learning rate of $0.1$, and $20$ local steps in each communication round. Thank you for the question and we will gladly clarify this in the subsequent version. \n\n\n**Comparison with other approach on highly heterogeneous dataset**\nTheoretically, as we discussed in Section 2, when clients' data distribution are extremely heterogeneous, e.g., first half clients and second half clients have the opposite labels for the same input, then any existing works that relies on learning a global model then personalizing will fail as the global model may transfer *negative* knowledge from other highly irrelevant data sources. However, in PERM, this can be avoided by preperly estimating the mixing parameters (almost zero contribution from highly irrelevant data sources thus entailing better accuracy). This is also verified by our experiments on synthetic dataset, where meta-learning based algorithms and fine-tuned FedAvg model fail to achieve good accuracy. In the attached PDF (Figure 2) we included result on an extremely non-IID setting where the best model for each client is just the model trained solely based on its local data without collaboration, which further demonstrates the effectiveness of PERM compared to other methods that rely on learning a single global model to transfer knowledge among data sources (which might harm the performance due to negative effect of unrelated data sources which are eliminated in PERM by estimating the relatedness of data sources). \n\n**What is the scalability of the proposed paradigm with respect to the number of participating devices?**\n\nThanks for the question. Please refer to Global comment for a detailed discussion on scalablity of PERM.  To summarize, the compute burden on clients and server is roughly same as existing methods thanks to shuffling (except extra overhead due to estimating mixing parameters which is same as running FedAvg in two-stage approach and an extra communication in interleaved approach). The only hurdle would be the required memory *at server* to maintain mixing parameters at server which scales at $O(N^2)$ which can be resolved by clustering devices or using approximate methods such as low rank approximation  which we leave as a future work.  To evaluate this empirically, we reported the **wall-clock time** of PERM and other methods in Figure.1 of attached PDF, demonstrating that the run-time is slightly worse due to overhead of estimating mixing parameters.\n\n\n**The practical usefulness of Theorem 2**\nThe purpose of Theorem 2 is to show that PERM can provably converges  but in practice, we totally agree with your comment that some of the constants are usually unknown  which holds for most of theoretical analysis of SGD based methods, unless we utilize parameter-free or adaptive methods which is not the focus of work. We note that in implementing the PERM in conducting our experiments  we do not need the parameters of the function/oracle as discussed in your question about details of implementation. Our experiments validate the convergence of PERM as well."}]

Author rebuttal:
We would like to thank all reviewers for your careful and constructive comments. We appreciate the compliments from reviewers on our algorithmic novelty and theory and  will gladly incorporate the suggestions. 

**Empirical evaluation on real data** We understand the main concern raised is that our experiments were only conducted on synthetic datasets. Importantly, we direct your attention to additional experiments *already included* in the **Appendix C** of main submission on EMNIST and MNIST using MLP models, which further validate our theoretical results.  We included additional results in an *extremely* non-IID setting in **Figure 2**  in attached PDF on CIFAR10 dataset using a 2-layer CNN model, which further demonstrates the effectiveness of PERM. To further address this concern more comprehensively, we are currently working to supplement our experiments with results on the CIFAR10 dataset using CNN models with varying heterogeneity among data sources and will include in the subsequent version. This will provide further evidence to underscore our theoretical findings.  We appreciate you highlighting this opportunity to strengthen our empirical evaluations. 

**Scalability of PERM** Another common question is the scalability of our method. Since in PERM we have to solve $N$ weighted ERM problems to yields $N$ personalized model, naively solution will make client $i$ update $N$ models at the same time. When $N$ is a large number, it is not affordable. Our proposed method only requires each client to update one model at the same time thanks to model shuffling/swapping by server, and the convergence rate and the experimental results show that it can still converge very fast. 
 
To better illustrate this, for simplicity, let us focus on the two-stage method. The first stage is simply running FedAvg to estimate the mixing parameters. In the second stage,  the compute burden of each client is exactly the same as FedAvg that optimizes a single ERM objective (without sampling), with the only difference  that in PERM the model each client receives at each communication round varies. On the server side, the averaging is replaced with model swapping which enjoys the same scalability as existing methods. The only issue would be the required *memory* at the server to maintain mixing parameters for all clients which can be reduced by low-rank approximation or clustering methods (which could be an interesting future work). A piece of empirical evidence to support this is included in the attached PDF (**Figure 1**) which is slightly worse due to the overhead of estimating mixing parameters in parallel to learning personalized models (i.e., two times more communication needed compared to FedAvg and same as  personalized methods such as personalization with interpolation).



Meta-review: This paper studies "personalized ERM" for learning from heterogeneous datasets. Instead of pursuing bilevel optimization proposed in prior work, the authors focus on a new formulation based on the mixing of local losses personalized for each client. The convergence of the proposed algorithm is analyzed and the experiments demonstrate the practicality of the method in certain applications. The reviewers found the topic interesting, the presentation clear, and the contributions solid. The authors' rebuttal addressed the concerns of multiple reviewers. I recommend the acceptance of the paper. In your revision, please include discussions on prior work (e.g. Fallah et al) that was brought up in the rebuttal period. Also, please include additional rebuttal experiments and discussions related to the scalability/limitations of the proposed method.
Decision: Accept (poster)