Global Convergence of Federated Learning for Mixed Regression
Anonymous Author(s)
AfﬁliationAddress email
Abstract
This paper studies the problem of model training under Federated Learning whenclients exhibit cluster structure.
We contextualize this problem in mixed regression, where each client has limited local data generated from one of k unknownregression models.
We design an algorithm that achieves global convergence fromany initialization, and works even when local data volume is highly unbalanced –there could exist clients that contain O (1) data points only.
Our algorithm ﬁrst runsmoment descent on a few anchor clients (each with ˜ ⌦ ( k ) data points) to obtaincoarse model estimates.
Then each client alternately estimates its cluster labelsand reﬁnes the model estimates based on FedAvg or FedProx.
A key innovationin our analysis is a uniform estimate on the clustering errors, which we prove bybounding the VC dimension of general polynomial concept classes based on thetheory of algebraic geometry.
Introduction
Federated learning (FL) [MMR + 17] enables a massive number of clients to collaboratively trainmodels without disclosing raw data.
Heterogeneity in clients renders local data non-IID and highlyunbalanced.
For example, smartphone users have different preferences in article categories (e.g.politics, sports or entertainment) and have a wide range of reading frequencies.
In fact, distributionof the local dataset sizes is often heavy-tailed [DDPM13, FAC, FHK12].
Existing methods to deal with data heterogeneity can be roughly classiﬁed into three categories: acommon model, fully personalized models, and clustered personalized models; see Section 2 fordetailed discussion.
Using a common model to serve highly heterogeneous clients has fundamentaldrawbacks; recent work [SXY21] rigorously quantiﬁed the heterogeneity level that the commonmodel can tolerate with.
Fully personalized models [SCST17, MNB + 21] often do not come withperformance guarantees as the underlying optimization problem is generally hard to solve.
In thiswork, we focus on clustered personalized models [SMS20], i.e., clients within the same cluster sharethe same underlying model and clients across clusters have relatively different underlying models.
The main challenge is that the cluster identities of the clients are unknown.
We target at designingalgorithms that simultaneously learn the clusters and train models for each cluster.
A handful of simultaneous clustering and training algorithms have been proposed [MMRS20,
XLS + 21, LLV21, GCYR20], mostly of which are heuristic and lack of convergence guarantees[MMRS20, XLS + 21, LLV21].
Towards formal assurance, [GCYR20] studied this problem throughthe lens of statistical learning yet postulated a number of strong assumptions such as the initial modelestimates are very close to the true ones, linear models, strong convexity, balanced and high-volumeof local data, and sample splitting across iterations.
Their numerical results [GCYR20] suggestedthat many random initializations would lead to at least one good realization satisfying the required closeness assumption.
However, the necessary number of random initializations scales exponentially in both the input dimension and the number of clusters.
Besides, in practice, it ishard to recognize which initialization is good and to winnow out high quality trained model.
In thiswork, following [GCYR20], we adopt a statistical learning setup.
In particular, we contextualizeour problem as the canonical mixed regression, where each client has a set of local data generatedfrom one of k unknown regression models.
Departing from standard mixed regression, in which eachclient keeps one data point only [MN19, LL18], in our problem the sizes of local datasets can varysigniﬁcantly across clients.
We only make a mild assumption that there exist a few anchor clients(each with ˜ ⌦ ( k ) data points) and sufﬁciently many clients with at least two data points.
Similar mixedregression setup with data heterogeneity has been considered in [KSS + 20] in a different context ofmeta-learning; the focus there is on exploring structural similarities among a large number of tasks incentralized learning.
On the technical side, their analysis only works when the covariance matrices ofall the clusters are identical and each client has ⌦ (log k ) data points.
Please refer to Remark 1 formore detailed technical comparisons.
Contributions The main contributions of this work are summarized as follows:
• We design a two-phase federated learning algorithm to learn clustered personalized models
in the context of mixed regression problems.
In Phase 1, the parameter server runs afederated moment descent on a few anchor clients to obtain coarse model estimates based
on subspace estimation.
In each global iteration of Phase 2, each client alternately estimates
its cluster label and reﬁnes the model estimates based on FedAvg or FedProx.
The algorithm
works even when local data volume is highly unbalanced – there could exist clients that
contain O (1) data points only.
• We prove the global convergence of our algorithm from any initialization.
The proof is
built upon two key ingredients: 1) We develop a novel eigengap-free bound to control the
projection errors in subspace estimation; 2) To deal with the sophisticated interdependence
between the two phases and across iterations, we develop a novel uniform
estimate on theclustering errors, which we derive by bounding the VC dimension of general polynomialconcept classes based on the theory of algebraic geometry.
Our analysis reveals that theﬁnal estimation error is dominated by the uniform deviation of the clustering errors, whichis largely overlooked by the previous work.
Related Work
FedAvg [MMR + 17] is a widely adopted FL algorithm due to its simplicity and low communicationcost.
However, severe data heterogeneity could lead to unstable training trajectories and land insuboptimal models [LSZ + 20, ZLL + 18, KKM + 20].
Based on the number of models trained, existingmethods to deal with data heterogeneity can be roughly classiﬁed into three categories.
A Common Model: To limit the negative impacts of data heterogeneity on the obtained commonmodel, a variety of techniques based on variance reduction [LSZ + 20, LXC + 19, KKM + 20] andnormalization [WLL + 20] have been introduced.
Their convergence results mostly are derived understrong technical assumptions such as bounded gradient and/or bounded Hessian dissimilarity whichdo not hold when the underlying truth in the data generation is taken into account [LSZ + 20, LXC + 19,
KKM + 20].
In fact, none of them strictly outperform others in different instances of data heterogeneity[LDCH21].
Besides, the generalization errors of the common model with respect to local data aremostly overlooked except for a recent work [SXY21], which shows that the common model cantolerate only a moderate level of model heterogeneity.
Fully Personalized Models: [SCST17] proposed Federated Multi-Task Learning (MTL) whereindifferent models are learned for each of the massive population of clients [SCST17, MNB + 21].
Though conceptually attractive, the convergence behaviors of Federated MTL is far from wellunderstood because the objective is not jointly convex in the model parameters and the modelrelationships [SCST17, AEP08, AZB05].
Speciﬁcally, [SCST17] focused on solving the subproblemof updating the model parameters only.
Even in the centralized setting, convergence is only shownunder rather restricted assumptions such as equal dataset sizes for different tasks (i.e. balanced localdata)
[AZB05] and small number of common features [AEP08].
Moreover, the average excess errorrather than the error of individual tasks is shown to decay with the dominating term O
(1 / p n ) , wheren is the homogeneous local dataset size [AZB05].
Despite recent progress [TJJ21, DHK + 20], their results are mainly for linear representation learning and for balanced local data.
Parallel to Federated
MTL, model personalization is also studied under the Model-Agnostic Meta-Learning (MAML)framework [FMO20, JKRK19] where the global objective is modiﬁed to account for the cost ofﬁne-tuning a global model at individual clients.
However, they focused on studying the convergencein training errors only– no characterization of the generalization error is given.
Clustered Personalized Models: Clustered Federated Learning (CFL) [SMS20, GCYR20, GHYR19,
XLS + 21, MMRS20, LLV21] can be viewed as a special case of Federated MTL where tasks acrossclients form cluster structures.
In addition to the algorithms mentioned in Section 1, i.e., whichsimultaneously learn clusters and models for each cluster, other attempts have been made to integrateclustering with model training.
[SMS20] hierarchically clustered the clients in a post-processingfashion.
To recover the k clusters, ⌦ ( k ) empirical risk minimization problems need to be solvedsequentially – which is time-consuming.
[GHYR19] proposed a modular algorithm that contains oneshot clustering stage, followed by k individual adversary-resilient model training.
Their algorithmscales poorly in the input dimension, and requires local datasets to be balanced and sufﬁciently large(i.e.
, n  d 2 ).
Moreover, each client sequentially solves two empirical risk minimization problems.
To utilize information across different clusters, [LLV21] proposed soft clustering.
Unfortunately,the proposed algorithm is only tested on the simple MNIST and Fashion-MNIST datasets, and notheoretical justiﬁcations are given.
Problem Formulation
A FL system consists of a parameter server (PS) and M clients.
Each client i
[ M ] keeps adataset D i
= { ( x ij , y ij ) } n i j =1 that are generated from one of k unknown regression models.
Let
N =
P
Mi =1 n i .
The local datasets are highly unbalanced with varying n i across clients.
If n i = ˜ ⌦ ( k ) ,we refer to client i as anchor client, which corresponds to active user in practice.
Anchor clients playa crucial rule in our algorithm design.
We consider the challenging yet practical scenario wherein anon-anchor client may have O (1) data points only.
We adopt a canonical mixture model setup:
For each client i 2 [ M ] ,
[Equation]
where z i 2 [ k ] is the hidden local cluster label, ✓ ⇤ 1 , · · · , ✓ ⇤ k are the true models of the clusters, ( x i ) 2 R n i ⇥ d is the feature matrix with rows given by  ( x ij ) , y i =
( y ij ) 2 R n i is the responsevector, and ⇣ i =
( ⇣ ij ) 2 R n i is the noise vector.
Note that the feature map  can be non-linear(e.g.
polynomials).
The cluster label of client i is randomly generated from one of the k componentsfrom some unknown p =
( p 1 , . . .
, p k ) in probability simplex  k  1 .
That is, P { z i = j } = p
j forj 2 [ k ] .
In addition, k ✓ ⇤ j k 2  R for each component.
The feature covariate  ( x ij ) is independent andsub-Gaussian ↵ I d  E
[  ( x ij )  ( x ij ) > ]   I d .
We assume that the covariance matrix is identicalwithin the same cluster but may vary across different clusters, i.e., E [  ( x ij )  ( x ij ) > ] =
⌃ j if z i =
j.
The noise ⇣ ij is independent and sub-Gaussian with E [ ⇣ ij ] = 0 and E [ ⇣ 2 ij ]
  2 .
Our formulation accommodates statistical heterogeneity in feature covariates, underlying local models,and observation noises [KMA + 21].
For the identiﬁability of the true cluster models ✓ ⇤ j ’s, we assumea minimum proportion and a pairwise separation of clusters.
Formally, let  = min
j 6 = j 0 k ✓ ⇤ j  ✓ ⇤ j 0 k 2and p min = min
j 2 [ k ] p j .
For ease of presentation, we assume the parameters ↵ ,  = ⇥ (1) ,  /  =
O (1) , and R/  = O (1) , while our main results can be extended to show more explicit dependencieson these parameters with careful bookkeeping calculations.
Notations: Let [ n ] , { 1 , . . . , n } .
For two sets A and B , let A
B denote the symmetric difference( A  B ) [ ( B  A ) .
We use standard asymptotic notation: for two positive sequences { a n } and{ b n } , we write a n = O ( b n ) (or a n .
b n ) if a n  Cb n for some constant C and sufﬁciently large n ; 133 a n = ⌦ ( b n ) (or a n & b n ) if b n = O ( a n ) ; a n = ⇥ ( b n ) (or a n ⇣ b n ) if a n = O ( b n ) and a n = ⌦ ( b n ) ;
Poly-logarithmic factors are hidden in ˜ ⌦ .
Given a matrix A 2 R n ⇥ d , let A =
P ri =1  i u
i v >
i denote 135 its singular value decomposition, where r = min { n, d } ,  1  · · ·   r  0 are the singular values, 136 and u i ( v i ) are the corresponding left (right) singular vectors.
We call U =
[ u 1 , u 2 , . . .
, u k ] asthe top- k left singular matrix of A. Let span( U ) =
span { u 1 , . .
, u k } denote the k -dimensionalsubspace spanned by { u 1 , . . .
, u k } .
Main Results
We propose a two-phase FL algorithm that enables clients to learn the model parameters ✓ ⇤ 1 , . . .
, ✓ ⇤ kand their clusters simultaneously:
(i) Coarse estimation via FedMD.
Run the federated moment descent algorithm to obtain coarseestimates of model parameters ✓ ⇤ i ’s.
(ii) Fine-tuning via iterative FedX+clustering.
In each iteration, each client ﬁrst estimates its clusterlabel and then reﬁnes its local model estimate via either FedAvg or FedProx (which we refer to asFedX)
[MMR + 17, LSZ + 20].
The details of FedMD and FedX+clustering are speciﬁed in Phase 1 and Phase 2, respectively.
4.1 Federated moment descent
The key idea of the ﬁrst phase of our algorithm is to leverage the existence of anchor clients.
Speciﬁcally, the PS chooses a set H of n H anchor clients uniformly at random.
Each selectedanchor client i
H maintains a sequence of estimators { ✓ i,t } that approaches
✓ ⇤ z
i , achievingk ✓ i,t  ✓ ⇤ z
i k 2  ✏  for some small constant ✏ > 0 when t is sufﬁciently large.
At high-level, we hope to have ✓ i,t move along a well calibrated direction r i,t that decreases theresidual estimation error
k
⌃ z
i ( ✓ ⇤ z i  ✓ i,t ) k 2 , i.e., the variance of the residual h  ( x ij ) , ✓ ⇤  ✓ i,t i .
Assuch, we like to choose r i,t to be positively correlated with ⌃
z i ( ✓ ⇤ z i  ✓ i,t ) .
However, to estimate⌃ z i ( ✓ ⇤ z i  ✓ i,t ) solely based on the local data of anchor client i , it requires n i = ˜
⌦ ( d ) , which isunaffordable in typical FL systems with high model dimension d and limited local data.
To resolvethe curse of dimensionality, we decompose the estimation task at each chosen anchor client into twosubtasks: we ﬁrst estimate a k -dimensional subspace that ⌃
z i ( ✓ ⇤ z i  ✓ i,t ) lies in by pooling localdatasets across sufﬁciently many non-anchor clients; then we project the local data of anchor client ionto the estimated subspace and reduce the estimation problem from d -dimension to k -dimension.
The precise description of our Phase 1 procedure is given as below.
For ease of notation, let" ( x, y, ✓ ) , ( y  h  ( x ) , ✓ i )  ( x ) .
[Equation]
[Equation]
/* described in Algorithm 3 */
Each
calls power
iteration
i
t A > i,t
T
to output (ˆ  i
t
ˆ
i
t
with A i
t deﬁned in (
); if ˆ  i,t > ✏  then ✓ i,t + ✓ i,t + r i,t ⌘ i,t and reports ✓ i,t +1 , where r i,t
= ˆ
U i,t ˆ
i,t and ⌘ i,t = ↵ ˆ
i,t / (2  2 ) ;
In Step 9, PS estimates the subspace that the residual estimation errors { ⌃ j ( ✓ ⇤ j  ✓ i,t ) } kj =1 lie in, in collaboration with clients in
S t .
In particular, for each anchor client i 2 H , deﬁne
[Equation]
We approximate the subspace spanned by { ⌃ j ( ✓ ⇤ j  ✓ i,t ) } kj =1 via that spanned by the top- k left 165 singular vectors of Y i,t .
To compute the latter, we adopt the following multi-dimensional generalization of the power method, known as orthogonal-iteration [GVL13, Section 8.2.4].
In general,given a symmetric matrix Y 2 R d ⇥ d , the orthogonal iteration generates a sequence of matrices
Q t 2 R d ⇥ k as follows:
Q 0 2 R d ⇥ k is initialized as a random orthogonal matrix Q > 0 Q
=
I and
Y Q t =
Q t +1 R t +1 with QR factorization.
When t is large, Q t approximates the top- k left singualrmatrix of Y , provided the existence of an eigen-gap  k >  k
+1 .
When k
= 1 , this is just the thepower-iteration and we can further approximate the leading eigenvalue of Y by the Raleigh quotient
Q > t Y Q t .
When Y is asymmetric, by running the orthogonal iteration on Y Y > , we can computethe top- k left singular matrix of Y .
In our setting, the orthogonal iteration can be implemented in adistributed manner in FL systems as shown in Algorithm 3 in the Appendix.
In Step 11, each anchor client i estimates the residual error ⌃ z
i ( ✓ ⇤ z i  ✓ i,t ) by projecting " ( x ij , y ij , ✓ i,t )onto the previously estimated subspace, that is, ˆ U > i,t " ( x ij , y ij , ✓ i,t ) .
This reduces the estimation fromd -dimension to k -dimension and hence ˜ ⌦ ( k ) local data points sufﬁce.
Speciﬁcally, deﬁne
[Equation]
where D i,t consists of 2 ` local data points ( x ij , y ij ) and (˜ x ij , ˜ y ij ) freshly drawn from D i at iterationt .
Client i runs the power-iteration to output ˆ  i,t and ˆ  2 i,t as approximations of the leading leftsingular vector and singular value of A i,t , Then anchor client i updates a new estimate ✓ i,t +1 bymoving along the direction of the estimated residual error r i,t with an appropriately chosen step size⌘ i,t .
We show that ✓ i,T is close to ✓ ⇤
z i for every anchor client i 2 H and the outputs ˆ ✓ j are close to ✓ ⇤ j upto a permutation of cluster indices.
Theorem 1.
Let ✏ 2 (0 , 1 / 4) be a small but ﬁxed constant.
Suppose that
[Equation]
With probability at least 1  O ( N  9 ) , for all initialization ✓ 0 with k ✓ 0 k 2  R,
[Equation] [Equation]
Furthermore, when n H  log N/p min , with probability at least 1  O ( N  9 ) ,
[Equation]
(5) whered (ˆ ✓ , ✓ ⇤ ) , min
[Equation]
max j
[Equation]
2 , where ⇡ is permutation over [ k ] .
[Equation] [Equation]
Note that in (6) we take a minimization over permutation, as the cluster indices are unidentiﬁable.
Phase 1 uses fresh data at every iteration.
In total we need p  2min ˜ ⌦
( d ) clients with at least two datapoints and ˜ ⌦ (1 /p min ) anchor clients.
This requirement is relatively mild, as typical FL systems havea large number of clients with O (1) data points and a few anchor clients with moderate data volume.
We defer the detailed proof of Theorem 1 to Appendix D. A key step in our proof is to show theresidual estimation errors { ⌃ j ( ✓ ⇤ j  ✓ i,t ) } kj =1 approximately lie in span( ˆ U i,t ) .
Unfortunately, the 195 eigengap of Y i,t could be small, especially when ✓ i,t gets close to ✓
⇤ z i ; and hence the standard
Davis-Kahan theorem [DK70] cannot be be applied.
This issue is further exaggerated by the fact thatthe convergence rate of the orthogonal iteration also crucially depends on the eigengaps [GVL13].
For these reasons, span( ˆ U i,t ) may not be close to span { ⌃ j ( ✓ ⇤ j  ✓ i,t ) } kj =1 at all.
To resolve thisissue, we develop a novel gap-free bound to show that projection errors ˆ U > i,t ⌃ j ( ✓ ⇤ j  ✓ i,t ) are smallfor every j
[ k ] (cf.
Lemma 6).
Remark 1 (Comparison to previous work [LL18, KSS + 20]) .
Our algorithm is partly inspiredby [LL18] which focuses on the noiseless mixed linear regression, but deviates in a number of crucialaspects.
First, our algorithm crucially utilizes the fact that each client chosen in S t has at least twodata points and hence the space of the singular vectors of E [ Y i,t ] is spanned by { ⌃ j ( ✓ ⇤ j  ✓ i,t ) } kj =1 .
In 205 contrast, [LL18] relies on the sophisticated method of moments which only works under the Gaussianfeatures and requires exponential in k 2 many data points.
Second, our algorithm crucially exploitsthe existence of anchor clients and greatly simpliﬁes the moment descent algorithm in [LL18].
Our algorithm also bears similarities with the meta-learning algorithm in [KSS + 20], which also usesclients collectively for subspace estimation and anchor clients for estimating cluster centers.
However,there are several key differences.
First, [KSS + 20] focuses on the centralized setting and relies onone-shot estimation, under the additional assumption that the covariance matrix of features across allclusters are identical.
Instead, our moment descent algorithm is iterative, is amenable to a distributedimplementation in FL systems, and allows for covariance matrices varying across clusters.
Second,in the ﬁne-tuning phase, [KSS + 20] uses the centralized least squares to reﬁne the clusters estimatedwith anchor clients, under the additional assumption that ⌦ (log k ) data points for every client.
Incontrast, as we will show later, we use the FedX+clustering to iteratively cluster clients and reﬁnecluster center estimation.
4.2 FedX+clustering
At the end of Phase 1, only the selected anchor clients in H obtained coarse estimates of their localtrue models (characterized in (4)).
In Phase 2, both anchor clients in H and all the other clients(anchor or not) will participate and update their local model estimates.
Phase 2 is stated in a generic form for any loss function L ( ✓ ,  ; D ) , where ✓ = ( ✓ 1 , . . .
, ✓ k ) 2 R dk isthe cluster parameters,  2  k  1 represents the likelihood of the cluster identity of a client, and Ddenotes the client’s dataset.
This generic structure covers the idea of soft clustering [LLV21].
Notethat unlike Phase 1 where each anchor client i only maintains an estimate ✓ i,t of its own model,in Phase 2, each client i maintains model estimates ✓ i · ,t = ( ✓ i 1 ,t , . . .
, ✓ ik,t ) for all clusters.
1 Input the output of
Phase
⌘
2 Output:
=
) 3 PS sets
✓ do
PS broadcasts
1 to all clients
Each client i estimates the
likelihood
of its
local cluster
label by
[Equation]
Each client i reﬁnes its local model based on either FedAvg or FedProx with L i ( ✓ ) =
L ( ✓ ,  i,t ; D i ) , and reports the updated local parameters ✓ i · ,t = ( ✓ i 1 ,t , . .
. , ✓ ik,t ) .
• FedAvg-based: it runs s steps of local gradient descent: ✓ i · ,t G si ( ✓ t  1 ) , where G i ( ✓ ) =
✓  ⌘ r L i ( ✓ ) • FedProx-based: it solves the local proximal optimization:
[Equation]
i ·
[Equation]
PS updates the global model as ✓ t
P
Mi =1 w i ✓ i · ,t , where w i = n i /N .
In Phase 2, the local estimation at each client has a ﬂavor of alternating minimization:
It ﬁrst runs aminimization step to estimate its cluster, and then runs a FedAvg or FedProx update to reﬁne modelestimates.
To allow the participation of clients with O (1) data points only, at every iteration theclients are allowed to reuse all local data, including those used in the ﬁrst phase.
Similar alternatingupdate is analyzed in [GCYR20] yet under the strong assumption that the update in each round is over fresh data with Gaussian distribution.
Moreover, the analysis therein is restricted to the settingwhere the model reﬁnement at each client is via running a single gradient step, which is barely usedin practice but much simpler to analyze than FedAvg or FedProx update.
In our analysis, we consider the square loss
[Equation]
In this context, (7) yields a vertex of the probability simplex  ij,t = { j = z i,t } , where
[Equation]
The estimate z i,t provides a hard clustering label.
Hence, in each round, only one regression modelwill be updated per client.
To capture the tradeoff between communication cost and statistical accuracy using FedAvg or FedProx,we introduce the following quantities from [SXY21]:
[Equation]
FedAvg
[Equation]
FedProx
[Equation]
We choose a properly small learning rate ⌘ such that  < 1 .
Here,   1 quantiﬁes the stability oflocal updates.
Notably,  ⇡ 1 using a relatively small ⌘ .
For the learnability of model parameters, we assume that collectively there are sufﬁcient data ineach cluster.
In particular, we assume N j & d , where N j =
P i : z i = j n i denotes the numberof data points in cluster j .
To further characterize the quantity skewness (i.e., the imbalance ofdata partition n = ( n 1 , . .
. , n M ) across clients), we adopt the  2 -divergence, which is deﬁned as
[Equation]
Q
[Equation]
=
[Equation]
dQ for a distribution P absolutely continuous with respect to a distribution Q.
Let  2 ( n ) be the chi-squared divergence between data partition p n over the clients p n ( i )
=
n i /Nand
the uniform distribution over [ M ] .
Note that when data partition is balanced (i.e., n i = N/M forall i ), it holds that  2 ( n ) = 0 .
We have the following theoretical guarantee of Phase 2, where s is the number of local steps in
FedAvg.
Notably, s is an algorithmic parameter for FedAvg only.
To recover the results for FedProx,we only need to set s = 1 .
Theorem 2.
Suppose that k  2 , ⌘ . 1 /s , and N j & d .
Let ⇢ =
min j N j /N .
If ⌫ log( e/ ⌫ ) .
⇢ /  ,then with probability 1  Cke  d , for all t  T + 1 and all ✓ T such that d ( ✓ T , ✓ ⇤ )  ✏  , where ✏ issome constant, it holds that
[Equation]
where⌫
[Equation] [Equation]
and p e ( n i ) = 4 ke  cn i ⇣ 1 ^  2  2 ⌘ 2 .
Furthermore, if t  T + 1 , for each client i , with probability1  p e ( n i ) , it is true that
[Equation]
where ˆ ✓ i,t is client
i ’s estimate of its own model parameter at time t.
Notably, ˆ ✓ i,t is the z i,t -th entry of ✓ i · ,t .
Theorem 2 shows that the model estimation errors decaygeometrically starting from any realization that is within a small neighborhood of ✓ ⇤ .
The parameter⌫ captures the additional clustering errors injected at each iteration.
It consists of two parts: the ﬁrstterm of (10) bounds the clustering error in expectation which diminishes exponentially in the localdata size and the signal-to-noise ratio  /  ; the second term bounds the uniform deviation of the clustering error across all initialization and iterations.
Note that if the cluster structure were knownexactly, we would get an model estimation error of ✓ ⇤ j scaling as p d/N j .
However, it turns out thatthis estimation error is dominated by our uniform deviation bound of the clustering error and hence isnot explicitly shown in our bound (9).
In comparison, the previous work [GCYR20] assumes freshsamples at each iteration by sample-splitting and good initialization independent of everything elseprovided a prior; hence their analysis fails to capture the inﬂuence of the uniform deviation of theclustering error.
4.2.1 Analysis of global iterations
Without loss of generality, assume the optimal permutation in (6) is identity.
In this case, if z i,t = j ,then client i will reﬁne ✓ j,t  1 .
To prove Theorem 2, we need to analyze the global iteration of ✓ t .
Following a similar argument to [SXY21] with a careful examination of cluster labels, we obtain thefollowing lemma.
The proof is deferred to Appendix E.1.
Lemma 1.
Let  ( x ) be the matrix that stacks all  ( x i ) vertically, and similarly for y .
It holds that
[Equation]
where B = 1 N  ( x ) > P , P is a block diagonal matrix with i th block P i of size n i ⇥ n i given by
[Equation]
for FedAvg
[Equation]
for FedProx , and ⇤ j,t is another block diagonal matrix with i th block being  ij,t I n i .
Lemma 1 immediately yields the evolution of estiamtion error.
Let ⇤ j be the matrix with
i th blockbeing { z i
=
j } I n i representing the true client identities.
Plugging model (1), the estimation errorevolves as
[Equation]
[Equation]
where K j = B ⇤ j  ( x ) and E j,t = ⇤ j,t  ⇤ j .
The estimation error is decomposed into three terms:1) the main contribution to the decrease of estimation error; 2) the clustering error; and 3) the noisyperturbation.
Let I
j = { i : z i = j } be the clients belonging to j th cluster, and I j,t =
{ i :
z i,t = j } 287 be the clients with estimated label j .
The indices of nonzero blocks of
E j,t are I j
I j,t indicating the 288 clustering errors pertaining to j th cluster.
For the ease of presentation, we introduce a few additional notations for the collective data over asubset of clients.
Given a subset I ✓ [ M ] of clients, let  ( x I ) denote the matrix that vertically stacks ( x i ) for i 2 I , and we similarly use notations y I and ⇣ I ; let P I be the matrix with diagonal blocks
P i for i 2 I .
Using those notations, we have K j = 1 N  ( x I j ) >
P
I j  ( x I j ) , which differs from theusual covariance matrix by an additional matrix P
I j .
Therefore, the analysis of the ﬁrst and thirdterms on the right-hand side of (12) follows from standard concentration inequalities for randommatrices.
In the remaining of this subsection, we focus on the second term, which is a major challengein the analysis.
The proof details are all deferred to Appendix E.2.
Lemma 2.
There exists a universal constant C such that, with probability 1  Ce  d ,
[Equation]
Lemma 2 aims to upper bound the error of
[Equation]
where S j,t = I j
I j,t .
The technical difﬁculty arises from the involved dependency between theclustering error S j,t and the estimated parameter ✓ j,t  1 as estimating label
z i,t and updating ✓ j,t  1 301 use a common set of local data.
Proof Sketch of Lemma 2.
It follows from the deﬁnition of z i,t in (8) that
[Equation]
Then,
[Equation]
Hence, it sufﬁces to upper bound k  ( x S j,t ) k 2 and k ⇣ S j,t k 2 given a small estimation error d ( ✓ t  1 , ✓ ⇤ )
305 from the last iteration.
To this end, we show a uniform upper bound of the total clustering error
P i 2 S j,t n i by analyzing a weighted empirical process.
Using the decision rule (8), the set S j,t canbe written as a function S j ( ✓ t  1 ) with
[Equation]
where
[Equation]
Then we derive the following uniform deviation of the incorrectly clustered data points
[Equation]
This is proved via upper bounds on the Vapnik–Chervonenkis (VC) dimensions of the binary functionclasses
[Equation]
Using classical results of VC dimensions, those functions are equivalently intersections of hyperplanesin ambient dimension
O ( d 2 ) , which yields an upper bound O ( d 2 ) .
However, the hyperplanes arecrucially rank-restricted as the total number of parameters in ✓ is dk .
We prove that the VC dimensionsare at most O ( dk log k ) using the algebraic geometry of polynomials given by the celebrated MilnorThom theorem (see, e.g., [Mat13, Theorem 6.2.1]).
Consequently, k  ( x S j,t ) k 2 , k ⇣ S j,t k 2 and thus (15) 317 can be uniformly upper bounded using sub-Gaussian concentration and the union bound, concludingthe proof of Lemma 2.
4.3 Global convergence
Combining Theorem 1 and Theorem 2, we immediately deduce the global convergence from anyinitialization within the ` 2 ball of radius R.
Theorem 3.
Suppose the conditions of Theorem 1 and Theorem 2 hold.
Let ˆ ✓ be the output of ourtwo-phase algorithm by running Phase 1 with T = ⇥
(1) iterations starting from any initialization✓ 0 with k ✓ 0 k 2  R , followed by Phase 2 with T 0 = ⇥ (  s ⌘⇢ log  ⌫ ) iterations.
Then with probability1  N  9  cke  d , it is true that
[Equation]
Furthermore, for each client i , with probability 1  p e ( n i ) , it holds that k ˆ ✓ i,T + T 0  ✓ ⇤ z
i k 2 .⇢ ⌫ log e ⌫ .
To the best of our knowledge, this is the ﬁrst result that proves the global convergence of clusteredfederated learning from any initialization.
Our bound (18) reveals that the ﬁnal estimation erroris dominated by the clustering error captured by ⌫ , and scales linearly in  which characterizesthe stability of local updates under FedAvg or FedProx.
Moreover, Theorem 3 shows that Phase 1converges very fast with only ⇥ (1) iterations and hence is relatively inexpensive in both computationand communication.
Instead, the number of iterations needed for Phase 2 grows logarithmically in / ⌫ and linearly in  / ( s ⌘⇢ ) .
Thus, by choosing s relatively large while keeping  close to 1, FedAvgenjoys a saving of the total communication cost.
B Broad Impact
Federated Learning has gained tremendous popularity over the past few years and has been widelydeployed in products such as Apple’s Siri and Google’s Gboard.
It opens up a world of newopportunities for training machine learning models without compromising data privacy.
Our papersigniﬁcantly advances the theory and algorithms for Federated Learning, providing important designinsights and key enabling technologies for efﬁcient and personalized model training under FL.
Ourstudy is highly interdisciplinary, bringing together ideas from statistics, optimization, and distributedcomputation.
We are unaware of any potential negative societal impacts of our work.