Global Convergence of Federated Learning for Mixed Regression
Authors: Lili Su
Jiaming Xu
Pengkun Yang
Venue: NeurIPS 2022 Accept (2022)
Keywords: Federated Learning
clustering
global convergence
FedAvg
FedProx
empirical process

Abstract: This paper studies the problem of model training under Federated Learning when clients exhibit cluster structure. 
We contextualize this problem in mixed regression, where each client has limited local data generated from one of $k$ unknown regression models. We design an algorithm that achieves global convergence from any initialization, and works even when local data volume is highly unbalanced  -- there could exist clients that contain $O(1)$ data points only. Our algorithm first runs moment descent on a few anchor clients (each with $\tilde{\Omega}(k)$ data points) to obtain coarse model estimates. Then each client alternately estimates its cluster labels  and refines the model estimates based on FedAvg or FedProx. A key innovation in our analysis is a uniform estimate on the clustering errors, which we prove by bounding the VC dimension of general polynomial concept classes based on the theory of algebraic geometry.

Global Convergence of Federated Learning for Mixed Regression
Lili Su
Jiaming Xu
Electrical and Computer Engineering
The Fuqua School of Business
Northeastemn University I.su @ northeastternn
Duke University
jiaming.xu868 @ duke.edu
Pengkun Yang
Center for Statistical Science
Tsinghua University
yangpengkun @ tsinghua.edu.cm
Abstract
This paper studies the problem of model training under Federated Learning when clients exhibit cluster structure. We contextualize this problem in mixed regres- sion, where each client has limited local data generated from one of & unknown regression models. We design algorithm that achieves global convergence from any initialization, and works even when local data volume is highly unbalanced there could exist clients that contain (1) data points only. Our algorithm first runs moment descent on few anchor clients (each with f2(k) data points) obtain coarse model estimates. Then each client alternately estimates its cluster labels and refines the model estimates based on FedAvg or FedProx. key innovation in our analysis is uniform estimate on the clustering errors, which we prove by bounding the VC dimension of general polynomial concept classes based on the theory of algebraic geometry.
Introduction
Federated learning (FL) MMR + 17 enables a massive number of clients collaboratively train models without disclosing raw data. Heterogeneity in clients renders local data non-IID and highly unbalanced. For example, smartphone users have different preferences in article categories (e.g. politics, sports entertainment) and have wide range of reading frequencies. In fact, distribution of the local dataset sizes often heavy-tailed [DPMI]3 FAC 1HK
Existing methods to deal with data heterogeneity can be roughly classified into three categories: a common model, fully personalized models, and clustered personalized models; see Section - for detailed discussion. Using a common model to serve highly heterogeneous clients has fundamental drawbacks; recent work [SXY2] rigorously quantified the heterogeneity level that the common model can tolerate with. Fully personalized models [SCSTIZ MNB 21 often do not come with performance guarantees as the underlying optimization problem generally hard solve. In this work, we focus on clustered personalized models SMSS0, i.c., clients within the same cluster share the 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 designing algorithms that simultaneously learn the clusters and train models for each cluster.
A handful of simultancous clustering and training algorithms have been proposed [MRS20
X1 LLV21 GCYR20) mostly of which are heuristic and lack of convergence guarantees
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

MMRS20 XLS 21 LLV21 Towards formal assurance, [GCYR2 studied this problem through
the lens of statistical learning yet postulated number of strong assumptions such as the initial model estimates are very close to the true ones, linear models, strong convexity, balanced and high-volume of local data. Their numerical results [GCYR20] suggested that sufficiently 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 is hard to recognize and winnow out good initialization. In this work, following [GCYR20), we adopt a statistical learning setup. In particular, we contextualize our problem the canonical mixed regression, where each client has set of local data generated from one of k unknown regression models. Departing from standard mixed regression, in which each client keeps one data point only [MN19 LLTS in our problem the sizes of local datasets can vary significantly across clients. We only make a mild assumption that there exist a few anchor clients (each with (A)) data points) and sufficiently many clients with at least two data points. Similar mixed regression setup with data heterogeneity has been considered in [KSS+20) in a different context of meta-learning: the focus there is on exploring structural similarities among a large number of tasks in centralized learning. On the technical side, their analysis only works when the covariance matrices of all the clusters are identical and each client has Mlog k) data points. Please refer to Remark T for more 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 federated 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 refines 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 (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- 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 the clustering errors, which we derive by bounding the VC dimension of general polynomial concept classes based on the theory of algebraic geometry. Our analysis reveals that the final estimation error is dominated by the uniform deviation of the clustering errors, which is largely overlooked by the previous work.
2 Related Work
FedAvg 1M is a widely adopted FL algorithm due to its simplicity and low communication cost. However, severe data heterogeneity could lead to unstable training trajectories and land in suboptimal models 20 ZLL K0M+ Based on the number of models trained, existing methods to deal with data heterogeneity can be roughly classified into three categories.
Common Model: To limit the negative impacts of data heterogeneity on the obtained common model, variety of techniques based variance reduction L LXC+19 19 KKM+20) and normalization [W]] have been introduced. Their convergence results mostly are derived under strong technical assumptions such as bounded gradient and/or bounded Hessian dissimilarity which do not hold when the underlying truth in the data generation is taken into account LLZ LXC+ 19 2K In fact, none of them strictly outperform others in different instances of data heterogeneity LDCH2I] Besides, the generalization errors of the common model with respect to local data are mostly overlooked except for a recent work [SXY2I], which shows that the common model can tolerate only= a moderate level of model heterogeneity.
Fully Personalized Models: [SCSTIZ] proposed Federated Multi-Tare Learning (MTL) wherein different models are learned for each of the massive population of clients (SCSTI7 MNB+211. Though conceptually attractive, the convergence behaviors of Federated MTL. is far from well- understood because the objective is not jointly convex in the model parameters and the model relationships (SCST17 AEP08 AZBOS]. Specifically, [SCSTI7 focused solving the subproblem

---

of updating the model parameters only. Even in the centralized setting, convergence is only shown under rather restricted assumptions such as equal dataset sizes for different tasks (i.e. balanced local data) [AZB05] and small number of common features [APPOS] Moreover, the average excess error rather than the error of individual tasks shown to decay with the dominating term 0(1/vn) /n), where
is the homogeneous local dataset size [AZB05] Despite recent progress TH22 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-Agnostie Meta-Learning (MAML) framework (EM02 [K where the global objective is modified to account for the cost of fine-tuning a global model at individual clients. However, they focused studying the convergence in training errors only- no characterization of the generalization error is given.
Clustered Personalized Models: Clustered Federated Learning (CFL) [SMS20] GCYR20 GHYRI9
XL MMRS20 LLV2I can be viewed as a special case of Federated MTI where tasks across clients form cluster structures. In addition to the algorithms mentioned in Section [T] i.e., which simultaneously learn clusters and models for each cluster, other attempts have been made to integrate clustering with model training.
[SMS20] hierarchically clustered the clients in a post-processing fashion. To recover the & clusters, (k(k) empirical risk minimization problems need to be solved sequentially which is time-consuming. [GHYRT9] proposed a modular algorithm that contains one- shot clustering stage, followed by k individual uudversary-resiliees
model training. Their algorithm scales poorly in the input dimension, and requires local datasets to be balanced and sufficiently large (i.e., n 2 d²). Morcover, each client sequentially solves two empirical risk minimization problems. To utilize information across different clusters, [LLV2] proposed soft clustering. Unfortunately, the proposed algorithm is only tested on the simple MNIST and Fashion-MNIST datasets, and no theoretical justifications are given.
3 Problem Formulation
A FL. system consists of a parameter server (PS) and M clients. Each client i E [M] keeps a dataset DD (Trg ((ig))i)))] (=))))= that are generated from of k unknown regression models. Let M N - i= we refer to client as anchor client, which corresponds to active user in practice. Anchor clients play crucial rule in our algorithm design. We consider the challenging yet practical scenario wherein non-anchor client may have ((1) data points only.
n2, The local datasets are highly unbalanced with varying n, across clients. If (2(k).
We adopt a canonical mixture model setup: For each client i E [M].
y/ c(x;)0 + Si+
where 2 E [k] is the hidden local cluster label, "i. 0f are the true models of the clusters, ((x;)) E Rnd is the feature matrix with rows given by o(z,z), y/ (wij) E I is the response vector, and S = (Se)) E R is the noise vector. Note that the feature map & can be non- non-linear (e.g. polynomials). The cluster label of client is randomly generated from one of the ic components from some unknown p = (pi..... Pk) in probability simplex A --1 That is, P {z+ = j} p, for j E [k). In addition, ||0] 2 S R for each component. The feature covariate (xx;)) is independent and sub-Gaussian ala < (((())))))
3 Bla. We assume that the covariance matrix is identical within the same cluster but may vary across different clusters, i.e., E( ()(g)) The noise Cij is independent and sub-Gaussian with E[[yy] and E[G S
]== E, if =
Our formulation accommodates statistical heterogeneity in feature covariates, local models, and observation noises [KMA For the identifiability of the true cluster models fj's, 's we assume minimum proportion and pairwise separation of clusters. Formally, let A min, 110 and Puiin = [i) Pj- For ease of presentation, we assume the parameters a, 3 O o o 17 O(),, and R/A O(1), while our main results be extended to show more explicit dependencies on these parameters with careful bookkceping calculations. Note that even under these assumptions, we still allow R, A, to scale with model dimension d. Also, the assumption R/ A O(1) basically requires the radius of 0f's 's is on the same scale as their pairwise separation. rules out the extreme setting where 077 themseeles are extremely large while their pairwise separations are tiny.
Notations: Let [r]] A (1...., n). For two sets A and B, let A () denote the symmetric difference (A - B) U (B- A). We use standard asymptotic notation: for two positive sequences } and
3

---

). we write an = (b,,) (or on S if an I Cbn, for some constant C' and sufficiently large n;; a, = [bb,,, (or an 2 b,) b, = (;,)): an (b,,) (or an b,, if a, = (b,,,) and an (b,b,,) Poly-logarithmic factors are hidden in S2. Given a matrix A E guxd let = Ei=1 o,u;t; denote its singular value decomposition, where = min(n, d}, 0ı 2 2 or 2 0 are the singular values, and uy (v)) are the corresponding left (right) singular vectors. We call = [ui, ug...., uv] as the top-k left singular matrix of A. Let span(U) spanfu1,.... ug denote the k-dimensional subspace spanned by fui...., 3.
4 Main Results
We propose a two-phase FL algorithm that enables clients to learn the model parameters .i..... .2 and their clusters simultaneously:
(i) Coarse estimation via FedMI). Run the federated moment descent algorithm to obtain coarse estimates of model parameters ST'.
(ii) Fine-tuning via iterative FedX+clustering, In each iteration, each client first estimates its cluster label and then refines its local model estimate via either FedAvg or FedProx (which we refer to as FedX) [MR 17 LSZ 20
FedMD and FedX+clustering are detailed in the pseudocode in Phase and Phase respectively.
4.1 Federated moment descent
With multiple clusters and aub-Gaussian features, simple procedures such as power method will no longer provide reasonably good coarse estimation. The reasons are two-fold: 1) With uauGGausian features, it is difficult to construct a matrix whose leading eigenspace approximately aligns with the space spanned by the true model parameters (0i,.... Of): 2) Even this is achievable, there still remains significant ambiguity in determining the model parameters from their spanned subspace.
The key idea of the first phase of our algorithm is to leverage the existence of anchor clients. Specifically, the PS chooses a set H of nH anchor clients uniformly at random. Each selected anchor client E И maintains a sequence of estimators
{0..) that approaches 0% achieving ||0,,, 0% ll2 cA for some small constant > O when t is sufficiently large.
At high-level, we hope to have 0s,, move along a well calibrated direction vi, that decreases the residual estimation error ||>]// - 0,,0)l.,... i.e., the variance of the residual (Ø(I))) 0 0s,x) As such, we like to choose Ti,, to be positively correlated with Σ,, (8%, - 0,,,). However, to estimate Er, (0: 0..z) solely based on the local data of anchor client i, requires n, (dd.. which unaffordable in typical FL systems with high model dimension d and limited local data. To resolve the curse of dimensionality, we decompose the estimation task at each chosen anchor client into two subtasks: we first estimate a k-dimensional subspace that Σ,, (#" 0,,)) lies in by pooling local datasets across sufficiently many non-anchor clients; then we project the local data of anchor client i onto the estimated subspace and reduce the estimation problem from d-dimension to k-dimension.
The precise description of our Phase procedure is given as below. For ease of notation, let E(x, y, 0) (y (0(x), ())(().
In Step 9, PS estimates the subspace that the residual estimation errors (E, - 0r,)) lie in, in collaboration with clients in S,. In particular, for each anchor client i H, define
1 m
Yr,
a(0v1, (i.)),,,)))0.... w0 Oc.))
PES.
We approximate the subspace spanned by {E,(0) - 0.,))} via that spanned by the top-k left singular vectors of Y.,. . To compute the latter, we adopt the following multi-dimeeniooal general- ization of the power method, known as orthogonalliteratio
GVLK Section 8.2.4]. In general, given symmetric matrix Y E Rdddd the orthogonal iteration generates a sequence of matrices Q, E dx as follows: Qo Rdxk is initialized as random orthogonal matrix Q Qo / and YQ: = Q++ Ret1 with QR factorization. When / large, Q: approximates the top-k left singualr matrix of Y, provided the existence of eigen-gap d Ast1- When k 1, this is just the the

---

Phase 1: Federated Moment Descent (FedMD)
1 Input: RH; k, m, E, T, N. a, c, E R. 0g E R' with 1100112 _ R
Output: 01....,0;, ,,
3 PS chooses set H of nn anchor clients uniformly at random;
4 for each anchor client E H do
L 0r,0 0p:
for & 0, T do
PS selects set St of clients from [M] U
PS broadcasts {0,.,. i H to all clients i' in S /* where ---0S, =0
*/; vrz, 0,, (s.e) k, Ti) output Ü for each anchor client i H; /* described in Algori thm
PS calls decerratd--rthogonalleee
(S,, {ε(œv1, ([(111,,11,
9/ (,0), 1
*/
PS sends Ü, each anchor client E H;
Each anchor client calls (ererritttaaiitA,AAA... T2) to output (Bi,, with Ait defined in 2
11
12
if δ., > é then
0,,+1 + 04, + r. and reports 0,,441, where = Ug and
13
Th.t a0.51
14
else
01,411 O2, and reports 0i.e+::
PS computes the pairwise distance ||0t,r rlly for every pair of anchor clients i, H, assigns them in the same cluster when the pairwise distance is smaller than 4/2, and outputs 0j to be the center of the estimated j-th cluster for j E [A]]
power-iteration and we can further approximate the leading eigenvalue of Y by the Raleigh quotient Q, Y yQ., When Y is asymmetric, by running the orthogonal iteration on YY we can compute the top-k left singular matrix of Y. In our setting, the orthogonal iteration can be implemented in distributed manner in FL systems as shown in Algorithm 3 in the Appendix D.I
In Step 11, each anchor client estimates the residual error Σ., (0%, -.,.) by projecting (Zij, Bijs 0,,e) onto the previously estimated subspace, that is, UT Mhj. 0,,)). This reduces the estimation from d-dimension to k-dimension and hence (2(k) local data points suffice. Specifically, define
1 Σ (TTT(aas Vig. Os.) (UTTe(2i) Øij. Br.. Ar. 8.
jED,
where Di,, consists of 28 local data points (Iij. Vig) and (æ.j. Ris) freshly drawn from D, at iteration t. Client runs the power-iteration to output Br. and 07, as approximations of the leading left singular vector and singular value of Ai,,, Then anchor client updates 0,,41 by moving along the direction of the estimated residual error Tir with an appropriately chosen step size T7.
We show that OgT is close to 0" for every anchor client E H and the outputs 0, are close to 0j up permutation of cluster indices.
Theorem 1. Let (0, 1/4) be small but fixed constant. Suppose that
m ス 1 (d), f ((k), T (1), Ti Q (og())))). Tr (log(Nd)).
(3)
With probability at least OnnnT/N10), for all initialization Og with [10ol22 I R
sup ieff
0,.T 02,112 5ell.ad. <eA.
(4)
Furthermore, when nu 2 log 1/pmin with probability at least OnnTT/N10)
d(0,0") < cA,
(5)
5

---

where
[Qx)]) - 09 0;|l.- where x is permutation over [k].
d(0,6") 0*) _ min max
(6)
Note that in 6 we take a minimization over permutation, as the cluster indices are unidentifiable.
Phase uses fresh data at every iteration. In total we need Q(d) clients with least two data points and (11/mmin)) anchor clients. This requirement is relatively mild, as typical FL systems have a large number of clients with O(1) data points and few anchor clients with moderate data volume.
We defer the detailed proof of Theorem D to Appendix D A key step in our proof is show the residual estimation errors {2; (r,))== approximately lie in span(U,.). Unfortunately, the eigengap of Yi, could be small, especially when 0s., gets close to 6%. and hence the standard Davis-Kahan theorem [K70 cannot be be applied. This issue is further exaggerated by the fact that the convergence rate of the orthogonal iteration also crucially depends on the eigengaps [GVL13] For these reasons, span(U...) may not be close to span (Ea
0,,)) at all. To resolve this issue, we develop a novel gap-fret free bound to show that projection errors U,T B,,1) are small for every j E [k]] (cf. Lemma 5
2 Our algorithm is partly inspired by [L which focuses on the noiseless mixed linear regression, but deviates in number of crucial aspects. First, our algorithm crucially utilizes the fact that each client chosen in S, has at least two data points and hence the space of the singular vectors of IE [Yr,]] is spanned by (0;000)0;,)) (3 In contrast, [LLI8] relies on the sophisticated method of moments which only works under the Gaussian features and requires exponential in 2 many data points. Second, our algorithm crucially exploits the existence of anchor clients and greatly simplifies the moment descent algorithm in LLT..
Remark 1 (Comparison to previous work LLIS
Our algorithm also bears similarities with the meta-learning algorithm in KSS+20 which also uses clients collectively for subspace estimation and anchor clients for estimating cluster centers. However, there are several key differences. First, KSS+20 20) focuses on the centralized setting and relies on one-shot estimation, under the additional assumption that the covariance matrix of features across all clusters are identical. Instead, our moment descent algorithm is iterative, amenable distributed implementation in FL systems, and allows for covariance matrices varying across clusters. Second, in the fine-tuning phase, IKSS+20) uses the centralized least squares to refine the clusters estimated with anchor clients. under the additional assumption that allog k) data points for every client. In contrast, as we will show later, we use the FedX+clustering to iteratively cluster clients and refine cluster center estimation.
(Data privacy risk) Compared to the standard FedAvg algorithm wherein only aggregated local /graainn/g are broadcasted by the parameter server, the major step of our wwo-phase algorithm that may leak additional privacy is Step 8 wherein the local model estimates of the anchor clients are broadcasted to many other non-anchor clients. However, this privacy leakage is minor and can be further mitigated by simple considerations.
Remark 2.
mechanism according to the following
First, in our algorithm, each chosen non-anchor client only receives a collection of local model esti- mates (without ID for anchor clients) from the parameter server, it does not know which broadcasted model corresponds to which anchor client and hence cannot directly identify each individual anchor client's local true model. Second, we only choose very few number of anchor clients (roughly on the order of the number of clusters) and in practice these anchor clients often specially recruited by the PS; hence they can be made less concerned about privacy leakage through some incentivizing schemes. Last but not least, we can better preserve the privacy of anchor clients by broadcasting perturbed versions of their local models to each client. Specifically, fix any anchor client i, each non non-anchor client i' receives 0p and x. that are equal to or,, subject to two independent noise perturbations. Then for the subspace estimation in Step 9, can replace one 0z,, by 0g i, and the other by 0yr,... in the definition of Yi,. Crucially, Y,, involves an average over m non-anchor clients; hence these independent noise perturbations for different i' will be averaged out. Since m is large, this implies that the injected random noises can be made large without deteriorating too much the accuracy of the subspace estimation, in a similar spirit as privatizing the model averaging step in FedAvg. This gives a promising pathway to maintain anchor clients' privacy; we leave rigorously analyzing its privacy guarantee future work.

---

4.2 FedX+clustering
At the end of Phase 1 only the selected anchor clients in H obtained coarse estimates of their local true models (characterized in D. In Phase 2 both anchor clients in and all the other clients (anchor not) will participate and update their local model estimates.
Phase is stated in a generic form for any loss function L(B, X; D), where ........ 0k)) E IR is the cluster parameters,
A k-1 represents the likelihood of the cluster identity of client, and D denotes the client's dataset. This generic structure covers the idea of soft clustering [LV2I Note that unlike Phase where each anchor client only maintains an estimate o₁, of own model, in Phase each client maintains model estimates d,, = (011.s..... Bsk,e) for all clusters.
Phase 2: FedX+clustering
Input: 0 == (0).-. 4x) from the output of Phase I 17, T"
2 Output: 0 (0,, 02)
3 PS sets 0r f 0
4 for f +
1..,,.,,,++ do
PS broadcasts - to all clients; Each client i estimates the likelihood of its local cluster label by
AS, arg min AEA
A;;
Each client i refines its local model based on either FedAvg or FedProx with (0) L0,, xi, D;), and reports the updated local parameters Oz.t (0.1,ty.... Oik,t).
FedAvg-basee it runs s steps of local gradient descent:
Og. (0:-1), where G:(0) 0 1
FedProx-baseed it solves the local proximal optimization:
Og, f- arg min [4(0)
110 0.-1
21j
PS updates the global model as 0 + EN Ei=1 w,Gee where we n./N.
In Phase 2 the local estimation at each client has a flavor of alternating minimization: It first runs a minimization step to estimate its cluster, and then runs a FedAvg FedProx update to refine model estimates. To allow the participation
clients with O(1) data points only, at every iteration the clients allowed reuse all local data, including those used in the first phase. Similar alternating update 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 setting where the model refinement at each client is via running a single gradient step, which is barely used in practice but much simpler analyze than FedAvg FedProx update.
our analysis, we consider the square loss
A 11 ((x,)0,12
2n
In this context,
i yields a vertex of the probability simplex Aej.1 = O{1 = 21, where
Zi = arg min llv.
((x,)0 ((.1))
(8)
jE(k]
The estimate i provides hard clustering label. Hence, in each round, only one regression model will be updated per client.

---

To capture the tradeoff between communication cost and statistical accuracy using FedAvg or FedProx, we introduce the following quantities from [SXY2I]:
1 [|o(2;)11
for FedAvg,
y - 1 max
dE[M] n4
1 +
for FedProx.
We choose a properly small learning rate such that y < 1. Here, 2 quantifies the stability of local updates. Notably, K s 1 using a relatively small 17.
For the learnability of model parameters, we assume that collectively there are sufficient data in each cluster. In particular, we assume N, Z d. where N, - E...= " denotes the number of data points in cluster j. To further characterize the quantity skewness (i.e., the imbalance of data partition n. = (M1,..., (2M) across clients), we adopt the <--divergence, which is defined as (P(I))) = (dP-dQ)² for a distribution P absolutely continuous with respect to a distribution Q- Let (n) be the chi-squared divergence between data partition p., the clients pa. (n()) n//N and the uniform distribution over [M]. Note that when data partition is balanced (i.e., n2 N/M M for all i), it holds that (n) 0.
dQ
We have the following theoretical guarantee of Phase 2 where & is the number of local steps in FedAvg. Notably, is algorithmic parameter for FedAvg only. To recover the results for FedProx, we only need to set s = 1.
Theorem 2. Suppose thut k: 2 2, 1) S 1/8, and N₂ Z d. Let p= min; N5/N. N. If log(e/v) S p/k, then with probability 1 Oke for all 2 T + 1 and all OT such that d(Or, 8*) S cA, where is some constant, it holds that
d(Or, 0*) s (1 Cismp/ k) d(Or-1, 0*) + Cysnov log
(9)
where
M mapa(ne) + CV
dk log
A
Σ
(x²(n) + +1,,
(10)
N
M
(In A2
and (m() 4ke -nn
Furthermore, ift T + 1, for each client i, with probability
1 - 1-pe(ng) (nt),, ), it is true that
0 02 0all2= S..l, A.e Cymppt- (T/ σK e log
v2
where 0z is client i's estimate of its own model parameter at time
Notably, Oz is the -t entry of Oz.,-- Theorem 2 shows that the model estimation errors decay geometrically starting from any realization that is within a small neighborhood of d". The parameter captures the additional clustering errors injected at each iteration. It consists of two parts: the first term of (TOp bounds the clustering error in expectation which diminishes exponentially in the local data size and the signal-to-noise ratio A/r; the second term bounds the uniform deviation of the clustering error across all initialization and iterations. Note that if the cluster structure were known exactly, we would get an model estimation eтror of 07 scaling as Vd/Nz. However, it turns out that this estimation error is dominated by our uniform deviation bound of the clustering error and hence is not explicitly shown in our bound ( In comparison, the previous work GCYR20) assumes fresh samples at each iteration by sample-splittin and good initialization independent of everything else provided prior; hence their analysis fails to capture the influence of the uniform deviation of the clustering error.
In passing, we briefly comment on the key assumption log(e/v) S p/KK. K. As aforementioned, the clustering eтror v consists of two parts shown in (T): the first term decays exponentially in the local dataset size and the signal-to- ggnaa--o-noise ratio and hence is very small in most typical scenarios; the second term is on the order of dk log k/ M when the quantity skewness (imbalance of data partition) is of constant order. Finally, captures imbalance of cluster sizes, which is typically of a constant order, and we can choose a small enough step size to ensure x is close 1. Thus the key assumption logíe/v) Σ P/N roughly translates to being subconstant, which further means that M (the number of clients) needs to be larger than (the model dimension) by polylog factors. This is often satisfied in the typical FL applications which involve a very large collection of clients.
8

---

4.2.1 Analysis of global iterations
Without loss of generality, assume the optimal permutation in 66 is identity. In this case, if 2i,t then client i will refine -, To prove Theorem we need to analyze the global iteration of 0t- Following a similar argument to [SXY2I] with a careful examination of cluster labels, we obtain the following lemma. The proof is deferred to Appendix E.
Lemma 1. Let ((x) be the matrix that stacks all (x;;) vertically, and similarly for y. holds that (11)
03 03,1-1 BBA
y), j E [k],
where B ( P. P is block diagonal matrix with ith block Pa of size ni ny given by
nob=)=n)((;))(0 /ns)² for FedAvg, for FedProx,
P1 [ no(z))((a)) /nn]
and Agaz is another block diagonal matrix with ith block being Xij,ehm.
Lemma M limmediately yields the evolution of estiamtion error. Let Aj be the matrix with ith block being 1 = -.tn, representing the true client identities. Plugging model (T). the estimation error evolves as
0,e 0 03 = (T 11 m),
( - y) BAyS, Vj E [b].
(12)
where K2- ,0,0((z) ()) and E₁ Ay,t- AA 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 noisy perturbation. Let 1; {i 2, j} be the clients belonging jth cluster, and IJ {i 24, = j} be the clients with estimated label j. The indices of nonzero blocks of Ej, are 1 θ J, indicating the clustering emross pertaining jth cluster.
For the ease presentation, we introduce a few additional notations for the collective data over a subset of clients. Given subset / [M] of clients, let ((ær) denote the matrix that vertically stacks ((x:) for i I, and we similarly use notations y/r and Çr; let Py be the matrix with diagonal blocks P. for i I. Using those notations. we have K₃ t 6 Pr, p(zx,), which differs from the usual covariance matrix by an additional matrix Ptş- Therefore, the analysis of the first and third terms on the right-hand side of (12) follows from standard concentration inequalities for random matrices. In the remaining of this subsection, we focus on the second term, which is major challenge in the analysis. The proof details are all deferred to Appendix E.2
Lemma 2. There exists a universal constant C such that, with probability Ce-
[ 2)ll2 x s(d(O.-1,0)) + to) E.
Vj E [k]]
(13)
Lemma 2 aims to upper bound the error of
BE BE,,( y) H TPs,, (6(I5))) - Usy.)),
(14)
where Sj.r = 1 e J5.t- The technical difficulty arises from the involved dependency between the clustering error Sy, and the estimated parameter 0,,-- as estimating label 21, and updating 0,,-1 use a common set of local data.
Proof Sketch of Lemma 2 It follows from the definition of 24,4 in that
[10((;)) villa S (10((x))0... vella. Vi E Sy
Then,
[1o( )0,0-1 US 12==
Σ
-
vill2 Σ (o: -
sES,,
&ES,
<
Σ (Ilo(x;) -000-11--1111
02,)11 + lk)l
< 2 (d(0 8*) 1|0(Es,)) + I[Ss),
(15)

---

Hence. it suffices to upper bound | lll and I|Ss,, ll given small estimation error 2(0t-1,6% 8*) from the last iteration. To this end, we show a uniform upper bound of the total clustering error Eies,, ni by analyzing weighted empirical process. Using the decision rule the set SJ can be written function Sj(0,-1) with
max E
[e vi)(0) 2 03 _ f/ f) W.), (2i,
i 1,,
1i E S,(0))
(16)
1 mi Vi (0) 2 0) - ( Yil, i i I
where
Par e. vi) _ |1: 0x;))0,1 - 11 - (0,))
Then we derive the following uniform deviation of the incorrectly clustered data points
M
miE 5,(0)) [iP((EE & CNV
M
dk log
sup
n; n; S,(())
(x² + 1).
M
PER
This is proved via upper bounds on the classes
(VC) dimensions of the binary function
Fj U3e0 8 E A U RR
(17)
Using classical results of VC dimensions, those functions equivalently intersections of hyper- planes in ambient dimension O(d)), which yields an upper bound (cd²). However, the hyperplanes are crucially aunkreesticcted as the total number of parameters in is dk. We prove that the VC dimen- sions are at most (dk log k) using the algebraic geometry of polynomials given by the celebrated Milnor- -hooo theorem (see, c.g., [Mat]] Theorem 6.2.1J) Consequently, [|o )ll2. llCs,, ll2 and thus 15) can be uniformly upper bounded using sub- sub-Gaussian concentration and the union bound, concluding the proof of Lemma
4.3 Global convergence
Combining Theorem I and Theorem 2 we immediately deduce the global convergence from any initialization within the l ball of radius R.
Theorem 3. Suppose the conditions of Theorem and Theorem hold. Let Ô be the oulput of our
two-phuse algorithm by running Phase 7 with T (1) iterations starting from any initialization
ig with [10ol|2 S R. followed by Phase with T O x log +) iterations. Then with probability anp
1 N-9 cke d. it is true that
a(0, ((0 $ log e 4'
(18)
Furthermore, for each client i, with probability pe(n)), holds that 1: 8% ll2 Z log y
To the best of our knowledge, this the first result that proves the global convergence clustered federated learning from any initialization. Our bound (8) reveals that the final estimation error dominated by the clustering error captured by v, and scales linearly in k which characterizes the stability of local updates under FedAvg or FedProx. Moreover, Theorem 3 shows that Phase 1 converges very fast with only (11) iterations and hence is relatively inexpensive in both computation and communication. Instead. the number of iterations needed for Phase 2 grows logarithmically in A /vv and linearly in k/ (smp). Thus, by choosing & relatively large while keeping x close to 1, FedAvg enjoys saving of the total communication cost.
Acknowledgement
J. Xu supported in part by the NSF Grant CCF- and an NSF CAREER award CCF-2144593, P. Yang is supported in part by the NSFC Grant 12101353 and Tsinghua University Initiative Scientific Research Program. The authors would like thank Philippe Rigollet for pointing out the related literature computing VC dimensions using Milnor- Thom theorem.
Similar applications of the -hom theorem have been known in the literature (see e.g.
GJ93 Theorem
2.2] and [2022 Theorem 2)).
10

---

References
[AEPO8]
Andreas Argyriou, Theodoros Evgeniou, and Massimiliano Pontil. Convex multi-task feature learning, Machine learning, 73(3):243-272, 2008
[AZB05] Rie Kubota Ando, Tong Zhang, and Peter Bartlett. A framework for learning predic- tive structures from multiple tasks and unlabeled data. Journal of Machine Learning Research, 6(11), 2005.
[DDPM13] Simon Dooms, Toon De Pessemier, and Lue Martens. Movietweetings: movie rating dataset collected from twitter. In Workshop on Crowdsourcing and human computation for recommender systems, CrowdRec RecSys, volume 2013, page 43, 2013.
[DHK*20] Simon S Du, Wei Hu, Sham M Kakade, Jason D Lee, and Qi Lei. Few-shot learning via learning the representation, provably. arXiv preprini arXiv:2002.09453 099444 2020.
[DK700 C. Davis and W.M. Kahan. The rotation of eigenvectors by perturbation. III. SIAM Journal on Numerical Analysis, 7(1):1-46, 1970.
[FAC] Frequency of facebook use in the united states of 3rd quar-
ter
http:://www.staaioen com http://swt..cossssttttttii/////a ffreuuenc-- of-oof-uuss use among- facebook- th- www.ed--itedee states/
2020.
Accessed: 2022-05-01.
[FHK12] Andrey Feuerverger, Yu He, and Shashi Khatri. Statistical significance of the netflix challenge. Statistical Science, 27(2):202-25 2012.
[FMO20] Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach. Advances in Neural Information Prucessing Systems, 33:3557-3568, 2020.
[GCYR20] Avishek Ghosh, Jichan Chung, Dong Yin, and Kannan Ramchandran. An efficient framework for clustered federated learning. Advances in Neural Information Processing Systems, 33:19586-19597 2020.
[GHYR19] Avishek Ghosh, Justin Hong, Dong Yin, and Kannan Ramchandran. Robust federated learning in heterogeneous environment. arXiv preprint arXiv::906.0000 2019.
[GJ93] Paul Goldberg and Mark Jerrum. Bounding the vapmik-chervonenkiis dimension of concept classes parameterized by real numbers. In Proceedings of the sixth annual conference Computational learning theory, pages 361-369 1993.
[GVL13] Gene H Golub and Charles F Van Loan. Matrix computations. Johns Hopkins University Press, edition, 2013.
[JKK119] Yihan Jiang, Jakub Konečny, Keith Rush, and Sreeram Kannan. Improving fed- erated learning personalization via model agnostic meta learning. arXiv preprint arXiv: 12488, 2019.
[KKM+20] Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sehastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International Conference on Machine Learning, pages 5132-5143, PMLR, 2020.
[MM+ Peter Kairouz, H. Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, Rafael G. L. D'Olierira, Hubert Eichner, 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 Konecný, Aleksandra Korolova, Farinaz Koushanfar, Sanmi Koyejo, Tancrède Lepoint, Yang Liu, Prateek Mittal, Mehryar Mohri, Richard Nock, Ayfer Özgür, Rasmus Pagh, Hang Qi, Daniel Ramage, Ramesh Raskar, Mariana Raykova, Dawn Song, Weikang Song, Sebastian U. Stich, Ziteng Sun, Ananda Theertha Suresh, Florian Tramèr, Pranecth

---

Vepakomma, Jianyu Wang, Li Xiong, Zheng Xu, Qiang Yang, Felix X. Yu, Han Yu, and Sen Zhao. Advances and open problems federated learning. Foundations and Trends@ in Machine Learning, 14(1-2):1-210, 2021.
[KSS+20] Weihao Kong, Raghav Somani, Zhao Song, Sham Kakade, and Sewoong Oh. Meta- learning for mixed linear regression. In International Conference on Machine Learning, pages 5394-5404. PMLR, 2020.
[LDCH21] Qinbin Li, Yiqun Diao, Quan Chen, and Bingsheng He. Federated learning on non-iid data silos: An experimental study. arXiv preprint arXhv:2102.0207, 2021.
[LL18] Yuanzhi Li and Yingyu Liang. Learning mixtures of linear regressions with nearly optimal complexity. In Conference On Learning Theory, pages 1125-1144, PMLR, 2018.
[LLV21] Chengxi Li, Gang Li, and Pramod K Varshney. Federated learning with soft clustering. IEEE Internet of Things Journal, 2021.
[LSZ*20] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Systems, 2:429-450, 2020.
[9X Liping Li, Wei Xu, Tianyi Chen, Georgios B Giannakis, and Qing Ling, Rsa: Byzantine- robust stochastic aggregation methods for distributed learning from heterogeneous datasets. In Proceedings of the AAAI Conference Artificial Intelligence, volume 33, pages 1544-1551, 2019.
[Mat13] Jiri Matousek. Lectures discrete geomelry, volume 212. Springer Science & Business Media, 2013.
[MMR 17] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication--fficceen learning of deep networks from decentralized data. In Artificial intelligence and statistics, pages 1273-1282. PMLR, 2017.
[MMRS20] Yishay Mansour, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh. Three ap- proaches for personalization with applications to federated learning. arXiv preprint arXiv:2002.10619 2020.
[MN19] Peter McCullagh and John A Nelder. Generalized linear models. Routledge, 2019.
[MNB+21] Othmane Marfoq, Giovanni Neglia, Aurélien Bellet, Lactitia Kameni, and Richard Vidal. Federated multi-task learning under mixture of distributions. Advances in Neural Information Processing Systems, 34, 2021.
[RV+13] Mark Rudelson, Roman Vershynin, al. Hanson-wright inequality and sub-gaussian concentration. Electronic Communications in Probability, 18, 2013.
[SCST17] Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi. and Ameet S Talwalkar. Federated multi-task learning. Advances in neural information processing systems, 30, 2017.
[SMS20) Felix Sattler, Klaus-Robert Müller, and Wojciech Samek. Clustered federated learning: Model-agnostic distributed multitask optimization under privacy constraints. IEEE transactions on neural networks and learning systems, 32(8):3710 33222,,, 2020.
[SXY21] Lili Su, Jiaming Xu, and Pengkun Yang. A non-parametric view of FedAvg and FedProx: Beyond stationary points. arXiv preprint arXiv:2106.15211 2021.
[Sza91] Stanislaw J Szarek. Condition numbers of random matrices. Journal of Complexity, 7(2):131-149, 1991.
[TJJ21] Nilesh Tripuraneni, Chi Jin, and Michael Jordan. Provable meta-learning of linear representations. International Conference on Machine Learning, pages 10434-10443. PMLR, 2021.
12

---

[Tro15] Joel A Tropp. An introduction to matrix concentration inequalities. arXiv preprint 1rXiv:1501.01571, 2015.
[Ver10] Roman Vershynin. Introduction to the non-asymptotic analysis of random matrices. arXiv preprini 1rX1v:1011.3027, 2010.
[Ver18] Roman Vershynin. High-dimensional probability: An introduction with applications in data science, volume 47. Cambridge university press, 2018.
[VR02] M Vidyasagar and J Maurice Rojas. An improved bound on the vc-dimension of neural networks with polynomial activation functions. Math ArXiV preprint math. OC/0112206 submitted for publication, 2002.
[WLL 20] Jianyu Wang, Qinghua Liu. Hao Liang. Gauri Joshi, and H Vincent Poor. Tackling the objective inconsistency problem heterogeneous federated optimization. Advances in neural information processing systems, 33:7611-76223 2020.
[XLS*21] Ming Xie, Guodong Long, Tao Shen, Tianyi Zhou, Xianzhi Wang, Jing Jiang, and Chengqi Zhang. Multi-center federated learning. arXiv preprint arXiv:2108.08644 2021.
[ZLL+18] Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data, arXiv preprint arXiv: 00582, 2018.
13

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See Appendix A
(c) Did you discuss any potential negative societal impacts of your work? [Ys] See Ap-
pendix B
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes]
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or URL)? [N/A]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [N/A]
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [N/A]
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, cloud provider)? [N/A]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [N/A]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any assets either in the supplemental material as URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you're 're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "zsejYJZ2p1K", "reviewer": "Reviewer_SAAF", "summary": "The following points summarize the paper:\n\n1. The authors propose a two-phase federated learning algorithm to learn cluster label and mixed regression parameter vectors. \n2. In Phase 1, few anchor clients (with $\\tilde{\\Omega(k)}$ data) obtain a course model estimate.\n3. In Phase 2, each client alternately estimates its cluster label and refines the model estimates based on FedAvg or FedProx.\n4. Their method works with unbalanced data across clients.\n5. They also provide global convergence of their algorithm with any initialization.  \n", "questions": "- Is $R$ (and $\\Delta, \\sigma$) treated as a constant in the proofs? If not, does it change the theoretical results (concentration rate etc)?  ", "limitations": "The limitations are adequately addressed by the authors.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\nAuthors have done a commendable job of providing theoretical gurantees for their claims. The paper is theoretical in nature and like any theoretical paper, seems a bit tedious to parse at times. However, on a high level all the proofs seem to follow naturally and support the claim. Overall, it is a well written paper with strong theoretical claims.\n\nWeaknesses:\n- The assumptions such $\\| \\theta_j^* \\|  \\leq R$ seem quite strong as they essentially remove $O(\\sqrt{d})$ dependency (assuming $R$ is treated as constant) from the problem. This combined with conditions like $\\frac{R}{\\Delta} = \\mathcal{O}(1)$ seem like quite stringent conditions on the entries of $\\theta_j^*$. \n- Although, it is a theoretical paper, it would have been nice to add some experiments (on synthetic or real data) to validate the theorems.     ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Zf42SgzYnNm", "reviewer": "Reviewer_777o", "summary": "The paper studies the mixture linear regression problem under the federated learning setting. One major contribution is that the paper studies the situation where when clients’ local data volume is highly unbalanced. The authors introduce a two-phase estimation procedure. The Phase 1 aims to estimate the parameters of anchor clients (clients with large data volume), and the Phase 2 then gets the estimations of all clients by using the anchor clients’ estimation results obtained from the Phase 1. Theoretical guarantees are developed to support their methodologies.", "questions": "May the authors propose some experimental design to justify their algorithms?", "limitations": "The authors have adequately addressed the limitation concerns. ", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Originality: The method and theory are both novel to me. Typically, in Phase 1, the author utilizes clients with low data volume to help reduce the sample complexity of estimating the parameters of anchor clients, which is clever. Besides, the theoretical results are highly nontrivial.\n\nQuality: The mythology and theoretical analysis are in good quality. Especially, the development of the theoretical results is highly nontrivial. However, I am not familiar with relevant literature on mixture linear regression, thus I cannot judge the significance of the contribution made by the paper on theoretical perspectives. \n\nOne major concern, however, is that there is no experiment in the paper. Given that the paper is proposing new algorithms, it is not ideal to miss the empirical evidence. I understand that the main contribution of the paper is on theory; but I personally think that empirical evidence is still necessary.\n\nClarity: The paper is in general clear. I can follow the main ideas of the paper. However, I think there should be more discussion on the motivation of the design of Federated Moment Descent (FedMD) algorithm. The current description is not clear enough for me to follow easily.\n\nSignificance: Both methodologies and theories are significant. However, the lack of experiments hurts the significance of the paper in all.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Ql_n2KoIuil", "reviewer": "Reviewer_DfZZ", "summary": "This paper generally studies the problem of clustered federated learning. Specifically, the paper proposes a method for clustered federated learning under a mixed regression setting. At its core, the paper explains the method (dividing it into an initial estimation phase and a fine-tuning phase) and gives convergence guarantees about the statistical error of incurred by the method. The theoretical proofs are the bulk of the paper's contribution, and involve eigengap-free bounds on subspace estimation and VC dimension analyses of certain classes of polynomials.", "questions": "1. The application of the Milnor-Thom theorem in the proof of Lemma 10 is not quite correct. The bound quoted on the number of sign patterns in L684 of the appendix is only valid if $m(k-1) \\geq dk$. However, if this does not hold then the resulting bound on $m$ follows anyways, so the lemma is still true. However, this is worth cleaning up.\n\n2. The assumption that $\\phi(x_{ij})$ is sub-Gaussian means that the assumption that $\\phi(\\cdot)$ can be a non-linear mapping is not particularly useful, and simply adds unnecessary notation. I think the paper would be clearer if $\\phi$ was omitted altogether, and the authors assume a mixed linear regression setting with sub-Gaussian data.\n\n3. Intuitively, what does the bound on $\\nu$ in the assumptions of Theorem 1 impose?", "limitations": "There are no negative societal impacts of the work. My only concern is whether or not the method proposed can be made private (in some formal notion) and this is discussed above.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Core Strengths\n\nThe core contribution of the paper (as it is claimed in the paper) is novel and exciting: Rigorous guarantees on clustered federated learning are few and far between, and I believe that trying to attain this is a significant strength of the paper. Moreover, I think the authors do a good job of motivating this type of approach and giving a rigorous taxonomy of the related literature.\n\nAt a high level, I believe that the paper's theoretical results are significant and enough to warrant acceptance. It is clear that there is a lot of interesting facets to the analysis, and the authors mainly do a good job of blending together a variety of disparate analytical viewpoints. In particular, the strategy outlined to prove Theorems 1 & 2 seems promising and mainly correct. I have not verified every detail though, but the portions I have looked at all seem correct. In particular, I especially appreciated the application of the Milnor-Thom theorem (though there is a slight, fixable issue with that, see the suggestions below).\n\n## Weaknesses\n\nThe biggest weakness of the paper (and the method proposed by the authors more generally) is its questionable pertinence to federated learning. In particular, as the authors state in their \"Broad Impact\" section, one of the primary tenets of federated learning is the ability to perform machine learning without compromising data privacy. In the original FedAvg paper, for example, the models trained by clients are never shared directly with one another or the server, but are instead all averaged together (which paves the way for things like differential privacy and secure aggregation, which have formal privacy guarantees).\n\nHowever, Algorithm 1 repeatedly sends models that are held by a client (specifically an anchor client) to many other clients. While this is not the same as broadcasting data directly, it seems questionable as to whether this method is preserving data privacy. This is particularly concerning due to the linear structure of the problem, where clients could conceivably use an anchor client's model to learn about the client's data. Note that I do not believe that all FL algorithms need to incorporate privacy from the start. However, I think that a lack of any reasonable pathway to privacy integration is a drawback to a method that should be discussed and considered. As a naive reference point, an algorithm that had clients send all their data directly to one another would be effectively useless as an FL method.\n\nThis is all to say that I believe that the paper should 1) be up front about what kind of privacy model it is actually envisioning and 2) discuss possible pathways to improving the privacy of the algorithm (if there are any). If the authors believe that broadcasting cluster model estimates (as in Algorithm 1, Line 9) directly to other clients is acceptable, then I believe that this warrants discussion.\n\nA more minor weakness of the work is that some of the notation becomes cumbersome, and relies on quantities/assumptions that could use extra explanation. For example, Theorem 2 makes an assumption on the \"clustering error\" term $\\nu$ (namely, k$\\nu\\log(e/\\nu) \\leq \\rho/\\kappa$). It is not immediately clear how strong or weak this assumption is. If I understand correctly, then this actually imposes a restriction on the dimension $d$, which seems like an important thing to note.\n\n## Summary\n\nI think this is a theoretically strong paper that studies an important topic in federated learning. However, I would prefer that the authors are more up front about the actual privacy model underlying their method, and whether or not that has a pathway to integration with privacy mechanisms like differential privacy or secure aggregation.\n\n### Review Score\n\nI have opted to only give the work a weak accept (6) at the time of writing this review. While I believe the theory is strong, this score is tied to the privacy concerns above. I would be happy to revisit this score, contingent upon the work including a useful, grounded discussion of the privacy risks of the method and potential pathways forward.\n\n\n# Concerning the Author Feedback and Discussion\n\nThis section has been added after reading the other reviews and author feedback. The authors have addressed my feedback well, and have even gone so far as to add entirely new experiments (as per the request of another reviewer). I do not believe that these were necessary for the purposes of acceptance, but cement my belief that this paper should be accepted. I have increased my score from a 6 to a 7 correspondingly.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept