Title: Distributed Personalized Empirical Risk Minimization
Venue: NeurIPS 2023
Accepted: yes
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

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 a distinct model for each client by learning who to learn with and 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 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 of individual clients. We rigorously analyze the convergence of the proposed algorithm and conduct experiments that corroborate the effectiveness of the proposed paradigm.

1 Introduction

Recently federated learning (FL) has emerged as an alternative paradigm to centralized learning to encourage federated model sharing and create a framework to support edge intelligence by shifting model training and inference from data centers to potentially scattered-and perhaps self-interestedsystems where data is generated [1]. While undoubtedly being a better paradigm than centralized learning, enabling the widespread adoption of FL necessitates foundational advances in the efficient use of statistical and computational resources to encourage a large pool of individuals or corporations to share their private data and resources. Specifically, due to heterogeneity of data and compute resources among participants, it is necessary, if not imperative, to develop distributed algorithms that are i) cognizant of statistical heterogeneity ( data-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 ( system-awareness ).

To mitigate the negative effect of data heterogeneity (non-IIDness), two common approaches are clustering and personalization. The key idea behind the clustering-based methods [2, 3, 4, 5] is to partition the devices into clusters (coalitions) of similar data distributions and then learn a single shared model for all clients within each cluster. While appealing, the partitioning methods are limited to heuristic ideas such as clustering based on the geographical distribution of devices without taking the actual data distributions into account and lack theoretical guarantees or postulate strong assumptions on initial models or data distributions [4, 5]. In personalization-based methods [6, 7, 8, 2, 9, 10, 11], the idea is to learn a distinct personalized model for each device alongside the global model, which

Pouria Mahdavinia

can be unified as minimizing a bi-level optimization problem [12]. Personalization aims to learn a 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 a 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 a 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 architecture to be deployed across the clients and server (model-homogeneity) [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 distillation 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 a 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 a 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 heterogeneous data sources to achieve optimal accuracy w.r.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-aware 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 a 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 multitask 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 a 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 a single global model can be trained in a model-heterogeneous 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 i = { ( x i,j , y i,j ) } n i j =1 with n i training samples that are realized by a local source distribution D i over instance space Ξ = X × Y . The data distributions across the devices are not independently and identically distributed (non-IID or heterogeneous ), i.e., D 1 = D 2 =

̸

̸

̸

. . . = D N , and each distribution corresponds to a local generalization error or true risk L i ( h ) = E ( x ,y ) ∼D i [ ℓ ( h ( x ) , y )] , i = 1 , 2 , . . . , N on unseen samples for any model h ∈ H , where H is the hypothesis set (e.g., a linear model or a deep neural network) and ℓ : Y × Y → R + is a given convex or non-convex loss function. We use ̂ L i ( h ) = (1 /n i ) ∑ ( x ,y ) ∈S i ℓ ( h ( x ) , y )) to denote the local empirical risk or training loss at i th data shard S i with n i 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 = 1 , . . . , N for all data sources (all-forall [33]). A simple non-personalized solution, particularly in FL, aims to minimize a (weighted) empirical risk over all data shards in a communication-efficient manner [34]:

where ∆ N = { p ∈ R N + | ∑ N i =1 p ( i ) = 1 } denotes the simplex set.

It has been shown that a single model learned by WERM, for example by using fixed mixing weights p ( i ) = n i /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 a poor generalization error on individual datasets by increasing the diversity among distributions [37, 38, 39, 40]. 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]):

where ⊕ denotes the mixing operation to combine local and global models, and ̂ F i is a modified local loss which is not necessarily same as local risk ̂ L i . By carefully designing the local loss ̂ F i and mixing operation ⊕ , we can develop different penalization schemes for FL including existing methods such as linearly interpolating global and local models [11, 2], multi-task learning [10] and metalearning [9] as special cases. For example, BERM reduces to zero-personalization objective WERM when h i ⊕ h ∗ = h ∗ , and ̂ F i = ̂ L i . At the other end of the spectrum lies the zero-collaboration where the i th client trains its own model without any influence from other clients by setting h i ⊕ h ∗ = h i , ̂ F i = ̂ L i . The personalized model with interpolation of global and local models can be recovered by setting h i ⊕ h ∗ = αh i +(1 -α ) h ∗ , and ̂ F i = ̂ L i . 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) a same model space is shared across servers and clients.

To motivate our proposal, let us consider the empirical loss ∑ N i =1 α ( i ) ̂ L i ( h ) in WERM (or the inner level objective in BERM) with fixed mixing weights α ∈ ∆ N , and denote the optimal solution by ̂ h α . The excess risk of the learned model ̂ h α on i th local distribution D i w.r.t. the optimal local model h ∗ i = arg min h ∈H L D i ( h ) (i.e. all-for-one ) can be bounded by (informal) [31]

where R j ( H ) is the empirical Rademacher complexity H w.r.t. S j , and disc H ( D i , D j ) is a pseudodistance on the set of probability measures on Ξ to assess the discrepancy between the distributions D i and D j 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 Ξ = X × Y ,

Intuitively, the discrepancy between the two distributions is 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

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. L i with ̂ L i ).

From GEN, it can be observed that a mismatch between pairs of distributions limits the benefits of ERM on all distributions. Indeed, the generalization risk w.r.t. D i will significantly increase when the distribution divergence terms disc H ( D j , D i ) are large. It leads to an ideal sample complexity 1 / √ n where n = n 1 + n 2 + . . . + n N is the total number of samples, which could have been obtained in the IID setting with α ( j ) = 1 /N when the divergence is 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 (e.g., over parametrized setting) and can entail a 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 , i = 1 , 2 , . . . , N . 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 a 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 D i by personalizing the WERM, i.e., (i) first estimating α i , i = 1 , 2 , . . . , N for each client individually, then (ii) solving a variant of WERM for each client with obtained mixing parameters:

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 a different aggregated empirical loss for each client, we can pick a different model space/model architecture H i 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 i , however, making it scalable requires overcoming two issues. First, estimating the statistical discrepancies between each pair of data sources (i.e., α i , i = 1 , . . . , N ) is a computing burden as it requires solving O ( N 2 ) difference of (non)-convex functions in a distributed manner and requires 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 learning). In the next section, we propose a simple yet effective idea to overcome these issues in a computationally efficient manner.

3 PERM at Scale via Model Shuffling

In this section, we propose a method to efficiently estimate the empirical discrepancies among data sources followed by a model shuffling idea to simultaneously 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 as 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 ⊆ R d and use f i ( w ) := ̂ L i ( w ) = ∑ ( x ,y ) ∈S i ℓ ( w ; ( x , y )) to denote the empirical loss at i th data shard.

3.1 Warmup: a two-stage algorithm

We start by proposing a 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 to solve all personalized empirical losses in parallel.

Stage 1: Mixing parameters estimation. In the first stage we aim to efficiently estimate the pairwise discrepancy among local distributions to construct mixing parameters α i , i = 1 , 2 , . . . N . From

generalization bound GEN and Definition 1, a direct solution to estimate α i is to solve the following convex-nonconcave minimax problem for each client:

where we estimate the true risks in pairwise discrepancy terms with their empirical counterparts and drop the complexity term as it becomes identical for all sources by fixing the hypothesis space H and bounding it with a computable distribution-independent quantity such as VC dimension [43], or it can be controlled by choice of H or 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, so most of the existing minimax algorithms will fail on this problem. To our best knowledge, the only provable deterministic algorithm is [44], and it is hard to generalize it to stochastic and distributed fashion. Moreover, since we have N clients, we need to solve N variants of (1), 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 w | f i ( w ) -f j ( w ) | in terms of gradient dissimilarity between local objectives ∥∇ f i ( w ) -∇ f j ( w ) ∥ [45], which 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 min w ∈W (1 /N ) ∑ N i =1 f i ( 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:

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 a communication-reduced 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 w R we optimize over α in g i ( w R , α ) using T α iterations of GD to get ˆ α i . Actually, we will show that as long as w R converge to w ∗ , ˆ α i , i = 1 , . . . , N converges to solution of (2) very fast. Our proof idea is based on the following Lipschitzness observation:

where α ∗ g i ( w ) := arg min α ∈ ∆ N g i ( w , α ) and κ g := n max / (2 λ ) is the condition number of g i ( w , · ) where n max = max i ∈ [ N ] n i . The Lipschitz constant mainly depends on gradient dissimilarity at optimum . As w R tends to w ∗ , the α ∗ g i ( · ) becomes more Lipschitz continuous, i.e., the coefficient in front of ∥ ∥ w R -w ∗ ∥ ∥ 2 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 f i ( x ) 's are L-smooth and µ -strongly convex, i.e.,

We denote the condition number by κ = L/µ .

Assumption 2 ( Bounded variance) . The variance of stochastic gradients computed at each local function is bounded, i.e., ∀ i ∈ [ N ] , ∀ w ∈ W , E [ ∥∇ f i ( w ; ξ ) -∇ f i ( w ) ∥ 2 ] ≤ δ 2 .

Assumption 3 ( Bounded domain) . The domain W ⊂ R d is a bounded convex set, with diameter D under ℓ 2 metric, i.e., ∀ w , w ′ ∈ W , ∥ w -w ′ ∥ ≤ D .

Definition 2 ( Gradient dissimilarity) . We define the following quantities to measure the gradient dissimilarity among local functions:

The following theorem gives the convergence rate of estimated discrepancies to optimal counterparts.

Theorem 1. Under Assumptions 1-3, if we run Algorithm A1 on F ( w ) := 1 N ∑ N j =1 f j ( w ) with γ = Θ ( log( RK ) µRK ) for R rounds with synchronization gap K , for κ g = 1 / ( λn min ) , it holds that

An immediate implication of Theorem 1 is that even we solve (2) at w R , 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 a parameter within a small region centered at w ∗ , the function α ∗ g i ( w ) becomes 'more Lipschitz'. The rigorous characterization of this property is captured by Lemma 3 in appendix.

Stage 2: Scalable personalized optimization with model shuffling. After obtaining the per client mixing parameters, in the second stage we aim at solving N different personalized variants of PERM denoted by Φ(ˆ α 1 , v ) , Φ(ˆ α 2 , v ) , . . . Φ(ˆ α N , v ) to learn local models where

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, as shown in Algorithm 1 the algorithm proceeds for R epochs where each epoch runs for N communication rounds. At the beginning of each epoch r the server generates a random permutation σ r over N clients. At each communication round j within the epoch, the server sends the model of client i to client i j = ( i + j ) mod N in the permutation σ r along with α i ( i j ) . After receiving a 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 Φ(ˆ α i , v ) , i = 1 , 2 , . . . , N during an epoch is equivalent to sequentially processing individual losses in (3) which can be considered as permutation-based SGD but with the different that each component now is updated for K steps. 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 is replaced with model shuffling, and the algorithms run over a 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.

Assumption 4 (Bounded Gradient) . The variance of stochastic gradients computed at each local function is bounded, i.e., ∀ i ∈ [ N ] , sup v ∈W ∥∇ f i ( v ) ∥ ≤ G .

We note that the Assumption 4 can be realized since we work with a bounded domain W .

Theorem 2. Let Assumptions 1- 4 hold. Assume α ∗ i is the solution of (2) . Then if we run Algorithm 1 on the ˆ α i obtained from Algorithm A1, then Algorithm 1 with η = Θ ( log( NKR 3 ) µR ) will output the solution ˆ v i , ∀ i ∈ [ N ] , such that with probability at least 1 -p , the following statement holds:

where κ = L µ , κ g = n max 2 λ and κ Φ = √ NG µ , and the expectation is taken over randomness of Algorithm A1. That is, to guarantee E [Φ( α ∗ i , ˆ v i ) -Φ( α ∗ i , v ∗ i )] ≤ ϵ , we choose R = O ( max { Lδ 2 µϵ , κ 2 Φ κ 2 g L 3 ¯ ζ i ( w ∗ ) D 2 ϵ }) and T α = O ( κ g log ( Lκ 2 Φ ϵ )) .

The above theorem shows that even though we run the optimization on Φ(ˆ α i , v ) , our obtained model ˆ v i will still converge to the optimal solution of Φ( α ∗ i , v ) . The convergence rate is contributed from two parts: convergence of ˆ α i (Algorithm A1) and convergence of personalized model ˆ v i (Algorithm 1). Notice that, for the convergence rate of ˆ v i , we roughly recover the optimal rate of shuffling SGD [55], which is O (1 /R 2 ) . However, we suffer from a O ( δ 2 /R ) term since each client runs vanilla SGD on their local data (the SGD-Update 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 two 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 a theoretical perspective, we are mostly interested in investigating whether the algorithm can converge to the true optimal solution of Φ( α ∗ i , v ) if we only optimize on a surrogate function Φ(ˆ α i , v ) .

One drawback of Algorithm 1 is that we have to wait for Algorithm A1 to finish and output ˆ α i , so that we can proceed with Algorithm 1. However, if we are not satisfied with the precision of ˆ α i , we 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 a 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 α i s and v i s as depicted in Algorithm 2 by intertwining the two stages in Algorithm A1 and Algorithm 1 in a 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 1 to 4 to be satisfied. Assume α ∗ i is the solution of (2) . Then if we run Algorithm 2 with η = Θ ( log( NKR 3 ) µR ) and γ = Θ ( log( NKR 3 ) µR ) , it will output the solution ˆ v i ,

∀ i ∈ [ N ] , such that with probability at least 1 -p , the following statement holds:

where κ = L µ , κ g = n max 2 λ , κ Φ = √ NG µ and the expectation is taken over the randomness of stochastic samples in Algorithm 2. That is, to guarantee E [Φ( α ∗ i , ˆ v i ) -Φ( α ∗ i , v ∗ i )] ≤ ϵ , we choose M = O ( L 2 κ 2 κ 2 g κ 2 Φ ¯ ζ i ( w ∗ ) δ 2 µ 2 ϵ ) , R = O ( max { Lδ 2 µϵ , κ 2 Φ κ 2 κ 2 g L 3 ¯ ζ i ( w ∗ ) DG ϵ }) and T α = O ( κ g log ( LR 2 ϵ )) .

Compared to Theorem 2, we achieve a slightly worse rate, since we need a large mini-batch when we update global model w . However, the advantages of the single-loop algorithm are two-fold. First, as we mentioned in the previous subsection, we have the freedom to optimize ˆ α i to arbitrary accuracy, while in double loop algorithm (Algorithm A1 + Algorithm 1), once we get ˆ α i , we do not have the chance to further refine it. Second, in practice, a 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 a 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 to support learning in model-heterogeneous 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 a 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 submodel updates sent from the clients to update the global model. We can consider the complexity of











































(a) Personalized Accuracy









































(b) Personalized Loss

models used by clients when estimating mixing parameters by solving a modified version of (2) as:

where we simply upper bounded the Rademacher complexity w.r.t. each data source in (GEN) with VC dimension [57]. Here m i is the masking operator to 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 i ⊙ w ∗ versus m j ⊙ 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 a distinct aggregated empirical loss for each client by interleaving permutations and shuffling models, we can utilize different model spaces W i , 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 it notably outshone other renowned methods as evident in Figure 1. Our experiments concluded with the CIFAR10 dataset, employing a 2-layer convolutional neural network, where PERM, despite a warmup 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 a total of N clients, where we draw samples from the distribution N ( µ 1 , Σ i ) for half of the clients, denoted by i ∈ [1 , N 2 ] , and from N ( µ 2 , Σ i ) for the remaining clients, denoted by i ∈ ( N 2 , N ] . Following the approach outlined in [58], we adopt a uniform variance for all samples, with Σ k,k = k -1 . 2 . Subsequently, we generate a labeling model using the distribution N ( µ w , Σ w ) .

Given a data sample x ∈ R d , the labels are generated as follows: clients 1 , ..., N 2 assign labels based on y = sign( w ⊤ x ) , while clients N 2 +1 , ..., N assign labels based on y = sign( -w ⊤ x ) . For this specific experiment, we set µ 1 = 0 . 2 , µ 2 = -0 . 2 , and µ w = 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 a logistic regression model on each client's data.

To demonstrate the superiority of our PERM algorithm, we conducted a 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 in highly heterogeneous scenarios, as seen in personalized accuracy and loss. This showcases PERM's exceptional ability to leverage relevant client learning.

0

50

100

150

200

250

300

Number of Communications

70.0

72.5

75.0

77.5

80.0

82.5

85.0

87.5

90.0

Personal Validation Accuracy

PERM

Localized FedAvg

PerFedAvg

pFedME

(a) Personalized Accuracy

0

50

100

150

200

250

300

Number of Communications

0.25

0.30

0.35

0.40

0.45

0.50

0.55

0.60

0.65

Personal Validation Accuracy

PERM

Localized FedAvg

PerFedAvg

pFedME

(b) Personalized Loss

Experiment on CIFAR10 dataset. We extend our experimentation to the CIFAR10 dataset using a 2-layer convolutional neural network. During this test, 50 clients participate, each limited to data from just 2 classes, resulting in a pronounced heterogeneous data distribution. We benchmark our algorithm against PerFedAvg, PFedMe, and the localized FedAvg. As illustrated in Figure 2, PERM demonstrates superior convergence performance compared to other personalized strategies. It's noteworthy that PERM's initial personalized validation is significantly lower than that of approaches like PerFedAvg and PFedMe. This discrepancy stems from our choice to implement 10 communica-

Time to run one round of communication (seconds)

1.0 -

0.8

0.6

0.4

0.2

0°0

PERM

Localized FedAvg

PerFedAvg

pFedMe

tion rounds as a warm-up phase before initiating personalization, whereas other models embark on personalization right from the outset.

Computational overhead. In demonstrating the computational efficiency of the proposed PERM algorithm, we present a comparison of wall-clock time of 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 a notably minimal computational overhead. The run-time is slightly worse due to overhead of estimating mixing parameters.

5 Discussion & Conclusion

This paper introduces a new data&system-aware 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, we propose a 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 to 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 is the same as running FedAvg in a two-stage approach and an extra communication in an interleaved approach). The only hurdle would be the required memory at server to maintain mixing parameters, which scales proportionally to the square of the number of clients, which can be alleviated by clustering devices which we leave as a future work.

References

Peter Kairouz, H. Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Keith Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, Rafael G. L. D'Oliveira, Salim El 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ˇ cný, 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, Praneeth Vepakomma, Jianyu Wang, Li Xiong, Zheng Xu, Qiang Yang, Felix X. Yu, Han Yu, and Sen Zhao. Advances and open problems in federated learning. Foundations and Trends® in Machine Learning , 2021.

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.

Chengxi Li, Gang Li, and Pramod K Varshney. Federated learning with soft clustering. IEEE Internet of Things Journal , 9(10):7773-7782, 2021.

Avishek Ghosh, Jichan Chung, Dong Yin, and Kannan Ramchandran. An efficient framework for clustered federated learning. Advances in Neural Information Processing Systems , 33:1958619597, 2020.

Jie Ma, Guodong Long, Tianyi Zhou, Jing Jiang, and Chengqi Zhang. On the convergence of clustered federated learning. arXiv preprint arXiv:2202.06187 , 2022.

Hubert Eichner, Tomer Koren, Brendan Mcmahan, Nathan Srebro, and Kunal Talwar. Semicyclic stochastic gradient descent. In Proceedings of the 36th International Conference on Machine Learning, PMLR , volume 97, 2019.

Canh T Dinh, Nguyen Tran, and Tuan Dung Nguyen. Personalized federated learning with moreau envelopes. Advances in Neural Information Processing Systems , 33, 2020.

Yutao Huang, Lingyang Chu, Zirui Zhou, Lanjun Wang, Jiangchuan Liu, Jian Pei, and Yong Zhang. Personalized federated learning: An attentive collaboration approach. arXiv preprint arXiv:2007.03797 , 2020.

Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning: A meta-learning approach. Advances in Neural Information Processing Systems , 2020.

Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet S Talwalkar. Federated multi-task learning. In Advances in Neural Information Processing Systems , pages 4424-4434, 2017.

Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Adaptive personalized federated learning. arXiv preprint arXiv:2003.13461 , 2020.

Filip Hanzely, Boxin Zhao, and Mladen Kolar. Personalized federated learning: A unified framework and universal optimization techniques. arXiv preprint arXiv:2102.09743 , 2021.

Tony Lancaster. The incidental parameter problem since 1948. Journal of econometrics , 95(2):391-413, 2000.

Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efficient learning of deep networks from decentralized data. In Artificial intelligence and statistics , pages 1273-1282. PMLR, 2017.

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 on Machine Learning , pages 5132-5143. PMLR, 2020.

Jenny Hamer, Mehryar Mohri, and Ananda Theertha Suresh. Fedboost: A communicationefficient algorithm for federated learning. In International Conference on Machine Learning , pages 3973-3983. PMLR, 2020.

Farzin Haddadpour, Mohammad Mahdi Kamani, Aryan Mokhtari, and Mehrdad Mahdavi. Federated learning with compression: Unified analysis and sharp guarantees. In International Conference on Artificial Intelligence and Statistics , pages 2350-2358. PMLR, 2021.

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.

Haibo Yang, Xin Zhang, Prashant Khanduri, and Jia Liu. Anarchic federated learning. In International Conference on Machine Learning , pages 25331-25363. PMLR, 2022.

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.

Chuhan Wu, Fangzhao Wu, Lingjuan Lyu, Yongfeng Huang, and Xing Xie. Communicationefficient federated learning via knowledge distillation. Nature communications , 13(1):1-8, 2022.

Chaoyang He, Murali Annavaram, and Salman Avestimehr. Group knowledge transfer: Federated learning of large cnns at the edge. Advances in Neural Information Processing Systems , 33:14068-14080, 2020.

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.

Sohei Itahara, Takayuki Nishio, Yusuke Koda, Masahiro Morikura, and Koji Yamamoto. Distillation-based semi-supervised federated learning for communication-efficient collaborative training with non-iid private data. IEEE Transactions on Mobile Computing , 22(1):191-205, 2021.

Enmao Diao, Jie Ding, and Vahid Tarokh. Heterofl: Computation and communication efficient federated learning for heterogeneous clients. arXiv preprint arXiv:2010.01264 , 2020.

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:12876-12889, 2021.

Sebastian Caldas, Jakub Koneˇ cny, H Brendan McMahan, and Ameet Talwalkar. Expanding the reach of federated learning by reducing client resource requirements. arXiv preprint arXiv:1812.07210 , 2018.

Samiul Alam, Luyang Liu, Ming Yan, and Mi Zhang. Fedrolex: Model-heterogeneous federated learning with rolling sub-model extraction. Advances in Neural Information Processing Systems , 35:29677-29690, 2022.

Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jennifer Wortman Vaughan. A theory of learning from different domains. Machine learning , 79:151-175, 2010.

Yishay Mansour and Mariano Schain. Robust domain adaptation. Annals of Mathematics and Artificial Intelligence , 71(4):365-380, 2014.

Nikola Konstantinov and Christoph Lampert. Robust learning from untrusted sources. In International conference on machine learning , pages 3488-3498. PMLR, 2019.

Koby Crammer, Michael Kearns, and Jennifer Wortman. Learning from multiple sources. Journal of Machine Learning Research , 9(8), 2008.

Mathieu Even, Laurent Massoulié, and Kevin Scaman. On sample optimality in personalized collaborative and federated learning. In NeurIPS 2022-36th Conference on Neural Information Processing System , 2022.

Jakub Koneˇ cn` y, 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.05492 , 2016.

Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh. Agnostic federated learning. In International Conference on Machine Learning , pages 4615-4625. PMLR, 2019.

Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Distributionally robust federated averaging. Advances in neural information processing systems , 33:15111-15122, 2020.

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.

Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank J Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for on-device federated learning. International Conference on Machine Learning , 119:5132-5143, 2020.

Farzin Haddadpour and Mehrdad Mahdavi. On the convergence of local descent methods in federated learning. arXiv preprint arXiv:1910.14425 , 2019.

Tao Yu, Eugene Bagdasaryan, and Vitaly Shmatikov. Salvaging federated learning by local adaptation. arXiv preprint arXiv:2002.04758 , 2020.

Bharath K Sriperumbudur, Kenji Fukumizu, Arthur Gretton, Bernhard Schölkopf, and Gert RG Lanckriet. On integral probability metrics, \ phi-divergences and binary classification. arXiv preprint arXiv:0901.2698 , 2009.

Steve Hanneke and Samory Kpotufe. A no-free-lunch theorem for multitask learning. The Annals of Statistics , 50(6):3119-3143, 2022.

Shai Shalev-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algorithms . Cambridge university press, 2014.

Zi Xu, Huiling Zhang, Yang Xu, and Guanghui Lan. A unified single-loop alternating gradient projection algorithm for nonconvex-concave and convex-nonconcave minimax problems. Mathematical Programming , pages 1-72, 2023.

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 6454-6462, 2022.

Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data. arXiv preprint arXiv:1806.00582 , 2018.

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.04723 , 2022.

Sebastian U Stich. Local sgd converges fast and communicates little. In International Conference on Learning Representations , 2018.

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.

Blake E Woodworth, Kumar Kshitij Patel, and Nati Srebro. Minibatch vs local sgd for heterogeneous distributed learning. Advances in Neural Information Processing Systems , 33:6281-6292, 2020.

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. PMLR, 2022.

Honglin Yuan and Tengyu Ma. Federated accelerated stochastic gradient descent. Advances in Neural Information Processing Systems , 33:5332-5344, 2020.

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-11092, 2019.

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.

Kwangjun Ahn, Chulhee Yun, and Suvrit Sra. Sgd with shuffling: optimal rates without component convexity and large epoch requirements. Advances in Neural Information Processing Systems , 33:17526-17535, 2020.

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.03109 , 2023.

Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar. Foundations of machine learning . MIT press, 2018.

Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. arXiv preprint arXiv:1812.06127 , 2018.

Sebastian Caldas, Peter Wu, Tian Li, Jakub Koneˇ cn` y, H Brendan McMahan, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for federated settings. arXiv preprint arXiv:1812.01097 , 2018.

Tianyi Lin, Chi Jin, and Michael I Jordan. On gradient descent ascent for nonconvex-concave minimax problems. arXiv preprint arXiv:1906.00331 , 2019.

Markus Schneider. Probability inequalities for kernel embeddings in sampling without replacement. In Artificial Intelligence and Statistics , pages 66-74. PMLR, 2016.

A Additional Experiments

In addition to experiments on synthetic and CIFAR10 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 our 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 CIFAR10 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 MLP 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 data according to the similarity of the data distribution between clients. The learned values of α , 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 algorithms. As it can be concluded, in terms of performance, PERM consistently excels beyond its peers, demonstrating exemplary results on various benchmark datasets.







































(a) Personalized Accuracy











































(b) Personalized Loss







             

























































Figure 5: The heat map of the learned α 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 data, which also highlight that the distribution of data is not significantly heterogeneous in this dataset.

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 α 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 1 class per client (highly heterogeneous). As it can be seen from Figure 6, when the data distribution is homogeneous the learned values of α 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 a subset of other clients. This discernible pattern reinforces the inherent confidence that each client is effectively learning from a limited but strategically chosen group of clients.







             

























































(a) Homogeneous distribution







             























































(b) Highly heterogeneous distribution

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 permutationbased variant of distributed SGD with periodic communication).

B.1 Technical Lemmas

Lemma 1. Define v ∗ ( α ) := arg min v ∈W Φ( α , v ) , and assume Φ( α , · ) is µ -strongly convex and ∇ v Φ( α , v ) is L Lipschitz in α . Then, v ∗ ( · ) is κ -Lipschitz where κ = L/µ .

Proof. The proof is similar to Lin et al's result on minimax objective [60]. First, according to optimality conditions we have:

Substituting v with v ∗ ( α ′ ) and v ∗ ( α ) in the above first and second inequalities respectively yields:

Adding up the above two inequalities yields:

Since Φ( α , · ) is µ strongly convex, we have:

Adding up (4) and (5) yields:

Finally, using L smoothness of Φ will conclude the proof:

Lemma 2 (Optimality Gap) . Let Φ( α , v ) be defined in (3). Let ˆ v = P W (˜ v -1 L ∇ v Φ(ˆ α , ˜ v )) . If we assume each f i is L -smooth, µ -strongly convex and with gradient bounded by G , then the following statement holds true:

where κ Φ = √ NG µ , v ∗ = arg min v ∈W Φ( α ∗ , v ) .

Proof. First we show that ∇ v Φ( α , v ) is √ NG Lipschitz in α . To see this:

√

Hence due to Lemma 1, we know v ∗ ( α ) is κ Φ := NG µ Lipschitz. According to property of projection, we have:

For T 1 , we notice:

where at last step we used Young's inequality. To bound ♠ , we apply the L smoothness and µ strongly convexity of Φ( α , · ) :

Putting above bound back yields:

Now we switch to bounding T 2 . Applying Cauchy-Schwartz yields:

where at last step we apply √ NG smoothness of Φ( · , v ) . Putting pieces together yields:

Re-arranging terms and setting v = v ∗ ( α ∗ ) = arg min v ∈W Φ( α ∗ , v ) yields:

At last, due to the κ Φ -Lipschitzness property of of v ∗ ( · ) as shown in Lemma 1, it follows that:

as desired.

Input: Number of clients N , number of local steps K , number of communications rounds R for r = 0 , . . . , R -1 do parallel for client i = 0 , ..., N -1 do

r +1

Server broadcasts w to all clients. Server computes ˆ α i , i = 1 , 2 , . . . , N by running T α steps of GD on g i ( w R , α ) . Output: ˆ α 1 , . . . , ˆ α N .

B.2 Proof of Convergence of Theorem 1

In this section we are going to prove the result in Theorem 1. To this end, we need to 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 A1) converges to optimal values. Notice that in Algorithm A1 we do not solve g i ( w ∗ , α ) directly, but optimize g i ( w R , α ) on α for T α iterations of GD. Hence, firstly we need to show that optimizing the surrogate function will also guarantee the convergence of output of algorithm ̂ α to α ∗ by deriving a property of the objective in (2). Formally the property is captured by the following lemma.

Lemma 3. Let g ( w , α ) := ∑ N j =1 α j ∥∇ f i ( w ) -∇ f j ( w ) ∥ 2 + λ ∑ N j =1 α 2 j /n j and α ∗ g ( w ) = arg min α ∈ ∆ N g ( w , α ) . Let w R be the output of Algorithm A1. Then the following statement holds:

∥ ∥ where κ g := n max 2 λ .

Proof. Define function

Apparently, W ( z , α ) is linear in z and 2 λ n max strongly convex in z . Next we show that ∇ α W ( z , α ) is Lipschitz in w . To see this,

Then, according to Proposition 1, α ∗ W ( z ) := arg min α ∈ ∆ N W ( z , α ) is κ g lipschitz in z where κ g = n max 2 λ , i.e., ∥ α ∗ W ( z ) -α ∗ W ( z ′ ) ∥ ≤ κ g ∥ z -z ′ ∥ . Now, let us consider the objective (2):

We define α ∗ g ( w ) = arg min α ∈ ∆ N g ( w , α ) .

We set

Then we know that

Since ∥ ∥ ∇ f i ( w R ) -∇ f j ( w R ) ∥ ∥ ≤ ∥∇ f i ( w ∗ ) -∇ f j ( w ∗ ) ∥ + 2 L ∥ ∥ w R -w ∗ ∥ ∥ , we can conclude that

With above lemma, to show the convergence of ˆ α to α ∗ , we do the following decomposition

Now it remains to show the convergence of Local SGD last iterate w R to optimal solution w ∗ . By convention, we use w t = 1 N ∑ N i =1 w t i 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 ∈ [ T ] :

Proof. According to updating rule in Algorithm A1, we have the following identity:

For T 1 , since each f j is L smooth and µ strongly convex, we have:

Due to Jensen's inequality we know: -1 N ∑ N i =1 µ 2 ∥ w t i -w ∗ ∥ 2 ≤ -µ 2 ∥ w t -w ∗ ∥ 2 . Hence we know:

For T

Now, plugging T and T

Lemma 5. [50, Lemma 8] For the iterates { w r,t i } generated in Algorithm A1, the following statement holds true:

Lemma 6 (Last iterate convergence of Local SGD) . Under the conditions of Theorem 1, the following statement holds true for the iterates in Algorithm A1:

Proof. We first unroll the recursion in Lemma 4 from t = K to 0 , within one communication round:

Since we choose γ ≤ 1 2 L , we know ∑ K -1 t =0 (1 -µγ ) K -t (2 γ -4 γ 2 L ) E ( F ( w ∗ ) -F ( w r,t )) ≥ 0 . Plugging in Lemma 5 yields:

Algorithm A2: Shuffling Local SGD ( One Client)

Input: Clients 0 , ..., N -1 , Number of Local Steps K , Number of Epoch R , Mixing parameter ˆ α

Plugging in γ = log( RK ) µRK gives the convergence rate:

Equipped with above results, we are now ready to provide the convergence of main theorem.

Proof of Theorem 1. The proof simply follows from Lemma 3:

Plugging in the convergence of ∥ ∥ w R -w ∗ ∥ ∥ 2 from Lemma 6, and the stepsize η α = 1 L g for α yields:

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 j once, while here we have to take K steps of SGD update on each component function. Second, we are considering a weighted sum objective in contrary to averaged objective in [56], which means we need to rescale the objective when we apply without-replacement 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.

Proposition 1. Assume a sequence { w t } K t =1 is obtained by

then we have

where δ t := ∇ f ( w t ; ξ t ) -∇ f ( w t ) , and by convention, we define ∏ b j = a A j = I if a < b .

Proof. According to updating rule, we have:

Since f is L smooth, and according to Mean Value Theorem, there is a matrix H t satisfying µ I ⪯ H t ⪯ L I , such that ∇ f ( w t ) -∇ f ( w 0 ) = H t ( w t -w 0 ) . Hence we have:

Unrolling the recursion from t to 0 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 no confusion, we drop the superscript r in σ r .

Lemma 7 (One epoch updating rule) . Let v r and v r +1 be two iterates generated by Shuffling Local SGD (Algorithm A2), then the following updating rule holds:

where

by convention, we define ∏ b j = a A j = I if a < b .

Proof. According to Proposition 1, we have

Plugging our definition of Q j and δ j yields:

Following the same reasoning in the proof of Proposition 1 will conclude the proof.

Lemma 8 (Summation by parts) . Let A j and B j be complex valued matrices. Then the following fact holds:

Proposition 2 (Spectral bound of polynomial expansion) . Given a collection of matrices { A t } and { B t } , such that A t ⪯ L I and B t ⪯ L I , the following bound hold:

Proof. We start with proving the first statement. Expanding the product yields:

Hence we have:

According to the upper bound for binomial coefficients: ( h -l m ) ≤ ( e ( h -l ) m ) m , we have:

Then we switch to the second one. Using the same expanding product yields:

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 ([61, Theorem 2]) . Suppose n ≥ 2 . Let g 1 , g 2 , . . . , g n ∈ R d satisfy ∥ g j ∥ ≤ G for all j . Let ¯ g = 1 n ∑ n j =1 g j . Let σ ∈ S n be a uniform random permutation of n elements. Then, for i ≤ n , with probability at least 1 -p , we have

Lemma 10 (Concentration of partial sum of gradients) . Given a uniformly randomly generated permutation σ , and simplex vector α , if we assume each sup v ∈W ∥∇ f j ( v ) ∥ ≤ G , then the following statement holds true:

Proof. The proof works by re-writing weighted sum of vectors to average of the these vectors:

Proposition 3 (Spectral norm bound of Q ) . Let Q j be defined in (11). Then the following bound for the spectral norm of Q j holds true for all j ∈ [ N ] :

Proof. The proof can be completed by writing down the definitin of Q j and applying CauchySchwartz inequality:

The last step is due to we choose η such that ηNL ≤ 1 K .

The following lemma establishes the bound regarding cumulative update between two epochs, namely, v r +1 -v r . In particular, Lemma 11 below shows that: (a) in shuffling Local SGD, our update from v r to v r +1 approximates performing NK times of gradient descent with ˆ α ( j ) N ∇ f σ ( j ) ( v r ) , namely, the bias is controlled, and (b) the update itself is bounded, and can be related to the norm of full gradient.

Lemma 11. During the dynamic of Algorithm A2, the following statements hold true with probability at least 1 -p :

(a)

(b) for any N ′ such that 0 ≤ N ′ < N

where

Proof. We start with proving statement (a). Let A j = Q j ˆ α ( σ ( j )) and B j = ˆ α ( σ ( j )) ∇ f σ ( j ) ( v r ) , applying the identity of summation by parts yields:

According to Proposition 2, we have:

Since we choose η ≤ 1 4 NKRL , we have:

where we use the fact that ∑ K -2 -τ m =1 ( e 4 Rm ) m ≤ ∑ K -2 -τ m =1 ( e 4 R ) m ≤ e 4 R 1 1 -e/ 4 R . Hence we know:

Thus we have:

For T 2 , we first examine the bound of Q n +1 ˆ α ( σ ( n +1)) -Q n ˆ α ( σ ( n )) :

where we evoke Proposition 2 at last step. Given that η ≤ 1 4 NKRL we have:

where we use the reasoining in (12). Hence for T 2 :

where at last step we evoke Lemma 10. So we can conclude T 2 ≤ 2 η 2 N 2 K 2 ( 2 e 4 R -e ) 2 ( G 2 8 N log(1 /p ) + ∥∇ Φ(ˆ α , v r ) ∥ 2 ) . Putting the bounds of T 1 and T 2 together will conclude the proof for (a).

Now we switch to proving (b). Once again by the summation of parts identity we have:

Taking the norm of both side yields:

Plugging our developed bound for ∥ Q N ′ ∥ and ∥ ∥ ∥ ∑ N ′ -1 n =1 ( Q n +1 ˆ α ( σ ( n +1)) -Q n ˆ α ( σ ( n )) )∥ ∥ ∥ yields:

where at last step we evoke Lemma 10. And for C, we use the similar reasoning:

Putting these pieces together yields:

Lemma 12. During the dynamic of Algorithm A2, the following statements hold true with probability at least 1 -p :

Proof. We first apply Cauchy-Schwartz inequality:

We proceed by applying the bound from Lemma 11 (b):

Therefore, it follows that:

Lemma 13 (Noise bound) . During the dynamic of Algorithm A2, the following statement for gradient noises holds true with probability at least 1 -p :

where

Proof. According to triangle and Cauchy-Schwartz inequalities we have:

B.4 Proof of Theorem 2

Proof. For notational convenience, let us define

Then we recall the updating rule of v (Lemma 7):

Hence we have:

Now, applying strongly convexity of Φ(ˆ α , · ) and Cauchy-Schwartz inequality yields:

Since Φ(ˆ α , · ) is L smooth, we have: E ∥∇ Φ(ˆ α , v r ) ∥ 2 ≤ 2 L E [Φ( ˆ α , v r ) -Φ(ˆ α , v ∗ ( ˆ α ))] . Therefore, we have:

Now, we examine the term ∥ g r -ηNK ∇ Φ(ˆ α , v r ) ∥ 2 . First according to summation by part (Lemma 8) by letting A j := ∏ j +1 j ′ = N ( I -Q j ′ H j ′ ) and B j = Q j ∇ f σ ( j ) ( v r ) , we have:

Hence we have:

where in (1) we apply Jensen's inequality, in (2) we plug in Lemma 11 (a), and Lemma 12, and in (3) we use the L -smoothness of Φ . Plugging above bound back in (19) yields:

Since we choose η = 4 log( NKR ) µNKR , and large enough epoch number:

we know that T 1 ≤ 0 . We thus have:

Unrolling the recursion from r = R to 0 :

Plugging in our choice of η will conclude the proof:

Finally, according to Lemma 2 we can complete the proof:

where we plug in the convergence result from Theorem 1 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 w t to be the t th iterate of minibatch SGD on smooth and strongly convex function F , with minibatch size M and learning rate γ . Also assume the variance is bounded by δ . Then the following statements hold true after T iterations of SGD:

Lemma 14 (Bounded iterates difference of α ) . Let { α r i } be iterates generated by Algorithm 2, then under conditions of Theorem 3, the following statement holds:

Proof. Define

According to updating rule of α in Algorithm 2 and Lemma 3 we have:

r r

where at last step we plug in Proposition 4 (16).

Lemma 15 (Convergence of α ) . Let { ˆ α i } N i =1 be the mixing parameters generated by Algorithm 2. Then under the conditions of Theorem 3, the following statement holds:

Proof. We notice the following decomposition:

where in the second inequality we apply Lemma 3, and in the third inequality we use Proposition 4 (17).

C.2 Proof of Theorem 3

Proof. According to Lemma 2, we have:

We first examine the convergence of ∥ ∥ v R -v ∗ ( ˆ α i ) ∥ ∥ 2 . Applying Cauchy-Schwartz inequality yields:

where a = 1 µηNK , and last step is due to that v ∗ ( α ) is κ Φ := √ NG µ Lipschitz, as proven in Lemma 2 . Similar to the proof of Theorem 2, we first define

Then we recall the updating rule of v :

Hence we have:

Now, applying strongly convexity of Φ( α r , · ) and Cauchy-Schwartz inequality yields:

where in the first inequality we applied Cauchy-Schwartz inequality and strongly convexity. Since Φ( α r , · ) is L smooth, we have: E ∥∇ Φ( α r , v r ) ∥ 2 ≤ 2 L E [Φ( α r , v r ) -Φ( α r , v ∗ ( α r ))] . Therefore, it follows that:

Now, we examine the term ∥ g r -ηNK ∇ Φ( α r , v r ) ∥ 2 in the right hand side of abovee inequality. First according to summation by part (Lemma 8): we let A j := ∏ j +1 j ′ = N -1 ( I -Q j ′ H j ′ ) and B j = Q j ∇ f σ ( j ) ( v r ) , then we have:

Hence we have:

where in (1) we apply Jensen's inequality, in (2) we plug in Lemma 11 (a), and Lemma 12, and in (3) we use the L -smoothness of Φ . Plugging above bound back in (19) yields:

Since we choose η = 4 log( √ NKR ) µNKR , and

hence we have:

Putting above inequality back to (18) yields:

where at second inequality we plug in Lemma 14. Unrolling the recursion from r = R to 0 , and plugging in η = 4 log( NKR 3 ) µNKR yields:

Plugging in γ = log( NKR 3 ) µR yields:

Since ˆ v i = v R and ˆ α i = α R , we have the convergence of ∥ ˆ v i -v ∗ ( ˆ α i ) ∥ 2 . Plugging this convergence rate together with the convergence of ∥ ˆ α i -α ∗ ∥ 2 from Lemma 15:

together with applying Lemma 2 leads to:

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\n1: 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.\n\nQUESTIONS:\n1: 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?\n\nLIMITATIONS:\nSee weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

comments: 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.

STRENGTHS:
The overall derivation of the algorithm is theory-driven and reasonable.

Good 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.
Does your algorithm work for a more general class of mixing ops?

The 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？

QUESTIONS:
See above

LIMITATIONS:
See above
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n+ 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).\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nTraining 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.\n\nLIMITATIONS:\nAs described in the Weakness section, I suspect the proposed method may lack scalability.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\n- 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).", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-31"}