# Provably Personalized and Robust Federated Learning

Mariel Werner1 Lie He2 Michael Jordan1 Martin Jaggi2 Sai Praneeth Karimireddy1

( 1Department of Electrical Engineering and Computer Sciences,
University of California, Berkeley
2Machine Learning and Optimization Laboratory (MLO),
EPFL, Switzerland )

###### Abstract

Identifying clients with similar objectives and learning a model-per-cluster
is an intuitive and interpretable approach to personalization in federated
learning. However, doing so with provable and optimal guarantees has remained
an open challenge. We formalize this problem as a stochastic optimization
problem, achieving optimal convergence rates for a large class of loss
functions. We propose simple iterative algorithms which identify clusters of
similar clients and train a personalized model-per-cluster, using local client
gradients and flexible constraints on the clusters. The convergence rates of
our algorithms asymptotically match those obtained if we knew the true
underlying clustering of the clients and are provably robust in the Byzantine
setting where some fraction of the clients are malicious.

##  1 Introduction

We consider the federated learning setting in which there are $N$ clients with
individual loss functions $\\{f_{i}\\}_{i\in[N]}$ who seek to jointly train a
model or multiple models. The defacto algorithm for problems in this setting
is FedAvg [1] which has an objective of the form

$\displaystyle
x^{*}_{\text{FedAvg}}=\operatorname*{arg\,min}_{x\in\mathcal{X}}\frac{1}{N}\sum_{i\in[N]}f_{i}(x).$
(1)

From (1), we see that FedAvg optimizes the average of the client losses. In
many real-world cases however, clients’ data distributions are heterogeneous,
making such an approach unsuitable since the global optimum (1) may be very
far from the optima of individual clients. Rather, we want algorithms which
identify clusters of the clients that have relevant data for each other and
that only perform training within each cluster. However, this is a challenging
exercise since 1) it is unclear what it means for data distributions of two
clients to be useful for each other, or 2) how to automatically identify such
subsets without expensive multiple retraining [2]. In this work we propose
algorithms which iteratively and simultaneously 1) identify $K$ clusters
amongst the clients by clustering their gradients and 2) optimize the clients’
losses within each cluster.

###  1.1 Related Work

##### Personalization via Clustering.

Personalization in federated learning has recently enjoyed tremendous
attention (see [3, 4] for surveys). We focus on gradient-based clustering
methods for personalized federated learning. Several recent works propose and
analyze clustering methods. [5] alternately train a global model with FedAvg
and partition clients into smaller clusters based on the global model’s
performance on their local data. [6] and [7] instead train personalized models
from the start (as we do) without maintaining a global model. They iteratively
update $K$ models and, using empirical risk minimization, assign each of $N$
clients one of the models at every step. In Section 2 we analyze these
algorithms on constructed examples and in Section 3.1.3 compare them to our
method.

Since our work is closest to [7], we highlight key similarities and
differences. Similarities: 1) We both design stochastic gradient descent- and
clustering-based algorithms for personalized federated learning. 2) We both
assume sufficient intra-cluster closeness and inter-cluster separation of
clients for the clustering task (Assumptions 1 and 2 in their work;
Assumptions 4 and 5 in ours). 3) Our convergence rates both scale inversely
with the number of clients and the inter-cluster separation parameter
$\Delta$. Differences: 1) They assume strong convexity of the clients’ loss
objectives, while our guarantees hold for all smooth (convex and non-convex)
functions. 2) They cluster clients based on similarity of loss-function values
whereas we cluster clients based on similarity of gradients. We show that
clustering based on loss-function values instead of gradients can be overly
sensitive to model initialization (see Fig. 1(b)). 3) Since we determine
clusters based on distances in gradient space, we are able to apply an
aggregation rule which makes our algorithm robust to some fraction of
malicious clients. They determine cluster identity based on loss-function
value and do not provide robustness guarantees.

Recently, [8] established lower bounds showing that the optimal strategy is to
cluster clients who share the same optimum. Our algorithms and theoretical
analysis are inspired by this lower-bound, and our gradient-based clustering
approach makes our algorithms amenable to analysis à la their framework.

##### Multitask learning.

Our work is closely related to multitask learning, which simultaneously trains
separate models for different-but-related tasks. [9] and [10] both cast
personalized federated learning as a multitask learning problem. In the first,
the per-task models jointly minimize an objective that encodes relationships
between the tasks. In the second, models are trained locally (for
personalization) but regularized to be close to an optimal global model (for
task-relatedness). These settings are quite similar to our setting. However,
we use assumptions on gradient (dis)similarity across the domain space to
encode relationships between tasks, and we do not maintain a global model.

##### Robustness.

Our methods are provably robust in the Byzantine [11, 12] setting, where
clients can make arbitrary updates to their gradients to corrupt the training
process. Several works on Byzantine robust distributed optimization [12, 13,
14, 15, 16] propose aggregation rules in lieu of averaging as a step towards
robustness. However, [17, 18] show that these rules are not in fact robust and
perform poorly in practice. [19] are the first to provide a provably
Byzantine-robust distributed optimization framework by combining a novel
aggregation rule with momentum-based stochastic gradient descent. We use a
version of their centered-clipping aggregation rule to update client
gradients. Due to this overlap in aggregation rule, components of our
convergence results are similar to their Theorem 6. However, our analysis is
significantly complicated by our personalization and clustering structure. In
particular, all non-malicious clients in [19, 20] have the same optimum and
therefore can be viewed as comprising a single cluster, whereas we consider
multiple clusters of clients (without necessarily assuming clients are i.i.d.
within a cluster). The personalization algorithm in [10] also has robustness
properties, but they are only demonstrated empirically and analyzed on toy
examples.

##### Recent Empirical Approaches.

Two recent works [21, 22] examine the setting in which clients’ marginal
distributions $p(x)$ differ, whereas most prior work only allows their
conditional distributions $p(y|x)$ to differ. One of our experiments (Section
4.1) assumes heterogeneity between clients’ marginal distributions, while the
others (Section 4.2) assume heterogeneity only between their conditional
distributions. In [21], the server maintains a global pool of modules (neural
networks) from which clients, via a routing algorithm, efficiently select and
combine sub-modules to create personalized models that perform well on their
individual distributions. Extending the work of [23], [22] model each client’s
joint distribution as a mixture of Gaussian distributions, with the weights of
the mixture personalized to each client. They then propose a federated
Expectation-Maximization algorithm to optimize the parameters of the mixture
model. In general, the contributions and style of our work and these others
differ significantly. We focus on achieving and proving optimal theoretical
convergence rates which we verify empirically, whereas [21] and [22] emphasize
empirical application over theoretical analysis.

###  1.2 Our Contributions

To address the shortcomings in current approaches, we propose two personalized
federated learning algorithms, which simultaneously cluster similar clients
and optimize their loss objectives in a personalized manner. In each round of
the procedure, we examine the client gradients to identify the cluster
structure as well as to update the model parameters. Importantly, ours is the
first method with theoretical guarantees for general non-convex loss
functions, and not just restrictive toy settings. We show that our method
enjoys both nearly optimal convergence, while also being robust to some
malicious (Byzantine) client updates. This is again the first theoretical
proof of the utility of personalization for Byzantine robustness. Specifically
in this work,

  * •

We show that existing or naive clustering methods for personalized learning,
with stronger assumptions than ours, can fail in simple settings (Fig. 1).

  * •

We design a robust clustering subroutine (Algorithm 3) whose performance
improves with the separation between the cluster means and the number of data
points being clustered. We prove nearly matching lower bounds showing its
near-optimality (Theorem 3.2), and we show that the error due to malicious
clients scales smoothly with the fraction of such clients (Theorem 3.1).

  * •

We propose two personalized learning algorithms (Algorithm 2 and Algorithm 4)
which converge at the optimal $\mathcal{O}(1/\sqrt{n_{i}T})$ rate in $T$ for
stochastic gradient descent for smooth non-convex functions and linearly scale
with $n_{i}$, the number of clients in client $i$’s cluster.

  * •

We empirically verify our theoretical assumptions and demonstrate
experimentally that our learning algorithms benefit from collaboration, scale
with the number of collaborators, are competitive with SOTA personalized
federated learning algorithms, and are not sensitive to the model’s initial
weights (Section 4).

##  2 Existing Clustering Methods for Personalized Federated Learning

Our task at hand in this work is to simultaneously learn the clustering
structure amongst clients and minimize their losses. Current methods do not
rigorously check similarity of clients throughout the training process.
Therefore they are not able to correct for early-on erroneous clustering (e.g.
due to gradient stochasticity, model initialization, or the form of loss-
functions far from their optima). In the next section we demonstrate such
failure modes of existing algorithms.

###  2.1 Failure Modes of Existing Methods

The first algorithm we discuss, Myopic-Clustering, does not appear in the
existing literature, but we create it in order to motivate the design of our
method (Algorithm 2). In particular, it is a natural first step towards our
method, but has limitations which we correct when designing our algorithm.

##### Myopic-Clustering (Algorithm 1).

At every step, each client computes their gradient at their current model and
sends the gradient to a central server. The server clusters the gradients and
sends each cluster center to the clients assigned to that cluster. Each client
then performs a gradient descent update on their model with their received
cluster center. This is a natural federated clustering procedure and it is
communication-efficient ($\mathcal{O}(N)$). However, it has two issues: 1) If
it makes a clustering mistake at one step, models will be updated with the
wrong set of clients. This can cause models to diverge from their optima,
gradients of clients in the same cluster to drift apart, and gradients of
clients in different clusters to drift together, thus obscuring the correct
clustering going forward. Furthermore, these errors can compound over rounds.
2) Even if Myopic-Clustering clusters clients perfectly at each step, the
clients’ gradients will approach zero as the models converge to their optima.
This means that clients from different clusters will appear to belong to the
same cluster as the algorithm converges and all clients will collapse into a
single cluster.

(a) Myopic-Clustering ($\eta=0.5$). Correct clustering: $\\{1,2\\}$ and
$\\{3\\}$. Client $\\{2\\}$ gets stuck at $x=1$, not reaching its optimum, and
clients $\\{1,3\\}$ converge to their optima. All gradients being $0$ at this
point, the clients are incorrectly clustered together: $\\{1,2,3\\}$.

(b) IFCA/HypCluster. Correct clustering: $\\{1\\},\\{2\\}$. Both clients’
function values are smaller at initialization point $x_{2,0}$ than $x_{1,0}$
causing IFCA/HypCluster to initially cluster them together. Since the average
of the clients’ gradients at $x_{2,0}$ is $0$, the models never update and the
algorithm thinks the initial erroneous clustering, $\\{1,2\\}$ is correct.

(c) Clustered Federated Learning. Correct clustering: $\\{1\\}$, $\\{2,3\\}$
(client 3 not drawn due to its stochastic gradient – details on pg. 5).
Clustered FL averages gradients of clients $\\{1\\}$ and $\\{2\\}$ to $0$,
clustering them together, and with non-$0$ probability clusters $\\{3\\}$
separately due to its stochastic gradient. Based on this initial erroneous,
the algorithm partitions the clients $\\{1,2\\}$ and $\\{3\\}$ and recursively
runs on each group, never recovering the correct clustering.

Figure 1: We show how existing personalized FL algorithms miscluster and fail
to converge on constructed examples.

The following example (Fig. 1(a)) demonstrates these failure modes of Myopic-
Clustering.

Let $N=3$ and $K=2$, with client loss functions

$\displaystyle f_{1}(x)$ $\displaystyle=\frac{1}{6\eta}x^{2}$ (2)
$\displaystyle f_{2}(x)$
$\displaystyle=\begin{cases}4(x-1)^{3}+3(x-1)^{4}+1&x<1\\\
\frac{1}{2\eta}(x-1)^{2}+1&x\geq 1,\end{cases}$ (3) $\displaystyle f_{3}(x)$
$\displaystyle=\frac{1}{2\eta}(x-2)^{2},$ (4)

where $\eta$ is the learning rate of the algorithm. With this structure,
clients $\\{1,2\\}$ share the same global minimum and belong to the same
cluster, and client $\\{3\\}$ belongs to its own cluster. Suppose Myopic-
Clustering is initialized at $x_{0}=1.5$. At step 1, the client gradients
computed at $x_{0}=1.5$ are $\nicefrac{{1}}{{2\eta}}$,
$\nicefrac{{1}}{{2\eta}}$, and $\nicefrac{{-1}}{{2\eta}}$ respectively.
Therefore, clients $\\{1,2\\}$ are correctly clustered together and client
$\\{3\\}$ alone at this step. After updates, the clients’ parameters will next
be $x_{1,1}=1,x_{2,1}=1,x_{3,1}=2$ respectively. At this point, clients
$\\{2,3\\}$ will be incorrectly clustered together since their gradients will
both be $0$, while client $\\{1\\}$ will be clustered alone. As the algorithm
proceeds, clients $\\{2,3\\}$ will always be clustered together and will
remain at $x=1$ and $x=2$ respectively, while client $\\{1\\}$ will converge
to its optimum at $x=0$. Consequently, two undesirable things happen: 1)
Client $\\{2\\}$ gets stuck at the saddle point at $x=1$ which occurred when
it was incorrectly clustered with client $\\{3\\}$ at $t=1$ and subsequently
did not recover. 2) All gradients converge to $0$, so at the end of the
algorithm all clients are clustered together.

Algorithm 1 Myopic-Clustering

Input Learning rate: $\eta$. Initial parameters:
$\\{x_{1,0}=...=x_{N,0}=x_{0}\\}$.

1:for round $t\in[T]$ do

2: for client $i$ in [N] do

3: Client $i$ sends $g_{i}(x_{i,t-1})$ to server.

4: Server clusters $\\{g_{i}(x_{i,t-1})\\}_{i\in[N]}$, generating cluster
centers $\\{v_{k,t}\\}_{k\in[K]}$.

5: Server sends $v_{k_{i},t}$ to client $i$, where $k_{i}$ denotes the cluster
to which client $i$ is assigned.

6: Client $i$ computes update: $x_{i,t}=x_{i,t-1}-\eta v_{k_{i},t}$.

7:Output: Personalized parameters: $\\{x_{1,T},...,x_{N,T}\\}$.

To further motivate the design choice for our algorithms, we now discuss three
clustering-based algorithms in the literature on personalized federated
learning. In particular, we generate counter-examples on which they fail and
show how our algorithm avoids such pitfalls.

The first two algorithms IFCA [7] and HypCluster [6] are closely related. They
both cluster loss function values rather than gradients, and like our
algorithm they avoid the myopic nature of Myopic-Clustering by, at each step,
computing all client losses at all current cluster parameters to determine the
clustering. However, as we show in the next example (Fig. 1(b)), they are
brittle and sensitive to initialization.

##### IFCA [7].

Let $N=2$, $K=2$ with loss functions

$\displaystyle f_{1}(x)$ $\displaystyle=(x+0.5)^{2}$ (5) $\displaystyle
f_{2}(x)$ $\displaystyle=(x-0.5)^{2},$ (6)

and initialize clusters 1 and 2 at $x_{1,0}=-1.5$ and $x_{2,0}=0$
respectively. Given this setup, both clients initially select cluster 2 since
their losses at $x_{2,0}$ are smaller than at $x_{1,0}$.

Option I: At $x_{2,0}=0$, the client gradients will average to 0. Consequently
the models will remain stuck at their initializations, and both clients will
be incorrectly assigned to cluster 2.

Option II: Both clients individually run $\tau$ steps of gradient descent
starting at their selected model $x_{2,0}$ (i.e. perform the LocalUpdate
function in line 18 of IFCA). Since the clients’ individually updated models
will be symmetric around $0$ after this process, the server will compute
cluster 2’s model update in line 15 of IFCA as: $x_{2,1}\leftarrow 0=x_{2,0}$.
Consequently, the outcome is the same as in Option I: the models never update
and both clients are incorrectly assigned to cluster 2.

##### HypCluster [6].

This algorithm is a centralized version of Option II of IFCA. The server
alternately clusters clients by loss function value and runs stochastic
gradient descent per-cluster using the clients’ data. It performs as Option II
of IFCA on the example above.

Finally, we discuss Clustered Federated Learning, the algorithm proposed in
[5], which runs the risk of clustering too finely, as in the next example
(Fig. 1(c)).

##### Clustered Federated Learning [5].

Clustered Federated Learning operates by recursively bi-partitioning the set
of clients based on the clients’ gradient values at the FedAvg optimum.
Consider the following example. Let $N=3$ and $K=2$ with client gradients

$\displaystyle g_{1}(x)$ $\displaystyle=x$ (7) $\displaystyle g_{2}(x)$
$\displaystyle=x-\nicefrac{{1}}{{2}}$ (8) $\displaystyle g_{3}(x)$
$\displaystyle=\begin{cases}x&\text{w.p. }\nicefrac{{1}}{{2}}\\\
x-1&\text{w.p. }\nicefrac{{1}}{{2}}.\end{cases}$ (9)

Therefore the correct clustering here is $\\{1\\}$ and $\\{2,3\\}$. The FedAvg
optimum is $x_{\text{FedAvg}}^{*}=\nicefrac{{1}}{{4}}$, at which the clients’
gradient values are $g_{1}(\nicefrac{{1}}{{4}})=\nicefrac{{1}}{{4}}$,
$g_{2}(\nicefrac{{1}}{{4}})=-\nicefrac{{1}}{{4}}$ and
$g_{3}=\nicefrac{{1}}{{4}}$ w.p. $\nicefrac{{1}}{{2}}$. Based on this
computation, Clustered Federated Learning partitions the client set into
$\\{1,2\\}$ and $\\{3\\}$ w.p. $\nicefrac{{1}}{{2}}$ and then proceeds to run
the algorithm separately on each sub-cluster. Therefore, the algorithm never
corrects its initial error in separating clients $\\{2\\}$ and $\\{3\\}$.

The behaviour of these algorithms motivates our method Federated-Clustering,
which by rigorously checking client similarity at every step of the training
process can recover from past clustering errors.

##  3 Proposed Method: Federated-Clustering (Algorithm 2)

At a high level, Federated-Clustering works as follows. Each client $i$
maintains a personalized model which, at every step, it broadcasts to the
other clients $j\neq i$. Then each client $j$ computes its gradient on clients
$i$’s model parameters and sends the gradient to client $i$. Finally, client
$i$ runs a clustering procedure on the received gradients, determines which
other clients have gradients closest to its own at its current model, and
updates its current model by averaging the gradients of these similar clients.
By the end of the algorithm, ideally each client has a model which has been
trained only on the data of similar clients.

The core of Federated-Clustering is a clustering procedure, Threshold-
Clustering (Algorithm 3), which identifies clients with similar gradients at
each step. This clustering procedure, which we discuss in the next section,
has two important properties: it is robust and its error rate is near-optimal.

##### Notation.

For an arbitrary integer $N$, we let $[N]=\\{1,...,N\\}$. We take $a\gtrsim b$
to mean there is a sufficiently large constant $c$ such that $a\geq cb$,
$a\lesssim b$ to mean there is a sufficiently small constant $c$ such that
$a\leq cb$, and $a\approx b$ to mean there is a constant $c$ such that $a=cb$.
We write $i\sim j$ if clients $i$ and $j$ belong to the same cluster,
$i\overset{i.i.d.}{\sim}j$ if they belong to the same cluster and their data
is drawn independently from identical distributions (we will sometimes
equivalently write $z_{i}\overset{i.i.d.}{\sim}z_{j}$, where $z_{i}$ and
$z_{j}$ are arbitrary points drawn from clients $i$’s and $j$’s
distributions), and $i\not\sim j$ if they belong to different clusters. For
two different clients $i$ and $j$, same cluster or not, we write $i\neq j$.
Finally, $n_{i}$ denotes the number of clients in client $i$’s cluster,
$\delta_{i}=\nicefrac{{n_{i}}}{{N}}$ denotes the fraction of clients in client
$i$’s cluster, and $\beta_{i}$ denotes the fraction of clients that are
malicious from client $i$’s perspective.

###  3.1 Analysis of Clustering Procedure

Given the task of clustering $N$ points into $K$ clusters, at step $l$ our
clustering procedure has current estimates of the $K$ cluster-centers,
$v_{1,l},...,v_{K,l}$. To update each estimate $v_{k,l+1}\leftarrow v_{k,l}$,
it constructs a ball of radius $\tau_{k,l}$ around $v_{k,l}$. If a point falls
inside the ball, the point retains its value; if it falls outside the ball,
its value is mapped to the current cluster-center estimate. The values of all
the points are then averaged to set $v_{k,l+1}$ (update rule (24)). The
advantage of this rule is that it is very conservative. If our algorithm is
confident that its current cluster-center estimate is close to the true
cluster mean (i.e. there are many points nearby), it will confidently improve
its estimate by taking a large step in the right direction (where the step
size and direction are determined mainly by the nearby points). If our
algorithm is not confident about being close to the cluster mean, it will
tentatively improve its estimate by taking a small step in the right direction
(where the step size and direction are small since the majority of points are
far away and thus do not change the current estimate).

To analyze the theoretical properties of this procedure, we look at a natural
setting in which clients within the same cluster have i.i.d. data (for
analysis of our federated learning algorithm, we will relax this strong notion
of intra-cluster similarity). In particular, in our setting there are $N$
points $\\{z_{1},...,z_{N}\\}$ which can be partitioned into $K$ clusters
within which points are i.i.d.. We assume the following.

  * •

Assumption 1 (Intra-cluster Similarity): For all $i\sim j$,

$\displaystyle z_{i}\overset{i.i.d.}{\sim}z_{j}.$ (10)

  * •

Assumption 2 (Inter-cluster Separation): For all $i\not\sim j$,

$\displaystyle\|\mathbb{E}z_{i}-\mathbb{E}z_{j}\|^{2}$
$\displaystyle\geq\Delta^{2}.$ (11)

  * •

Assumption 3 (Bounded Variance): For all $z_{i}$,

$\displaystyle\mathbb{E}\|z_{i}-\mathbb{E}z_{i}\|^{2}$
$\displaystyle\leq\sigma^{2}.$ (12)

######  Theorem 3.1.

Suppose there $N$ points $\\{z_{i}\\}_{i\in[N]}$ for which Assumptions [1-3]
hold with inter-cluster separation parameter
$\Delta\gtrsim\nicefrac{{\sigma}}{{\delta_{i}}}$. Running Algorithm 3 for

$\displaystyle l$
$\displaystyle\gtrsim\max\bigg{\\{}1,\max_{i\in[N]}\frac{\log(\nicefrac{{\sigma}}{{\Delta}})}{\log(1-\nicefrac{{\delta_{i}}}{{2}})}\bigg{\\}}$
(13)

steps with fraction of malicious clients $\beta_{i}\lesssim\delta_{i}$ and
thresholding radius $\tau\approx\sqrt{\delta_{i}\sigma\Delta}$ guarantees that

$\displaystyle\mathbb{E}\|v_{k_{i},l}-\mathbb{E}z_{i}\|^{2}$
$\displaystyle\lesssim\frac{\sigma^{2}}{n_{i}}+\frac{\sigma^{3}}{\Delta}+\beta_{i}\sigma\Delta.$
(14)

###### Proof.

See A.1. ∎

Supposing $\beta_{i}=0$, if we knew the identity of all points within
$z_{i}$’s cluster, we would simply take their mean as the cluster-center
estimate, incurring estimation error of $\nicefrac{{\sigma^{2}}}{{n_{i}}}$
(i.e. the sample-mean’s variance). Since we don’t know the identity of points
within clusters, the additional factor of $\nicefrac{{\sigma^{3}}}{{\Delta}}$
in (14) is the price we pay to learn the clusters. This additional term scales
with the difficulty of the clustering problem. If true clusters are well-
separated and/or the variance of the points within each cluster is small (i.e.
$\Delta$ is large, $\sigma^{2}$ is small), then the clustering problem is
easier and our bound is tighter. If clusters are less-well-separated and/or
the variance of the points within each cluster is large, accurate clustering
is more difficult and our bound weakens.

##### Setting $\tau$.

To achieve the rate in (14), we set $\tau\approx\sqrt{\sigma\Delta}$, which is
the geometric mean of the standard deviation, $\sigma$, of points belonging to
the same cluster and the distance, $\Delta$, to a different cluster. The
intuition for this choice is that we want the radius for each cluster to be at
least as large as the standard deviation of the points belonging to that
cluster in order to capture in-cluster points. The radius could be
significantly larger than the standard deviation if $\Delta$ is large, thus
capturing many non-cluster points as well. However, the conservative nature of
our update rule (24) offsets this risk. By only updating the center with a
step-size proportional to the _fraction_ of points inside the ball, it limits
the influence of any mistakenly captured points.

Threshold-Clustering has two important properties which we now discuss: it is
Byzantine robust and has a near-optimal error rate.

####  3.1.1 Robustness

We construct the following definition to characterize the robustness of
Algorithm 3.

######  Definition 1 (Robustness).

An algorithm $\mathcal{A}$ is _robust_ if the error introduced by bad clients
can be bounded i.e. malicious clients do not have an arbitrarily large effect
on the convergence. Specifically, for a specific objective, let
$\mathcal{E}_{1}$ be the base error of $\mathcal{A}$ with no bad clients, let
$\beta$ be the fraction of bad clients, and let $\mathcal{E}_{2}$ be some
bounded error added by the bad points. Then $\mathcal{A}$ is _robust_ if

$\displaystyle Err(\mathcal{A})$
$\displaystyle\leq\mathcal{E}_{1}+\beta\mathcal{E}_{2}.$

##### Threat model.

Our clustering procedure first estimates the centers of the $K$ clusters from
the $N$ points and constructs a ball of radius $\tau_{k}$ around the estimated
center of each cluster $k$. If a point falls inside the ball, the point
retains its value; if it falls outside the ball, its value is mapped to the
current cluster-center estimate. Following the update rule (24), the values of
all the points are averaged to update the cluster-center estimate. Therefore,
a bad point that wants to distort the estimate of the $k$’th cluster’s center
has the most influence by placing itself just within the boundary of the ball
around that cluster-center, i.e. at $\tau_{k}$-distance from the cluster-
center.

From (14) we see that the base squared-error of Algorithm 3 in estimating
$z_{i}$’s cluster-center is
$\lesssim\nicefrac{{\sigma^{2}}}{{n_{i}}}+\nicefrac{{\sigma^{3}}}{{\Delta}}$,
and that the bad points introduce extra squared-error of order $\sigma\Delta$.
Given our threat model, this is exactly expected. The radius around $z_{i}$’s
cluster-center is order $\sqrt{\sigma\Delta}$. Therefore, bad points placing
themselves at the edge of the ball around $z_{i}$’s cluster-center estimate
will be able to distort the estimate by order $\sigma\Delta$. The scaling of
this extra error by $\beta_{i}$ satisfies our definition of robustness, and
the error smoothly vanishes as $\beta_{i}\rightarrow 0$.

####  3.1.2 Near-Optimality

The next result shows that the upper bound (14) on the estimation error of
Algorithm 3 nearly matches the best-achievable lower bound. In particular, it
is tight within a factor of $\nicefrac{{\sigma}}{{\Delta}}$.

######  Theorem 3.2 (Near-optimality of Threshold-Clustering).

For any algorithm $\mathcal{A}$, there exists a mixture of distributions
$\mathcal{D}_{1}=(\mu_{1},\sigma^{2})$ and
$\mathcal{D}_{2}=(\mu_{2},\sigma^{2})$ with $\|\mu_{1}-\mu_{2}\|\geq\Delta$
such that the estimator $\hat{\mu}_{1}$ produced by $\mathcal{A}$ has error

$\displaystyle\mathbb{E}\|\hat{\mu}_{1}-\mu_{1}\|^{2}\geq\Omega\bigg{(}\frac{\sigma^{4}}{\Delta^{2}}+\frac{\sigma^{2}}{n_{i}}\bigg{)}.$
(15)

###### Proof.

See A.2. ∎

####  3.1.3 Federated-Clustering on Examples in Section 2.1

We describe how Federated-Clustering successfully handles the examples in
Section 2.1.

##### Example 1: Fig. 1(a).

Federated-Clustering checks at every step the gradient values of all $N$
clients at the current parameters of all $K$ clusters. This verification
process avoids the type of errors made by Myopic-Clustering. For instance, at
$t=1$ when Myopic-Clustering makes its error, Federated-Clustering computes
the gradients of all clients at client $\\{1\\}$’s current parameters:
$g_{1}(1)=\nicefrac{{1}}{{3\eta}}$, $g_{2}(1)=0$, and
$g_{3}(1)=-\nicefrac{{1}}{{\eta}}$. Therefore it correctly clusters
$\\{1,2\\}$ together at this point, and client $\\{2\\}$’s parameters update
beyond the saddle-point and converge to the global minimum at $x=0$.

##### Example 2: Fig. 1(b).

By clustering clients based on gradient instead of loss value, Federated-
Clustering initially computes the clients’ gradients of $+1$ and $-1$
respectively at $x_{2,0}=0$, and given the continued separation of their
gradients around $0$ as the algorithm converges, correctly identifies that
they belong to different clusters.

##### Example 3: Fig. 1(c).

Recall how Clustered FL fails on this example. Based on an initial clustering
error, it partitions the clients incorrectly early on and then evaluates each
subset separately going forward, thus never recovering the correct clustering.
Our algorithm avoids this type of mistake by considering all clients during
each clustering at every step.

###  3.2 Analysis of Federated-Clustering

We now proceed with the analysis of Federated-Clustering. First, we establish
necessary assumptions: intra-cluster similarity, inter-cluster separation,
bounded variance of stochastic gradients, and smoothness of loss objectives.

  * •

Assumption 4 (Intra-cluster Similarity): For all $x$, $i\sim j$, and some
constant $A\geq 0$,

$\displaystyle\|\nabla f_{i}(x)-\nabla\bar{f}_{i}(x)\|^{2}\leq
A^{2}\|\nabla\bar{f}_{i}(x)\|^{2},$ (16)

where $\bar{f}_{i}(x)\triangleq\frac{1}{n_{i}}\sum_{j\sim i}f_{j}(x)$.

  * •

Assumption 5 (Inter-cluster Separation): For all $x$, $i\not\sim j$, and some
constant $D\geq 0$,

$\displaystyle\|\nabla f_{i}(x)-\nabla
f_{j}(x)\|^{2}\geq\Delta^{2}-D^{2}\|\nabla f_{i}(x)\|^{2}.$ (17)

This formulation is motivated by the information theoretic lower-bounds of [8]
who show that the optimal clustering strategy is to group all clients with the
same optimum (even if they are non-iid). Assumptions 4 and 5 are in fact a
slight strengthening of this very statement. To see this, note that for a
client with loss function $f_{i}(x)$, belonging to cluster $\bar{f}_{i}(x)$
with a first-order stationary points ${\bar{x}}^{\ast}$, Assumption 4 implies
that if $\nabla\bar{f}_{i}({\bar{x}}^{\ast})=0\Rightarrow\nabla
f_{i}({\bar{x}}^{\ast})=0$, and so ${\bar{x}}^{\ast}$ is also a stationary
point for client $i$. Thus, all clients within a cluster have shared
stationary points. Assumption 4 further implies that the gradient difference
elsewhere away from the optima is also bounded. This latter strengthening is
motivated by the fact the the loss functions are smooth, and so the gradients
cannot diverge arbitrarily as we move away from the shared optima. In fact, it
is closely related to the strong growth condition (equation (1) in [24]),
which is shown to be a very useful notion in practical deep learning. We also
empirically verify its validity in Fig. 2 (left).

Similarly, Assumption 5 is a strengthening of the condition that clients
across different clusters need to have different optima. For two clients $i$
and $j$ who belong to different clusters and with first-order stationary
points $x_{i}^{\ast}$ and $x_{j}^{\ast}$, Assumption 5 implies that $\|\nabla
f_{i}(x_{j}^{\ast})\|^{2}\geq\Delta$ and $\|\nabla
f_{j}(x_{i}^{\ast})\|^{2}\geq\Delta$. Thus, they do not share any common
optimum. Similar to the Assumption 4, Assumption 5 also describes what happens
elsewhere away from the optima - it allows for the difference between the
gradients to be smaller than $\Delta$ as we move away from the stationary
points. Again, this specific formulation is motivated by smoothness of the
loss function, and empirical validation (Fig. 2, right).

In the following lemma, we give a specific setting in which Assumptions 4 and
5 hold.

######  Lemma 1.

Suppose losses $f_{i}$ are $L$-smooth and $\mu$-strongly convex and clients in
the same cluster have the same optima. Then for all clients $i$

$\displaystyle\|\nabla f_{i}(x)-\nabla\bar{f}_{i}(x)\|^{2}$
$\displaystyle\leq(\nicefrac{{2L}}{{\mu}})^{2}\|\nabla\bar{f}_{i}(x)\|^{2},$
(18)

and for clients $i\not\sim j$ in different clusters

$\displaystyle\|\nabla f_{i}(x)-\nabla f_{j}(x)\|^{2}$
$\displaystyle\geq\frac{1}{2}(\max_{j\not\sim i}\|\nabla
f_{j}(x_{i}^{*})\|)^{2}-2(1+(\nicefrac{{L}}{{\mu}})^{2})\|\nabla
f_{i}(x)\|^{2}$ (19)

for all $x$.

###### Proof.

See A.3. ∎

  * •

Assumption 6 (Bounded Variance of Stochastic Gradients): For all $x$,

$\displaystyle\mathbb{E}\|g_{i}(x)-\nabla f_{i}(x)\|^{2}\leq\sigma^{2},$ (20)

where $\mathbb{E}[g_{i}(x)|x]=\nabla f_{i}(x)$ and each client $i$’s
stochastic gradients $g_{i}(x)$ are independent.

  * •

Assumption 7 (Smoothness of Loss Functions): For any $x,y$,

$\displaystyle\|\nabla f_{i}(x)-\nabla f_{i}(y)\|\leq L\|x-y\|.$ (21)

######  Theorem 3.3.

Let _Assumptions [4-7]_ hold with inter-cluster separation parameter
$\Delta\gtrsim\nicefrac{{\max(1,A^{4})\max(1,D^{2})\sigma}}{{\delta_{i}}}$.
Under these conditions, suppose we run Algorithm 2 for $T$ rounds with
learning rate $\eta\leq\nicefrac{{1}}{{L}}$, fraction of malicious clients
$\beta_{i}\lesssim\delta_{i}$, and batch size
$|B_{i}|\gtrsim\min(\sqrt{\max(1,A^{2})(\nicefrac{{\sigma^{2}}}{{n_{i}}}+\nicefrac{{\sigma^{3}}}{{\Delta}}+\beta_{i}\sigma\Delta)m_{i}},m_{i})$,
where $m_{i}$ is the size of client $i$’s training dataset. If, in each round
$t\in[T]$, we cluster with radius $\tau\approx\sqrt{\delta_{i}\sigma\Delta}$
for

$\displaystyle l$ $\displaystyle\geq\max\bigg{\\{}1,\max_{i\in
N]}\frac{\log(\nicefrac{{\sigma}}{{\sqrt{|B_{i}|}\Delta}})}{\log(1-\nicefrac{{\delta_{i}}}{{2}})}\bigg{\\}}$
(22)

steps, then

$\displaystyle\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}$
$\displaystyle\lesssim\sqrt{\frac{\max(1,A^{2})(\nicefrac{{\sigma^{2}}}{{n_{i}}}+\nicefrac{{\sigma^{3}}}{{\Delta}}+\beta_{i}\sigma\Delta)}{T}}.$
(23)

###### Proof.

See A.4. ∎

We note a few things. 1) The rate in (23) is the optimal rate in $T$ for
stochastic gradient descent on non-convex functions [25]. 2) The dependence on
$\sqrt{\sigma^{2}/n_{i}}$ is intuitive, since convergence error should
increase as the variance of points in the cluster increases and decrease as
the number of points in the cluster increases. It is also optimal as shown in
[8]. 3) The dependence on $\sqrt{\beta_{i}\sigma\Delta}$ is also expected. We
choose a radius $\tau\approx\sqrt{\sigma\Delta}$ for clustering. Given our
threat model, the most adversarial behavior of the bad clients from client
$i$’s perspective is to place themselves at the edge of the ball surrounding
the estimated location of $i$’s gradient, thus adding error of order
$\sqrt{\beta_{i}\sigma\Delta}$. When there are no malicious clients, this
extra error vanishes. 4) If the constraint on batch-size in Theorem 3.3
requires $|B_{i}|=m_{i}$, then the variance of stochastic gradients vanishes
and the standard $\mathcal{O}(\nicefrac{{1}}{{T}})$ rate for deterministic
gradient descent is recovered (see equation (118) in proof). We also note that
as long as there are no malicious clients (i.e. $\beta_{i}\sigma\Delta$ term
is $0$), there are a large enough number of clients $n_{i}$ in the cluster,
and inter-cluster separation $\Delta$ is sufficiently larger than the inter-
cluster-variance $\sigma^{2}$, then minimum-batch size will likely be less
than $m_{i}$.

If losses are smooth and strongly convex, our proof techniques for Theorem 3.3
get an $\mathcal{O}(\nicefrac{{1}}{{T}})$ convergence rate with the specific
constants $A$, $D$, and $\Delta$ stated in Lemma 1.

##### Privacy.

Since Federated-Clustering requires clients to compute distances between
gradients, they must share their models and gradients which compromises
privacy. The focus of our work is not on optimizing privacy, so we accommodate
only the lightest layer of privacy for federated learning: sharing of models
and gradients rather than raw data. Applying more robust privacy techniques is
a direction for future work. In the meantime, we refer the reader to the
extensive literature on differential privacy, multi-party computation, and
homomorphic encryption in federated learning.

##### Communication Overhead.

At each step, Federated-Clustering requires $\mathcal{O}(N^{2})$ rounds of
communication since each client sends its model to every other client,
evaluates its own gradient at every other client’s model and then sends this
gradient back to the client who owns the model. We pay this communication
price to mitigate the effect of past clustering mistakes. For example, say at
one round a client mis-clusters itself and updates its model incorrectly. At
the next step, due to communication with all other clients, it can check the
gradients of all other clients at its current model, have a chance to cluster
correctly at this step, and update its model towards the optimum, regardless
of the previous clustering error. Recall on the other hand that an algorithm
like Myopic-Clustering (Alg. 1), while communication efficient ($N$ rounds per
step), may not recover from past clustering mistakes since it doesn’t check
gradients rigorously in the same way. In the next section, we propose a more
communication-efficient algorithm, Momentum-Clustering (Algorithm 4), which
clusters momentums instead of gradients (reducing variance and thus clustering
error) and requires only $\mathcal{O}(N)$ communication rounds per step.

Algorithm 2 Federated-Clustering

Input Learning rate: $\eta$. Initial parameters for each client:
$\\{x_{1,0},...,x_{N,0}\\}$. Batch-size $|B_{i}|$ (see Theorem 3.3 for a lower
bound on this quantity)111The batch-size constraint reduces variance of the
stochastic gradients (Lemma 7). In Section 3.3 we propose another algorithm
Momentum-Clustering for which there is no batch-size restriction and which
reduces variance by clustering momentums instead of gradients.

1:for client $i\in[N]$ do

2: Send $x_{i,0}$ to all clients $j\neq i$.

3:for round $t\in[T]$ do

4: for client $i$ in [N] do

5: Compute $g_{i}(x_{j,t-1})$ with batch-size
$|B_{i}|$222$g_{i}(x_{j,t-1})=\frac{1}{|B_{i}|}\sum_{b}g_{i}(x_{j,t-1};b)$,
where $g_{i}(x_{j,t-1};b)$ is the gradient computed using sample $b$ in the
batch.and send to client $j$ for all $j\neq i\in[N]$.

6: Compute $v_{i,t}\leftarrow\texttt{Threshold-
Clustering}(\\{g_{j}(x_{i,t-1})\\}_{j\in[N]};1\text{
cluster};g_{i}(x_{i,t-1}))$.

7: Update parameter: $x_{i,t}=x_{i,t-1}-\eta v_{i,t}$.

8: Send $x_{i,t}$ to all clients $j\neq i$.

9:Output: Personalized parameters: $\\{x_{1,T},...,x_{N,T}\\}$.

Algorithm 3 Threshold-Clustering

Input Points to be clustered: $\\{z_{1},...,z_{N}\\}$. Number of clusters:
$K$. Cluster-center initializations: $\\{v_{1,0},...,v_{K,0}\\}$.

1:for round $l\in[M]$ do

2: for cluster $k$ in [K] do

3: Set radius $\tau_{k,l}$.

4: Update cluster-center estimate:  $\displaystyle
v_{k,l}=\frac{1}{N}\sum_{i=1}^{N}\bigg{(}z_{i}\mathbbm{1}(\|z_{i}-v_{k,l-1}\|\leq\tau_{k,l})+v_{k,l-1}\mathbbm{1}(\|z_{i}-v_{k,l-1}\|>\tau_{k,l})\bigg{)}.$
(24)

5:Output: Cluster-center estimates $\\{v_{1}=v_{1,M},...,v_{K}=v_{K,M}\\}$.

###  3.3 Improving Communication Overhead with Momentum

Federated-Clustering is inefficient, requiring $N^{2}$ rounds of communication
between clients at each step (each client computes their gradient at every
other client’s parameter). Since momentums change much more slowly from round-
to-round than gradients, a past clustering mistake will not have as much of a
harmful impact on future correct clustering and convergence as when clustering
gradients.

In Algorithm 4, at each step each client computes their momentum and sends it
to the server. The server clusters the $N$ momentums, computes an update per-
cluster, and sends the update to the clients in each cluster. Therefore,
communication is limited to $N$ rounds per step.

####  3.3.1 Analysis of Momentum-Clustering

The analysis of the momentum based method requires adapting the intra-cluster
similarity and inter-cluster separation assumptions from before.

  * •

Assumption 8 (Intra-cluster Similarity): For all $i\sim j$ and $t\ \in[T]$,

$\displaystyle m_{i,t}\overset{i.i.d.}{\sim}m_{j,t},$ (25)

where $m_{i,t}$ is defined as in (29).

  * •

Assumption 9 (Inter-cluster Separation): For all $i\not\sim j$ and $t\in[T]$,

$\displaystyle\|\mathbb{E}m_{i,t}-\mathbb{E}m_{j,t}\|^{2}$
$\displaystyle\geq\Delta^{2}.$ (26)

Note that the intra-cluster similarity assumption in this momentum setting is
stronger than in the gradient setting (Assumption 4): namely we require that
the momentum of clients in the same cluster be i.i.d. at all points. This
stronger assumption is the price we pay for a simpler and more practical
algorithm. Finally, due to the fact that momentums are low-variance
counterparts of gradients (Lemma 14), we can eliminate constraints on the
batch size and still achieve the same rate.

######  Theorem 3.4.

Let _Assumptions [6-9]_ hold with inter-cluster separation parameter
$\Delta\gtrsim\nicefrac{{\sigma}}{{\delta_{i}}}$. Under these conditions,
suppose we run Algorithm 4 for $T$ rounds with learning rate
$\eta\lesssim\min\bigg{\\{}\frac{1}{L},\sqrt{\frac{\mathbb{E}(f_{i}(x_{i,0})-f_{i}^{*})}{LT(\nicefrac{{\sigma^{2}}}{{n_{i}}}+\nicefrac{{\sigma^{3}}}{{\Delta}})}}\bigg{\\}}$
($f_{i}^{*}$ is the global minimum of $f_{i}$), fraction of bad clients
$\beta_{i}\lesssim\delta_{i}$, and momentum parameter $\alpha\gtrsim L\eta$.
If, in each round $t\in[T]$, we cluster with radius
$\tau\approx\sqrt{\delta_{i}\sigma\Delta}$ for

$\displaystyle l$ $\displaystyle\geq\max\bigg{\\{}1,\max_{i\in
N]}\frac{\log(\nicefrac{{\sqrt{\alpha}\sigma}}{{\Delta}})}{\log(1-\nicefrac{{\delta_{i}}}{{2}})}\bigg{\\}}$
(27)

steps, then for all $i\in[N]$

$\displaystyle\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}$
$\displaystyle\lesssim\sqrt{\frac{\nicefrac{{\sigma^{2}}}{{n_{i}}}+\nicefrac{{\sigma^{3}}}{{\Delta}}}{T}}+\frac{\beta_{i}\sigma\Delta}{T^{\frac{1}{4}}(\nicefrac{{\sigma^{2}}}{{n_{i}}}+\nicefrac{{\sigma^{3}}}{{\Delta}})^{\frac{1}{4}}}.$
(28)

###### Proof.

See A.5. ∎

We see from (28) that when there are no malicious clients ($\beta_{i}=0$),
Momentum-Clustering achieves the same
$\sqrt{\nicefrac{{\sigma^{2}}}{{n_{i}T}}}$ convergence rate observed in (23),
with no restrictions on the batch size.

Algorithm 4 Momentum-Clustering

Input Learning rate: $\eta$. Initial parameters: $\\{x_{1,0}=...=x_{N,0}\\}$.

1:for round $t\in[T]$ do

2: for client $i$ in [N] do

3: Client $i$ sends  $\displaystyle m_{i,t}=\alpha
g_{i}(x_{i,t-1})+(1-\alpha)m_{i,t-1}$ (29) to server.

4: Server generates cluster centers
$\displaystyle\\{v_{k,t}\\}_{k\in[K]}\leftarrow\texttt{Threshold-
Clustering}(\\{m_{i,t}\\};K\text{ clusters};\\{v_{k,t-1}\\}_{k\in[K]})$ (30)
and sends $v_{k_{i},t}$ to client $i$, where $k_{i}$ denotes the cluster to
which $i$ is assigned in this step.

5: Client $i$ computes update: $x_{i,t}=x_{i,t-1}-\eta v_{k_{i},t}$.

6:Output: Personalized parameters: $\\{x_{1,T},...,x_{N,T}\\}$.

##  4 Experiments

In this section, we first use a synthetic dataset to verify the assumptions
and rates claimed in our theoretical analysis in the previous section; and
second, we use the MNIST dataset [26] and CIFAR dataset [27] to compare our
proposed algorithm, Federated-Clustering, with existing state-of-the-art
federated learning algorithms. All algorithms are implemented with PyTorch
[28], and code for all experiments is at this github repo.

###  4.1 Synthetic dataset

Figure 2: Here, we show empirically on a synthetic dataset that the intra-
cluster variance ratio (31) is upper-bounded by a constant (left subplot) and
the inter-cluster variance (32) that is lower-bounded by a constant (right
subplot).

##### Construction of synthetic dataset.

We consider a synthetic linear regression task with squared loss for which we
construct $K=4$ clusters, each with $n_{i}=75$ clients. Clients in cluster
$k\in[K]$ share the same minimizer $x_{k}^{\star}\in\mathbb{R}^{d}$. For each
client $i$ in cluster $k$, we generate a sample matrix
$A_{i}\in\mathbb{R}^{d\times n}$ from $\mathcal{N}(k,\mathbf{1}_{d\times n})$
and compute the associated target as
$y_{i}=A_{i}^{\top}x_{k}^{\star}\in\mathbb{R}^{n}$. We choose the model
dimension $d=10$ to be greater than the number of local samples $n=9$ such
that the local linear system $y_{i}=A_{i}^{\top}x$ is overdetermined and the
error $\lVert x^{\star}-x\rVert_{2}^{2}$ is large. A desired federated
clustering algorithm determines the minimizer by incorporating information
from other clients $j\sim i$ in the same cluster.

##### Estimating constants in Assumptions 4 and 5.

In Fig. 2, using the above synthetic dataset

  * •

we estimate the intra-cluster variance ratio $A^{2}$ by finding the upper
bound of (31)

$\displaystyle\frac{\lVert\nabla
f_{i}(x)-\nabla\bar{f}_{i}(x)\rVert_{2}^{2}}{\lVert\nabla\bar{f}_{i}(x)\rVert_{2}^{2}};$
(31)

  * •

we estimate the inter-cluster variance $\Delta^{2}$ by setting $D=0$ and
computing the lower bound of (32)

$\displaystyle\lVert\nabla f_{i}(x)-\nabla f_{j}(x)\rVert_{2}^{2}.$ (32)

We run Federated-Clustering with perfect clustering assignments and estimate
these bounds over time. The result is shown in Fig. 2 where grey lines are the
quantities in (31) and (32) for individual clients, black lines are those
quantities averaged within clusters, and red dashed lines are empirical bounds
on the quantities. The left figure demonstrates that the intra-cluster
variance ratio does not grow with time and can therefore reasonably be upper
bounded by a constant $A^{2}$. Similarly, the right figure shows that the
inter-cluster variance can be reasonably lower bounded by a positive constant
$\Delta^{2}$. These figures empirically demonstrate that Assumptions 4 and 5
are realistic in practice.

Figure 3: The performance of our algorithm vs. baselines on a synthetic
dataset. When $n_{i}$ is small, the ground-truth outperforms our algorithm,
but this difference vanishes with increasing $n_{i}$. This behavior is
consistent with the dependence of the convergence rate on $n_{i}$ in Theorem
3.3: increasing $n_{i}$ improves convergence.

##### Performance.

In Fig. 3, we compare the performance of our algorithm Federated-Clustering
(FC) with several baselines: standalone training (Local), IFCA [7], FedAvg
(Global) [1], and distributed training with ground truth (GT) cluster
information. We consider the synthetic dataset from before, starting with
cluster parameters $(K,n_{i})=(4,4)$ and observe performance when increasing
parameters to $n_{i}=16$ and $K=16$ separately. In each step of optimization,
we run Threshold-Clustering for $l=10$ rounds so that heuristically the
outputs are close enough to cluster centers, cf. Fig. 7. We tune the learning
rate separately for each algorithm through grid search, but preserve all other
algorithmic setups. Our algorithm outperforms the non-oracle baselines in all
cases. While Federated-Clustering is slightly worse than ground truth when
$(K,n_{i})=(4,4)$, their performances are almost identical in the middle
subplot for $n_{i}=16$. This observation is consistent with the
$n_{i}$-scaling observed in (23): as the number of clients-per-cluster
increases, convergence improves.

###  4.2 MNIST experiment

Table 1: Comparison of test losses and accuracies for federated personalization algorithms on MNIST. FC outperforms all non-oracle baselines on two learning tasks. | Rotation | Private label
---|---|---
| Acc.(%) | Loss | Acc.(%) | Loss
Local | 71.3 | 0.517 | 75.2 | 0.489
Global | 46.6 | 0.631 | 22.2 | 0.803
Ditto | 62.0 | 0.576 | 61.7 | 0.578
IFCA | 54.6 | 0.588 | 65.4 | 0.531
KNN | 52.1 | 2.395 | 63.2 | 1.411
FC (ours) | 75.4 | 0.475 | 77.0 | 0.468
GT (oracle) | 84.7 | 0.432 | 85.1 | 0.430

In this section, we compare Federated-Clustering to existing federated
learning baselines on the MNIST dataset. The dataset is constructed as
follows, similar to [7]. The data samples are randomly shuffled and split into
$K=4$ clusters with $n_{i}=75$ clients in each cluster. We consider two
different tasks: 1) the rotation task transforms images in cluster $k$ by
$k*90$ degrees; 2) the private label task transforms labels in cluster $k$
with $T_{k}(y):y\mapsto(y+k\mod 10)$, such that the same image may have
different labels from cluster-to-cluster.

Algorithm hyperparameters. For these two experimental tasks, in addition to
the baselines from our synthetic experiment, we include the KNN-
personalization [23] and Ditto [10] algorithms which both interpolate between
a local and global model. The KNN-personalization is a linear combination of a
global model, trained with FedAvg [1], and a local model which is the
aggregation of nearest-neighbor predictions in the client’s local dataset to
the global model’s prediction. We set the coefficients of this linear
combination to be $\lambda_{\text{knn}}=0.5$ and $\lambda_{\text{knn}}=0.9$
for the _rotation_ and _private label_ tasks, respectively. The Ditto
objective is a personalized loss with an added regularization term that
encourages closeness between the personalized and global models. Since tuning
this regularization parameter $\lambda_{\text{ditto}}$ leads to a degenerated
”Local” training where $\lambda_{\text{ditto}}=0$, we fix
$\lambda_{\text{ditto}}=1$ for both _rotation_ and _private label_ tasks. To
reduce the computation cost of our algorithm, in each iteration we randomly
divide the $N$ clients into $16$ subgroups and apply Federated-Clustering to
each subgroup simultaneously. The clipping radius $\tau_{k,l}$ for each
cluster $k$ is adaptively chosen to be the $20$th-percentile of distances to
the cluster-center.

Performance. The experimental results are listed in Table 1. Since an image
can have different labels across clusters in the _private label_ task, a model
trained over the pool of all datasets only admits inferior performance.
Therefore, distributed training algorithms that maintain a global model, such
as FedAvg, Ditto, and KNN, perform poorly compared to training alone. On the
other hand, our algorithm Federated-Clustering outperforms standalone training
and all personalization baselines. This experiment suggests that our algorithm
successfully explored the cluster structure and benefited from collaborative
training.

###  4.3 CIFAR experiment

In this section, we evaluate the efficacy of various clustering algorithms on
the CIFAR-10 and CIFAR-100 datasets [27].

####  4.3.1 CIFAR-10

For the CIFAR-10 experiment, we create 4 clusters, each containing 5 clients
and transform the labels in each cluster such that different clusters can have
different labels for the same image (the _private label_ task in Section 4.2).
We train a VGG-16 model [29] with batch size 32, learning rate 0.1, and
momentum 0.9. The outcomes are presented in Fig. 4. The left subplot
illustrates that collaborative clustering algorithms designed for global model
training (e.g., Ditto, IFCA, Global) yield suboptimal models, as not all
participating clients benefit from each other. On the other hand, Local and
GroundTruth training are not influenced by the conflicting labels from other
clusters so they significantly outperform Ditto and IFCA. Our Federated-
Clustering (FC) algorithm also excludes such adversarial influence and, more
importantly, outperforms Local training, showing that FC benefits from
collaboration.

In the middle subplot, we examine the impact on accuracy of varying the
thresholding radius $\tau$ (i.e. $\tau$ is set as the percentile of gradient
distances from the cluster-center, so smaller percentile corresponds to
smaller $\tau$). Our findings indicate that adopting a more conservative value
for $\tau$ (lower percentile) does not substantially compromise accuracy.

The right subplot demonstrates the behavior of Federated-Clustering when the
clustering oracle is invoked intermittently. The results suggest that
increasing the number of local iterations boosts the learning curve early in
training when gradients from different clusters are close together and
clusters are ill-defined. However, this improvement plateaus when gradients
become separated and clusters become well-defined.

Figure 4:  Performance of algorithms on CIFAR-10 dataset with the _private
labels_ task. Left: Relative accuracy of clustering algorithms. Algorithms
optimized for global model performance, such as Ditto, IFCA, Global (FedAvg),
perform poorly on personalization. FC outperforms Local training, showing that
it benefits from collaboration between clients, and is competitive with
GroundTruth. Middle: Impact of thresholding radius $\tau$ on accuracy. $\tau$
is the percentile of gradients distances from the cluster-center. Right:
Impact of local gradient steps between two clustering calls. Early on in
training, clusters are less identifiable so local optimization helps but these
gains lessen later on when gradients from different clusters drift apart and
clusters are better defined.

Figure 5: Performance of Federated-Clustering (FC) on the CIFAR-100 dataset.
The Global (FedAvg) accuracy plateaus while FC continually improves.

Figure 6: Performance of Federated-Clustering (FC) against a large gradient
attack (left) and bit-flipping attack (right). FC is robust to these attacks
and significantly outperforms Global (FedAvg) performance.

####  4.3.2 CIFAR-100

We consider the CIFAR-100 dataset distributed over 10 clusters so that each
cluster contains 10 unique labels. In each cluster, we set 10 clients with IID
data. We use a VGG-8 model for training and the same hyperparameters as those
in the CIFAR-10 experiment (Section 4.3.1) and report the results in Fig. 6.
While clients’ data within each cluster share similar features, a small model
like VGG-8 cannot sufficiently benefit from intra-cluster collaboration.
Therefore the performance of Global training plateaus at a very low level
while in contrast Federated-Clustering (FC) still benefits from collaboration
and continues to improve over time.

###  4.4 Defense against Byzantine attacks

Byzantine attacks, in which attackers have full knowledge of the system and
can deviate from the prescribed algorithm, are prevalent in distributed
environments [11]. There are many forms of Byzantine attacks. For example, our
_private label_ setting in Sections 4.2 and 4.3.1 corresponds to the _label-
flipping attack_ in the Byzantine-robustness literature, since a malicious
client can try to corrupt the model by assigning the wrong label to an image
in training data.

In this section, we investigate two other attacks: Byzantine workers send
either very large gradients or gradients with opposite signs. Using the MNIST
dataset with the _private label_ task, we set 4 clusters with 50 non-malicious
clients each (so non-malicious clients from different clusters can have
private labels) and add 50 Byzantine works to each cluster. We demonstrate the
robustness of Federated-Clustering (FC) in Fig. 6. In both cases, Global
training suffers from serious model degradation while FC successfully reaches
high accuracy under these attacks, demonstrating its robustness.

###  4.5 Empirical Study on Clipping Iterations in Algorithm 3

We employ Algorithm 3 (Threshold-Clustering) on datasets to discern the
effectiveness of the clipping iterations in identifying optimal cluster
centers. These datasets share the same groundtruth cluster centers, and thus
the same $\Delta$, but vary in their inter-cluster standard deviations, with
$\sigma$ values of $0.5,1,2,4$. Each dataset is made up of 90 ten-dimensional
samples from 10 clusters, generated using the scikit-learn package [30]. For
each iteration $l$ within cluster $k$, the clipping radius $\tau_{k,l}$ is
defined as the $10$th percentile of gradients’ distances from the cluster-
center. We repeat this experimental setup ten times for consistency.

The outcomes, presented in Fig. 7, show that the average distances initially
decrease rapidly, then steadily approach convergence. To identify the elbow of
a given curve $f$, we use the formula
$\tfrac{f(l)-f(l-1)}{f(l-1)-\min_{l}f(l)}$, where curves post-elbow are
notably flat. These elbows elucidate the correlation between $\sigma$ and $l$,
indicating that for a fixed dataset (and its corresponding $\sigma$), one can
pinpoint the minimal iterations $l$ needed for convergence. Notably, this
observation appears to align with the $l\gtrsim\log\sigma$ lower bound stated
in Theorem 3.1.

Figure 7: The average distance to cluster centers as a function of the number
of clipping iterations $l$. The distance between cluster-centers ($\Delta$) is
fixed while inter-cluster standard deviation $\sigma$ differs.

##  5 Conclusion

We develop gradient-based clustering algorithms to achieve personalization in
federated learning. Our algorithms have optimal convergence guarantees. They
asymptotically match the achievable rates when the true clustering of clients
is known, and our analysis holds under light assumptions (e.g., for all smooth
convex and non-convex losses). Furthermore, our algorithms are provably robust
in the Byzantine setting where some fraction of the clients can arbitrarily
corrupt their gradients. Future directions involve developing bespoke analysis
for the convex-loss case and developing more communication-efficient versions
of our algorithms. Further, our analysis can be used to show that our
algorithms are incentive-compatible and lead to _stable coalitions_ as in
[31]. This would form a strong argument towards encouraging participants in a
federated learning system. Investigating such incentives and fairness concerns
is another promising future direction.

## References

  * [1] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson and Blaise Aguera Arcas  “Communication-efficient learning of deep networks from decentralized data”  In _Artificial intelligence and statistics_ , 2017, pp. 1273–1282  PMLR
  * [2] Amir R Zamir, Alexander Sax, William Shen, Leonidas J Guibas, Jitendra Malik and Silvio Savarese  “Taskonomy: Disentangling task transfer learning”  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2018, pp. 3712–3722
  * [3] Alysa Ziying Tan, Han Yu, Lizhen Cui and Qiang Yang  “Towards personalized federated learning”  In _IEEE Transactions on Neural Networks and Learning Systems_ IEEE, 2022
  * [4] Viraj Kulkarni, Milind Kulkarni and Aniruddha Pant  “Survey of personalization techniques for federated learning”  In _2020 Fourth World Conference on Smart Trends in Systems, Security and Sustainability (WorldS4)_ , 2020, pp. 794–797  IEEE
  * [5] Felix Sattler, Klaus-Robert Muller and Wojciech Samek  “Clustered Federated Learning: Model-Agnostic Distributed Multi-Task Optimization under Privacy Constraints”  In _IEEE Transactions on Neural Networks and Learning Systems_ 32(8), 2021
  * [6] Yishay Mansour, Mehryar Mohri, Jae Ro and Ananda Theertha Suresh  “Three Approaches for Personalization with Applications to Federated Learning”, 2021  arXiv:2002.10619 [cs.LG]
  * [7] Avishek Ghosh, Jichan Chung, Dong Yin and Kannan Ramchandran  “An efficient framework for clustered federated learning”  In _Advances in Neural Information Processing Systems_ 33, 2020, pp. 19586–19597
  * [8] Mathieu Even, Laurent Massoulié and Kevin Scaman  “On sample optimality in personalized collaborative and federated learning”  In _Advances in Neural Information Processing Systems_ 35, 2022, pp. 212–225
  * [9] V. Smith, Chao-Kai Chiang, Maziar Sanjabi and Ameet Talwalkar  “Federated Multi-Task Learning”  In _31st Conference on Neural Information Processing Systems_ , 2017
  * [10] T. Li, Shengyuan Hu, Ahmad Beirami and Virginia Smith  “Fair and Robust Federated Learning through Personalization”  In _38th International Conference on Machine Learning_ , 2021
  * [11] L. Lamport, R. Shostak and M. Pease  “The Byzantine Generals Problem”  In _In Concurrency: the Works of Leslie Lamport._ , 2019
  * [12] Peva Blanchard, El Mahdi El Mhamdi, Rachid Guerraoui and Julien Stainer  “Machine learning with adversaries: Byzantine tolerant gradient descent”  In _Advances in Neural Information Processing Systems_ 30, 2017
  * [13] D. Yin, Y. Chen, R. Kannan and P. Bartlett  “Byzantine-robust distributed learning: Towards optimal statistical rates”  In _International Conference on Machine Learning (PMLR)_ , 2018
  * [14] G. Damaskinos, E.M. El-Mhamdi, R. Guerraoui, A. Guirguis and S. Rouault  “Aggregator: Byzantine machine learning via robust gradient aggregation”  In _Proceedings of Machine Learning and Systems_ 1, 2018, pp. 81–106
  * [15] R. Guerraoui and S. Rouault  “The hidden vulnerability of distributed learning in byzantium”  In _International Conference on Machine Learning (PMLR)_ , 2018, pp. 35213530
  * [16] K. Pillutla, S. Kakade and Z. Harchaoui  “Robust aggregation for federated learning”  In _IEEE Transactions on Signal Processing_ 70, 2022, pp. 1142–1154
  * [17] B. Baruch, M. Baruch and Y. Goldberg  “A little is enough: Circumventing defenses for distributed learning”  In _Advances in Neural Information Processing Systems_ 32, 2019
  * [18] C. Xie, O. Koyejo and I. Gupta  “Fall of empires: Breaking byzantine-tolerant sgd by inner product manipulation”  In _Uncertainty in Artificial Intelligence (PMLR)_ , 2020, pp. 261–270
  * [19] Sai Praneeth Karimireddy, Lie He and Martin Jaggi  “Learning from History for Byzantine Robust Optimization”  In _Proceedings of the 38th International Conference on Machine Learning_ 129, 2021, pp. 5311–5319
  * [20] Sai Praneeth Karimireddy, Lie He and Martin Jaggi  “Byzantine-Robust Learning on Heterogeneous Datasets via Bucketing”  In _International Conference on Learning Representations_ , 2022
  * [21] Tianchun Wang, Wei Cheng, Dongsheng Luo, Wenchao Yu, Jingchao Ni, Liang Tong, Haifeng Chen and Xiang Zhang  “Personalized Federated Learning via Heterogeneous Modular Networks”  In _IEEE International Conference on Data Mining_ , 2022, pp. 1197–1202
  * [22] Yue Wu, Shuaicheng Zhang, Wenchao Yu, Yanchi Liu, Quanquan Gu, Dawei Zhou, Haifeng Chen and Wei Cheng  “Personalized Federated Learning under Mixture of Distributions”  In _Proceedings of the 40th International Conference on Machine Learning (PMLR)_ , 2023
  * [23] Othmane Marfoq, Giovanni Neglia, Laetitia Kameni and Richard Vidal  “Federated multi-task learning under a mixture of distributions”  In _Proceedings of the 35th International Conference on Machine Learning_ 34, 2021
  * [24] Sharan Vaswani, Francis Bach and Mark Schmidt  “Fast and Faster Convergence of SGD for Over-Parameterized models and an Accelerated Perceptron”  In _Proceedings of the 22nd International Conference on Artificial Intelligence and Statistics (AISTATS)_ 89, 2019
  * [25] Yossi Arjevani, Yair Carmon, John C. Duchi, Dylan J. Foster, Nathan Srebro and Blake Woodworth  “Lower Bounds for Non-Convex Stochastic Optimization”  In _Mathematical Programming_ 199, 2023, pp. 165–214
  * [26] Yann LeCun, Corinna Cortes and CJ Burges  “MNIST handwritten digit database”  In _ATT Labs [Online]. Available: http://yann.lecun.com/exdb/mnist_ 2, 2010
  * [27] Alex Krizhevsky  “Learning Multiple Layers of Features from Tiny Images”, 2009  URL: https://api.semanticscholar.org/CorpusID:18268744
  * [28] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga and Adam Lerer  “Automatic differentiation in PyTorch”  In _NIPS-W_ , 2017
  * [29] Karen Simonyan and Andrew Zisserman  “Very deep convolutional networks for large-scale image recognition”  In _International Conference on Learning Representations_ , 2015
  * [30] F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot and E. Duchesnay  “Scikit-learn: Machine Learning in Python”  In _Journal of Machine Learning Research_ 12, 2011, pp. 2825–2830
  * [31] Kate Donahue and Jon Kleinberg  “Optimality and stability in federated learning: A game-theoretic approach”  In _Advances in Neural Information Processing Systems_ 34, 2021, pp. 1287–1298

##  Appendix A Proofs

###  A.1 Proof of Theorem 3.1

First we establish some notation.

##### Notation.

  * •

$\mathcal{G}_{i}$ are the good points and $\mathcal{B}_{i}$ the bad points
from point $z_{i}$’s perspective. Therefore
$|\mathcal{G}_{i}|+|\mathcal{B}_{i}|=N$.

  * •

$k_{i}$ denotes the cluster to which client $i$ is assigned at the end of
Threshold-Clustering.

  * •

To facilitate the proof, we introduce a variable $c_{k_{i},l}^{2}$ that
quantifies the distance from the cluster-center-estimates to the true cluster
means at each step of thresholding. Specifically, for client $i$’s cluster
$k_{i}$ at round $l$ of Threshold-Clustering, we set

$\displaystyle
c_{k_{i},l}^{2}=\mathbb{E}\|v_{k_{i},l-1}-\mathbb{E}z_{i}\|^{2}.$ (33)

  * •

For client $i$’s cluster $k_{i}$ at round $l$ of Threshold-Clustering, we use
thresholding radius

$\displaystyle\tau_{k_{i},l}^{2}\approx
c_{k_{i},l}^{2}+\delta_{i}\sigma\Delta.$ (34)

  * •

We introduce a variable $y_{j,l}$ to denote the points clipped by Threshold-
Clustering:

$\displaystyle
v_{k,l}=\frac{1}{N}\sum_{j\in[N]}\underbrace{z_{j}\mathbbm{1}(\|z_{j}-v_{k,l-1}\|\leq\tau_{k,l})+v_{k,l-1}\mathbbm{1}(\|z_{j}-v_{k,l-1}\|>\tau_{k,l})}_{y_{j,l}}.$
(35)

###### Proof of Theorem 3.1.

We prove the main result with the following sequence of inequalities, and then
justify the labeled steps afterwards.

$\displaystyle\mathbb{E}\|v_{k_{i},l}-\mathbb{E}z_{i}\|^{2}$
$\displaystyle=\mathbb{E}\bigg{\|}\frac{1}{N}\sum_{j\in[N]}y_{j,l}-\mathbb{E}z_{i}\bigg{\|}^{2}$
(36)
$\displaystyle=\mathbb{E}\bigg{\|}(1-\beta_{i})\bigg{(}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}}y_{j,l}-\mathbb{E}z_{i}\bigg{)}+\beta_{i}\bigg{(}\frac{1}{|\mathcal{B}_{i}|}\sum_{j\in\mathcal{B}_{i}}y_{j,l}-\mathbb{E}z_{i}\bigg{)}\bigg{\|}^{2}$
(37)
$\displaystyle\overset{(i)}{\leq}(1+\beta_{i})(1-\beta_{i})^{2}\mathbb{E}\bigg{\|}\bigg{(}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}}y_{j,l}\bigg{)}-\mathbb{E}z_{i}\bigg{\|}^{2}+\bigg{(}1+\frac{1}{\beta_{i}}\bigg{)}\beta_{i}^{2}\mathbb{E}\bigg{\|}\bigg{(}\frac{1}{|\mathcal{B}_{i}|}\sum_{j\in\mathcal{B}_{i}}y_{j,l}\bigg{)}-\mathbb{E}x_{i}\bigg{\|}^{2}$
(38)
$\displaystyle\lesssim\underbrace{\mathbb{E}\bigg{\|}\bigg{(}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}}y_{j,l}\bigg{)}-\mathbb{E}z_{i}\bigg{\|}^{2}}_{\mathcal{E}_{1}}+\beta_{i}\underbrace{\mathbb{E}\bigg{\|}\bigg{(}\frac{1}{|\mathcal{B}_{i}|}\sum_{j\in\mathcal{B}_{i}}y_{j,l}\bigg{)}-\mathbb{E}z_{i}\bigg{\|}^{2}}_{\mathcal{E}_{2}}$
(39)
$\displaystyle\overset{(ii)}{\lesssim}\bigg{(}(1-\delta_{i})c_{k_{i},l}^{2}+\frac{\sigma^{2}}{n_{i}}+\frac{\sigma^{3}}{\Delta}\bigg{)}+\beta_{i}(c_{k_{i},l}^{2}+\delta_{i}\sigma\Delta)$
(40)
$\displaystyle\overset{(iii)}{\lesssim}(1-\nicefrac{{\delta_{i}}}{{2}})c_{k_{i},l}^{2}+\bigg{(}\frac{\sigma^{2}}{n_{i}}+\frac{\sigma^{3}}{\Delta}+\beta_{i}\sigma\Delta\bigg{)}$
(41)
$\displaystyle\overset{(iv)}{\lesssim}(1-\nicefrac{{\delta_{i}}}{{2}})^{l}c_{k_{i},1}^{2}+\bigg{(}\frac{\sigma^{2}}{n_{i}}+\frac{\sigma^{3}}{\Delta}+\beta_{i}\sigma\Delta\bigg{)}\sum_{q=0}^{l-1}(1-\nicefrac{{\delta_{i}}}{{2}})^{q}$
(42)
$\displaystyle\overset{(v)}{\lesssim}(1-\nicefrac{{\delta_{i}}}{{2}})^{l}\sigma^{2}+\bigg{(}\frac{\sigma^{2}}{n_{i}}+\frac{\sigma^{3}}{\Delta}+\beta_{i}\sigma\Delta\bigg{)}$
(43)
$\displaystyle\overset{(vi)}{\lesssim}\frac{\sigma^{2}}{n_{i}}+\frac{\sigma^{3}}{\Delta}+\beta_{i}\sigma\Delta.$
(44)

Justifications for the labeled steps are:

  * •

(i) Young’s inequality:
$\|x+y\|^{2}\leq(1+\epsilon)\|x\|^{2}+(1+\nicefrac{{1}}{{\epsilon}})\|y\|^{2}$
for any $\epsilon>0$.

  * •

(ii) We prove this bound in Lemmas 2 and 6. Importantly, it shows that the
clustering error is composed of two quantities: $\mathcal{E}_{1}$, the error
contributed by good points from the cluster’s perspective, and
$\mathcal{E}_{2}$, the error contributed by the bad points from the cluster’s
perspective.

  * •

(iii) Assumption that $\beta_{i}\lesssim\delta_{i}$

  * •

(iv) Since $\mathbb{E}\|v_{k_{i},l}-\mathbb{E}z_{i}\|^{2}=c_{k_{i},l+1}^{2}$,
the inequality forms a recursion which we unroll over $l$ steps.

  * •

(v) Assumption that
$c_{k_{i},1}^{2}=\mathbb{E}\|v_{k_{i},0}-\mathbb{E}z_{i}\|^{2}\leq\sigma^{2}$.
Also, the partial sum in the second term can be upper-bounded by a large-
enough constant.

  * •

(vi) Assumption that
$l\geq\max\bigg{\\{}1,\frac{\log(\nicefrac{{\sigma}}{{\Delta}})}{\log(1-\nicefrac{{\delta_{i}}}{{2}})}\bigg{\\}}$

From (44), we see that
$c_{k_{i},l}^{2}\lesssim\frac{\sigma^{2}}{n_{i}}+\frac{\sigma^{3}}{\Delta}+\beta_{i}\sigma\Delta$.
Plugging this into the expression for $\tau_{k_{i},l}^{2}$ gives
$\tau_{k_{i},l}^{2}\approx\frac{\sigma^{2}}{n_{i}}+\frac{\sigma^{3}}{\Delta}+\delta_{i}\sigma\Delta\approx\delta_{i}\sigma\Delta$
for large $n_{i}$ and $\Delta$. ∎

######  Lemma 2 (Clustering Error due to Good Points).

$\displaystyle\mathbb{E}\bigg{\|}\bigg{(}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}}y_{j,l}\bigg{)}-\mathbb{E}z_{i}\bigg{\|}^{2}$
$\displaystyle\lesssim(1-\delta_{i})c_{k_{i},l}^{2}+\frac{\sigma^{2}}{n_{i}}+\frac{\sigma^{3}}{\Delta}$
(45)

###### Proof of Lemma 2.

We prove the main result with the following sequence of inequalities and
justify the labeled steps afterward.

$\displaystyle\mathbb{E}\bigg{\|}\bigg{(}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}}y_{j,l}\bigg{)}-\mathbb{E}z_{i}\bigg{\|}^{2}$
$\displaystyle=\mathbb{E}\bigg{\|}\bigg{(}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\sim
i}(y_{j,l}-\mathbb{E}z_{j})\bigg{)}+\bigg{(}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\not\sim
i}(y_{j,l}-\mathbb{E}z_{i})\bigg{)}\bigg{\|}^{2}$ (46)
$\displaystyle\overset{(i)}{\leq}\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\mathbb{E}\bigg{\|}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\sim
i}(y_{j,l}-\mathbb{E}z_{j})\bigg{\|}^{2}+\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}\mathbb{E}\bigg{\|}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\not\sim
i}(y_{j,l}-\mathbb{E}z_{i})\bigg{\|}^{2}$ (47)
$\displaystyle\overset{(ii)}{\lesssim}\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\mathbb{E}\bigg{\|}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\sim
i}(\mathbb{E}y_{j,l}-\mathbb{E}z_{j})\bigg{\|}^{2}+\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\mathbb{E}\bigg{\|}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\sim
i}(y_{j,l}-\mathbb{E}y_{j,l})\bigg{\|}^{2}$ (48)
$\displaystyle\phantom{{}=1}+\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}\mathbb{E}\bigg{\|}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\not\sim
i}(y_{j,l}-\mathbb{E}z_{i})\bigg{\|}^{2}$ (49)
$\displaystyle\leq\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\delta_{i}^{2}\|\mathbb{E}_{j\in\mathcal{G}_{i}:j\sim
i}(y_{j,l}-z_{j})\|^{2}+\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\mathbb{E}\bigg{\|}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\sim
i}(y_{j,l}-\mathbb{E}y_{j,l})\bigg{\|}^{2}$ (50)
$\displaystyle\phantom{{}=1}+\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}(1-\delta_{i})^{2}\mathbb{E}_{j\in\mathcal{G}_{i}:j\not\sim
i}\|y_{j,l}-\mathbb{E}z_{i}\|^{2}$ (51)
$\displaystyle\lesssim\delta_{i}\underbrace{\|\mathbb{E}_{j\in\mathcal{G}_{i}:j\sim
i}(y_{j,l}-z_{j})\|^{2}}_{\mathcal{T}_{i}}+\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\underbrace{\mathbb{E}\bigg{\|}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\sim
i}(y_{j,l}-\mathbb{E}y_{j,l})\bigg{\|}^{2}}_{\mathcal{T}_{2}}$ (52)
$\displaystyle\phantom{{}=1}+\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}(1-\delta_{i})^{2}\underbrace{\mathbb{E}_{j\in\mathcal{G}_{i}:j\not\sim
i}\|y_{j,l}-\mathbb{E}z_{i}\|^{2}}_{\mathcal{T}_{3}}$ (53)
$\displaystyle\overset{(iii)}{\lesssim}\delta_{i}\bigg{(}c_{k_{i},l}^{2}+\frac{(c_{k_{i},l}^{2}+\sigma^{2})\sigma}{\delta_{i}\Delta}\bigg{)}+\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\frac{n_{i}}{|\mathcal{G}_{i}|^{2}}\sigma^{2}$
(54)
$\displaystyle\phantom{{}=1}+(1-\delta_{i})^{2}\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}\bigg{(}\bigg{(}1+\frac{\delta_{i}}{2}+\frac{\sigma^{2}}{\delta_{i}\Delta^{2}}\bigg{)}c_{k_{i},l}^{2}+\frac{\sigma^{3}}{\Delta}\bigg{)}$
(55)
$\displaystyle\lesssim\bigg{(}1-\delta_{i}+\frac{\sigma}{\Delta}+\frac{\sigma^{2}}{\delta_{i}\Delta^{2}}\bigg{)}c_{k_{i},l}^{2}+\bigg{(}\frac{\sigma^{2}}{n_{i}}+\frac{\sigma^{3}}{\Delta}\bigg{)}$
(56)
$\displaystyle\overset{(iv)}{\lesssim}(1-\delta_{i})c_{k_{i},l}^{2}+\bigg{(}\frac{\sigma^{2}}{n_{i}}+\frac{\sigma^{3}}{\Delta}\bigg{)}.$
(57)

  * •

(i), (ii) Young’s inequality

  * •

(iii) We prove this bound in Lemmas 3, 4, and 5. Importantly, it shows that,
from point $i$’s perspective, the error of its cluster-center-estimate is
composed of three quantities: $\mathcal{T}_{1}$, the error introduced by our
thresholding procedure on the good points which belong to $i$’s cluster (and
therefore ideally are included within the thresholding radius);
$\mathcal{T}_{2}$, which accounts for the variance of the points in $i$’s
cluster; and $\mathcal{T}_{3}$, the error due to the good points which don’t
belong to $i$’s cluster (and therefore ideally are forced outside the
thresholding radius).

  * •

(iv) Assumption that $\Delta\gtrsim\nicefrac{{\sigma}}{{\delta_{i}}}$.

∎

######  Lemma 3 (Bound $\mathcal{T}_{1}$: Error due to In-Cluster Good
Points).

$\displaystyle\|\mathbb{E}_{j\in\mathcal{G}_{i}:j\sim i}(y_{j,l}-z_{j})\|^{2}$
$\displaystyle\lesssim
c_{k_{i},l}^{2}+\frac{(c_{k_{i},l}^{2}+\sigma^{2})\sigma}{\delta_{i}\Delta}.$
(58)

###### Proof of Lemma 3.

By definition of $y_{j,l}$,

$\displaystyle\mathbb{E}_{j\in\mathcal{G}_{i}:j\sim i}\|y_{j,l}-z_{j}\|$
$\displaystyle=\mathbb{E}[\|v_{k_{i},l-1}-z_{j}\|\mathbbm{1}(\|v_{k_{i},l-1}-z_{j}\|>\tau_{k_{i},l})]$
(59)
$\displaystyle\leq\frac{\mathbb{E}[\|v_{k_{i},l-1}-z_{j}\|^{2}\mathbbm{1}(\|v_{k_{i},l-1}-z_{j}\|>\tau_{k_{i},l})]}{\tau_{k_{i},l}}$
(60)
$\displaystyle\leq\frac{\mathbb{E}\|v_{k_{i},l-1}-z_{j}\|^{2}}{\tau_{k_{i},l}}$
(61)
$\displaystyle\lesssim\frac{\mathbb{E}\|v_{k_{i},l-1}-\mathbb{E}z_{i}\|^{2}+\mathbb{E}\|\mathbb{E}z_{j}-z_{j}\|^{2}}{\tau_{k_{i},l}}$
(62) $\displaystyle\leq\frac{c_{k_{i},l}^{2}+\sigma^{2}}{\tau_{k_{i},l}}$ (63)

Finally, by Jensen’s inequality and plugging in the value for
$\tau_{k_{i},l}$,

$\displaystyle\|\mathbb{E}(y_{j,l}-z_{j})\|^{2}$
$\displaystyle\leq(\mathbb{E}\|y_{j,l}-z_{j}\|)^{2}$ (64)
$\displaystyle\lesssim\frac{(c_{k_{i},l}^{2}+\sigma^{2})^{2}}{\tau_{k_{i},l}^{2}}$
(65)
$\displaystyle\lesssim\frac{c_{k_{i},l}^{4}}{c_{k_{i},l}^{2}}+\frac{c_{k_{i},l}^{2}\sigma^{2}}{\delta_{i}\sigma\Delta}+\frac{\sigma^{4}}{\delta_{i}\sigma\Delta}$
(66)
$\displaystyle=c_{k_{i},l}^{2}+\frac{(c_{k_{i},l}^{2}+\sigma^{2})\sigma}{\delta_{i}\Delta}.$
(67)

∎

######  Lemma 4 (Bound $\mathcal{T}_{2}$: Variance of Clipped Points).

$\displaystyle\mathbb{E}\bigg{\|}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\sim
i}(y_{j,l}-\mathbb{E}y_{j,l})\bigg{\|}^{2}$
$\displaystyle\leq\frac{n_{i}}{|\mathcal{G}_{i}|^{2}}\sigma^{2}.$ (68)

###### Proof of Lemma 4.

Note that the elements in the sum $\sum_{j\in\mathcal{G}_{i}:j\sim
i}(y_{j,l}-\mathbb{E}y_{j,l})$ are not independent. Therefore, we cannot get
rid of the cross terms when expanding the squared-norm. However, if for each
round of thresholding we were sample a fresh batch of points to set the new
cluster-center estimate, then the terms would be independent. With this
resampling strategy, our bounds would only change by a constant factor.
Therefore, for ease of analysis, we will assume the terms in the sum are
independent. In that case,

$\displaystyle\mathbb{E}\bigg{\|}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\sim
i}(y_{j,l}-\mathbb{E}y_{j,l})\bigg{\|}^{2}$
$\displaystyle\leq\frac{n_{i}}{|\mathcal{G}_{i}|^{2}}\mathbb{E}\|y_{j,l}-\mathbb{E}y_{j,l}\|^{2}$
(69)
$\displaystyle\leq\frac{n_{i}}{|\mathcal{G}_{i}|^{2}}\mathbb{E}\|z_{j}-\mathbb{E}z_{j}\|^{2}$
(70) $\displaystyle\leq\frac{n_{i}}{|\mathcal{G}_{i}|^{2}}\sigma^{2},$ (71)

where the second-to-last inequality follows from the contractivity of the
thresholding procedure. ∎

######  Lemma 5 (Bound $\mathcal{T}_{3}$: Error due to Out-of-Cluster Good
Points).

$\displaystyle\mathbb{E}_{j\in\mathcal{G}_{i}:j\not\sim
i}\|y_{j,l}-\mathbb{E}z_{i}\|^{2}$
$\displaystyle\lesssim\bigg{(}1+\frac{\delta_{i}}{2}+\frac{\sigma^{2}}{\delta_{i}\Delta^{2}}\bigg{)}c_{k_{i},l}^{2}+\frac{\sigma^{3}}{\Delta}.$
(72)

###### Proof of Lemma 5.

By Young’s inequality,

$\displaystyle\mathbb{E}_{j\in\mathcal{G}_{i}:j\not\sim
i}\|y_{j,l}-\mathbb{E}z_{i}\|^{2}$
$\displaystyle\leq\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}\mathbb{E}\|v_{k_{i},l-1}-\mathbb{E}z_{i}\|^{2}+\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\mathbb{E}_{j\in\mathcal{G}_{i}:j\not\sim
i}\|y_{j,l}-v_{k_{i},l-1}\|^{2}$ (73)
$\displaystyle\leq\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}c_{k_{i},l}^{2}+\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\mathbb{E}_{j\in\mathcal{G}_{i}:j\not\sim
i}\|y_{j,l}-v_{k_{i},l-1}\|^{2}$ (74)
$\displaystyle=\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}c_{k_{i},l}^{2}+\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\mathbb{E}_{j\in\mathcal{G}_{i}:j\not\sim
i}[\|z_{j}-v_{k_{i},l-1}\|^{2}\mathbbm{1}\\{\|z_{j}-v_{k_{i},l-1}\|\leq\tau_{k_{i},l}\\}]$
(75)
$\displaystyle\leq\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}c_{k_{i},l}^{2}+\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\tau_{k_{i},l}^{2}\mathbb{P}_{j\in\mathcal{G}_{i}:j\not\sim
i}(\|z_{j}-v_{k_{i},l-1}\|\leq\tau_{k_{i},l}).$ (76)

We now have to bound the probability in the expression above. Note that if
$\|v_{k_{i},l-1}-z_{j}\|\leq\tau_{k_{i},l}$, then

$\displaystyle\|\mathbb{E}z_{j}-\mathbb{E}z_{i}\|^{2}$
$\displaystyle\lesssim\|z_{j}-\mathbb{E}z_{j}\|^{2}+\|z_{j}-\mathbb{E}v_{k_{i},l-1}\|^{2}+\|\mathbb{E}v_{k_{i},l-1}-\mathbb{E}z_{i}\|^{2}$
(77)
$\displaystyle\lesssim\|z_{j}-\mathbb{E}z_{j}\|^{2}+\|z_{j}-\mathbb{E}z_{j}\|^{2}+\|\mathbb{E}z_{j}-\mathbb{E}v_{k_{i},l-1}\|^{2}+\mathbb{E}\|v_{k_{i},l-1}-\mathbb{E}z_{i}\|^{2}+\mathbb{E}\|z_{i}-\mathbb{E}z_{i}\|^{2}$
(78)
$\displaystyle\lesssim\|z_{j}-\mathbb{E}z_{j}\|^{2}+\tau_{k_{i},l}^{2}+c_{k_{i},l}^{2}+\sigma^{2}.$
(79)

By Assumption 2, this implies that

$\displaystyle\Delta^{2}$
$\displaystyle\lesssim\|z_{j}-\mathbb{E}z_{j}\|^{2}+\tau_{k_{i},l}^{2}+c_{k_{i},l}^{2}+\sigma^{2}$
(80)

which means that

$\displaystyle\|z_{j}-\mathbb{E}z_{j}\|^{2}$
$\displaystyle\gtrsim\Delta^{2}-(\tau_{k_{i},l}^{2}+c_{k_{i},l}^{2}+\sigma^{2}).$
(81)

By Markov’s inequality,

$\displaystyle\mathbb{P}(\|z_{j}-\mathbb{E}z_{j}\|^{2}$
$\displaystyle\gtrsim\Delta^{2}-(\tau_{k_{i},l}^{2}+c_{k_{i},l}^{2}+\sigma^{2}))\leq\frac{\sigma^{2}}{\Delta^{2}-(\tau_{k_{i},l}^{2}+c_{k_{i},l}^{2}+\sigma^{2})}\lesssim\frac{\sigma^{2}}{\Delta^{2}}$
(82)

as long as

$\displaystyle\Delta^{2}$
$\displaystyle\gtrsim\tau_{k_{i},l}^{2}+c_{k_{i},l}^{2}+\sigma^{2},$ (83)

which holds due to the constraint on $\Delta$ in the theorem statement.
Therefore

$\displaystyle\mathbb{E}_{j\in\mathcal{G}_{i}:j\not\sim
i}\|y_{j,l}-\mathbb{E}z_{i}\|^{2}$
$\displaystyle\lesssim\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}c_{k_{i},l}^{2}+\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\frac{(c_{k_{i},l}^{2}+\delta_{i}\sigma\Delta)\sigma^{2}}{\Delta^{2}}$
(84)
$\displaystyle\leq\bigg{(}1+\frac{\delta_{i}}{2}+\frac{\sigma^{2}}{\delta_{i}\Delta^{2}}\bigg{)}c_{k_{i},l}^{2}+\frac{\sigma^{3}}{\Delta}.$
(85)

∎

######  Lemma 6 (Clustering Error due to Bad Points).

$\displaystyle\mathbb{E}\bigg{\|}\bigg{(}\frac{1}{|\mathcal{B}_{i}|}\sum_{j\in\mathcal{B}_{i}}y_{j,l}\bigg{)}-\mathbb{E}z_{i}\bigg{\|}^{2}$
$\displaystyle\lesssim c_{k_{i},l}^{2}+\delta_{i}\sigma\Delta$ (86)

###### Proof of Lemma 6.

$\displaystyle\mathbb{E}\bigg{\|}\bigg{(}\frac{1}{|\mathcal{B}_{i}|}\sum_{j\in\mathcal{B}_{i}}y_{j,l}\bigg{)}-\mathbb{E}z_{i}\bigg{\|}^{2}$
$\displaystyle\leq\mathbb{E}_{j\in\mathcal{B}_{i}}\|y_{j,l}-\mathbb{E}z_{i}\|^{2}$
(87)
$\displaystyle\lesssim\mathbb{E}_{j\in\mathcal{B}_{i}}\|y_{j,l}-v_{k_{i},l-1}\|^{2}+\mathbb{E}\|v_{k_{i},l-1}-\mathbb{E}z_{i}\|^{2}$
(88) $\displaystyle\lesssim c_{k_{i},l}^{2}+\delta_{i}\sigma\Delta.$ (89)

The last inequality follows from the intuition that bad points will position
themselves at the edge of the thresholding ball, a distance $\tau_{k_{i},l}$
away from the current center-estimate $v_{k_{i},l-1}$. Therefore we cannot do
better than upper-bounding
$\mathbb{E}_{j\in\mathcal{B}_{i}}\|y_{j,l}-v_{k_{i},l-1}\|^{2}$ by
$\tau_{k_{i},l}^{2}\approx c_{k_{i},l}^{2}+\delta_{i}\sigma\Delta$, the
squared-radius of the ball. ∎

###  A.2 Proof of Theorem 3.2

###### Proof of Theorem 3.2.

Let

$\mathcal{D}_{1}=\left\\{\begin{array}[]{<EMAIL_ADDRESS>}p\\\ 0&\textnormal{w.p. }1-p\end{array}\right.$

and

$\mathcal{D}_{2}=\left\\{\begin{array}[]{<EMAIL_ADDRESS>}1-p\\\ 0&\textnormal{w.p. }p\end{array}\right.$

and define the mixture
$\mathcal{M}=\frac{1}{2}\mathcal{D}_{1}+\frac{1}{2}\mathcal{D}_{2}$. Also
consider the mixture
$\tilde{\mathcal{M}}=\frac{1}{2}\tilde{\mathcal{D}}_{1}+\frac{1}{2}\tilde{\mathcal{D}}_{2}$,
where $\tilde{\mathcal{D}}_{1}=0$ and $\tilde{\mathcal{D}}_{2}=\delta$. It is
impossible to distinguish whether a sample comes from $\mathcal{M}$ or
$\tilde{\mathcal{M}}$. Therefore, if you at least know a sample came from
either $\mathcal{M}$ or $\tilde{\mathcal{M}}$ but not which one, the best you
can do is to estimate $\mu_{1}$ with $\hat{\mu}_{1}=\frac{\delta p}{2}$, half-
way between the mean of $\mathcal{D}_{1}$, which is $\delta p$, and the mean
of $\tilde{\mathcal{D}}_{1}$, which is $0$. In this case

$\mathbb{E}\|\hat{\mu}_{1}-\mu_{1}\|^{2}=\frac{\delta^{2}p^{2}}{4}.$

If $p\leq\frac{1}{2}$, then

$\Delta=(1-p)\delta-p\delta=(1-2p)\delta.$ (90)

Also,

$\sigma^{2}=\delta^{2}p(1-p).$ (91)

Equating $\delta^{2}$ in (90) and (91),

$\displaystyle\frac{\Delta^{2}}{(1-2p)^{2}}=\frac{\sigma^{2}}{p(1-p)},$ (92)

which can be rearranged to

$(4\sigma^{2}+\Delta^{2})p^{2}-(4\sigma^{2}+\Delta^{2})p+\sigma^{2}=0.$

Solving for $p$,

$p=\frac{1}{2}-\frac{\Delta}{2\sqrt{4\sigma^{2}+\Delta^{2}}}.$ (93)

Note that,

$\displaystyle\frac{\delta^{2}p^{2}}{4}$
$\displaystyle=\frac{\sigma^{2}p^{2}}{4p(1-p)}$
$\displaystyle=\frac{\sigma^{2}p}{4(1-p)}.$ (94)

Plugging the expression for $p$ from (93) into (94), we can see that

$\displaystyle\frac{\delta^{2}p^{2}}{4}$
$\displaystyle=\frac{\sigma^{2}}{4}\bigg{(}\frac{\sqrt{4\sigma^{2}+\Delta^{2}}-\Delta}{\Delta}\bigg{)}=\frac{\sigma^{2}}{4}\bigg{(}\sqrt{1+\frac{4\sigma^{2}}{\Delta^{2}}}-1\bigg{)}\geq\frac{\sigma^{2}}{4}\bigg{(}\frac{2\sigma^{2}}{\Delta^{2}}-\frac{2\sigma^{4}}{\Delta^{4}}\bigg{)}.$
(95)

The last step used an immediately verifiable inequality that $\sqrt{1+x}\geq
1+\frac{x}{2}-\frac{x^{2}}{8}$ for all $x\in[0,8]$. Finally, we can choose
$\Delta^{2}\geq 2\sigma^{2}$ to give the result that

$\mathbb{E}\|\hat{\mu}_{1}-\mu_{1}\|^{2}\geq\frac{\delta^{2}p^{2}}{4}\geq\frac{\sigma^{4}}{4\Delta^{2}}\,.$

Finally, suppose that there is only a single cluster with $K=1$. Then, given
$n$ stochastic samples. standard information theoretic lower bounds show that
we will have an error at least

$\mathbb{E}\|\hat{\mu}_{1}-\mu_{1}\|^{2}\geq\frac{\sigma^{2}}{4n}\,.$

Combining these two lower bounds yields the proof of the theorem. ∎

###  A.3 Proof of Lemma 1

###### Proof.

For $h$ an $L$-smooth function and $g$ a $\mu$-strongly-convex function with
shared optimum $x^{*}$, the following inequality holds for all $x$:

$\displaystyle\|\nabla h(x)\|^{2}$
$\displaystyle\leq\bigg{(}\frac{L}{\mu}\bigg{)}^{2}\|\nabla g(x)\|^{2}.$ (96)

To see this, note that by $L$-smoothness of $h$

$\displaystyle\|\nabla h(x)\|^{2}$ $\displaystyle=\|\nabla h(x)-\nabla
h(x^{*})\|^{2}\leq L^{2}\|x-x^{*}\|^{2}.$ (97)

By $\mu$-strong-convexity of $g$ and Cauchy-Schwarz inequality,

$\displaystyle\mu\|x-x^{*}\|^{2}$ $\displaystyle\leq\langle\nabla g(x)-\nabla
g(x^{*}),x-x^{*}\rangle\leq\|\nabla g(x)\|\|x-x^{*}\|.$ (98)

Rearranging terms in (98), squaring both sides, and combining it with (97)
gives (96).

We can now apply (96) to show that Assumptions 4 and 5 hold.

For Assumption 4, let $h(x)=f_{i}(x)-\nabla\bar{f}_{i}(x)$ and
$g(x)=\nabla\bar{f}_{i}(x)$. Thus $h$ and $g$ have the same optimum. Since the
average of $\mu$-strongly-convex functions is $\mu$-strongly-convex, $g$ is
$\mu$-strongly-convex. By $L$-smoothness of $f_{i}$,

$\displaystyle\|(\nabla f_{i}(x)-\nabla\bar{f}_{i}(x))-(\nabla
f_{i}(y)-\nabla\bar{f}_{i}(y))\|$ $\displaystyle\leq\|\nabla f_{i}(x)-\nabla
f_{i}(y)\|+\|\nabla\bar{f}_{i}(x))-\nabla\bar{f}_{i}(y)\|$ (99)
$\displaystyle\leq 2L\|x-y\|,$ (100)

showing that $h$ is $2L$-smooth. Therefore, by (96)

$\displaystyle\|\nabla f_{i}(x)-\nabla\bar{f}_{i}(x)\|^{2}$
$\displaystyle\leq\bigg{(}\frac{2L}{\mu}\bigg{)}^{2}\|\nabla\bar{f}_{i}(x)\|^{2},$
(101)

which shows that Assumption 4 is satisfied with $A=\nicefrac{{2L}}{{\mu}}$.

For Assumption 5, let $x_{i}^{*}$ be client $i$’s optimum (equivalently the
optimum of all clients in client $i$’s cluster).

$\displaystyle\|\nabla f_{i}(x)-\nabla f_{j}(x)\|^{2}$ $\displaystyle=\|\nabla
f_{i}(x)-(\nabla f_{j}(x)-\nabla f_{j}(x_{i}^{*}))-(\nabla
f_{j}(x_{i}^{*})-\nabla f_{i}(x_{i}^{*})))\|^{2}$ (102)
$\displaystyle\overset{(i)}{\geq}\frac{1}{2}\|\nabla f_{j}(x_{i}^{*})-\nabla
f_{i}(x_{i}^{*})\|^{2}-\|\nabla f_{i}(x)-(\nabla f_{j}(x)-\nabla
f_{j}(x_{i}^{*}))\|^{2}$ (103)
$\displaystyle\overset{(ii)}{\geq}\frac{1}{2}\|\nabla f_{j}(x_{i}^{*})-\nabla
f_{i}(x_{i}^{*})\|^{2}-2\|\nabla f_{i}(x)\|^{2}-2\|\nabla f_{j}(x)-\nabla
f_{j}(x_{i}^{*})\|^{2}$ (104)
$\displaystyle\overset{(iii)}{\geq}\frac{1}{2}\|\nabla f_{j}(x_{i}^{*})-\nabla
f_{i}(x_{i}^{*})\|^{2}-2\|\nabla
f_{i}(x)\|^{2}-2(\nicefrac{{L}}{{\mu}})^{2}\|\nabla f_{i}(x)\|^{2}$ (105)
$\displaystyle=\frac{1}{2}\|\nabla f_{j}(x_{i}^{*})-\nabla
f_{i}(x_{i}^{*})\|^{2}-2(1+(\nicefrac{{L}}{{\mu}})^{2})\|\nabla
f_{i}(x)\|^{2}$ (106) $\displaystyle=\frac{1}{2}\|\nabla
f_{j}(x_{i}^{*})\|^{2}-2(1+(\nicefrac{{L}}{{\mu}})^{2})\|\nabla
f_{i}(x)\|^{2},$ (107)

where justification for the steps are:

  * •

(i) For all $a$, $b$, it holds that $(a-b)^{2}\geq\frac{1}{2}b^{2}-a^{2}$,
since this inequality can be rearranged to state
$(\nicefrac{{b}}{{\sqrt{2}}}-\sqrt{2}a)^{2}\geq 0$.

  * •

(ii) Young’s inequality.

  * •

(iii) Set $h(x)=f_{j}(x)-\langle f_{j}(x_{i}^{*}),x\rangle$ and
$g(x)=f_{i}(x)$. Then $\nabla h(x)=\nabla f_{j}(x)-\nabla f_{j}(x_{i}^{*})$,
from which we see that $h$ and $g$ have the same optimum $x_{i}^{*}$ and $h$
is $L$-smooth (since $\|\nabla h(x)-\nabla h(y)\|=\|(\nabla f_{j}(x)-\nabla
f_{j}(x_{i}^{*}))-(\nabla f_{j}(y)-\nabla f_{j}(x_{i}^{*}))\|=L\|x-y\|$).
Applying (96) gives the desired result.

Therefore (107) shows that Assumption 5 is satisfied with
$\Delta=\frac{1}{\sqrt{2}}\max_{j\not\sim i}\|\nabla f_{j}(x_{i}^{*})\|$ and
$D=\sqrt{2(1+(\nicefrac{{L}}{{\mu}})^{2})}$. ∎

###  A.4 Proof of Theorem 3.3

First we establish some notation.

##### Notation.

  * •

$\mathcal{G}_{i}$ are the good clients and $\mathcal{B}_{i}$ the bad clients
from client $i$’s perspective. Therefore
$|\mathcal{G}_{i}|+|\mathcal{B}_{i}|=N$.

  * •

$\mathbb{E}_{x}$ denotes conditional expectation given the parameter, e.g.
$\mathbb{E}_{x}g(x)=\mathbb{E}[g(x)|x]$. $\mathbb{E}$ denotes expectation over
all randomness.

  * •

$\overline{X_{t}}\triangleq\frac{1}{T}\sum_{t=1}^{T}X_{t}$ for a general
variable $X_{t}$ indexed by $t$.

  * •

We use $f_{i}(x)$ to denote average loss on a general batch $B$ of samples.
That is, if $f_{i}(x;b)$ is the loss on a single sample $b$, we define
$f_{i}(x)=\frac{1}{|B|}\sum_{b\in B}f_{i}(x;b)$.

  * •

$\bar{f}_{i}(x)\triangleq\frac{1}{n_{i}}\sum_{j\in\mathcal{G}_{i}:j\sim
i}f_{j}(x)$

  * •

We introduce a variable $\rho^{2}$ to bound the variance of the gradients

$\displaystyle\mathbb{E}\|g_{i}(x)-\mathbb{E}_{x}g_{i}(x)\|^{2}$
$\displaystyle\leq\rho^{2},$ (108)

and show in Lemma 7 how this can be written in terms of the variance of the
gradients computed over a batch size of 1.

  * •

$l_{t}$ is the number of rounds that Threshold-Clustering is run in round $t$
of Federated-Clustering.

  * •

$k_{i}$ denotes the cluster to which client $i$ is assigned.

  * •

$v_{i,l,t}$ denotes the gradient update for client $i$ in round $t$ of
Federated-Clustering and round $l$ of Threshold-Clustering. That is,
$v_{i,l_{t},t}$ corresponds to the quantity returned in Step 6 of Algorithm 2.

  * •

To facilitate the proof, we introduce a variable $c_{k_{i},l,t}$ that
quantifies the distance from the cluster-center-estimates to the true cluster
means. Specifically, for client $i$’s cluster $k_{i}$ at round $t$ of
Federated-Clustering and round $l$ of Threshold-Clustering we set

$\displaystyle
c_{k_{i},l,t}^{2}=\mathbb{E}\|v_{i,l-1,t}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\|^{2}.$
(109)

  * •

For client $i$’s cluster $k_{i}$ at round $t$ of Federated-Clustering and
round $l$ of Threshold-Clustering, we use thresholding radius

$\displaystyle\tau_{k_{i},l,t}^{2}\approx
c_{k_{i},l,t}^{2}+A^{4}\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}+\delta_{i}\rho\Delta.$
(110)

  * •

Finally, we introduce a variable $y_{j,l,t}$ to denote the points clipped by
Threshold-Clustering:

$\displaystyle
v_{i,l,t}=\frac{1}{N}\sum_{j\in[N]}\underbrace{\mathbbm{1}(\|g_{j}(x_{i,t-1})-v_{i,l-1,t}\|\leq\tau_{k_{i}^{t},l})+v_{i,l-1,t}\mathbbm{1}(\|g_{j}(x_{i,t-1})-v_{i,l-1,t}\|>\tau_{k_{i}^{t},l}}_{y_{j,l,t}}).$
(111)

###### Proof of Theorem 3.3.

In this proof, our goal is to bound $\overline{\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}}$ for each client $i$, thus showing convergence. Recall
that our thresholding procedure clusters the gradients of clients at each
round and estimates the center of each cluster. These estimates are then used
to update the parameters of the clusters. Therefore, we expect
$\overline{\mathbb{E}\|\nabla f_{i}(x_{i,t-1})\|^{2}}$ to be bounded in terms
of the error of this estimation procedure. The following sequence of
inequalities shows this.

By $L$-smoothness of $f_{i}$ and setting $\eta\leq\nicefrac{{1}}{{L}}$,

$\displaystyle f_{i}(x_{i,t})$ $\displaystyle\leq
f_{i}(x_{i,t-1})+\langle\nabla
f_{i}(x_{i,t-1}),x_{i,t}-x_{i,t-1}\rangle+\frac{L}{2}\|x_{i,t}-x_{i,t-1}\|^{2}$
$\displaystyle=f_{i}(x_{i,t-1})-\eta\langle\nabla
f_{i}(x_{i,t-1}),v_{i,l_{t},t}\rangle+\frac{L\eta^{2}}{2}\|v_{i,l_{t},t}\|^{2}$
$\displaystyle=f_{i}(x_{i,t-1})+\frac{\eta}{2}\|v_{i,l_{t},t}-\nabla
f_{i}(x_{i,t-1})\|^{2}-\frac{\eta}{2}\|\nabla
f_{i}(x_{i,t-1})\|^{2}-\frac{\eta}{2}(1-L\eta)\|v_{i,l_{t},t}\|^{2}$
$\displaystyle\leq
f_{i}(x_{i,t-1})+\eta\|v_{i,l_{t},t}-\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}+\eta
A^{2}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}-\frac{\eta}{2}\|\nabla
f_{i}(x_{i,t-1})\|^{2}-\frac{\eta}{2}(1-L\eta)\|v_{i,l_{t},t}\|^{2}.$ (112)

Recall that $v_{i,l_{t},t}$ is client $i$’s cluster-center estimate at round
$t$ of optimization, so the second term on the right side of (112) is the
error due to the clustering procedure. In Lemma 8, we show that, in
expectation, this error is bounded by

$\displaystyle\delta_{i}\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+\frac{\rho}{\Delta}D^{2}\overline{\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}}+\bigg{(}\frac{\rho^{2}}{n_{i}}+\frac{\rho^{3}}{\Delta}+\beta_{i}\rho\Delta\bigg{)}.$
(113)

Therefore, subtracting $f_{i}^{*}$ from both sides, summing (112) over $t$,
dividing by $T$, taking expectations, applying Lemma 8 to (112), and applying
the constraint on $\Delta$ from the theorem statement, we have

$\displaystyle\eta\overline{\mathbb{E}\|\nabla f_{i}(x_{i,t-1})\|^{2}}$
$\displaystyle\lesssim\frac{\mathbb{E}(f_{i}(x_{i,0})-f_{i}^{*})}{T}+\eta
A^{2}\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+\eta\bigg{(}\frac{\rho^{2}}{n_{i}}+\frac{\rho^{3}}{\Delta}+\beta_{i}\rho\Delta\bigg{)}.$
(114)

The third term on the right side of (112) reflects the fact that clients in
the same cluster may have different loss objectives. Far from their optima,
these loss objectives may look very different and therefore be hard to cluster
together.

In order to bound this term, we use a similar argument as above. By
$L$-smoothness of $f_{i}$’s and setting $\eta\leq\nicefrac{{1}}{{L}}$,

$\displaystyle\bar{f}_{i}(x_{i,t})$
$\displaystyle\leq\bar{f}_{i}(x_{i,t-1})+\langle\nabla\bar{f}_{i}(x_{i,t-1}),x_{i,t}-x_{i,t-1}\rangle+\frac{L}{2}\|x_{i,t}-x_{i,t-1}\|^{2}$
$\displaystyle=\bar{f}_{i}(x_{i,t-1})-\eta\langle\nabla\bar{f}_{i}(x_{i,t-1}),v_{i,l_{t},t}\rangle+\frac{L\eta^{2}}{2}\|v_{i,l_{t},t}\|^{2}$
$\displaystyle=\bar{f}_{i}(x_{i,t-1})+\frac{\eta}{2}\|v_{i,l_{t},t}-\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}-\frac{\eta}{2}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}-\frac{\eta}{2}(1-L\eta)\|v_{i,l_{t},t}\|^{2}.$
(115)

Subtracting $\bar{f}_{i}^{*}$ from both sides, summing over $t$, dividing by
$T$, taking expectations, and applying Lemma 8,

$\displaystyle\eta\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}$
$\displaystyle\lesssim\frac{\mathbb{E}(\bar{f}_{i}(x_{i,0})-\bar{f}_{i}^{*})}{T}+\frac{\eta\rho}{\Delta}D^{2}\overline{\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}}+\eta\bigg{(}\frac{\rho^{2}}{n_{i}}+\frac{\rho^{3}}{\Delta}+\beta_{i}\rho\Delta\bigg{)}.$
(116)

Combining (114) and (116), and applying the constraint on $\Delta$ from the
theorem statement, we have that

$\displaystyle\eta\overline{\mathbb{E}\|\nabla f_{i}(x_{i,t-1})\|^{2}}$
$\displaystyle\lesssim\frac{\mathbb{E}(f_{i}(x_{i,0})-f_{i}^{*})+A^{2}\mathbb{E}(\bar{f}_{i}(x_{0})-\bar{f}_{i}^{*})}{T}+\eta\max(1,A^{2})\bigg{(}\frac{\rho^{2}}{n_{i}}+\frac{\rho^{3}}{\Delta}+\beta_{i}\rho\Delta\bigg{)}.$
(117)

Dividing both sides of (117) by $\eta=\nicefrac{{1}}{{L}}$ and noting that
$\rho^{2}=\nicefrac{{\sigma^{2}}}{{|B|}}$ from Lemma 7,

$\displaystyle\overline{\mathbb{E}\|\nabla f_{i}(x_{i,t-1})\|^{2}}$
$\displaystyle\lesssim\frac{\mathbb{E}(f_{i}(x_{i,0})-f_{i}^{*})+A^{2}\mathbb{E}(\bar{f}_{i}(x_{i,0})-\bar{f}_{i}^{*})}{\eta
T}+\max(1,A^{2})\bigg{(}\frac{\rho^{2}}{n_{i}}+\frac{\rho^{3}}{\Delta}+\beta_{i}\rho\Delta\bigg{)}$
(118)
$\displaystyle\leq\frac{L(\mathbb{E}(f_{i}(x_{i,0})-f_{i}^{*})+A^{2}\mathbb{E}(\bar{f}_{i}(x_{i,0})-\bar{f}_{i}^{*}))}{T}+\frac{\max(1,A^{2})(\nicefrac{{\sigma^{2}}}{{n_{i}}}+\nicefrac{{\sigma^{3}}}{{\Delta}}+\beta_{i}\sigma\Delta)}{\sqrt{|B|}}$
(119)
$\displaystyle\lesssim\sqrt{\frac{\max(1,A^{2})(\nicefrac{{\sigma^{2}}}{{n_{i}}}+\nicefrac{{\sigma^{3}}}{{\Delta}}+\beta_{i}\sigma\Delta)}{T}},$
(120)

where the last inequality follows from setting

$\displaystyle|B|\gtrsim\max(1,A^{2})\bigg{(}\frac{\sigma^{2}}{n_{i}}+\frac{\sigma^{3}}{\Delta}+\beta_{i}\sigma\Delta\bigg{)}T.$
(121)

Since $T=\nicefrac{{m_{i}}}{{|B|}}$, this is equivalent to setting

$\displaystyle|B|\gtrsim\sqrt{\max(1,A^{2})\bigg{(}\frac{\sigma^{2}}{n_{i}}+\frac{\sigma^{3}}{\Delta}+\beta_{i}\sigma\Delta\bigg{)}m_{i}}.$
(122)

∎

######  Lemma 7 (Variance reduction using batches).

If, for a single sample $b$,

$\displaystyle\mathbb{E}_{x}\|g_{i}(x;b)-\mathbb{E}_{x}g_{i}(x;b)\|^{2}\leq\sigma^{2},$
(123)

then for a batch $B$ of samples,

$\displaystyle\mathbb{E}_{x}\|g_{i}(x)-\mathbb{E}_{x}g_{i}(x)\|^{2}$
$\displaystyle\leq\frac{\sigma^{2}}{|B|}.$ (124)

###### Proof of Lemma 7.

Due to the independence and unbiasedness of stochastic gradients,

$\displaystyle\mathbb{E}_{x}\|g_{i}(x)-\mathbb{E}_{x}g_{i}(x)\|^{2}$
$\displaystyle=\frac{1}{|B|^{2}}\sum_{b\in
B}\mathbb{E}_{x}\|g_{i}(x;b)-\mathbb{E}_{x}g_{i}(x;b)\|^{2}$ (125)
$\displaystyle\leq\frac{\sigma^{2}}{|B|}.$ (126)

∎

######  Lemma 8 (Bound on Clustering Error).

$\displaystyle\overline{\mathbb{E}\|v_{i,l_{t},t}-\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}$
$\displaystyle\lesssim\delta_{i}\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+\frac{\rho}{\Delta}D^{2}\overline{\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}}+\bigg{(}\frac{\rho^{2}}{n_{i}}+\frac{\rho^{3}}{\Delta}+\beta_{i}\rho\Delta\bigg{)}$
(127)

###### Proof of Lemma 8.

We prove the main result with the following sequence of inequalities and
justify the labeled steps afterwards.

$\displaystyle\overline{\mathbb{E}\|v_{i,l_{t},t}-\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}$
$\displaystyle=\overline{\mathbb{E}\|v_{i,l_{t},t}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\|^{2}}$
(128)
$\displaystyle=\overline{\mathbb{E}\bigg{\|}\frac{1}{N}\sum_{j\in[N]}y_{j,l_{t},t}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\bigg{\|}^{2}}$
(129)
$\displaystyle=\overline{\mathbb{E}\bigg{\|}(1-\beta_{i})\bigg{(}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}}y_{j,l_{t},t}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\bigg{)}+\beta_{i}\bigg{(}\frac{1}{|\mathcal{B}_{i}|}\sum_{j\in\mathcal{B}_{i}}y_{j,t,l}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\bigg{)}\bigg{\|}^{2}}$
(130)
$\displaystyle\overset{(i)}{\leq}(1+\beta_{i})(1-\beta_{i})^{2}\overline{\mathbb{E}\bigg{\|}\bigg{(}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}}y_{j,l_{t},t}\bigg{)}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\bigg{\|}^{2}}$
(131)
$\displaystyle\phantom{{}=1}+\bigg{(}1+\frac{1}{\beta_{i}}\bigg{)}\beta_{i}^{2}\overline{\mathbb{E}\bigg{\|}\bigg{(}\frac{1}{|\mathcal{B}_{i}|}\sum_{j\in\mathcal{B}_{i}}y_{j,l_{t},t}\bigg{)}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\bigg{\|}^{2}}$
(132)
$\displaystyle\leq\underbrace{\overline{\mathbb{E}\bigg{\|}\bigg{(}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}}y_{j,l_{t},t}\bigg{)}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\bigg{\|}^{2}}}_{\mathcal{E}_{1}}+\beta_{i}\underbrace{\overline{\mathbb{E}\bigg{\|}\bigg{(}\frac{1}{|\mathcal{B}_{i}|}\sum_{j\in\mathcal{B}_{i}}y_{j,l_{t},t}\bigg{)}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\bigg{\|}^{2}}}_{\mathcal{E}_{2}}$
(133)
$\displaystyle\overset{(ii)}{\lesssim}(1-\delta_{i}+\beta_{i})\overline{c_{k_{i},l_{t},t}^{2}}+(\delta_{i}+\beta_{i}A^{4})\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+\frac{\rho}{\Delta}D^{2}\overline{\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}}$ (134)
$\displaystyle\phantom{{}=}+\bigg{(}\frac{\rho^{2}}{n_{i}}+\frac{\rho^{3}}{\Delta}+\beta_{i}\rho\Delta\bigg{)}$
(135)
$\displaystyle\overset{(iii)}{\lesssim}(1-\nicefrac{{\delta_{i}}}{{2}})\overline{c_{k_{i},l_{t},t}^{2}}+\delta_{i}\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+\frac{\rho}{\Delta}D^{2}\overline{\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}}+\bigg{(}\frac{\rho^{2}}{n_{i}}+\frac{\rho^{3}}{\Delta}+\beta_{i}\rho\Delta\bigg{)}$
(136)
$\displaystyle\overset{(iv)}{\lesssim}\overline{(1-\nicefrac{{\delta_{i}}}{{2}})^{l_{t}}c_{k_{i},1,t}^{2}}+\delta_{i}\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+\frac{\rho}{\Delta}D^{2}\overline{\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}}+\bigg{(}\frac{\rho^{2}}{n_{i}}+\frac{\rho^{3}}{\Delta}+\beta_{i}\rho\Delta\bigg{)}$
(137)
$\displaystyle\overset{(v)}{\leq}\frac{\rho}{\Delta}\overline{c_{k_{i},1,t}^{2}}+\delta_{i}\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+\frac{\rho}{\Delta}D^{2}\overline{\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}}+\bigg{(}\frac{\rho^{2}}{n_{i}}+\frac{\rho^{3}}{\Delta}+\beta_{i}\rho\Delta\bigg{)}.$
(138)

Justifications for the labeled steps are:

  * •

(i) Young’s inequality:
$\|x+y\|^{2}\leq(1+\epsilon)x^{2}+(1+\nicefrac{{1}}{{\epsilon}})y^{2}$ for any
$\epsilon>0$.

  * •

(ii) We prove this bound in Lemmas 9 and 13. Importantly, it shows that the
clustering error is composed of two quantities: $\mathcal{E}_{1}$, the error
contributed by good points from the cluster’s perspective, and
$\mathcal{E}_{2}$, the error contributed by the bad points from the cluster’s
perspective.

  * •

(iii) Assumption that
$\beta_{i}\lesssim\min(\delta_{i},\nicefrac{{\delta_{i}}}{{A^{4}}})$

  * •

(iv) Since
$\mathbb{E}\|v_{i,l_{t},t}-\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}=c_{k_{i},l_{t}+1,t}^{2}$,
the inequality forms a recursion which we unroll over $l_{t}$ steps.

  * •

(v) Assumption that
$l_{t}\geq\max(1,\frac{\log(\nicefrac{{\rho}}{{\Delta}})}{\log(1-\nicefrac{{\delta_{i}}}{{2}})})$

Finally, we note that

$\displaystyle c_{k_{i},1,t}^{2}$
$\displaystyle=\mathbb{E}\|g_{i}(x_{i,t-1})-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\|^{2}$
(139)
$\displaystyle\lesssim\mathbb{E}\|g_{i}(x_{i,t-1})-\mathbb{E}_{x}g_{i}(x_{i,t-1})\|^{2}+\mathbb{E}\|\mathbb{E}_{x}g_{i}(x_{i,t-1})-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\|^{2}$
(140)
$\displaystyle\leq\rho^{2}+A^{2}\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}.$
(141)

Applying this bound to (138), and applying the bound on $\Delta$ from the
theorem statement, we have

$\displaystyle\overline{\mathbb{E}\|v_{i,l_{t},t}-\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}$
(142) $\displaystyle\lesssim\bigg{(}\delta_{i}+\frac{\rho
A^{2}}{\Delta}\bigg{)}\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+\frac{\rho}{\Delta}D^{2}\overline{\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}}+\bigg{(}\frac{\rho^{2}}{n_{i}}+\frac{\rho^{3}}{\Delta}+\beta_{i}\rho\Delta\bigg{)}$
(143)
$\displaystyle\lesssim\delta_{i}\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+\frac{\rho}{\Delta}D^{2}\overline{\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}}+\bigg{(}\frac{\rho^{2}}{n_{i}}+\frac{\rho^{3}}{\Delta}+\beta_{i}\rho\Delta\bigg{)}.$
(144)

∎

######  Lemma 9 (Clustering Error due to Good Points).

$\displaystyle\overline{\mathbb{E}\bigg{\|}\bigg{(}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}}y_{j,l_{t},t}\bigg{)}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\bigg{\|}^{2}}$
(145)
$\displaystyle\lesssim(1-\delta_{i})\overline{c_{k_{i},l_{t},t}^{2}}+\delta_{i}\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+\frac{\rho}{\Delta}D^{2}\overline{\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}}+\bigg{(}\frac{\rho^{2}}{n_{i}}+\frac{\rho^{3}}{\Delta}\bigg{)}$
(146)

###### Proof of Lemma 9.

We prove the main result in the sequence of inequalities below and then
justify the labeled steps.

$\displaystyle\overline{\mathbb{E}\bigg{\|}\bigg{(}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}}y_{j,l_{t},t}\bigg{)}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\bigg{\|}^{2}}$
(147)
$\displaystyle=\overline{\mathbb{E}\bigg{\|}\bigg{(}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}}y_{j,l_{t},t}\bigg{)}-\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\sim
i}\mathbb{E}_{x}g_{j}(x_{i,t-1})-\bigg{(}\frac{1}{n_{i}}-\frac{1}{|\mathcal{G}_{i}|}\bigg{)}\sum_{j\in\mathcal{G}_{i}:j\sim
i}\mathbb{E}_{x}g_{j}(x_{i,t-1})\bigg{\|}^{2}}$ (148)
$\displaystyle=\overline{\mathbb{E}\bigg{\|}\bigg{(}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\sim
i}(y_{j,l_{t},t}-\mathbb{E}_{x}g_{j}(x_{i,t-1}))\bigg{)}+\bigg{(}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\not\sim
i}(y_{j,l_{t},t}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1}))\bigg{)}\bigg{\|}^{2}}$
(149)
$\displaystyle\overset{(i)}{\leq}\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\overline{\mathbb{E}\bigg{\|}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\sim
i}(y_{j,l_{t},t}-\mathbb{E}_{x}g_{j}(x_{i,t-1}))\bigg{\|}^{2}}+\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}\overline{\mathbb{E}\bigg{\|}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\not\sim
i}(y_{j,l_{t},t}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1}))\bigg{\|}^{2}}$ (150)
$\displaystyle\overset{(ii)}{\lesssim}\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\overline{\mathbb{E}\bigg{\|}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\sim
i}(\mathbb{E}y_{j,l_{t},t}-\mathbb{E}_{x}g_{j}(x_{i,t-1}))\bigg{\|}^{2}}+\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\overline{\mathbb{E}\bigg{\|}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\sim
i}(y_{j,l_{t},t}-\mathbb{E}y_{j,l_{t},t})\bigg{\|}^{2}}$ (151)
$\displaystyle\phantom{{}=1}+\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}\overline{\mathbb{E}\bigg{\|}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\not\sim
i}(y_{j,l_{t},t}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1}))\bigg{\|}^{2}}$ (152)
$\displaystyle\overset{(iii)}{\lesssim}\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\overline{\mathbb{E}\bigg{\|}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\sim
i}\mathbb{E}(y_{j,l_{t},t}-g_{j}(x_{i,t-1}))\bigg{\|}^{2}}+\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\overline{\mathbb{E}\bigg{\|}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\sim
i}(\mathbb{E}_{x}g_{j}(x_{i,t-1})-\mathbb{E}g_{j}(x_{i,t-1}))\bigg{\|}^{2}}$
(153)
$\displaystyle\phantom{{}=1}+\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\overline{\mathbb{E}\bigg{\|}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\sim
i}(y_{j,l_{t},t}-\mathbb{E}y_{j,l_{t},t})\bigg{\|}^{2}}+\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}\overline{\mathbb{E}\bigg{\|}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\not\sim
i}(y_{j,l_{t},t}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1}))\bigg{\|}^{2}}$ (154)
$\displaystyle\overset{(iv)}{\lesssim}\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\delta_{i}^{2}\underbrace{\overline{\mathbb{E}_{j\in\mathcal{G}_{i}:j\sim
i}\|\mathbb{E}(y_{j,l_{t},t}-g_{j}(x_{i,t-1}))\|^{2}}}_{\mathcal{T}_{1}}+\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\frac{n_{i}}{|\mathcal{G}_{i}|^{2}}\rho^{2}$
(155)
$\displaystyle\phantom{{}=1}+\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\underbrace{\overline{\mathbb{E}\bigg{\|}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\sim
i}(y_{j,l_{t},t}-\mathbb{E}y_{j,l_{t},t})\bigg{\|}^{2}}}_{\mathcal{T}_{2}}+\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}(1-\delta_{i})^{2}\underbrace{\overline{\mathbb{E}_{j\in\mathcal{G}_{i}:j\not\sim
i}\|y_{j,l_{t},t}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\|^{2}}}_{\mathcal{T}_{3}}$
(156)
$\displaystyle\overset{(v)}{\lesssim}\delta_{i}\bigg{(}\overline{c_{k_{i},l_{t},t}^{2}}+\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+\frac{\rho^{3}}{\delta_{i}\Delta}\bigg{)}+\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\frac{n_{i}}{|\mathcal{G}_{i}|^{2}}\rho^{2}$
(157)
$\displaystyle\phantom{{}=1}+\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}(1-\delta_{i})^{2}\bigg{(}\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}\overline{c_{k_{i},l_{t},t}^{2}}+\delta_{i}\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+\frac{\rho}{\Delta}D^{2}\overline{\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}}+\frac{\rho^{3}}{\Delta}\bigg{)}$ (158)
$\displaystyle\overset{(vi)}{\lesssim}\delta_{i}\bigg{(}\overline{c_{k_{i},l_{t},t}^{2}}+\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+\frac{\rho^{3}}{\delta_{i}\Delta}\bigg{)}+\frac{\rho^{2}}{n_{i}}$
(159)
$\displaystyle\phantom{{}=1}+(1-\delta_{i})\bigg{(}\overline{c_{k_{i},l_{t},t}^{2}}+\delta_{i}\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+\frac{\rho}{\Delta}D^{2}\overline{\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}}+\frac{\rho^{3}}{\Delta}\bigg{)}$ (160)
$\displaystyle\lesssim(1-\delta_{i})\overline{c_{k_{i},l_{t},t}^{2}}+\delta_{i}\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+\frac{\rho}{\Delta}D^{2}\overline{\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}}+\bigg{(}\frac{\rho^{2}}{n_{i}}+\frac{\rho^{3}}{\Delta}\bigg{)}.$
(161)

Justifications for the labeled steps are:

  * •

(i),(ii),(iii) Young’s inequality

  * •

(iv) First, we can can interchange the sum and the norm due to independent
stochasticity of the gradients. Then by the Tower Property and Law of Total
Variance for the 1st and 3rd steps respectively,

$\displaystyle\mathbb{E}\|\mathbb{E}_{x}g_{j}(x_{i,t-1})-\mathbb{E}g_{j}(x_{i,t-1}))\|^{2}$
$\displaystyle=\mathbb{E}\|\mathbb{E}_{x}g_{j}(x_{i,t-1})-\mathbb{E}[\mathbb{E}_{x}g_{j}(x_{i,t-1})]\|^{2}$
(162) $\displaystyle=\text{Var}(\mathbb{E}_{x}(g_{j}(x_{i,t-1})))$ (163)
$\displaystyle=\text{Var}(g_{j}(x_{i,t-1}))-\mathbb{E}(\text{Var}_{x}(g_{j}(x_{i,t-1})))$
(164)
$\displaystyle\leq\text{Var}(g_{j}(x_{i,t-1}))-\mathbb{E}\|g_{j}(x_{i,t-1})-\mathbb{E}_{x}g_{j}(x_{i,t-1})\|^{2}$
(165) $\displaystyle\lesssim\rho^{2},$ (166)

where the last inequality follows since the two terms above it are both
bounded by $\rho^{2}$.

  * •

(v) We prove this bound in Lemmas 10, 11, and 12. It shows that, from point
$i$’s perspective, the error of its cluster-center-estimate is composed of
three quantities: $\mathcal{T}_{1}$, the error introduced by our thresholding
procedure on the good points which belong to $i$’s cluster (and therefore
ideally are included within the thresholding radius); $\mathcal{T}_{2}$, which
accounts for the variance of the clipped points in $i$’s cluster; and
$\mathcal{T}_{3}$, the error due to the good points which don’t belong to
$i$’s cluster (and therefore ideally are forced outside the thresholding
radius).

  * •

(vi) $(1+\nicefrac{{x}}{{2}})^{2}(1-x)^{2}\leq 1-x$ for all $x\in[0,1]$

∎

######  Lemma 10 (Bound $\mathcal{T}_{1}$: Error due to In-Cluster Good
Points).

$\displaystyle\overline{\mathbb{E}_{j\in\mathcal{G}_{i}:j\sim
i}\|\mathbb{E}(y_{j,l_{t},t}-g_{j}(x_{i,t-1}))\|^{2}}$
$\displaystyle\lesssim\overline{c_{k_{i},l_{t},t}^{2}}+\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+\frac{\rho^{3}}{\delta_{i}\Delta}.$
(167)

###### Proof of Lemma 10.

In this sequence of steps, we bound the clustering error due to good points
from client $i$’s cluster. By definition of $y_{j,l_{t},t}$,

$\displaystyle\mathbb{E}\|y_{j,l_{t},t}-g_{j}(x_{i,t-1})\|$
$\displaystyle=\mathbb{E}[\|v_{i,l_{t}-1,t}-g_{j}(x_{i,t-1})\|\mathbbm{1}(\|v_{i,l_{t}-1,t}-g_{j}(x_{i,t-1})\|>\tau_{k_{i},l_{t},t})]$
(168)
$\displaystyle\leq\frac{\mathbb{E}[\|v_{i,l_{t}-1,t}-g_{j}(x_{i,t-1})\|^{2}\mathbbm{1}(\|v_{i,l_{t}-1,t}-g_{j}(x_{i,t-1})\|>\tau_{k_{i},l_{t},t})]}{\tau_{k_{i}^{t},l_{t}}}$
(169)
$\displaystyle\leq\frac{\mathbb{E}\|v_{i,l_{t}-1,t}-g_{j}(x_{i,t-1})\|^{2}}{\tau_{k_{i},l_{t},t}}.$
(170)

Therefore, by Jensen’s inequality and plugging in the value for
$\tau_{k_{i},l_{t},t}$,

$\displaystyle\|\mathbb{E}(y_{j,l_{t},t}-g_{j}(x_{i,t-1}))\|^{2}$ (171)
$\displaystyle\leq(\mathbb{E}\|y_{j,l_{t},t}-g_{j}(x_{i,t-1})\|)^{2}$ (172)
$\displaystyle\leq\frac{(\mathbb{E}\|v_{i,l_{t}-1,t}-g_{j}(x_{i,t-1})\|^{2})^{2}}{\tau_{k_{i},l_{t},t}^{2}}$
(173)
$\displaystyle\lesssim\frac{(\mathbb{E}\|v_{i,l_{t}-1,t}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\|^{2}+\mathbb{E}\|\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})-\mathbb{E}_{x}g_{j}(x_{i,t-1})\|^{2}+\mathbb{E}\|g_{j}(x_{i,t-1})-\mathbb{E}_{x}g_{j}(x_{i,t-1})\|^{2})^{2}}{\tau_{k_{i},l_{t},t}^{2}}$
(174)
$\displaystyle\leq\frac{(c_{k_{i},l_{t},t}^{2}+A^{2}\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}+\rho^{2})^{2}}{\tau_{k_{i},l_{t},t}^{2}}$
(175)
$\displaystyle\lesssim\frac{(c_{k_{i},l_{t},t}^{2}+A^{2}\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}+\rho^{2})^{2}}{c_{k_{i},l_{t},t}^{2}+A^{4}\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}+\delta_{i}\rho\Delta}$
(176)
$\displaystyle\lesssim\bigg{(}1+\frac{\rho}{\delta_{i}\Delta}\bigg{)}c_{k_{i},l_{t},t}^{2}+\bigg{(}1+\frac{\rho
A^{2}}{\delta_{i}\Delta}\bigg{)}\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}+\frac{\rho^{3}}{\delta_{i}\Delta}$
(177) $\displaystyle\lesssim
c_{k_{i},l_{t},t}^{2}+\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}+\frac{\rho^{3}}{\delta_{i}\Delta},$
(178)

where the last inequality follows from constraints on $\Delta$. The second
inequality follows from Young’s inequality. The second-to-last inequality
follows by separating the fraction into a sum of fractions and selecting terms
from the denominator for each fraction that cancel with terms in the numerator
to achieve the desired rate.

Summing (178) over $t$ and dividing by $T$, we have

$\displaystyle\overline{\mathbb{E}_{j\in\mathcal{G}_{i}:j\sim
i}\|\mathbb{E}_{x}(y_{j,l_{t},t}-g_{j}(x_{i,t-1}))\|^{2}}$
$\displaystyle=\overline{\mathbb{E}\|y_{j,l_{t},t}-g_{j}(x_{i,t-1})\|^{2}}$
(179)
$\displaystyle\lesssim\overline{c_{k_{i},l_{t},t}^{2}}+\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+\frac{\rho^{3}}{\delta_{i}\Delta}.$
(180)

∎

######  Lemma 11 (Bound $\mathcal{T}_{2}$: Variance of Clipped Points).

$\displaystyle\mathbb{E}\bigg{\|}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\sim
i}(y_{j,l_{t},t}-\mathbb{E}y_{j,l_{t},t})\bigg{\|}^{2}$
$\displaystyle\leq\frac{n_{i}}{|\mathcal{G}_{i}|^{2}}\rho^{2}.$ (181)

###### Proof of Lemma 11.

The first thing to note is that the elements in the sum
$\sum_{j\in\mathcal{G}_{i}:j\sim i}(y_{j,l_{t},t}-\mathbb{E}y_{j,l_{t},t})$
are not independent. Therefore, we cannot get rid of the cross terms when
expanding the squared-norm. However, if for each round of thresholding we
sampled a fresh batch of points to set the new cluster-center estimate, then
the terms would be independent. With such a resampling strategy, our bounds in
these proofs only change by a constant factor. Therefore, for ease of
analysis, we will assume the terms in the sum are independent. In that case,

$\displaystyle\mathbb{E}\bigg{\|}\frac{1}{|\mathcal{G}_{i}|}\sum_{j\in\mathcal{G}_{i}:j\sim
i}(y_{j,l_{t},t}-\mathbb{E}y_{j,l_{t},t})\bigg{\|}^{2}$
$\displaystyle\leq\frac{n_{i}}{|\mathcal{G}_{i}|^{2}}\mathbb{E}\|y_{j,l_{t},t}-\mathbb{E}y_{j,l_{t},t}\|^{2}$
(182)
$\displaystyle\leq\frac{n_{i}}{|\mathcal{G}_{i}|^{2}}\mathbb{E}\|g_{j}(x_{i,t-1})-\mathbb{E}g_{j}(x_{i,t-1})\|^{2}$
(183) $\displaystyle\leq\frac{n_{i}}{|\mathcal{G}_{i}|^{2}}\rho^{2},$ (184)

where the second-to-last inequality follows from the contractivity of the
thresholding procedure. ∎

######  Lemma 12 (Bound $\mathcal{T}_{3}$: Error due to Out-of-Cluster Good
Points).

$\displaystyle\overline{\mathbb{E}_{j\in\mathcal{G}_{i}:j\not\sim
i}\|y_{j,l_{t},t}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\|^{2}}$
$\displaystyle\lesssim\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}\overline{c_{k_{i},l_{t},t}^{2}}+\delta_{i}\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+\frac{\rho}{\Delta}D^{2}\overline{\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}}+\frac{\rho^{3}}{\Delta}.$ (185)

###### Proof of Lemma 12.

This sequence of steps bounds the clustering error due to points not from
client $i$’s cluster. Using Young’s inequality for the first step,

$\displaystyle\mathbb{E}_{j\in\mathcal{G}_{i}:j\not\sim
i}\|y_{j,l_{t},t}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\|^{2}$ (186)
$\displaystyle\leq\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}\mathbb{E}\|v_{i,l_{t}-1,t}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\|^{2}+\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\mathbb{E}_{j\in\mathcal{G}_{i}:j\not\sim
i}\|y_{j,l_{t},t}-v_{i,l_{t}-1,t}\|^{2}$ (187)
$\displaystyle\leq\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}c_{k_{i},l_{t},t}^{2}+\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\mathbb{E}_{j\in\mathcal{G}_{i}:j\not\sim
i}\|y_{j,l_{t},t}-v_{i,l_{t}-1,t}\|^{2}$ (188)
$\displaystyle=\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}c_{k_{i},l_{t},t}^{2}+\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\mathbb{E}_{j\in\mathcal{G}_{i}:j\not\sim
i}[\|g_{j}(x_{i,t-1})-v_{i,l_{t}-1,t}\|^{2}\mathbbm{1}\\{\|g_{j}(x_{i,t-1})-v_{i,l_{t}-1,t}\|\leq\tau_{k_{i},l_{t},t}\\}]$
(189)
$\displaystyle\leq\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}c_{k_{i},l_{t},t}^{2}+\bigg{(}1+\frac{2}{\delta_{i}}\bigg{)}\tau_{k_{i},l_{t},t}^{2}\mathbb{P}_{j\in\mathcal{G}_{i}:j\not\sim
i}(\|g_{j}(x_{i,t-1})-v_{i,l_{t}-1,t}\|\leq\tau_{k_{i},l_{t},t})$ (190)
$\displaystyle\lesssim\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}c_{k_{i},l_{t},t}^{2}+\bigg{(}\frac{1}{\delta_{i}}\bigg{)}(c_{k_{i},l_{t},t}^{2}+A^{4}\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}+\delta_{i}\rho\Delta)\mathbb{P}_{j\in\mathcal{G}_{i}:j\not\sim
i}(\|g_{j}(x_{i,t-1})-v_{i,l_{t}-1,t}\|\leq\tau_{k_{i},l_{t},t}).$ (191)

The next step is to bound the probability term in the inequality above. Note
that if $\|v_{i,l_{t}-1,t}-g_{j}(x_{i,t-1})\|\leq\tau_{k_{i}^{t},l_{t}}$, then

$\displaystyle\|\mathbb{E}_{x}g_{j}(x_{i,t-1})-\mathbb{E}_{x}g_{i}(x_{i,t-1})\|^{2}$
$\displaystyle\lesssim\|g_{j}(x_{i,t-1})-\mathbb{E}_{x}g_{j}(x_{i,t-1})\|^{2}+\|g_{j}(x_{i,t-1})-v_{i,l_{t}-1,t}\|^{2}+\|v_{i,l_{t}-1,t}-\mathbb{E}_{x}g_{i}(x_{i,t-1})\|^{2}$
(192)
$\displaystyle\lesssim\|g_{j}(x_{i,t-1})-\mathbb{E}_{x}g_{j}(x_{i,t-1})\|^{2}+\tau_{k_{i},l_{t},t}^{2}$
(193)
$\displaystyle\phantom{{}=1}+\|v_{i,l_{t}-1,t}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\|^{2}+\|\mathbb{E}_{x}g_{i}(x_{i,t-1})-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\|^{2}$
(194)
$\displaystyle\lesssim\|g_{j}(x_{i,t-1})-\mathbb{E}_{x}g_{j}(x_{i,t-1})\|^{2}+\tau_{k_{i},l_{t},t}^{2}$
(195)
$\displaystyle\phantom{{}=1}+\|v_{i,l_{t}-1,t}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\|^{2}+A^{2}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}.$
(196)

By Assumption 4, the previous inequality implies

$\displaystyle\Delta^{2}-D^{2}\|\nabla f_{i}(x_{i,t-1})\|^{2}$
$\displaystyle\lesssim\|g_{j}(x_{i,t-1})-\mathbb{E}_{x}g_{j}(x_{i,t-1})\|^{2}+\tau_{k_{i},l_{t},t}^{2}+A^{2}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}+\|v_{i,l_{t}-1,t}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\|^{2}$
(197)

which, summing over $t$ and dividing by $T$, implies

$\displaystyle\overline{\|g_{j}(x_{i,t-1})-\mathbb{E}_{x}g_{j}(x_{i,t-1})\|^{2}}+\overline{\|v_{i,l_{t}-1,t}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\|^{2}}+A^{2}\overline{\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+D^{2}\overline{\|\nabla
f_{i}(x_{i,t-1})\|^{2}}$
$\displaystyle\gtrsim\Delta^{2}-\overline{\tau_{k_{i},l_{t},t}^{2}}.$ (198)

By Markov’s inequality, the probability of this event is upper-bounded by

$\displaystyle\frac{\rho^{2}+\overline{\mathbb{E}\|v_{i,l_{t}-1,t}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\|^{2}}+A^{2}\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+D^{2}\overline{\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}}}{\Delta^{2}-\overline{\tau_{k_{i},l_{t},t}^{2}}}$
(199)
$\displaystyle\lesssim\frac{\rho^{2}+\overline{c_{k_{i},l_{t},t}^{2}}+A^{2}\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+D^{2}\overline{\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}}}{\Delta^{2}},$ (200)

where the second inequality holds due to the constraint on $\Delta$ from the
theorem statement. Therefore,

$\displaystyle\overline{\mathbb{E}_{j\in\mathcal{G}_{i}:j\not\sim
i}\|y_{j,l_{t},t}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\|^{2}}$ (201)
$\displaystyle\lesssim\bigg{(}1+\frac{\delta_{i}}{2}+\frac{\rho}{\Delta}+\frac{\rho^{2}+\overline{c_{k_{i},l_{t},t}^{2}}+A^{2}\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+D^{2}\overline{\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}}}{\delta_{i}\Delta^{2}}\bigg{)}\overline{c_{k_{i},l_{t},t}^{2}}$
(202) $\displaystyle\phantom{{}=1}+\bigg{(}\frac{\rho
A^{2}}{\Delta}+\frac{A^{4}(\rho^{2}+\overline{c_{k_{i},l_{t},t}^{2}}+A^{2}\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+D^{2}\overline{\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}})}{\delta_{i}\Delta^{2}}\bigg{)}\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}$
(203)
$\displaystyle\phantom{{}=1}+\frac{\rho}{\Delta}D^{2}\overline{\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}}+\frac{\rho^{3}}{\Delta}$ (204)
$\displaystyle\lesssim\bigg{(}1+\frac{\delta_{i}}{2}\bigg{)}\overline{c_{k_{i},l_{t},t}^{2}}+\delta_{i}\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+\frac{\rho}{\Delta}D^{2}\overline{\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}}+\frac{\rho^{3}}{\Delta},$ (205)

where for the last inequality we again apply the constraint on $\Delta$. ∎

######  Lemma 13 (Clustering Error due to Bad Points).

$\displaystyle\overline{\mathbb{E}\bigg{\|}\bigg{(}\frac{1}{|\mathcal{B}_{i}|}\sum_{j\in\mathcal{B}_{i}}y_{j,l_{t},t}\bigg{)}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\bigg{\|}^{2}}$
$\displaystyle\lesssim\overline{c_{k_{i},l_{t},t}^{2}}+A^{4}\overline{\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}}+\delta_{i}\rho\Delta$
(206)

###### Proof of Lemma 13.

This lemma bounds the clustering error due to the bad clients from client
$i$’s perspective. The goal of such clients would be to corrupt the cluster-
center estimate of client $i$’s cluster as much as possible at each round.
They can have the maximum negative effect by setting their gradients to be
just inside the thresholding radius around client $i$’s cluster-center
estimate. This way, the gradients will keep their value (rather than be
assigned the value of the current cluster-center estimate per our update
rule), but they will have maximal effect in moving the cluster-center estimate
from its current position. Therefore, in step 3 of the inequalities below, we
can not do better than bounding the distance between these bad points and the
current cluster center estimate (i.e. $\|y_{j,l_{t},t}-v_{i,l_{t}-1,t}\|^{2}$)
by the thresholding radius ($\tau_{k_{i},l_{t},t}^{2}$).

$\displaystyle\mathbb{E}\bigg{\|}\bigg{(}\frac{1}{|\mathcal{B}_{i}|}\sum_{j\in\mathcal{B}_{i}}y_{j,l_{t},t}\bigg{)}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\bigg{\|}^{2}$
$\displaystyle\leq\mathbb{E}_{j\in\mathcal{B}_{i}}\|y_{j,l_{t},t}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\|^{2}$
(207)
$\displaystyle\lesssim\mathbb{E}_{j\in\mathcal{B}_{i}}\|y_{j,l_{t},t}-v_{i,l_{t}-1,t}\|^{2}+\mathbb{E}\|v_{i,l_{t}-1,t}-\mathbb{E}_{x}\bar{g}_{i}(x_{i,t-1})\|^{2}$
(208) $\displaystyle\leq\tau_{k_{i},l_{t},t}^{2}+c_{k_{i},l_{t},t}^{2}$ (209)
$\displaystyle\lesssim
c_{k_{i},l_{t},t}^{2}+A^{4}\mathbb{E}\|\nabla\bar{f}_{i}(x_{i,t-1})\|^{2}+\delta_{i}\rho\Delta.$
(210)

The last inequality applies the definition of $\tau_{k_{i},l_{t},t}$, and the
result of the lemma follows by summing this inequality over $t$ and dividing
by $T$. ∎

###  A.5 Proof of Theorem 3.4

First we establish some notation.

##### Notation.

  * •

$\mathcal{G}_{i}$ are the good clients and $\mathcal{B}_{i}$ the bad clients
from client $i$’s perspective.

  * •

$\mathbb{E}_{x}$ denotes conditional expectation given the parameter, e.g.
$\mathbb{E}_{x}g(x)=\mathbb{E}[g(x)|x]$. $\mathbb{E}$ denotes expectation over
all randomness.

  * •

$k_{i}^{t}$ is the cluster to which client $i$ is assigned at round $t$ of the
algorithm.

  * •

$\overline{X_{t}}\triangleq\frac{1}{T}\sum_{t=1}^{T}X_{t}$ for a general
variable $X_{t}$ indexed by $t$.

  * •

$\bar{f}_{i}(x)\triangleq\frac{1}{n_{i}}\sum_{j\in\mathcal{G}_{i}:j\sim
i}f_{j}(x)$

  * •

$\bar{m}_{i,t}\triangleq\frac{1}{n_{i}}\sum_{j\in\mathcal{G}_{i}:j\sim
i}m_{j,t}$

  * •

We introduce a variable $\rho^{2}$ to bound the variance of the momentums

$\displaystyle\mathbb{E}_{x}\|m_{i,t}-\mathbb{E}_{x}m_{i,t}\|^{2}$
$\displaystyle\leq\rho^{2},$ (211)

and show in Lemma 14 how this can be written in terms of the variance of the
gradients, $\sigma^{2}$.

  * •

$l_{t}$ is the number of rounds that Threshold-Clustering is run in round $t$
of Federated-Clustering.

  * •

$k_{i}$ denotes the cluster to which client $i$ is assigned.

  * •

$v_{k_{i},l,t}$ denotes the gradient update for client $i$ in round $t$ of
Momentum-Clustering and round $l$ of Threshold-Clustering. That is,
$v_{k_{i},l,t}$ corresponds to the quantity returned in Step 4 of Algorithm 4.

  * •

To facilitate the proof, we introduce a variable $c_{k_{i},l,t}$ that
quantifies the distance from the cluster-center-estimates to the true cluster
means. Specifically, for client $i$’s cluster $k_{i}$ at round $t$ of
Federated-Clustering and round $l$ of Threshold-Clustering we set

$\displaystyle
c_{k_{i},l,t}^{2}=\mathbb{E}\|v_{k_{i},l-1,t}-\mathbb{E}_{x}\bar{m}_{i,t}\|^{2}.$
(212)

  * •

For client $i$’s cluster $k_{i}$ at round $t$ of Federated-Clustering and
round $l$ of Threshold-Clustering, we use thresholding radius

$\displaystyle\tau_{k_{i},l,t}^{2}\approx
c_{k_{i},l,t}^{2}+\delta_{i}\rho\Delta.$ (213)

  * •

Finally, we introduce a variable, $y_{j,l,t}$, to denote the points clipped by
Threshold-Clustering:

$\displaystyle
v_{k_{i},l,t}=\frac{1}{N}\sum_{j\in[N]}\underbrace{\mathbbm{1}(\|m_{j,t}-v_{k_{i},l-1,t}\|\leq\tau_{k_{i},l,t})+v_{k_{i},l-1,t}\mathbbm{1}(\|m_{j,t}-v_{k_{i},l-1,t}\|>\tau_{k_{i},l,t}}_{y_{j,l,t}}).$
(214)

###### Proof of Theorem 3.4.

In this proof, our goal is to bound $\overline{\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}}$. We use $L$-smoothness of the loss objectives to get
started, and justify the non-trivial steps afterwards.

$\displaystyle\mathbb{E}f_{i}(x_{i,t})$
$\displaystyle\overset{(i)}{\leq}\mathbb{E}f_{i}(x_{i,t-1})+\mathbb{E}\langle\nabla
f_{i}(x_{i,t-1}),x_{i,t}-x_{i,t-1}\rangle+\frac{L}{2}\mathbb{E}\|x_{i,t}-x_{i,t-1}\|^{2}$
$\displaystyle=\mathbb{E}f_{i}(x_{i,t-1})-\eta\mathbb{E}\langle\nabla
f_{i}(x_{i,t-1}),v_{k_{i},l_{t},t}\rangle+\frac{L\eta^{2}}{2}\mathbb{E}\|v_{k_{i},l_{t},t}\|^{2}$
$\displaystyle=\mathbb{E}f_{i}(x_{i,t-1})+\frac{\eta}{2}\mathbb{E}\|v_{k_{i},l_{t},t}-\nabla
f_{i}(x_{i,t-1})\|^{2}-\frac{\eta}{2}\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}-\frac{\eta}{2}(1-L\eta)\mathbb{E}\|v_{k_{i},l_{t},t}\|^{2}$
$\displaystyle\overset{(ii)}{\lesssim}\mathbb{E}f_{i}(x_{i,t-1})+\eta\mathbb{E}\|v_{k_{i},l_{t},t}-\mathbb{E}_{x}\bar{m}_{i,t}\|^{2}+\eta\mathbb{E}\|\mathbb{E}_{x}\bar{m}_{i,t}-\nabla
f_{i}(x_{i,t-1})\|^{2}$ (215)
$\displaystyle\phantom{{}=1}-\frac{\eta}{2}\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}-\frac{\eta}{2}(1-L\eta)\mathbb{E}\|v_{k_{i},l_{t},t}\|^{2}$
(216)
$\displaystyle\overset{(iii)}{\lesssim}\mathbb{E}f_{i}(x_{i,t-1})+\eta\bigg{(}\frac{\rho^{2}}{n_{i}}+\frac{\rho^{3}}{\Delta}+\beta_{i}\rho\Delta\bigg{)}+\eta\mathbb{E}\|\mathbb{E}_{x}\bar{m}_{i,t}-\nabla
f_{i}(x_{i,t-1})\|^{2}$ (217)
$\displaystyle\phantom{{}=1}-\frac{\eta}{2}\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}-\frac{\eta}{2}(1-L\eta)\mathbb{E}\|v_{k_{i},l_{t},t}\|^{2}.$
(218)

Justifications for the labeled steps are:

  * •

(i) $L$-smoothness of $f_{i}$ and $\eta\leq\nicefrac{{1}}{{L}}$

  * •

(ii) Young’s inequality

  * •

(iii) Lemma 15

Now it remains to bound the $\mathbb{E}\|\mathbb{E}_{x}\bar{m}_{i,t}-\nabla
f_{i}(x_{i,t-1})\|^{2}$ term.

$\displaystyle\mathbb{E}\|\mathbb{E}_{x}\bar{m}_{i,t}-\nabla
f_{i}(x_{i,t-1})\|^{2}$
$\displaystyle\leq\mathbb{E}\|\mathbb{E}_{x}m_{i,t}-\nabla
f_{i}(x_{i,t-1})\|^{2}$ (219)
$\displaystyle=(1-\alpha)^{2}\mathbb{E}\|\mathbb{E}_{x}m_{i,t-1}-\nabla
f_{i}(x_{i,t-2})+\nabla f_{i}(x_{i,t-2})-\nabla f_{i}(x_{i,t-1})\|^{2}$ (220)
$\displaystyle\overset{(i)}{\lesssim}(1-\alpha)^{2}(1+\alpha)\mathbb{E}\|\mathbb{E}_{x}m_{i,t-1}-\nabla
f_{i}(x_{i,t-2})\|^{2}+(1-\alpha)^{2}\bigg{(}1+\frac{1}{\alpha}\bigg{)}L^{2}\eta^{2}\mathbb{E}\|v_{k_{i},l_{t-1},t-1}\|^{2}$
(221)
$\displaystyle\overset{(ii)}{\leq}\frac{1}{2}(1-L\eta)\sum_{t=2}^{T}\mathbb{E}\|v_{k_{i},l_{t-1},t-1}\|^{2},$
(222)

where justifications for the labeled steps are:

  * •

(i) Young’s inequality

  * •

(ii) Assumption that $\alpha\gtrsim L\eta$

Plugging this bound back into (218), summing over $t=1:T$, and dividing by $T$
gives

$\displaystyle\frac{\eta}{2T}\sum_{t=1}^{T}\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}$
$\displaystyle\lesssim\frac{\mathbb{E}(f_{i}(x_{i,0})-f_{i}^{*})}{T}+\eta\bigg{(}\frac{\rho^{2}}{n_{i}}+\frac{\rho^{3}}{\Delta}+\beta_{i}\rho\Delta\bigg{)}.$
(223)

By the variance reduction from momentum (Lemma 14) it follows from (223) that

$\displaystyle\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}\|\nabla
f_{i}(x_{i,t-1})\|^{2}$
$\displaystyle\lesssim\frac{\mathbb{E}(f_{i}(x_{i,0})-f_{i}^{*})}{\eta