An Improved Federated Clustering Algorithm with Model-based Clustering
Authors: Harsh Vardhan
Avishek Ghosh
Arya Mazumdar
Venue: Accepted by TMLR ()

Abstract: Federated learning (FL) is a distributed learning paradigm that allows multiple clients to collaboratively train a shared model via communications to a central server. However, optimal models of different clients often differ due to heterogeneity of data across clients. 
In this paper, we address the dichotomy between heterogeneous models and simultaneous training in FL via a clustering structure among the clients.  The clustering framework is one way to allow for   high heterogeneity level between clients, while clients with similar data can still train a shared model. We define a new clustering framework for FL based on the (optimal) local models of the clients: two clients belong to the same cluster if their local models are close. We propose an  algorithm, \emph{Successive Refine Federated Clustering Algorithm} (\texttt{SR-FCA}), that treats each client as a singleton cluster as an initialization, and then successively refine the cluster estimation via exploiting similarity with other clients. In any intermediate step, \texttt{SR-FCA} uses an {\em error-tolerant} federated learning algorithm within each cluster to exploit simultaneous training and to correct clustering errors. Unlike some prominent prior works \texttt{SR-FCA} does not require any  \emph{good} initialization (or warm start), both in theory and practice. We show that with proper choice of learning rate, \texttt{SR-FCA} incurs arbitrarily small clustering error. Additionally,  \texttt{SR-FCA} does not require the knowledge of the number of clusters apriori like some prior works. We  validate the performance of \texttt{SR-FCA} on real-world FL datasets including FEMNIST and Shakespeare in non-convex problems and show the benefits of \texttt{SR-FCA} over several baselines.

Published in Transactions on Machine Learning Research (02/2024)
An Improved Federated Clustering Algorithm with Model-based Clustering
Harsh Vardhan
Sharshoardhan@rcceee edu
Halicioğlu Data Science Instâtute University of California, San Diego
Avishek Ghosh
avishek ghosh whwsh in
Systems and Control Engg.,
Centre for Machine Intelligence and Data Sciences Indian Institute of Technology, Bombuy
Arya Mazumdar
arya Cryadu edu
Holicioğlu Data Science Institute University of Colifornio, San Diego
Reviewed on OpenReview: kttps: /tteee///ppnrreveeuu net/ forum? id= 12GA5mSkoB
Abstract
Federated learning (FL) is a distributed learning paradigm that allows multiple clients to collaboratively train shared model via commuications central server. However, optimal models of different clients often differ due to hetcrogeneity of data across clients. In this paper, we address the dichotomy between heterogeneous models and simultaneous training in FT. via a clustering structure among the clients. The clustering framework is one way to allow for high heterogeneity level between clients, while clients with similar data can still train shared model. We define new clustering framework for FL based on the (optimal) local models of the clients: two clients belong to the same cluster if their local models close. We propose algorithm, Successive Refine Federated Clustering Algorithm (SR-FCA), that treats each client singleton cluster
initialization, and then successively refine the cluster estimation via exploiting similarity with other clients. In any intermediate step, SR-FCA uses an error-tolerant federated learning algorithm within each cluster to exploit simultaneous training and to correct clustering errors. Unlike some prominent prior works SR-FCA does not require any good initialization (or warm start), both theory and practice. We show that with proper choice of learning rate, SR-FCA incurs arbitrarily small clustering error. Additionally, SR-FCA docs not require the knowledge of the number of clusters apriori like some prior works. We validate the performance of SR-FCA on real-world FL datasets including FEMNIST and Shakespeare in non-convex problems and show the benefits of SR-FCA over several baselines.
1 Introduction
Federated Learning (FL), introduced in McMahan et al. (2016); Konečný et al. (2016): McMahan Ramage (2017) is large scale distributed learning paradigm aimed to exploit the machine intelligence in users' local devices. Owing to its highly decentralized nature, several statistical and computational challenges arise in FL, and in this paper, we aim to address one such challenge: heterogeneity.
The issue of heterogeneity is crucial for FL, since the data resides users' own devices, and naturally no two devices have identical data distribution. There has been a rich body of literature in FL address this problem of non-iid data. We direct the readers two survey papers (and the references therein), Li et al. (2020a): Kairouz al. (019) for comprehensive list of papers on heterogeneity in FT. A line of research assumes

---

Published in Transactions on Machine Learning Research (02/202))
the degree of dissimilarity across users is small, and hence focuses on learning a single global model Zhao et al. (2018): Li et al. 2020b 2019); Sattler et al. (2019); Mohri et al. (201 Karimireddy et al. (2020). Along with
this, a line of research in FL focuses on obtaining models personalized to individual users. For example Li ct al. 2020b 2021) use a regularization to obtain individual models for users and the regularization ensures that the local models stay close to the global model. Another set of work poses the heterogeneous FL, as a meta. learning problem Chen et al. (2018); Jiang et al. (2019); Fallah et al. (2020bla). Here, the objective is to first obtain a single global model, and then each device run some local iterations (fine tune) the global model to obtain their local models. Furthermore Collins et al. (022)) exploits shared representation across users by running an alternating minimization algorithm and personalization. Note that all these personalization algorithms, including meta learning, work only when the local models of the users' are close to one another (see bounded heterogeneity terms YI and To lerms in Assumption 5 of Fallah et al. (2020b)).
On the other spectrum, when the local models of the users may not be close to one another, another,[Satller et al. (2021) )
Mansour ct al. (2020); Ghosh et al. (2022) propose a framework of Clustered Federated Learning. Here users with dissimilar data are put into different. clusters, and the objectiiv is to oblain individual models for each cluster; iee,,
joint training is performed within each cluster. Among these, /Sattler et all (2021) uses top-down approach using cosine similarity metric between gradient norm as optimization objective. However, it uses a centralized clustering scheme, where the center has a significant amount of compute load, which is not desirable for FL. Also, the theoretical guarantees of Sattler et al. (2021) are limited. Further, in Duan et al. similarity metric is used extending the cosine similarity and the framework of Sattler et al. (2021) h. Moreover, in Mansour et al.] (2020), ; the authors propose algorithms for both clustering and personalization However, they provide guarantees only on generalization, not iterate convergence. In /Smihh et al. (2017) the job of multi-task learning is framed as clustoring where a regularizer in the optimization problem defines clustering objective.
(2021) a. data-driven
Very recently, in Ghosh et al. [2022). an iterative method in the clustered federated learning framework called Iter- ative Federated Clustering Algorithm, or IFCA, was proposed and local convergence guarantee was obtained. The problem sctup for IFCA is somewhat restrictive-it it requires the model (or data distribution) of all the users in the same cluster to be (exactly) identical. In order to converge, IFCA necessarily requires suilable init ialization in clus- tering, which can be impractical. Furthermore, Gn/hhsh et al. (2022). all the users are partitioned into a fixed and known number of clusters, and it is discussod in the same paper that the knowledge about the number of clusters is quite non-trivial to obtain (see Section 6.3 in Ghosh et al. 20222) There are follow up works, such as Ruan & Joe Wong 2021) Xie et al. (020)), that extend IFCA in certain directions, but the crucial shortcomings, namely the re- quirements on good initialization and identical local models still remain unaddressed to the best of our knowledge.
In this paper, we address the above-mentioned whw introduce a new clustering algorithm, Successive Refinement Federated Clustering Algorithm or SR-FCA, which leverages pairwise distance based clustering and refines the estimates over multiple rounds. We show that SR-FCA does not require any specific initialization. More- over, we can allow the same users in cluster to have non-identical models (or data distributions); in Section 2 we define clustering structure (see Definition 2.1 that allows the models of the users in the same cluster to be differ- ent (we denote this discrepancy by parameter 61 (20) 0). Furthermore, SR-FCA works with different set of hyper- parameters which does not include the number of clusters and SR-FCA iteratively estimates this hyper-parameter
Classically, clustering is defined in terms of distribution from which the users sample data. However, in a fodcrated framework, it is common to define a heterogencous framework such as clustering in terms of other discrepancy metric; for example in Mansour al. (0220) a metric that depends on the local loss is used.
Clustering Framework and Distance Metric:
In this paper, we use a distance metric across users' local model as a discrepancy mcasure and define a clustering setup based on this. Our distance metric may in general include non-trivial metric like Wasserstein distance, by norm (with 721) that captures desired practical properties like permutation invariance and sparsity for (deep) neural-net training. For our theoretical results, we focus on strongly convex and smooth loss for which lg norm of iterates turns out to be the natural choice. However, for non-convex neural networks on which we run most of our experiments, we use a cross-cluster loss metric. For two clients i,j, we define their cross-cluster loss metric as the average of the loss of one client on the other's model, i.c., client i's loss on the model of j and the other way round. If this metric is low, we can use the model of client for client j and vice- versa, implying
2

---

Published in Transactions on Machine Learning Research (02/2024)
With the above discrepancy metric, we put the users in same cluster if their local models are close- - otherwise they are in different clusters.
that the clients are similar. We explain this in detail in Section 5
1.1 Our Contributions
Algorithms and Technical Contribution. We introduce a novel clustering framework based on local user models and propose an iterative clustering algorithm, SR-FCA. Note that, since the clustering is defined based on the optimal models of the users, we have no way to know the clustering at the beginning of the process. To mitigate this, we start with initially assigning a different cluster for each user, and run few local iterations of SGD/GT in parallel. We then form a clustering based on the pairwise distance between iterates. This clustering is refined (including merges/splits if necessary) over multiple rounds of our algorithm. We run federated training on each of the clusters to further improve the models. This step exploits collaboration across users in the same cluster. However, clustering based on iterates might lead to many mis-clustered clients, the iterates might be far from optimal models. The mis-clustered clients might lead to more errors from ordinary lederated training within clusters (error propagation) due to high heterogeneily between the original clusters.
To counter this we run a robust federated training (based on trimmed mean) within each clusters in the intermodiate rounds, instead of straightforwart federated lcarning. In particular, we use the first order gradient based robust FL algorithm of Yin et al. (2018) to handle the clustering error. Within a cluster, we treat the wrongly clustered users as lutliers. However, instead of throwing the outliers away like Yin et al. (2018) we reassign them to their closest cluster.
When the loss is strongly convex and smooth, and dist(...) is l2 norm we show that, the mis-chustering error in the first stage of SR-FCA is given by (Imdexp(-n/Val ( [Lemma 4.6), where m, n and d denote the number of users, the amount of data in each user and the dimensionality of the problem respectively. Moreover, successive stages of SR-FCA further reduce the mis- clustering error by factor of ((1/m) (Theorem { and hence yields arbitrarily small error. In practice we require very few refinement steps (we refine at most twice in experiments, see Scction 5). Comparing our results with IFCA Ghosh et al. (2022), we notice that the requirement on the separalion of clusters is quile mild for SR-FCA. We only need the separation to be ( On the other hand, in certain regimes, IFCA requires separation of M (,,,/)), ) which is much stronger requirement.
To summarize, key assumption in any clustering problem (which is non-convex) is suitable initialization. However, SR-FCA removes this requirement completely by the above technique, and allows the clients to start arbitrarily. For our results, we crucially leverage (a) sharp generalization guarantees for strongly convex losses with sub-exponential gradients and (b) robustness property of the trimmed mean estimator (of Yin et al. (018))
As a by-product of SR-FCA, we also obtain an appropriate loss minimizer for cach cluster ( (Theorem We notice that the statistical error we obtain here is ((L/V)). This statistical rate primarily comes from the usage of the robust estimator Yin et al. (208)).
Experiments. We implement SR-FCA on wide variety of simulated heterogencous datasets (rotated or inverted
MNIST, CIFAR10) and real federated datasets (FEMNIST and Shakespeare Caldas et al. 2018 ) With
cross-cluster loss distance metric, we compare the test performance of SR-FCA with five baselines (a) global (one model for all users) and (b) local (one model per user), IFCA, (d) CFL [Sattler [FL|Sattler et al. (2021) and (e)
(local models clustered by KMeans on model weights). On simulated datasets, SR-FCA obtains test. accuracy no worse than the best baseline and is able to recover the correct clustering. For CIFAR10, in particular, SR-FCA has 5% better test accuracy than IFCA. On real datasets, SR-FCA outperforms all baselines.
2 Federated Clustering and Our Setup
In this section, we formally define the clustering problem. Let, [n] We have users (or machines) that are partitioned into disjoint clusters, denoted by the clustering map C* : [m] + [K], where K is the (unknown) number of clusters. Let C²: [m]] + rg(C") denote any arbitrary clustering map C', where rg(C") is 12;
the range of the clustering. Each user i [m] contains n, 2 n data points { sampled from a distribution
D, For any clustering map C², let rg(C") denote the range of the map. We define f(::2) WW - R as the loss
Here, O and Ω hide logarithmic dependence.
3

---

Published in Transactions on Machine Learning Research (02/2024)
w
Co
62
E1
W
Cluster 1
Cluster 2
Clusters split
Impurities
Figure 1: The dots represent the population risk minimizers for two clusters in dist(...) space according to C*.
The dots represent the ERM in dist(...) space and the corresponding clustering Co obtained after ONE_SHOT
Figure 2:
function for the sample 2, where W WCR Here, W is closed and convex set with diameter D. We use C to denote different clustering maps on the set of clients im] encountered by our algorithm. We use rg(C) to denote the range of the clustering map, i.e,, the cluster indices. We define the population loss, F:W - -2d, and its minimizer, w/. ; for each user E [m]] (.(w)= = E,wD, [f(,w)] w = min F(().. The original clustering C is based on the population minimizers of users, wy- This is defined as:
Definition 2.1 (Clustering Structure). For distance metric dist(...), the local models satisfy
max dist(wt (uf (wf) Sel. C+)
min
dist( wt awf) 62.
(1)
where E1,62, are non-negative constants with 62 >1. This is illustrated in fig.
The above allows the population minimizers inside clusters to be close, but not necessarily equal.
Let G. = ié: [m]].c*
denote the set of users in cluster c according to the original clustering C*. We can then define the population loss and its minimizer, per cluster cE[K] as follows,
(1)) (i)=c) [
F(w), w argmminF.(w)) Fc(w)= |Gel IEG.
(2)
WEW
We use war for every client E [m] to denote the local model on client obtained after T local iterations of Our goal is to find a population loss minimizer for each cluster [K[K] [K]],ee, i.., w obtain this, we need to find the correct clustering C* and recover the minimizer of each cluster's population loss. Note that we have access neither F. nor w but only the sample mean variant of the loss, the empirical risk, fi(w) = 1.- , for cach user i E [m.. There are two major difficulties in this setting: (a) the number of clusters is not known beforehand. This prevents us from using most clustering algorithms like k-means, and (b) The clustering depends on = which we do not have access to. We can estimate wi by minimizing fú, however, when n, the minimum number of data points per user, is small, this estimate may be very far from w.--
(w,2z.))
The above difficulties can be overcome by utilizing federation. First, instead of estimating w¹ for single user, we can estimate w the population minimizer for each cluster, where users in the cluster collaborate to improve the estimate. Second, we can use these estimates of w improve the clustering, according to Definition 3.1
3 Algorithm: SR-FCA
In this section, formally present clustering algorithm, SR-FCA. We first run the subroutine ONE_SHOT to obtain an appropriate initial clustering, which can be further improved. SR-FCA then successively calls the REFINE() subroutine to improve this clustering. In cach step of REFINE() we first estimate the cluster models for each cluster. Then, based on these models we regroup all the users using RECLUSTER() () and, if required, we merge the resulting clusters, using MERGE(), The full algorithm along with its subroutines is provided Algorithms I ONE_SHOT and TrimmedMeanGD can be decomposed into server and client subroutines. For the clustering subrou- tines, RECLUSTER and MERGE, all steps are performed on the server, barring the computation of dist which might require communication bet ween the server and the clients. We now explain the different subroutines in detail.
4

---

Published in Transactions on Machine Learning Research (02/2024)
Algorithm SR-FCA
Algorithm TrimmedMeanGD
Input: Threshold A, Size parameter t
Output: Clustering CR
Input:
05 B< Clustering C- Output: Cluster models fwe,r herre(),)
Co - DNE_SHOT SHOT (A, D)
for to R do
Server:
C, f REFINE (C,-1,A)
for all clusters cErg(C.) in parallel do
end for
S. -iE (t)=
ONE_ (X,tt)
for f=0 1
1 do Send We,s to all clients i
Server:
for clients in parallel do
Receive ft (w(...)) from all clients iE
Send we to client i
g(we.)) TrMeang([V[ (w.,)
Receive WA.r from client end for
We,41 t - (we, [1g(w))))
end for end Return (we.r)eerg(e.)
G Graph with m nodes (one per user) and no edges
for all pairs of clients i [m],i do
Add edge (i.j) to the graph G if (is,(w);,,,,, end for
Client
Receive Wa,7 from Server. Send V f((we.) to Server.
Cu f Connected components from graph G with size 2
Client(i):
Receive we from Server.
Algorithm 4: MERGE ()
W Perform training iterations initialized from wn.
Input: Cluster models {we, )ecrg(e)) Clustering C,, Threshold A, Size parameter
Send Wayr Server.
Output: Merged Clustering C.+1; Cluster models
REF HEFINE (Cr-1.A)
(wr,T)eeng(c...)
for all clusters c do
Ge Graph with nodes rg(C.) and edges
We.7 TrimmedMeanGD(
for all pairs of clusters c, rrg(C)), do
end for
Add edge (c,c') to the graph G if dist(w,,,0e)) S end for
C, RECLUSTER C,-1)
C, + MERGE (,t)
Ctemp Connected components from graph of size 2t
Algorithm 2: RECLUSTER
each cluster in Cremp, merge the nodes of its component clusters to get Cr+1
Input: Cluster models [we,T)eerg(C.)) User models
fun H Clustering C,
for cETg(Ctemp) do
Output: Improved Clustering C, for all users éE [m] do + arrgmin destrricccddistwseeeee dist(w..ee.))
Ge Erg(C,) (cg(())) which merged into c} We.r + [G. Eeec, Weir end for return CeaunnCCruuuueeeeeeettc.. )eers(C.4))-
end for return Clustering C,
3.1 ONE_SHOTE ()
For our initial clustering, create edges between users based on the distance bet their locally trained models if dist(w,,wy (1;,,) X, for a threshold A. We obtain clusters from this graph by simply finding the connected compo- nents, which can be done in time linear in number of edges. We only keep the clusters which have at least users.
We use w to denote the model obtained after local iterations of any optimizer on client i E [m]. Tf our locally trained models, Wa,r, were close to their population minimizers, w for all users E [m], then choosing threshold À E (2), we obtain edges between only clients which in the same cluster in C*. However, if n, the number of local datapoints is small, then estimates of local models war might be far from their corresponding w and we will not be able to recover C*.
However, Co is still a good clustering if it satisfies these requirements: (a) if every cluster in the range of the clustering map rg(C*)= [K] has a good proxy (in the sense of definition 3.1 in rg(Co), and (b) each cluster in rg(C,) has most small fraction (< ()) of mis-clustered users such good clustering when C* is defined according to fig. D We can see that even though Co C**, the two green clusters and the single orange cluster in Co are mostly "pure' and are proxies of Cluster and Cluster 2 in fig.
E.g., fig. 2 provides an example of one
5

---

Published in Transactions on Machine Learning Research (02/2024)
To formally define the notion of "purity" and "proxy", we introduce the notion of cluster label for any arbitrary clustering C', which relates it to the original clustering C*.
Definition 3.1 (Cluster label). We define E [K], as the cluster label of cluster e' rg((C)) if the majority 1/2 fraction) of users in d' are originally from c.
This definition allows us to map each cluster c' Exg(C") to cluster e in C* and thus define the notion of "proxy". In fig. [ the cluster label of green clusters is Cluster 1 and that of orange cluster is Cluster 2. Further, using the cluster label we can define the impuritics in cluster c' as the users that did not come from c'. In fig. the green node in orange cluster is an impurity. Based on these definitions, we can see that if clusters in Co are mostly pure and can represent all clusters in C*, then Co is a good chustering.
3.2 REFINE ()
We iteratively refine the clustering obtained by ONE_ QNE_SHOT() () using REFINE(). We describe the subroutines of a single REFINE step below.
Subroutine TrimmedMeanGD (). The main issue with ONE_ QWE_SHOTGG O, namely, small n, can be mitigated if we use federation. Since Cy has atleast & users per cluster, training a single model for each cluster will utilize 2 In datapoints, making the estimation more accurate. However, from fig. 2 contain impurities, i.e, users from a different cluster. To handle them, we use a robust training algorithm, TrimmedMcar Yin al. (2018).
we can see that the clusters
This subroutine is similar to Fed McMahan et al. (2016), but instead of taking the average of local models, we take the coordinate-wise trimmed mean, referred to as TrMeanp where 8 (0,1/2) defines the trimming level.
Definition 3.2 (TrMeang)) For [0,5), and a set of vectors x' E j e [J], their coordinate-we wise trimmed mean g g= TrMeans (fx' _x² .....z")) is vector g e Rd, with each coordinate g (1-29)J Erre, 2, for each & [d], where Ug is subset of { (xl.= obtained by removing the smallest and largest /3 fraction of its clements.
1
The full algorithm for TrimmedMeanGD is provided in algorithm 3 Nole that TrMeans has been used to handle Byzantine users, achieving optimal statistical rates Yin et al. (0018). when < 3 fraction of the users are Byzantine. For our problem setting, there are no Byzantine users as such and we use TrMeang to handle users from different clusters as impurities.
Note the two requirements for good clustering Co from ONE_ WHE_SHOT: (a) every cluster in C* has a proxy in Co, then the TrimmedMeanGD obtains at least one cluster model for every cluster in C*, (b) if every cluster in Co has a small fraction (3 < 5) of impurities, then we can apply TrMeans operation can recover the correct cluster model for every cluster.
We end up with a trained model for each cluster as an output of this subroutine. Since these models are better estimates of their population risk minimizers than before, we can use them to improve Co.
Subroutine RECLUSTER(). The full algorithm for this subroutine is provided in algorithm 2 This subroutine reduces the impurity level of each cluster in Co by assigning each client i to its nearest cluster c in terms of dist(we, (.irr) Here, We,T refers to the model on cluster e obtained after TrimmedMeanGD with T itcrations. Since W are better estimates, we hope that the each impure user will go label. For instance, in fig. _ the impure green node should go to one of the green clusters. If some clusters in rg(C*) do not have good proxy in rg(Co), then the users of this cluster will always remain as impurities.
cluster with its actual cluster
Subroutine (). We provide the full algorithm for this subroutine in algorithm Even after removing all impuritics from each cluster, we can still end up with clusters in C* being split, for instance the green clusters in fig. 2 In C*, these form the same cluster, thus they should be merged. As these were originally from the same cluster in C*, their learned models should also be very close. Similar to ONE_SHOT we create a graph G but instcad with nodes being the clusters in C, Then, we add edges between clusters based on a threshold À and find all the clusters in the resultant graph by finding all connected components. Then, each of these clusters in G correspond to set of clusters in C,,, so we merge them into a single cluster to obtain the final clustering C,+1-
6

---

Published in Transactions on Machine Learning Research (02/2024
3.3 Discussion
SR-FCA uses a bottom-up approach to construct and refine clusters. The initialization in ONE_SHOT is obtained by distancc--bsed thresholding on local models. These local models are improper estimates of their population minimizers due to small n, causing C However, if Co is not very bad, i.e, each cluster has impurity Traction and all clusters C* are represented, we can refine it.
REFINE () is an alternating procedure, where first estimate cluster centers from impure clusters. Then, we RECLUSTER () to remove the impurities in each cluster and then MERGE the clusters which should be merged according to C*. Note that as these steps use more accurale cluster estimates, they should have smaller error. This iterative procedure should recover one cluster for each cluster in C*, thus obtaining the number of clusters and every cluster should be pure so that C* exactly recovered.
3.4 Computation and Communication Complexity
We now analyze the computation and communication complexity of SR-FCA. Note that the complexity of the REFINE step is the same as that of IFCA in terms of both computation time and communication since in each case, need to find the loss of every cluster model every client's 's data (all pairs). This requires O(m²) forward passes. Note that performing all pairwise comparisons is unavoidable if an initial clustering is not known, as we need to know which clients can be clustered together. (eg. see KMeann Lloyd (1982) v/s DBSCAN Ester et al. (1996) or Ward's ' algorithm).
The REFINE step is comparable to IFCA for which we provide a detailed analysis. Assume that we run IFCA on K clusters for T rounds with E local steps with a fraction of m clients participating each round. Then, we require O(omllE)) backward passes and (amll") communication for local training and aggregation. For cluster identity estimation, in each round, we need to compute the loss of each cluster model every participating client. To send cluster models lo each client, need O(amKT)) communication and to compute losses, we need ((mm(T) forward passes.
If We run REFINE steps with S clusters per step where each TrimmedMeanGD procodure runs for rounds with E local steps and a fraction of clients participating each round. We need ((amETT R) backward passes and O(aTTR) communication for aggregation and local training. For clustering in REFINE and MERGE steps, we need to compute the loss of every model on every client, which requires O(mK communication and O(m R) forward passes. We summarize these results in Table
Algorithm
Communication
Training Steps
Forward Passcs
IFCA
OoomTK) (omTRR+MKT
OamTE) O(oMTER)
O(amTK) O(m
SR-FCA
Table 1: Communication, training runtime and forward passes for loss computation in clustering for SR-FCA and IFCA. E is the number local steps, T' is the number rounds, R is the number of REFINE steps, K is the number clusters and a is the fraction of clients participating per round.
Comparing the two algorithms, where the number of REFINE steps R and the number of clusters C are assumod to be constants, for a constant a, both algorithms require the same communication and backward passes, but IFCA needs more forward passes for clustering, If f (,),,, which corresponds to selecting constant mumber of clients per round, then if m =TT),, IFCA is better than SR-FCA in terms of communication and forward passes, and vice-versa if T = 2(m).
As for the other baselines, Meeans is the most efficient in terms of computation and communication complexity as performs only a single round of commiication, at the end of training. CFL, on the other hand, adopts a top-down approach to split clusters into two parts based on the cosine similarity of gradients. Fach split calls bipartitioning subroutine with runtime O(m²) making it the slowest baseline.
7

---

Published in Transactions on Machine Learning Research (02/2024)
4 Theoretical Guarantees
In this section, we obtain the convergence guarantees of SR-FCA. For theoretical tractability, we impose additional conditions on SR-FCA. First, the dist(...) is the Fuclidean (f2). However, in experiments (see next, section), we remove this restriction and work with other dist(...) functions. Here, we show an example where l norm comes naturally as the dist(...) function.
Proposition 4.1. Suppose that there are m users, each with a local model w E Rd and its datapoint (I, E TRd x R is generated according to yi = (uet,a) + 4s. x ~ N(0, Ia) and & ~ N(0,,²), then d [[.pKK/ppp///e)lll(ssselll)))]l = 20 l|u; (ll;--wwk ww
iidd
Hence, we see that minimizing a natural measure (KT. divergence) between the distributions for different. users is equivalent to minimizing the lg distance of the underlying local models. This example only serves as a motivation, and our theoretical results hold for strictly larger class of functions, defined by our assumptions. Remark 4.2 (A range). For the guarantees of this section to hold, we require XE (c1,62) and t I Cuin, where Cuin is the minimum size of the cluster. We emphasize that, in practice (as shown in the experimenls), we treat À and as hyper-parametie and obtain them by tuning. Hence, we do not require the knowledge of 61,62 and Gmin-
The following assumptions specify the exact class of losses for which our analysis holds. Definitions provided in appendix E
Assumption 4.3 (Strong convexity). The loss per sample (w..) is p-strongly convex with respect to w.
Assumption 4.4 (Smoothness). The loss per sample f(w...) is also L-smooth with respect to w.
Assumption 4.5 (Lipschitz). The loss per sample f (w..) is Lę- Lx-Lipschizz for every coordinate ke [d]. Define L- L3
We want to emphasize that the above assumptions are standard and have appcared the previous litcrature. For example, the strong convexity and hnest conditions are often required oblain theoretical guarantees for clustering models (see Ghosh et al.
2022) Lu & Zhou (2016), including the classical k-means which assume
quadratic objective. The coordinate-we wise Li Lipschitz assumption is also not new and (equivalent assumptions) featured in previous works Yin et al.
2018 2019), with it being necessary to establish convergence of the
trimmed mean procedure. Throughout this section, we require Assumption 4.3 4.4 and 4.5 to hold.
Misclustering Error Since the goal of SR-FCA is to recover both the clustering and cluster models, we first quantify the probability of not recovering the original clustering, i.e., C, CC Here and subsequently, two clusters being not equal means they are not equal after relabeling (see definition 3.1). We are now ready to show the guarantees of several subroutines of SR-FCA. First, we show the probability of misclustering after the ONE_SHOT step.
Lemma 4.6 (Error after ONE_ SHOT). After running ONE SHOT with 7 7, for T iterations, for the threshold E and some conslant b2 > 0, the probabilily of ervor is CC* Sp= mnd exp(-n b24) 22/3 44/3
Lva), provided
1,73 2d, where A= min -mil (1- and n m
We would like to emphasize that the probability of error is exponential in n, yielding a reasonable good clustering after the ONE_ SHOT step. Note that the best probability of error is obtained when \= 51
Remark 4.7 (Separation). Tn order to obtain p < 1, we require A f2 logmm ). Since A < 2 - we require
(e2 (6-6e) e1) O logm = Ö(2). Note that we require a condition only on the separation 62 61, instead of just 62 or é individually.
Although obtain exponentially decreasing probability of error, we would like to improve the dependence m, the number of users. REFINE() () step does this job.
Theorem 4.8 (One step REFINE()). Let 3t (Cemiu), and REFINE() run with TrimmedMeanCD(3). Provided
m
2/3 44/8
LPlog(Cmun ( } Z d, with 0 < B < 1 where A = A aB A2
> 0 and B = 22Le1/N Then, for any
D4/3
2
8

---

Published in Transactions on Machine Learning Research (02/2024)
constant e E (1.2) and 1/2 (1,2 22,, such that after running 1 step of REFINE() with 17 1, have
&dd" (xp(-o=))((1-3mm)) P 10" +ml P ph +8 exp(-asm A' 21.
Pr(c+#C*s
exp(-a1Cmin)
Cain
m
where Cmin the minimun size of the cluster. Further for some small constants p1 00.2 (0,1), we can select 8,71 and Y/2 such that for large m,n and A', with B B<< 3'<""'' we have $
Remark 4.9 (Misclustering error improvement). Note that p2 can be made arbitrarily close to 0 by a proper choice of 71 and 72. So, one step of REFINE() brings down the mischustering error by (almost) factor of 1/m, where m is the number of users.
Remark 4.10 (Condition on B). Note that we require B<< 2A' for the above to hold. From the definition of B, it when the intra-cluster separation 6 is small, B is small. So, for setup like IFCA, where 61 =0, this condition is automatically satisfied.
With the notation, C₁ #C*, we can identify the main sources of error for this event as: (1) There evists rErg(C*) such that no cluster in Co has cluster label c. (2) Each cluster CE rErg (g(Co) should have < a fraction of impurities for some 3/ 3>3 > a. (3)MERGE error: Two clusters with same label are not mergod or two clusters with different labels are merged. (4)RECLUSTER() ervor: A client does not go to its correct cluster after (ERGE() and REFINE().
Proof Sketch for Theorem 4.8
For the first two sources of error, utilize the clustering structure obtained by sources of error we leverage strong concentration of coordinato-wise Lipschitz losses to bound the performance of improved cluster models obtained by TrimmedMeanGD. In appendix C we upper bound the probabilities of error for each of the above ment ioned cases and apply a union bound on these errors to obtain the final result.
while for the last two
Using single step improvement of REFINE, we obtain the improvement after R steps of REFINE.
Thcorem 4.11 (Multi-step REFINE()). ) If we run R steps of REFINE resampling n, points from DD and recompule wi as in ONEOOONSHOO SHOT for every step of REFINE() then the probability of error for SR-FCA with R steps of REFINE() is Pr(C m mn
02
Remark 4.12 (Re-sampling). Although the theoretical of of Multi-step REFINE() requires resampling of data points in each iteration of REFINE(), we experimentally validate (see section 5 that this is not required at all. Remark 4.13. Since each step of REFINE() reduces the probability of misclusteing by (almost) factor of 1/m, very few steps of REFINE() is often sufficient. In our experiments (section 5) we need 1-2 REFINE () steps.
Convergence of cluster models: SR-FCA also obtain appropriate loss minimizer for each cluster.
Theorem 4.14 (Cluster models). Under the conditions described in theorem 4.8 after running SR-FCA for (R+1) sleps of (EFINE(), we have CR+1_C" and
|| - will S (1- 1TT//2 1 wheee, A=
8 1-23 Vn
id
1 /ncmmi
log(nmLD))
Vcerg(C"), with probability 1-
Pg -,,,,-p (-M3)) Grsin (1 LDya for some constant >.
R
m
idu"
Remark 4.15 (Convergence rate matches IFCA). Note that the models converge exponentially fast to the truc cluster parameter we, which matches the convergence speed of IFCA.
Remark 4.16 (Comparison with IFCA in statistical error). Note that for IFCA, 6 0 and the statistical error rate of IFCA is ((1/n) (see Theorem 2 in Ghosh et al. (2022)) ). Looking at theorem 4.14 we see that under similar condition (61 = 0 and hence B B=0), SR-FCA obtains an error rate of (11/v/), which is weaker than IFCA. This can be thought of as the price of initialization. In fact for IFCA, a good initialization implies that only a very few users will be mis-clustered, which was crucially required to obtain the ((1/n) rate. But, for SR-FCA, do not have such guarantees which results in weaker statistical error.
5 Experiments
We compare the performance of SR-FCA against several baselines on simulated and real datasets.
9

---

Published in Transactions on Machine Learning Research (02/2024)
Table 2: Test Accuracy and standard deviations across 5 random seeds on simulated datasets. The highest accuracy is bold. SR-FCA is competitve with IFCA for MNIST and beats it for CIFAR1O.
MNIST
MNIST (ROTATED)
CIFAR (ROTATED)
CIFAR (LABEL)
BASELINE
(INVERTED)
SR-FCA
92.03 0.030
91.66 + 0.13
91.38 +: 0.27
93.06 + 0.20
LOCAL
76.52 =0.54
85.55 0.19
75.87 0.33
80.09 + 0.53
GLOBAL
88.61 + 0.77
80.88 =1.55
88.75 0.52
81.48 + 3.55
CFL
SATTLER ET AL.
2021
88.30 + 1.12
80.47 10.44
87.59 1 0.42
82.16 1 1.73
LOCAL-KI K MEANS GHOSH ET AL.
2019
10.56 + 1.31
10.35 = 0.71
10.00 10.20
72.89 + 0.61
IFCA
GHOSH ET AL.
[2022)
91.55+ 0.81
91.80 0.25
86.05 + 0.43
85.19 + 2.01
Simulated Datasets: We generate clustered FL datasets from MNIST LeCun & Cortes (2010) and CIFARIO Krizhevsky et al. by splitting them into disjoint sets, one per client. For MNIST, by inverting pixel value, we create 2 clusters (referred to as inverted in table p and by rotating the image by 90,180,270 degrees we get 4 clusters. Note that this is a common practice continual Icarning Lopez- Paz & Ranzato (2017) and FL Ghosh et. a (2022). We set m= 100,n = 600. We obtain Rotated CIFAR1O by creating 2 clusters with the images rotated by 180 degrees. To test with label heterogeneity,
create Label CIFAR10 with 2 clusters. The first cluster contains the first 7 of the 10 labels and the sccond cluster contains the last 7 of the 10 labels. We set m = 32 for both CIFARIO datasets and n= 3125 and n= 4375 for Rotated and Label CIFARIO respectively. To emulate practical FT, scenarios, we assume that only a fraction of the nodes participate in the learning procedure. For Rotated and Inverted MNIST, we assume that all the nodes participate, while for Rotated and Label CIFAR10, 50% of the nodes participate. For MNIST, we train 2-layer feedforward NN, while for CIFAR10, we train a 9esNet Page (019)) We train Rotated MNIST, Inverted MNIST, Rotated CIFARIO, and Label CIFAR1O for 250, 280, 2400 and 2400 iterations respectively with 2 refine steps for SR-FCA.
We use two real federated datasets from leaf database Caldas et al. (018)). We sample m= 50 machines from FEMNIST and Shakespeare. FEMNIST is a Federated version of FMNIST with data on each client being handwritten symbols from a different person. Shakospeare is a NLP dataset where the task is next character prediction. For FEMNIST, train a CNN for while for Shakespeare we train a 2-layer stacked LSTM. For clustered FL baselines, we tune K, the number of clusters, with K e (2,3,4,55 for FEMNIST and K (1,2,33) We run FEMNIST and Shakespeare for 1000 and 2400 iterations respectively and set number refine steps to be 1 for SR-FCA.
Real Datasets:
We compare with standard FL baselines Local (every client trains its own local model) and Global (a single model trained via FodAvg McMahan & Ramage 2017)) on all clients). The main bascline we compare to is IFCA. Among clustered FL baselines, we consider CFL Sattler et al. (2021), which uses a top-down approach with cosine distance metric, and Local-KMeanss Ghosh et al. (2019). which performs KMeans on the model weights of each client's local model. For real datasets, we compare with two additional baselines FedSoft Ruan & Joc- Wong (2021) and ONE_ SHOT-IFCA (initial clustering of IFCA obtained by ONE_SHOT] assess if these variants can fix the issues of initialization in IFCA.
For SR-FCA, we tune the parameters X and / for trimmed mean and set t and require at most 2 REFINE steps. We utilize the following metric based on cross-cluster loss which is bet suited to measure distances between clients' models as these are neural networks.
Definition 5.1 (Cross-Cluster distance). For any two clients i, j E [m], with corresponding local models wi and wy and local empirical losses Si and fi, we define the cross-cluster loss metric as distaross-cluster (w.,w,) = +f
To extend this definition to distances between cluster e and client j, such as those required by REFINE, we replace the client model W; and client loss fi by the cluster model w. and empirical cluster loss fe respectively. Similarly, to obtain distances between clusters and e', which are required by MERGE, we replace the client. models and losses with cluster models and losses respectively.
Note that no data is shared between clients to compute the cross-cluster distances, rather the models are shared between clients via the server. To compute the cross-cluster distance between two clients and j, we need to
10

---

Published in Transactions on Machine Learning Research (02/2024)
Test Accuracy and standard deviations across seeds on Real datasets. Highest accuracy is bold. SR-FCA consistently outperforms IFCA.
Table 4:
Table 3: Average Misclustering error of clustered FL algorithms test set across 5 random seeds for simulated datasets. The lowest error is bold. SR-FCA is competitve with IFCA and beats it for Rotated CIEARIO.
BASELINE
FEMNIST
SHAKESPEARE
MNIST (INVERTED)
MNIST
CIFAR
CIFAR
SR-FCA
83.831 1.49
48.54 + 0.69
BASELINE
(ROTATED) (ROTATED) (LABEL)
LOCAL
66.18 + 2.14
33.86 11.2
GLOBAL
80.00 3.02
45.28 = 0.78
SR-FCA
0.0
0.0
0.0
0.0
CFL
79.48 + 3.48
44.14 + 1.03
CFL
0.08
0.14
0.18
0.38
LOCAL- 0.36
IFCA
81.93# 1.56
46.12 + 1.22
KMEANS
0.28
0.38
0.26
FEDSOFT
78.74 2.61
46.98 + 1.25
IFCA
0.0
0.0
0.50
0.5
DNE_SHOT
81.62 + 2.29
45.56 + 1.15
-IFCA
compute fw;) and f:(wz). To compute (w;)), the server sends the model of client i, w;, to client j and client j computes its loss on the model wy and reports it back to the server. Js(wry) is computed similarly. To compute the cross-cluster distance between client i and cluster we need to compute fi(w)) and fe(w:) fi(w)) can be computed by sending the model we to client i, To compute the quantity fo(w:,, the server sends the model w, to each of the clients in cluster c to compute client losses on model wg. The cluster loss is simply the average of the losses of all clients in cluster We the same procedure to compute cluster loss, but instead of computing it client model, we compute it on the cluster model, to obtain cross-cluster distances between two clusters.
Further, for clustered FL baselines (IFCA, CFL, Local-KMeans, FedSoft, SWE_SHOT-IFCA, on real datasets, we tune the number of clusters. This is not required in SR-FCA. Note that for each baseline and dataset pair, we perform 20 trials of random search for hyperparameter tuning.
As our global model corresponds to Fed for a fair comparison, we use Fed McMahan Ramage [2017 inside the TrimmedMeanGD subroutine, bt applying T'r Many operation to the model updates after local steps from different clients. Note that we could have used different base foderated optimization algorithm for global model, for instance, FedProx Li et al. (2020b) or SCAFFOLD Karimireddy el al. 22020) In that case, would have to modify all our clustered FL algorithms baselines and SR-FCA, to use this base federated optimization algorithm for training cluster models.
Test Metrics:
The lest performance of any baseline is obtained by averaging over the clients, the test perfor- mance of each client on model trained by the bascline. For the local bascline, it is the client's local model and for the global baseline it is the single global model. For SR-FCA and clustered FL baselines, is the cluster model for the client. Note that we do not present convergence plots as different algorithms run in different number of stages.
For simulated datascts, the true clustering C* is known, so we report both the test accuracy and misclustering error in table and 3 respectively. For real datasets, the true clustering is not known, therefore, we report only the final test accuracy in table g Further, we provide the test accuracy after ONE_ SHOT and every REFINE step in table The total time to run all experiments including hyperparameter tuning on a single is 2.5 weeks and the code is provid ed here
5.1 Results
Across all datasets, we find that SR-FCA is competitive with or outperforms all other algorithms in terms of both misclustering error and test accuracy.
Comparison with Local and Global: cannot handle the heterogeneity. Hence, as seen in tables 2 and [ SR-FCA and other clustered FL baselines perform better as they find correct clusters with low heterogeneity inside each cluster.
The Local algorithu has access to little data, while the Global model
use the cosine distance between gradients and l2 distance between model weights which are not suitable for NN models. Local-KMeans performs the worst with 10% test accuracy for simulated datascts. For real datasots, it's accuracy is < <5% so we do not report it. SR-FCA and IFCA use cross-cluster loss and client loss respectively, which are better suited to
Comparison with CFL and Local-KMeans: CFL and
NN models, thus performminng these baselines (see tables 2 and
11

---

Published in Transactions on Machine Learning Research (02/2024)
Normelaed halogrum parsiee datunnee mettt
Normel and masiarom patte distannce ferrot
limdle
libbs
3.2
08 04
008
8.5
0.9
1.0
Figure 3: Histogram of pairwise distances after ONE_SHOT and in SR-FCA for Inverted MNIST (Teft) and FEMNiST((ighht)
Comparison with IFCA: On simulated datasets (tables 2 and 3 we find that IFCA recovers C* and outperforms SR-FCA marginally for MNIST datasets. This is duc to MNIST being simpler and casier to learn dataset, even after adding heterogeneity via rotations or inversions. In contrast, for CIFAR10 10, the learning task is much more difficult, and IFCA, without proper initialization, ends up with all clients in only a single cluster after a few rounds resulting in a misclustering of 0.5, as seen in table 3 Thus it's performance is comparable
to the global bascline in terms of test accuracy as seen in table 2 From table 3 we see that SR-FCA correctly
identifies C* and comprehensively beats IFCA in lerms of test. accuracy in table 2
On real datasets, for IFCA, need to find the correct number of clusters by tuning, For random sample of clients, the true number of clusters might not be the same. SR-FCA can compute both the correct clustering and cluster models for every random sample, allowing it to beat IFCA, which fits the same mumber of clusters to every random sample. The difference is more pronounced for the more difficult Shakespeare dataset than the easier FEMNIST dataset. Further, the variants of FedSoft and ONE_SHOT-IFCA, have similar test performance to IFCA on real datasets. For FedSoft, which is a soft-clustering version of IFCA, the issue of initialization remains unresolved. Running IFCA after ONE_ SHOT can only re-cluster the clients thus results in similar performance. In short, SR-FCA outperforms IFCA as well as its variants.
Intermediate Steps of SR-FCA: For MNIST, multiple REFINE steps are necessary for best performance, however, for CIFARIO and real datasets, only a single REFINE step achieves best performance (see table 5
Distribution of pairwise distances Note that À is an important hyperparameter for our algorithun and its choice depends on the distribution of pairwise distances. Tn Figure D we provide the histogram of the pairwise distances after ONE_SHOT for Inverted MNIST and FEMNIST respectively. For simulated datasets, there is clear separation between the pairwise distances clients in the same clusters have very low pairwise distance and clients in different clusters have very high pairwise distances. Therefore, any choice of À in the middle sufficos. However, for real datascts, there isn't 't clear scparation between clusters. Therefore, we use hyperparameter tuning to choose the which maximizes test accuracy.
6 Conclusion
We conclude with a polential scope of improvement in SR-FCA. SR-FCA trains models obtained via TrimmedMeanCD() which assumes / fraction of users inside each cluster are corrupted. Instead, if we run any federated optimization algorithm which can accommodate low heterogeneity, for instance FedProx Li et al. (2020b), inside each cluster in the final clustering CR, then we may obtain an improved error in theory (A in theorem 4.14 We leave this as a future work.
Acknowledgements. This research is supported in part by NSF awards 21 2112665, 2217058, and 2133484.
Table 5: Test Accuracies of intermediate steps of SR-FCA for all datasets
MNIST
MNIST
CIFAR10
CIFARIO (LABEL)
BASELINE
FEMNIST
SHAKESPEARE
(INVERTED) (ROTATED) (ROTATED)
AFTER SNE_SHOT 76.52 1 0.54
85.55 1 0.19 75.87 1 0.33 80.09 1 0.53 66.18 1 2.14
33.86 1 1.22
AFTER 1d REFINE
91.88 + 0.39
91.63 + 0.12 91.38 + 0.27 93.06 + 0.20 83.83 + 1.49
48.54 + 0.69
AFTER 2nd REFINE 92.03 1 0.30
91.66 = 0.13 91.38
0.27 93.06 1 0.20 83.83
1.49
48.54 1 0.69
12

---

Published in Transactions on Machine Learning Research (02/2024
References
Scbastian Caldas, Peter Wu, Tian Li, Jakub Koncčný, H. Brendan McMahan, Virginia Smith, and Ameet Talwalkar. LEAF: A benchmark for federated settings. CoRR, abs/1812.01097, 2018.
Fei Chen, Mi Luo, Zhenhua Dong, Zhengno Li, and Xiuqiang He. Federated meta-learning with fast convergence and efficient. communication. or or preprint. ar arXiv: 802.07876, 2018.
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized fedcrated lcarning. In Marina Meila and Tong Zhang (ods.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 2089 2099. PMLR, 18 18-24 Jul 2021.
Moming Duan, Duo Liu, Xinyuan Ji, Renping Liu, Liang Liang, Xianzhang Chen, and Yujuan Tan. Fedgroup: Efficient federated learning via decomposed similarity clustering. In 2021 IEEE Intl Conf on Parallel & Distribuled Processing with Applications, Big Dala & Cloud Compulinng, Sustainable Compuling & Communications, Social Computing & Networking Haaa///DDCoooSSoiaaCoonoo. 2021. doi: 10.1109/ISPA-BR
pp. 228-237,
2021.00042
Martin Ester, Pete- Kriegel, Jörg Sander, and Xiaowei Xu. A density-based algorithm for discovering clusters in large spatial databases with noise. In Proceedings of the Second International Conference Knowledge Discovery and Data Mining, KDD'96, pp. 226-231. AAAT Press, 1996.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. On the convergence theory of gradient-base based model-agnostic meta-lcarning algorithms. In International Conference on Artificial Intelligence and Statistics, pp.
PMLR, 2020a.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 3557-3568. Curran Associates, Inc., 2020b.
Avishek Ghosh, Justin Hong, Dong Yin, and Kannan Ramchandran. Robust federated learning in a heterogeneous environment. CoRR, abs/ hbs/1906.06629, 2019.
Avishek Ghosh, Jichan Chung, Dong Yin, and Kannan Ramchandran. An efficient framework for clustered foderated learning. IEEE Transactions on Information Theory, (8(12):8076-809)) 2022. shorter version in NeurIPS 2021.
Yihan Jiang, Jakub Konečný, Keith Rush, and Sreeram Kannan. Improving federated learning personalization via model agnostic meta learning. arXie preprint arXiv: 190999999888 2019.
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. arXie preprint arXiv:1112009970 2019.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic Controlled Averaging for Federated Learning, In Proceedings of the 37th International Conference Machine Leorning, pp. 5132-50 5143. PMLR, November 2020. ISSN: 26400-3498.
Jakub Konečný, H Brendan McMahan, Daniel Ramage, and Peter Richtárik. Federated optimization: distributed machine learning for on-device intelligence. arXiv preprint arXiv: 02527, 2016.
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar-10 (canadian institute for advanced research). URL
http://wwwh cs. www..torotoooooom www.- html
Yann LeCun and Corinna Cortes,
MNIST handwritten digit database.
2010.
URL
http: :
7/yann. l/yaan.lecun.com www.l.cwmmwe.c.....eeeeeeeceecce wwtp///www/mnnt////////////////t
13

---

Published in Transactions on Machine Learning Research (02/2024)
Liping Li, Wei Xu, Tianyi Chen, Georgios B. Giannakis, and Qing Ling. Rsa: Byzantinc-robust stochastic aggregation methods for distributed learning from heterogeneous datasets. In Proceedings of the Thirly- Third AAAI Conference on Artificial Intelligence and Thirty-First Innovative Applications of Artificial Intelligence Conference and Ninth AAAI Symposium on Educational Advances in Artificial Intelligence, AMAT19/TAATT1//AATTI AAAT Press, 2019. ISBN 978-1-57735-8094
doi: 10.1609/ama va3/..33ii.30 v33i01 1.30111..11111 1544.
Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. TEEE Signal Processing Magazine, 37(3):50-60) 2020a.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in hetcrogencous networks. In I. Dhillon, D. Papailiopoulos, and V. Sze (eds.), Proceedings of Machine Learning and Systems, volume 2, pp. 42944-9--44 450, 2020b.
Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In International Conference on Machine Learning, pp. 6357-63666 PMLR, 2021.
S. Lloyd. Least squares quantization in pem. IEEE Transactions on Information Theory, 28(2):129 137, 1982. doi: 11LL109/III..992.10000
David Lopez-Paz and Marc' Aurelio Ranzato. Gradient episodic memory for continual learning. In 1. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Yu Lu and Harrison H Zhou. Statistical and computational guarantees of lloyd's algorithm and its variants. arXiv preprint orXfc:1612.02099 2016.
Yishay Mansour, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh. Three approaches for personalization with applications to federated learning. arXie preprint arXiv:2002.10619 2020.
Brendan McMahan and Daniel Ramage. Federated leamnnng Collaborative machine learning without contralized training data. https://research, pow.g//agggggooooooooooooaaaaaaa
html 2017.
Brendan McMahan, Fider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas Communication- efficient learning of deep networks from decentralized data. arXiv preprint arXiv:1602.00 05629,, 2016.
Mehryar Mohri, Gary Sivek, and Ananda. Theertha Suresh. Agnostic federated learning. In International Conference on Machine Learning, pp. 4615 4625. PMLR, 2019.
David Page. How to Train your ResNet 4 : Architecture.
sttps://myrtlee.a/llean
2019.
Yichen Ruan and Carlee Joe- Wong. Fedsoft: Soft clustered federated learning with proximal local updating. CoRR, abs/21122.60555 2021.
Felix Sattler, Simon Wiedemann, Mabus- Müller, and Wojeiech Samek. Robust and communication- efficient federated learning from non-iid data. IEFE Transactions on Neural Networks and Learning Systems, 31(9):3400 3413, 2019.
Felix Sattler, Robert Müller, and Wojciech Samek. Clusterod federated lcarning: Model-agnostic distributed multitask optimization under privacy constraints. IEEE Transactions on Newral Networks and Learning Systems, 32(8):3710 3722, 2021. doi: 10.1109/ 10.1109/TNNT.S TNNLN..N...2 T...S..220..015500
Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet S Talwalkar. Federated multi-task learning. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (ods.), A Advances in Nearal Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Ming Xie, Guodong Long, Tao Shen, Tianyi Zhou, Xianzhi Wang, and Jing Jiang. Multi-center foderated learning. CoRR, abs abs/2005.01026, 2020.
14

---

Published in Transactions on Machine Learning Research (02/2024)
Dong Yin, Yudong Chen, Ramchandran Kannan, and Peter Bartlett. Byzantino-robust distributed learning: Towards optimal statistical rates. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 95th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. -659. PMLR, 10-15 Jul 2018
Dong Yin, Yudong Chen, Ramchandran Kannan, and Peter Bartlett. Defending against saddle point attack in Byzantine-robust distributed learning. In Kamalika Chaudhuri and Ruslan Salakhuldinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 7074-7084. PMLR, 09-15 Jun 2019.
Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data. ar Xiv preprint arXiv: 1006.0058, 2018.
15

---

Published in Transactions on Machine Learning Research (02/2024)
Appendix for "An Improved Federated Clustering Algorithm with Model-based Clustering"
A Proof of proposition 4.1
According the proposition, for two users i and j, the data is generated by first sampling each coordinate of x ERd from (0,1) iid and then computing y as
y= (x,wi) +e +
where 6; id (0, Then, the distribution of yı1z is N (((,) o2?). Therefore, the KL divergence between u:2z and Vj x is given by
[L(p(wl=)ll//e/]))) (w, - x
20²
Therefore, if we take expectation wrt æ, have
d|I W w3
20
B Proof of lemma 4.6
In ONE_SHOTE O) CO= if all the edges formed in the graph are correct. This means that if i are in the same cluster in C*, then ||wcrr-w,,r < X and if i.j are in different clusters, ||wir Wjrr >> A.
Note that,
W WjT (w) wff) + (wi,T 11:,r-w()) wx) - -W, wz)
Now, if apply triangle inequality, obtain
dist( (.,,,,,,,,,) 2 (is)(w(,,))))=ie))).) - - dist(ws,r,,,)) dist(w; ,wf) + E where E, - dist diss(www....)) This decomposition forms the key motivation for algorithm.
Therefore, if i.j are in the same cluster, then a sufficient condition for edge (i.j) to be incorrect is
A = (ww(wwtuuuu w5)) + +Ei = + =
Similarly, i.j are different clusters, then a sufficient condition for edge (i.j) to incorrect is
> dist(wffuu) w.)) E., - E,
Therefore, set A min A. 1 ) and then sufficient condition for any edge to be incorrect is Ea, Ax.
Thus,
Pr[C* <Pr(at least edge is incorrect]
<Pr[maxE,, Axl
<Pr[max (.) kmi.j
| - w 2 AN
A
<Pr(maxmax (amaaxx([ww)) (IwWx/
bn k=ij
<Pr[max w 1c/m]
33)
2
16

---

Published in Transactions on Machine Learning Research (02/2024)
The second and third incqualities are obtained by expanding the terms. The fourth inequality is obtained by Pr(a+ b 2 c) & [[]mmxxxa))) 2 c/2 For the fifth inequality, merge maxi,j maxk=i, into max4g[m]. As we can see in Equation equation 3 we need to bound ||wi, -w;] for each node The subsequent Lemma
allow us to bound this quantitios.
Lemma B.1 (Convergence of wir)). Let D2/9/273 S bıd, for some constant b, > 0. Then, after running ONE_ WHOTSOTTO with n t,, for some constant b2 > under assumption 4.3 assumptiion 4.4 and assumption 4.5 have
n2/3
Pr[lwwir--u w E2 - Sd exp( -1
L
where A = -1-))-(1-))) and n= mm]
This lemma follows from Yin al. 2018
The complete proof of this Lemma is present in appendix B.1
Now, can apply lemma B.1 in Eq equation 3
A
CC]
wi
kE[m]
A
<mmax iE(m]
wi ||
Smd exp(-n b
Lvdd
For the second incquality, we use Pr[max;e[m]a; ii/m from union bound.
Pr & Pr(a; c]. which follows
Note that for 1 we need the separation to be order of 0 logm
B.1 Proof of lemma B.1
We utilize results from Yin al. (2018); which hold for TrimmedMeanGD to analyze convergence for a single node as they yield stronger guarantees under the given assumptions.
Lemma B.2 (Convergence of wi,T). If assumption 4.3 assumplion 4-4 and assumption 4.5 hold, and S +, then
2 A viemm]
| (1-x
(4)
where k= L and A₃ is posilive random variable with
D ) 2L 2i2" 1)
x2
Pr(A
(5)
for some r,ō 0.
We provide the proof of this lemma in appendix C.8
Using the above Lemma, we can bound the probability Pr[] w A1
Pr[][w, w 2 2
1/Pp/22111 1/2
2
-
A
2
A>
<Pr(A, 24]. where A
(1 - 11T/21)
2
<Pr( 2dr 2v 22
A <dexp( -nbbp Wd
for some constants by.by, 00, where we set r= bybyymax iva' iva
and 8= 2 and for b 2/33 2/13
D2/3 1,483
such that 22r 22/2f.222222 A and amin 1T 27.3
,3
Dd in lemma
3
B.2
17

---

Published in Transactions on Machine Learning Research (02/2024)
C Proof of theorem 4.8
C.1
Preliminaries
First, we define certain random variables and their respective probabilities which we will use throughout this proof. Since the edge based analysis and corresponding clique identification involves a lot of dependent events, we try to decompose the dusence/presence of edge into a combination of independent events.
Define,
Xi { If the edge (i.j) in Co incorrect in C*
(6)
Otherwise
An edge (i.j) in Co incorrect in C* if either is present in C* and absent in Co or vice versa. We analyze the probability of this event for the case when C* contains the edge (i.j). The case when C* doesn't 't contain edge (i,j) and it is present in Co has exaclty same probability. When | -w] || Sei, then edge is present is C If it is absent in Co. then
Pr| 24
The analysis is similar to the proof of ONE_ [HE_SHOT() in appendix
Note that the random variables {Xir] are not independent. We now define independent random variables X; such that
IAA; Otherwise
X,
(7)
Thus, we can see that X,y XX, Additionally,
PH|X:=1]sP\MM2A5
(8)
This follows from analysis of ONE (NE_SHOT() () in appendix B
We can further generalize this notion the random variables defined as Yi.
Y {
(9)
Otherwise
Then,
Pr|y,, =|M,,,= <
74
P
Lvd
m
Note that the set of random variables (Yt,) are mutually independent random variables.
Further, we define the w for every cluster rg(Co). Let c E C* be the cluster label of node If G. fi:ie [m],C* =c)), which is the set of nodes in which were from d' in the original clustering, then we can define w and Fe(w) as
w = argminE wEW
Σ f((w)
(10)
Ge
lEG,
=argmin
=(ww)=agmin weW
(11)
Ge
wEW
We use this definition of w in the appendices C.5 and C.6
18

---

Published in Transactions on Machine Learning Research (02/2024)
C.2 Analysis of REFINE()
Our goal is to compute total probability of error for REFINE() to fail. If we define this error as Ci C then we can define the main sources of error for this event.
1. Be rg((") such that no cluster in Co has cluster label If the cluster CE rg(C*) is absent in Co, then subsequent steps of REFINE() will never be able to recover it, they only involve node reclustering and merging existing clusters. The lemma presented below gives an upper bound on the probability of this event.
Lemma C.1. Under the conditions of temma 4.6 and if (Cminn)) then there exists constant > 0 such that
Pr(3ce rg(C*) such that no cluster in Co has cluster label
exp(-a11min)
Cmin
The proof of this Lemma is presented in appendix
2. Each cluster E rg(C)o should have <a fraction of impurities for some 1 > >>>>> If some cluster has more than a-fraction of impure nodes, then we cannot expect convergence guarantees for TrimmedMeanGD
The below lemma bounds the probability of this error as
Lemma C.2. For some constants 0 0<a< 22 a2 m (1,2) and at aat=0(m), under the conditions in lemma 4.6. we have
PrDeerg(CCo) which has > fraction of impurities
exp((aam)+ m
P ym
The proof of this Lemma is presented in appendix
3. MERGE() error: We define this the error for the MERGE () to fail. Even though MERGE() operates after RECLUSTER(), O), RECLUSTER () does not change the cluster iterates. The goal of MERGE () is to ensure that all clusters in Co with the same cluster labels are merged. Therefore, we define MERGE error as the event when either two clusters with same cluster label are not merged or two clusters with different cluster labels are merged. The below lemma bounds this probability.
Lemma C.3. min 2//3 A²
D2/3 Lllog(emm)) und for some constants un 0, then for some constand
d/3 0, where =
>0, 0, where B = / 2he1 we have nB 2
4dm
A'
Pr Error]
exp( exp(-asn 2L
The proof of this Lemma is presented in appendix
RECLUSTER() () error: This event is defined as a node going to the wrong cluster after both MERGE and (EFINE() operations. After MERGE(), each cluster Co corresponds to a single cluster in C,. Therefore, we incur an error due to the RECLUSTER operation if any node does not go the cluster cEC, which has cluster label C* (i). The below lemma provides upper bound the probability of this error.
Lemma C.4. If min 2//3 A2 ) u2d for some constants un 0, then for some constants a > 0 and 72 e (1,2
1 aB
we have
24
A' < 4d exp(-a"n +m nm
Pr Pr|RECIUSTER
(12)
2i
19

---

Published in Transactions on Machine Learning Research (02/202))
The proof of this Lemma is presented in appendix C.6
The total probability of error after for single step of REFINE() is the sum of probability of errors for these events by the union bound. Therefore,
Pr[ expl-011uam 010mm
m Cmin
m exp(-aam)
+(1-8)m P y + 8d exp( -ssm
A' ml Ly yrz
2
where we set 0g minnaagaaj))
For some small constants P1 > 0, P2 E (0, 1), we can choose 71 E (1, E (0, 2) and 72 E 21, - ) such that (1 - B) m + ( ) < 2m sm/m and for large enough m, A and n, Omin exep( (ssp-_(1mmnn) + exp(-agm) + 8d exp(-a3n 2L S 2m ..p. This happens because we have terms of exp(-m),cxp) (xxp(_Cmin) and exp( (x))-nA'), which decrease much faster than m which has terms of O(mexp(-nA)) where A and A are of the same order. Therefore, the total probability of error can be bounded by
Pr(C
PI ml-p2! P
(13)
C.3 Proof of lemma C.1
Pr]cerg((C' such that no cluster in Co has cluster label cl
Σ
Pr(No cluster in Co has cluster label c]
(14)
OEC*
Here, we use union bound over the clusters for the second inequality, Now, we analyze the probability that no cluster in rg(Co) has cluster label for some CE rgg(C*). Consider cluster in rg(Cu). This cluster has cluster label e if a majority of its nodes are from cluster e rg(C*). Since the size of each cluster in rg(Co) is atleast and there are K clusters in rg(C*), if all clusters in rg(C,) have S K nodes from cluster c, then no cluster will have cluster label c.
Assume that the clique formed by nodes from cluster has r nodes. Then, every node in cluster e, must. S. odges absent, which correspond to the edges between a node of the clique and those outside it. Thus, we obtain,
No cluster Co has cluster label c) SPr 2 { Σ
(1)= (i),
<Pr( ΣΣ Ky>S.(S. t 11
(i)==
<Prl ΣΣ (X++X,) Sctts))S.(S.- t +
SPr(
Σ X,>1
Se C*
KS.
Scxp(
1
t.
p
2 Sc)
KS.
m
In the first step, we require each node i to have S.- R wrong edges. For the second inequality, we remove the intersection and thus, the total number of incorrect edges has to be S.(S. 4), since each node has S.- f incorrect edges. For the third incquality, we use X,, < X, + X, and collect the terms of X, for the fourth inequality. In the fifth incquality, we obtain a condition on the sum of independent Bernoulli random variables each with mean 2. Therefore,
K
apply Chernoff bound for their sum to obtain the fifth inequality.
20

---

Published in Transactions on Machine Learning Research (02/2024)
A necessary condition for us is 1 KS.
which translates to t< R.). If we select & Cmin -1, this
a1 for some positive constant a1 If we choose t= (m), which is possible if if t- (Cmin)) as we assume Cmin (m)), then this is satisfied. We use the lower bound and Se 2 Cmin to obtain the final inequality. Plugging this in Eq equation 14 we obtain our result.
inequality is always satisfied. Note that want the term
1 1-
KS.
m
C.4 Proof of lemma C.2
Pr(3eErg(Co) which has 2 fraction of impurities)
Σ
Prfcluster has 2 fraction of wrong nodes]
(15)
eErg(Co)
We use a simple union bound on clusters in Co for the above inequality. Let the set of nodes the cluster which are from same cluster of C* as the cluster label of c, ie., which are not impurities, be Re. Then let Qa=|RA Rel. Let Q% denote the number of impurities in cluster c.
Prfcluster e has Za fraction of wrong nodes] =Pr[Q' 1-a
a -0:00-2--1000 Qal
Pr(Q]
We use the fact that Q,+Q'2t, which is the minimum size of any cluster, for the second inequality.
Now, we analyze the probability of single node to be incorrect. A node is an impurity in cluster c if it has edge to each of nodes in R..
Pr[Node i is an impurity in cluster c| [Pr[min |u.r-
(16)
jER.
Pr Pr|mmin min (1=0; ww -3 een
PrrAA + maxA, 2
jER
Now, if ma An for 71 (1,2), then we need A, (2-27)1) A for error.
Using the definition of random variables in appendix C.1
att 2 at maxA, A Pr ImaxA, 2 A1 yeR.
JER.
m
<Pr
Y Y. 2 al + +Pr[maxx 271
JER.
For the first inequality, we use union bound over the value of maxjER A, and for the second inequality, we need atleast at impurities, atleast at of all Y, should be 1.
We now bound the two terms in the final inequality separately.
For the second term, if maxjeR, An.
Pr[maxA, 27 S)..Pt =
JER.
m
Here, use union bound over all elements in Rs for the first inequality and the second inequality is plugging in the value of Pr[Y,,,, = =1], which we have already computed.
Now, we nced to provide a bound on Qc. Note that if Qe denotes the correct number of nodes, which corresponds to the majority of nodes, then (1-x)_ where S. the size of the cluster c.
For the first term, we can use Chernoff bound as Y- are independent random variables with expectation
1
t Y, 1 2a -eeepp-(e t m
Pri
m -|Y 5
i=
21

---

Published in Transactions on Machine Learning Research (02/2024)
We need / 2 |, implies at 2 since Y12-m is a bernoulli random variable. Further, we require (m), so that we can bound the probability using constant 02 20. 0. If we choose constant. independent of m, then we are done.
Now, plugging all these inequalities into Eq equation we get
Pr[3cerg(Ch, which has 2a fraction of wrong nodes)
$ <rg(C)ccp((aa>m)))
Σ
(1
(1-a)S,(
)m
m
eErg(Co)
< (g(Co))oxp(-a2m) mm P
ym
m
m i
expl (oxp( +(--m))((1a))))) ym
For the second inequality, we use Eeec, S. =m and for the third inequality, we use |rg(Co))(5m m
C.5 Proof of lemma C.3
First, let. i,j E [m]] be node in cluster o,d' E rg((Co) respectively such that (j) and C+ are the cluster labels of clusters and d' respectively. Then, if we repeat our thresholding analysis for MERGE operation, we obtain
dist(wī mez) - -V Veret S dist(we, (T) dist(w" (05) Veoe
where We,e' = dist(w wy dist(wer w Σ dist(wk,r,wfg
A=c.
We obtain the above equations by simple application of triangle inequality. Here, w is as defined in appendix C.1
To analyze the above quantities, we need to bound || werll and ||w| - ww for some j jEG. The following Lemmas provide these bounds.
Lemma C.5 (Convergence w..). If assumption 4.5 assumption 4.4 and assumption 4. hold, and ns t: then
2A Aa Veerg(Co)
(17)
where k= L and Ag is positive random variable with
Pr[A. 2 2d
2 +338s
+ /
22 2(133))
1-2,3 8f.
1-23
,2
2dd(1+ (c(op(((111))))
2i 212 })
(18)
5²
+(1 +(1-a)S,ex((-nmil)
2i 20- 1))
for some r,s,8>>0 where S. is the size of cluster c.
Proof is presented in appendix
Lemma C.6 (Distance between cluster minima and node minima). If assumption 4.3 and assumption 4.5 are satisfied then, for all j [m), where a node in cluster CEC, where C* (j) is the chuster label of node e, we have
2Le1
(wt-wi) ||
=
(19)
1
22

---

Published in Transactions on Machine Learning Research (02/2024)
Proof is presented in appendix C.9
Now, that we have our required quantities, we are ready to analyze the probability of error after the merge and reclustering operations.
First, we analyze the probabilty of MERGE() operation. Note that correct nodes of and d' were from the same cluster C* then, l|us - -Wll e. Gg. If correct nodes of d' and e were from different clusters in C*,, then, |w] will 2 Ve, EG.. Therefore, the probability of MERGE () error is upper bounded by
Pr Pr(MERGE() Error] SPr(at least 1 edge is incorrect]
SPr[max o, 2 A
0,0'
2A
<Pr[max G,a"
221 -41-2(1-x~))722---
A
kmm.e!
A
< max oErg(Co)
Pr[Ae
2
(1-k' - 1)1
2
2
I max eErg(Co)
Pr(A. 21
(20)
A'
I max cErg(Co)
4dexp(- (-6:p
2L
A'
4dm
exp( dgy A'
4dexp( (00(-aa)) a'z71
)
(21)
21.
1
2f.
cErg(Co)
For the second incquality, we expand all the terms of & c.e'. We set
A 11
(11- K 1) TT/2 1-/2--BB BBB Then, weset
r= (i.max
Sevdi 1)s==)(Lmaax
2log(S.)
A'
2log( and 8 0 ( Now, if
5.vat
Seval.
Vevdi.
d=iQ(min
p2/3 ((llg((mma) 1), such that /2d +-3 +
1/3
-
sizA, then there exist some constant d3 > 0 We then use the union bound, followed by |rg(Co)) s 7
1-29
1-29
such that the second incquality is satisfied by lemma C.5
C.6 Proof of lemma C.4
We can apply our thresholding analysis to
(w,.r-W;rrll for rgr(CCo). First, let j be a node in cluster such that C**(j is the cluster label of
dist w3 w=)) + dist(distttwwss.o.. < _u ) Be.
where L.,; = dist a))) +2) dist(w" + dist w,r,,,;)
From appendix B and appendix C.5 we have bounds for all the terms involved. Note that after merging, each cluster in C should have only cluster in C1. Therefore, after we recluster according to || - (We.r-W;rTll we incur an error if i goes to the wrong cluster. Suppose that the corresponds to the correct cluster for and df is the cluster to which it is assigned, with c, ef Then,
Pr Pr[Reclustering Error] <P Pr[maxxaa We, T - w. -
ve(m) c'
<Pr[maxmaxe (E(m) d'cc
& S Se Po.
C2 -1
<Pr <Pr[maxmaxy t 0, vE[m]c' EC%
<Pr[max max (A + A,) > A+A
(22)
VE(m)c'EC'
<Pr[maxA, 2 A - (72 1) + Pr CEC
max 2 A]
3m/m)
< max Pr[A. A + 2
(23)
eErg(Co)
23

---

Published in Transactions on Machine Learning Research (02/2024)
the thresholding analysis (uier-Warrrll For the third inequality, we rearrange the terms and combine max over c' f with and use. For the fourth inequality, expand the terms of o. and substitute the values of A and A', using the inequality Ax 2 For the fifth inequality,
For the second inequality,
BB and break the terms using union bound such that A" = A' (72 1) O. consider some 7/2 E (1,2 2
Finally, we use the union bound CE rg(Co)' and iE [m].
Now, we bound the two terms in Eq equation 23 separately, The second term can be bounded in terms of Y2,ya Thus,
max maxXPrAA,2224 maxPr(X =1] ml 2 2 m
(24)
it/m]
We use expectation of Yi, calculated in appendix C.4 and then bound max by sum.
For the first term, our analysis is similar to that of MERGE() error. Assume that there is some constant U2 > 1 such that A" 2 uz We set 8 = O P
O(L max SevaL V
S.vaL H),
/A'
(l max {
A' 2log(Se)1), 1), and if d 22min D2/af2/3> (oolg( }), such that 2(1+3 si> A' then there exist some constant ag > such that the second inequality is satisfied 1-29
A'
2log((e)) S.var.
2//344/ n²A?
Sevar.
v 2d
+/2
1-28
by lemma
C,5
We then use the union bound, followed by |rg(Co)) m.
cErg(Ca)' max Pr[A.2A']| max 4dexp(- - (-ag A' cErg(Co)
(25)
2L
Σ
4dexp(- agag
(26)
21
cErg(Ca)'
4dm
$ t
exp( - az
(27)
2L
C.7 Proof of lemma C.5
First, intermediate Lemma from Yin et al. 2018 This characterizes the behavior of rimmeddean gradient estimator.
Lemma C.7 (TrimmedMean Estimator Variance). Let ge(w) be the output of TrMeany estimator for cluster cECo with size of cluster S.- Iʃ assumplion 4.5 holds, then
- (lec(w)- VFF(w))
where Pr|A> v 2d +333s
1-28 + V2 2(1+33)
1-2,3
D - (e00-((1a)Ssmminn)
,.2
(28)
<2d(1+
2i. 21.
s²
+(1 (11-a))eep(((ni))
2F. 2F.221.
for some r,8,5 O.
Proof. The proof of this Lemma follows from coordinale-wise sub-exponential distribution of F... Since loss per sample f(w,z) is Lipschitz in each its coordinates with Lipschitz constant Lg for be [d]. Thus, F((w) is also Lg-Lipschitz for each coordinate E [d] from corollary E.6 Now, every subgaussian variable with variance o² is o--uub exponential. Thus, each coordinate of V., w. (..2) is L-sub-exponential, since Î Vk [d]. The remainder of proof can be found in Appendix F. in Yin al. 2018
24

---

Published in Transactions on Machine Learning Research (02/2024)
Now, using the above Lemma, can bound the iterate error for cluster ceCo. Consider ||e,141 - w/
| - <
-w
< <|lo, - Vg((wo))
Wect 17 ()w,,. - w/ll (VFlo(())))-((uu.)))))
We,t - VV(,, - +n
Now, we bound (Ver-1//F(44 -(((w,rr)-333) wwll using u-strong convexity and L-smoothness of Fe. The analysis is similar the convergence analysis in appendix B.1 Thus, for n +
-w]1 m
Using this bound we can analyze the original term with ||
|4rr41-w5ll S -
T-
(1-11)++2
t=0
<(1=x-1)1/6 - will+nA1 ( 7/ )
2 1
5(1-x D A
For the second inequality, we use k L and unroll the recursion for steps. For the third inequality, we use v1-r<i 1 and upper bound the finite geometric sum by its infinite counterpart. Finally we use the boundedness of W and the sum of the geometric series to get our result.
C.8 Proof of lemma B.2
We present the proof for this lemma here as it is corollary of lemma C.5
Now, if we set a = 8 = 0 and S. = 1, we obtain the generalization Further, we do not need the terms of s as they appear with B, and thus, we can choose very large, so that we can ignore its contribution to error probability. The remainder of the proof follows that of lemma
We utilize the intermediate lemma C.7
guarantee for GD on a single node E m
25

---

Published in Transactions on Machine Learning Research (02/2024)
C.9 Proof of lemma C.6
Since Fg is I-Lipshccitz and j--strongly convex with minima w
) -F((wE)
F((w))-FFp) -F)(w²)
E
Qc
Qc
j+(CC,l)==
F((w;) F((wo)
) FF)(33)
Σ
Qc
Qc
i+i.Coli)=o
14 1/ - will
i 1 wIll
2Qc
Qc
jti,Col)=a
1|| w
(Qa-
w - w|ll 11
2Qc
Qc
W/ wwll? - u1 wlll
(Qc-
2Qc
Qc
2LE1 1
-wāll 112=
2t.11
||0] - =1
For the first equation, we expand F. into its component terms, where Qe denotes the number of correct nodes in cluster c. For the second inequality, we use the fact that w arg (w). For the third inequality, strong-convexity of F; and î-Láppchiizzess for Fj.i.. For the fourth inequality, ((T))=F((<u) using p-strong convexity of Fo. Finally, manipulate the remaining terms to obtain the final bound.
lower bound
D Proof of theorem 4.14
By theorem 4.8 CR with probability ( PTTP) For the (R+1)th step, we bound probability of error
by Therefore, with probability 1-exp(-g0 R)p. For the (R+1)th step, we optimize the cluster iterates from TrimmedMeanGD() to improve convergence instead of clustering error. Since CR+1 =Ca, each cluster in CR+1 maps to some cluster in C*. Without loss of generality, assume that cluster cE rg(CR+1) maps to the same cluster cEC. Now, if (C,,C2,....C4} are the clusters in Cn which merged to form cluster cE e rg(Cn+1).. Then, we can write
| will =
(W w:)
7
1=1
1
Wey -w
J=1
+ +|| wt)
Wey. r w6,
1=1
For the first incquality, we used the definition of war from MERGE(). For the sccond incquality, we used the triangle incquality for the elements. The third inequality is obtained by using triangle inequality and adding and subtracting w defined in appendix C.1
C.
26

---

Published in Transactions on Machine Learning Research (02/2024)
Now, consider the set of nodes f2..d S [m], such that iy EecVVjell and (*(1,) [1]. Therefore, we can split each term of | ww
We, w | - ww, brrs|+|luug ||| will)
| - 5
(
==1
Σ 7 ==1
We, w w +2B
From lemma C.6 since iz contributes to both clusters ej and e*, bound the difference from their minima by B. Further, we use lemma
C.5 and the lemma C.7 which is adapted from Theorem 4 in Yin et al.
(2018),t bound the convergence of Way, - wt,
|| If we set 8 nSe,
and
8d log(1 (gg(+nS.LL),
8d n.S.y
r= Lmax{ n.
log( (1InnnSSLD))
4d 8 Lmax
(dlog(1 + nSe, LD) logm). 4d (dlog(1+nSe, Í.D) logm))
n
where Sey is the size of cluster cj, obtain
||wrr - w ()
where
id
A' =
log log(nnmaxS LD)
1-23 V1 VTCCmin
JE(
We can further upper bound max by m. Now, the probability of error for each cluster e cErg(CR) for given values of and is
4d (1+ne DDya' therefore, can use union bound and multiply this probability of error by rg(CR) & m. Since t= OCCmi)), we can upper bound this by mu" for some positive constant Cmin. Omdo
E
Additional Definitions and Lemmas
We start with reviewing the standard definitions of strongly convex and smooth functions f:R ++ R. Definition E.1. f p-strongly convex if Vw,w', f(w') (V) -w) w)
I|w' ww|1?
Definition E.2. f L-smooth if Vw,w", V
Definition E.3. k-th coordinate of V fww). Lemma E.4. If J.g::Re Rd IR are two p-strongly conver functions on domain W. Then, conver on the same domain.
La Lipschitz for every coordinate E [4] if, |o, [0,[(w) < La, where (u](w) denotes the
{tg is also p-strongly
2
Proof. If f and are p-strongly convex on a domain W, then for any w1 W
f(w)) 2 f(m)) (/) / [ml - wol 2
g(=)) -wo) || wall
Adding the above equations, we gel
f(11++((w1) f(wo))
V V((wo) Vg(wo)
,01--10 - ww ) 2
1 - woll
Thus, 1t is also p-strongly convex.
27

---

Published in Transactions on Machine Learning Research (02/2024)
It is also L-smooth Lemma E.5. If f.g: Rd +R are two L-smooth functions domain W. Then, 2 the some domain.
Corollary E.6. If f.g. RR two L-Lipschitz functions domain W. Then, Ita is also L-Liscchitz on the same domain.
Proof. Consider the following term for any 11,40 E W
V Vwr ) (g((o))
Vi(wa) + VVg Vg(wn)
2
2
((V.(w)) (/(":/))))/())(((V(()(() (
[[V ( V/ Vy(w)))
2 (4||w)- wool + v - w
<b <
In the second inequality, use the triangle inequality of norms. For the third inequality, use the I-smoothness of f and g. Thus, f is also I-smooth The proof of the corollary is same as above, by replacing terms of V f and Vg by f and respectively.
Lemma E. If each coordinate of a function f: + R is La Lipschitz for ee [d] on the domain W, then f is -- - 1 -ipiscctt on the same domain W.
Proof. Consider w1,40 E W.Define a sequence of variables
{w[kk kk] = ((wi) (,(()))))))))))))))())) )2..)), )a, (wo) (wo))) 32 Then, wi = wrld] and wio [0
d
1f(wi)-f(wo)
f((lk)))f(uu( 1)))
La (a((w))--(w))))
The second inequality follows by using triangle rule. Then, ((w(A)) and f(w - 1]) differ only in the Ath coordinate, so we apply La coordinate-wise Lipschitzness. Now, consider a random variable ve IRd such that Uj: Lxx (wn (wooeel if (wo )a (wo) 00, else O. Then,
(wa)e (woo)
La((wr)e (wo))) (w, wo) k=
||ll|lw - wall
L ll w1 woll
Here, we use the Cauchy-Schwartz incquality for the second step. Then, note that each coordinate of v is bounded by Lg-
28


Reviews:
[{"review_id": "Asr9aqUOZe", "reviewer": "Reviewer_RdDU", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. Introduction of a reasonably intuitive framework.\n2. Theoretical analysis is comprehensive, albeit a bit derivative of prior works. \n\nWeaknesses:\n\n1. The paper's notation is all over the place. For example, on page 4, the usage of terms such as C_0 precedes before defining them (which they do later in the page 5 in the description of the routines). T (in w_{I,T}) is undefined similarly until page 5. Further terms like rg{C} are never defined and it took the reader a while to figure out what they meant. \n2. Apart from the framework, I feel the novelty is limited. Indeed the algorithm is a natural variant of the standard algorithm used in k-means. The analysis needs some work but uses standard apparatus (union bounds and probabilistic methods). \n3. The biggest let down, in my opinion, is the experimental analysis. It needs to be lot more comprehensive. Firstly, a lot of detail is missing from a reproducibility point of view: for example, the size of the neural network model used is not specified. Indeed this is important, as the similarity metric is defined based on the weight vector. Also, it is mentioned that only a fraction of clients participate in the model training. What fraction is used? If the fraction is too low, it could explain why the performance of the models on a simple dataset such as MNIST is so low (<90%) for the global model. Also, it would be instructive to see what the impact of model size would be on the performance of the algorithm w.r.t the baselines. Indeed, the whole setup seems to create an unnatural performance gradient from the baseline. For instance, simply taking a convolutional neural network, instead of feedforward network, could boost the model performance across the board and diminish the impact of the proposed algorithm. The best known models on MNIST yield a >99% accuracy.  \n4. One of the key claims that the paper makes is that it gets rid of the initialization problem. However, since they need to perform a rather expensive hyper parameter tuning on the $\\lambda$, this seems like a moot point. Indeed, one of the aspects missing from the results is the runtime analysis for different methods (including HPO).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces a \"novel\" clustering framework for training local models in a federated learning setup. The key contribution of the work is an iterative refinement algorithm (reminiscent of the EM-algorithm), its theoretical and experimental analysis. The key idea is to cluster \"similar clients\" in the federated setup  and train individual models for each cluster. The aim here is to address heterogeneity in client data. Similarity of clients is defined in terms of distance between the locally trained model weights on the client data. Prior works address the heterogeneity by training local models for each client. The key advantage here seems to be that they are able to leverage data across a few similar clients.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "As mentioned in \"Strengths and Weaknesses\", the following are the requested changes:\n\n1. Fix the terminology issues, ensure that the terms such as C_0, T, rg(.) etc are appropriately defined before usage.\n\n2. The experimental section needs to be expanded significantly. Firstly, the reproducibility issue must be fixed: give a detailed description of the all the parameters, hyper parameter ranges used. Include an empirical analysis of the algorithm runtime. More detailed experiments are requested analysis the impact of a) size of the model used b) change in the participation fraction of the clients c) number of clients (which seems to be too high for the 2-4 clusters used) and data size with each client. I think the datasets are either synthetic or semi-real (because the rotations based clusters are created). It would be great if some experiments can be included involving real heterogeneity (this needs some thinking through but perhaps images from different classes - cats vs dogs, could be one starting point).", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "asZl0gBTJX", "reviewer": "Reviewer_7Gsg", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n1) The proposed similarity metric across different clients (models) is interesting.\n2) Extensive theoretical results are provided for the proposed framework;\n3) The effectiveness of the proposed method is validated empirically;\n\nWeakness:\n1) According to my understanding, the proposed method would require $O(n^2)$ communication complexity to measure the similarity between all pairs of clients (models). I wonder if this can be improved via some heuristic algorithm.\n2) The experiments of the paper are mainly focus on domain heterogeneity (e.g., rotation, inversion, etc.). The empirical evidence can be enhanced with more results on label heterogeneity (e.g., Dirichlet heterogeneity as in [1-2]).\n\n[1] Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Nghia Hoang, and\nYasaman Khazaeni. Bayesian nonparametric federated learning of neural networks. In International Conference on Machine Learning, pp. 7252–7261. PMLR, 2019.\n\n[2] Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni.\nFederated learning with matched averaging. arXiv preprint arXiv:2002.06440, 2020a.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper targets the task of Federated clustering. The main insight of the proposed algorithm is to cluster clients according to the simlarity of the model parameters, and thus striking a balance between handle heterogeneous clients and simultaneous training. Different from the previous proposed approaches, this framework do not need to pre-define the number of clusters. Thorough theoretical analysis are also provided for the proposed framework.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. I suggest the author discuss the communication complexity of this framework and previous methods, as well as potential ways to improve this complexity.\n2. I suggest the author include more results on label heterogeneity.", "broader_impact_concerns": "Broader impact statemnent is not needed in my opinion.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "TCqxD35xAw", "reviewer": "Reviewer_Xk4R", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The paper proposes a novel clustered federated learning algorithm SR-FCA which doesn't require any specific initialization.\n2. The convergence guarantees for the proposed method are presented.\n3. The experiments are conducted on both simulated and real heterogeneous distributed datasets.\n4. All the relevant baselines are provided and the proposed algorithm either outperforms or is comparable with the baselines.\n5. The paper is well written and structured. However, there is room for improvement in Algorithm section (sec 3) and conclusion.\n6. It is appreciated that the authors have reported the total time spent for conducting all the experiments and hardware details. \n\nWeaknesses/questions:\n1. It is not clear to me why recluster subroutine needs to be before merge. Recluster doesn't effect the model parameters of cluster centers $w_{c,T}$ and merging only changes $w_{c,T}$. Is there any difference in running recluster module before merge and after.\n2. In the one_short method: do we have to train the model on each client till convergence? Do all the clients' models have to be initialized to the same point?\n3. For the experimental results, can you provide some insights on the accuracy improvements obtained by using TrimmedMean over FedAvg?\n4. \" Note that the complexity of the REFINE step is the same as that of IFCA in terms of both computation time and communication since in each case, we need to find the loss of every cluster model on every client’s data.\" It is not clear to me why this statement is true.\nFor IFCA on the workers end, the compute overhead as compared to traditional FL is because of the cluster estimation which requires $kT$ additional forward passes (referring to line 6 of algorithm 2 in [1]) and the communication overhead is $k \\times$ where $k$ is the number of clusters. Can you clarify why it is the same for SR-FCA's refine step?\n5. How do we calculate empirical cluster loss $f_c$? \n6. Do all the clients have same test dataset? \n\n[1] Ghosh, Avishek, et al. \"An efficient framework for clustered federated learning.\" Advances in Neural Information Processing Systems 33 (2020): 19586-19597.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a novel framework for clustered federated learning called *Successive Refine Federated Clustering Algorithm* to deal with data heterogeneity issue. \n1. The algorithm starts with a ONE_SHOT subroutine which trains a local model on each client independently for few iterations and then clusters them based on pair-wise distances.\n2. Now under REFINE subroutine, the algorithm runs robust federated learning (TrimmedMean) for each cluster.\n3. RECLUSTER subroutine refines the clusters by reallocating the cluster labels for each client based on the distance been local client model and the cluster center model. This aims to reduce the impurity level of each cluster.\n4. MERGE subroutine merges the clusters whose pairwise cluster centers distance is less than a set threshold.\n\nThe paper presents convergence guarantees of SR-FCA when the loss is strongly convex and smooth, and dist(.,.) is $L_2$ norm. The paper shows that the misclustering error in the first stage of SR-FCA is O(mdexp(−n/√d) ( Lemma 4.6) and the successive stages further reduce the misclustering error by a factor of O(1/m) (Theorem 4.8).\n\nThe experimental results are shown on both simulated and real heterogeneous distributed datasets. SR-FCA (with cross-cluster loss) is compared with multiple relevant baselines. On simulated datasets, SR-FCA obtains test accuracy no worse than the best baseline and is able to recover the correct clustering. For CIFAR10, in particular, SR-FCA has 5% better test accuracy than IFCA. On real datasets, SR-FCA outperforms all baselines.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Algorithms 1-4 are hard to follow. Rewrite the algorithms for better understanding and clearly differentiate the worker's computation from the server. Include the participation rate into algorithm 3. \n2. An ablation study on $\\lambda$ would be interesting.\n3. Add more discussion on computation and communication complexity of SR-FCA. This could help in better understanding of the benefits and limitations of the proposed method.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper was reviewed by three reviewers with expertise ranging from theory to more practical aspects of the clustering and federated learning. The reviewers gave extensive initial feedback which was successfully addressed by the authors through revisions and responses. All reviewers are satisfied with the revised paper and recommending the paper to be (conditionally) accepted modulo small minor remaining comments detailed below:

* Please include experiments on label heterogeneity with a larger dataset, such as CIFAR10 to showcase the scalability of the proposed framework.

* Please revise the paper for more clarity in response to Reviewer Xk4R's latest comments.

Congratulations!
Decision: Accept (poster)