One-Pass Distribution Sketch for Measuring Data Heterogeneity in Federated Learning
Authors: Zichang Liu
Zhaozhuo Xu
Benjamin Coleman
Anshumali Shrivastava
Venue: NeurIPS 2023 poster (2023)
Keywords: Distribution sketch
federated learning

Abstract: Federated learning (FL) is a machine learning paradigm where multiple client devices train models collaboratively without data exchange. Data heterogeneity problem is naturally inherited in FL since data in different clients follow diverse distributions.  To mitigate the negative influence of data heterogeneity, we need to start by measuring it across clients. However, the efficient measurement between distributions is a challenging problem, especially in high dimensionality. In this paper, we propose a one-pass distribution sketch to represent the client data distribution. Our sketching algorithm only requires a single pass of the client data, which is efficient in terms of time and memory. Moreover, we show in both theory and practice that the distance between two distribution sketches represents the divergence between their corresponding distributions. Furthermore, we demonstrate with extensive experiments that our distribution sketch improves the client selection in the FL training. We also showcase that our distribution sketch is an efficient solution to the cold start problem in FL for new clients with unlabeled data.

One-Pass Distribution Sketch for Measuring Data Heterogeneity in Federated Learning
Zichang Liu" Rice University zi changliuurice.ccmmmooo edu
Zhaozhuo Xu'
Stevens Institute of Technology 2xu79@stevens.eda
Benjamin Coleman
Anshumali Shrivastava Rice University & ThirdAI Corp. anshumali@rice. edu
Rice University
amin. benjamin.
benj
ray. cow.coleeaaaaoo wol.em@@gmmil...mm com
Abstract
Federated learning (FL.) is machine learning paradigm where multiple client devices train models collaboratively without data exchange. Data heterogeneity problem is naturally inherited in FL since data in different clients follow diverse distributions. To mitigate the negative influence of data heterogeneity, we need to start by measuring across clients. However, the efficient measurement between distributions is a challenging problem, especially in high dimensionality. In this paper, we propose one-pass distribution sketch to represent the client data dis- tribution. Our sketching algorithm only requires single pass of the client data, which efficient in terms of time and memory. Moreover, we show in both theory and practice that the distance between two distribution sketches represents the di- vergence between their corresponding distributions. Furthermore, we demonstrate with extensive experiments that our distribution sketch improves the client selection the FL training. We also showcase that our distribution sketch is an efficient solution to the cold start problem in FL for new clients with unlabeled data. |
Introduction
Federated learning (FL) has recently become the go-to solution for challenging on-device applications such as typing prediction (T. machine translation 2 and fraud detection 4 [5. Initially proposed by McMahan et. al. (6], the FL setting defines a learning problem where multiple clients collaborate construct a global model from local client data. Unlike centralized learning, where clients can freely exchange data, FL requires that private data never leave the client device. The FL framework is uniquely positioned to comply with recent data privacy legislation
9 and has correspondingly
seen surge of interest in machine learning (ML) research and deployment [10 11 12).
Data Heterogeneity in FL. Clients' local datasets may be drawn from significantly different statistical 8
21 22). For instance, in the typing prediction task, each client a cellphone user. Differences user interest, language, and geography will produce extremely diverse word distributions. Data heterogeneity is severe problem for FL because the models trained on each client may have many different behaviors. is challenging aggregate such different client models into a coherent and performant global model.
distributions (6] [3 14 15 16
17
19
20
'Equal contribution. The order of authors is determined by flipping coin. 'Now with Google DeepMind. Code available at
Distaace
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Benefits of Measuring Data Heterogeneity. We currently lack good methods to measure data heterogeneity but we expect that such methods could substantially improve the efficiency and We will use the term global distribution to denote the combination of client data distributions. If a client's data distribution is close to the global distribution, we can obtain better global model by focusing training resources on this client (improving efficiency). Data heterogeneity measures can also help implement an essential personalization procedure of FL [26] 27 In personalized FL, each client increases local predictive performance by finetuning the global model on their own data. This results in a personalized model for each client. Now, suppose that new client arrives with unlabeled data. We can search for client who has data that is similar to the new client and use its personalized model. To summarize, data heterogeneity measurement allows us to quantify the differences between clients in FL to improve FL algorithms and systems.
quality of FL 23 24 25
Challenges in Measuring Data Heterogeneity in FI.. The data heterogeneity measurement can be formulated as a statistical problem where we wish to quantify the degree to which probability distributions differ. However, this divergence estimation can be difficult
high-dimensiona space. Existing approaches, including earth mover distance 20 30 or Fréchet distance [3] 32] are compu- tationally expensive. Moreover, many classical algorithms fail for complex high-dimensional input distributions; learning is much more effective (32. However, learning introduces extra computational overheads. Since the clients may be mobile devices with limited hardware resources in the FL setting, expensive computation will introduce significant efficiency issue in practical deployment. Further, to keep with the privacy advantage of FL, data heterogeneity measurements must to private.
1.1 Our Proposal: One-Pass Distribution Sketch with Hashing
In this paper, we propose an efficient approach to represent client data distributions in FL. Motivated by the recent success of locality sensitive hashing in kernel density estimation, we develop a one-pass distribution sketch for measuring data heterogeneity on FL clients. We make the following specific contributions.
1. We propose an effective and efficient distribution sketching method that requires only one pass over the client data. We represent the data distribution with a matrix of counts, and we introduce a data heterogeneity measure which is the Euclidean distance between clients' sketch matrices.
2. We develop an efficient sietch-basedd sampling strategy for federated optimization. Given sketches for each FL client, we find a global distribution sketch by merging all client sketches. We propose form of client importance sampling where, in each training round, we preferentially select clients that have low distance to the global sketch. We observe a significant acceleration in FL, with faster convergence and better generalization than uniform sampling.
3. We develop efficient cold start approach for personalized FL where new client (with unlabeled data) selects the most appropriate model from the personalized finetuned models of other clients. To do this,
compare the distribution sketch of the new client with each existing client. We perform nearest neighbor search over clients under the sketch distance identify the client that has similar data distribution. Our cold-start approach provides simple yet effective way to find a suitable model for new clients.
2 Preliminary
Locality Sensitive Hashing (LSH)): LSH functions serve as a key ingredient in our distribution sketch. LSH is a randomized function family [33 34 35
36 37). Each function in this family maps an input vector into hash value, usually binary code or integer. If two vectors are close to each other, with high probability they will share the same hash value. Formally, we define the LSH function as:
Definition 2.1 (D,cD, Pı, PP.Pp))-ssse senssivve hash family). We define a function family H. to be (D, cD, pı, P2) sensitivee with respect to distance function d : Rd x RRd + R if for any two vector x,y E Rd. any h E H chosen uniformly at random satisfies:
If d(z y) < D then Pr[h(2)= = h((y)] 2 Pr
iff((, v) 2 cD then Pr(h(x) = h(y) < P2

---

Here D E R and C > 1. We denote Pr(h(z) = h(y]] as the collision probability of vector r and Here the term "collision" means I and y have the same hash value. In practice, we will use LSH functions where the collision probability between and y is monotonically increasing function with respect to a distance measure d TRd TRd + R. Formally, we denote this relationship as:
Pr[b(x)) = h(y) x f(d(2, y)).
We introduce the realization of the LSH function in Section A
of the Appendix.
Repeated Array of Count Estimators (RACE): RACE is an efficient Kernel Density Estimator
(KDE) with applications in outlier detection 38 39,,
graph compression [39,, genetic profiling [40]
and continue learning [1]. RACE builds two-dimensional matrix A E R Rx using R LSH
functions (see Definition 2.1 each LSH function h; maps a vector in Rd to an integer in [B]. Given a dataset D TR for every D, RACE maps
R different hash values, denoted as {h,(z)]² E [R]]. Then we increment entry (s) with 1 for every i E [R]. Following this manner, RACE maintains matrix where each column can be regarded as histogram of the dataset D. As result, we can use matrix A for KDE and its downstream tasks,
3
One-Pass Distribution Sketch
In this section, we introduce the one-pass distribution sketch for measuring data heterogeneity in FL..
3.1 One-Pass Sketch Construction
Intuition: We observe that the RACE matrix can be used as a representation for data distribution. Given a dataset DD C RR if we build a RACE matrix S by inserting every element in D, we build histogram of data samples with hashing-baa
projection. Next, we normalize the counts in the RACE matrix S with the dataset size [D], we will get representation of D.
Following this intuition, we introduce our distri- bution sketch algorithm in Algorithm ! Our dis- tribution sketch S is two-dimensional matrix with R rows and B columns. Given dataset D, we first initialize zero matrix S Rxx B and R LSH functions (h,,- ... p Here, each hash function h4 generates an integer hash value within [B). We suggest to be OV(ID). The choice of B depends on the LSH functions use. Next, for every x E D. we compute R hash values (h1) increment S,,).(z) for every i E [R]. In practice, for vision tasks with dense feature vectors, use SRP hash (See Definition A. as our LSH function. For language modeling tasks, tok- enize every sentence in the dataset into a bag of words characters 42 Next, sentence can be formed as a binary vector with nonzero index representing a token. In this case, we can use MinHash (see Definition A.2) as an LSH function.
Algorithm 1 One-Pass Distribution Sketch
Input: Dataset D C IR LSH function family
H. Hash range B, Rows R, Random Seed
Output: Sketch S E RRx Initialize: S 0
Generated R independent LSH functions h 1 hR from H with range B and random seed 8. for E do
(z)).. Finally, we
for i Ro
S.,.
end for end for S= 5/ return
Given dataset with vectors, the running time of Algorithm is O(nRd RB). Since the number LSH functions R and hash range B are usually smaller than the dataset size Algorithm[ i a linear passing over the dataset. Moreover, the space complexity of Algorithm I is ORBB), which is a much smaller than the memory cost of storing the dataset.
Complexity Analysis:
3.2 Theoretical Analysis
In this section, we describe how to measure the data heterogeneity using the one-pass distribution sketch in Algorithm I We will estimate the total variation distance (TVD), which roughly measures the largest pointwise difference between two distributions over their support.
Definition 3.1. Given two probability measures and Q on a space F).
TVD(P,Q) = sup P(A) Q(A))
(2)
AEF
3

---

If P and Q have density functions and q. this expression is equivalent to the following & norm.
1 TVD(p. q) = 2 //1 1 gldu
Our goal is to describe a method for consistent estimation of the TVD that is efficient to implement the distributed and private FL setting. At a high level, our strategy as follows.
1. Given a dataset D drawn from a distribution with probability density function p(r), construct sketch-based consistent estimator Sp(x) for the kernel density estimate (KDE) over D.
2. Argue that SD is a consistent estimator of p. because the KDE is a consistent estimator of and SD consistent estimator of the KDE.
3. Use Sp as a plug-in estimator for p in the TVD integral (Equation 3
Assumptions: We consider probability measures on R that have distribution function. We assume that the probability densities p and are continuous almost everywhere, so that the KDE - in step 2 [43]. We also assume that the support of P and Q is bounded subset of R so that our kernel estimator has the properties described in Section 3.2.1 Finally, our results are for convergence in probability rather than the (stronger) almost uniform convergence.
Notation: In a slight abuse of notation, we will use O(-) to denote both stochastic bounds and the more typical (deterministic) asymptotic bounds. More concretely, Xxx (ff(N) means that for any e > 0, we can find constants and No such that for all N > No. Pr[[ 2 cf(N)] S e. In the deterministic case, this still holds with densities p. 9 and two datasets D2 ~ p and Dg ~ We denote the sketch of Dp with Sp and Dj with S.. We write the distribution sketch S generated by Algorithm as functional S(z). We explicitly
O. We will assume two probability measures P, P,Q having
define this functional in Algorithm 4 Section
3.2.1 Consistent Estimation with Sketches
This section shows that our sketches provide consistent estimator of the data-generating distribution. We introduce the distribution sketch's 's query algorithm in Algorithm@i (See Section Din the appendix). For sketch constructed via Algorithm 1 given query q E Rdd we use the same LSH functions in construction sketch and get hash values. The äth hash value corresponds to location in the ith row. Next, we take the median over the average counts in the hashed locations. We show that the median-of-men -faa-ss---mmmss---- count Sgg) converges in probability to the kernel density of the LSH kernel. We begin with the definition of an LSH kernel.
Definition 3.2 (LSH Kernel), We define kernel k(x,, y) to be an LSH kernel for every x. y. k(x, v) represents the LSH collision probability (See Eq. (TT) of x and y. which is monotonically decreasing with distance d(z, y) 0. An LSH kernel k( y) is positive semi semi- definite kernel.
Next, we have the following statement.
Theorem 3.3 (Informal version of Theorem D.I. Let P(x) denote a probability density function.
Let D P(x) denote dataset. Let k(z, y) be LSH kernel (see Definition 3.2). Let S define the function implemented by Algorithm
We show that
S(x)
+
Σ k(
Lp. N
I.ED
with convergence rate Oog R/R).
Theorem 3.3 shows the consistent estimation of the KDE R co, which is itself consistent estimator of the distribution under some assumptions about k(x, y) and P. The convergence of the KDE the distribution as N + oo is well-studied problem under various conditions [44]. However, for the weak form of convergence which we consider here, it is sufficient for the kernel to integrate I, for the kernel bandwidth to decay with N, and for the bandwidth be bounded above by 1//d 45
We argue that these conditions hold for the L.SH kernels described by [46). The bandwidth require- ments are easy to satisfy because explicitly set the LSH kernel bandwidth by adjusting the

---

hash function parameters for example, by setting the number of bits in an SRP hash. The other requirement can be tricky to satisfy for LSH kemels,, such as the p-stable LSH probability, which is not absolutely integrable over RR. For this reason, some LSH kernels cannot be normalized to have unit integral. However, this is not a problem because we assume that P and Q have bounded support in Rd. Note that this type of input domain truncation is standard component of LSH design and analysis [47]. Under these conditions, Theorem 3.3 implies consistent estimation of P(x) by S(x) provided that R is allowed to weakly grow with N (i.c. R w()).
When the sketch required to preserve e-differential privacy, we must add independent Laplace noise to the sketch so that each counter is private (48]. This constrains our
Private Sketches:
choice of because the noise is distributed iid Laplace(Re"
Theorem 3.4 (Informal version of Theorem D.2p. Let S be e-differentially private distribution sketch of a dataset D = {21, } with iid P(x) and let k(x, y) be an LSH kernel (see Definition 3.2y. Let S define the function implemented by Algorithm 4 Then
S(z)
+
Σ k(z,,,) x)
Lp.
RED
Vooo with convergence rate O(v log R/R+ + v Vlg (N(Ne) when R = w(1), egg. R = log N.
In exchange for privacy, we pay a price in the convergence rate, which now has an additional - term. While this dependence can be improved 1/2 (by optimizing It for the fastest convergence), this requires linear memory, which is prohibitive in distributed settings.
3.2.2 Divergence Estimation
Under the conditions assumed by our analysis (continuous density Rd),, the TVD between KDEs a consistent estimator of the TVD [49 p1579]. Therefore, we can construct a consistent estimate of the TVD by simply plugging our estimators into Equation The integral itself be calculated using standard numerical integration techniques (e.g.. Monte Carlo estimation or sparse grids), giving us a consistent skatch-based estimator for the TVD.
This method is implementable and comes with error bounds, but it requires a costly numerical integration whose complexity scales poorly with dimension (O(d) per sample). Therefore, we justiffy a simpler method of computing the distance, albeit without formal guarantees.
As collection of histograms, our sketch is a piecewise constann approximation to the KDE over partition æ of events. This permits simpler expression for Equation 3 where we compute the value in each partition, weighed by the partition size. Here, suppose that the partition consists of cells (AI,.... We let x E be any point inside A and use X(A) to denote the Lebesgue measure A.
TVD(p. max X(A)-U TEAEE
lim
Σ -
To further simplify this expression, we make the following observations. Each cell of the partition, which we denote as A I,, is defined by the intersected buckets of the hash function. In other words, each cell contains the points that have the same hash mapping under all R partitions.
While the expression for max,, X(4) is not generally known for most LSH functions, it is clear that X(A) is non-increasing with R because each cell of the partition क is formed by intersecting R sub-partitions (one from each hash function). This operation almost surely decreases the size of A with each additional row in the sketch, since drawing two identical I.S functions is a zero-probability event. Also, (E{(A) A)] is the same for all the cells in the partition by symmetry arguments.
Therefore, it reasonable to consider an approximation where we compute entry-wise differences between sketches. Let Ar, be the difference between row and column (or bucket) 6 of the matrices Sp and S.y We can construct a TVD estimator as a sum over all possible intersections of hash buckets.
TVD(p, g) = B
Σ
A 1.b1
+ A An
bR=0
5

---

There are B buckets in each repetition and, therefore, B total "intersection paths" through the sketch. To circumvent the O(BR) complexity of computing each path. we consider the sum of absolute deviations 1 + bound on the TVD that can be computed in O(B R)) time because we can now distribute the sums. In particular, we can compute the upper bound by taking the Frobenius norm (Fuclidean distance) of Sp - S,. We denote ||S,, S,ll2 as sketch distance. Our empirical results show that this heuristic is an excellent indicator of distribution similarity.
n. instead of A1, + + A RRbaal This gives upper
4 One-Pass Distribution Sketch for FL
In this section, we introduce how to use our one-pass distribution sketch for effective and efficient FL. We will start with the client selection task with our sketch. Next, we will introduce how to perform an efficient cold start with our sketch in personalized FL shown in Section 6
4.1 Distribution-Awarc Client Selection
In our work, we study the setting where there are clients checking into the server. In every round of iterative federated optimization, clients are activated and ready for training. We assume that clients are activated uniformly random out of n clients. Our client selection algorithm presented Algorithm 3 in the supplemen- tary material. For each one of the clients, compute its sketch with Algorithm| D and send it to the server. As introduced in Theorem 3.4 our sketch can be c-differentially private with re- able
Algorithm 2 Cold Start with Distribution Sketch
Input: Clients C {cr,... .Cn). New Client
Cg. LSH function family H. Hash range B, Rows R. Random seed s.
Output: Personalized model
for i [m]] do
Get client data D; from client c;-
Get client model w; from client c;.
Compute sketch S; using Algorithm with
parameter D; H. B, R and 8.
Send S, to server
spect to the client data. As result, to communicate with the sketch between the server and clients. Next, the server takes the av- erage of all the client sketches towards a global sketch Sy- In each round of iterative federated optimization, the server looks up the sketch for every activated client. Next, compute the Eu- clidean distance between S,, and every activated client sketch Sj. Next, we use the reciprocal of the distance as sample probability Pj- Next, we normalize the p, across the L clients with Soft- max function [50]. Next, we sample K clients of the L activated clients. Each client is sam- pled with normalized probability pj. Finally, we perform local training on the selected clients and aggregate their model updates.
end for Get client data D2 from client Cg. Compute sketch S, using Algorithm with pa- rameter Dy. H. B, R and s. Send Sy to the server.
duni OC id = null
for i [n] do
d, = 1|s S:ll2
< dmin then
dmin + d;
id t
end if
end for return Wid
Our algorithm has three features: (1) We propose a way of estimating global data distribution across clients without data exchange. (2) We can select clients with data distribution closer to the global distribution for a better global model. (3) Our algorithm efficient. We only need one-shot, one- pass computation to generate a one me-pass distribution sketch of client data. Moreover, the sketch is mergeable so linear combination will give us sketch for global data distribution. Furthermore, we can efficiently estimate the divergence of the client data distribution to the global distribution by taking the Fuclidean distance of their sketches.
4.2 Nearest Neighbor Retrieval for Cold Start Client in Personalized FL
In this paper, we study the cold start problem in personalized FL. In this problem, many clients are connected to the server. Moreover, they collaborate together and produce global model through iterative federated optimization [12] Next, each client builds personalized model based on the global model that achieves decent predictive performance on its local data (See Section 6). Next, client connects to the server. The new client has local but unlabeled data. In this setting,

---

remains unknown which model, either global or local, should be applied to the new client as cold starter.
We propose Algorithm 2 for this cold start problem. For the clients that already connect to the server, compute its sketch with Algorithm Jnd send it the server. Next, given new client, we perform the sketching via Algorithm [ with the same parameters for existing clients. Next, we compute the Euclidean distance between the new client sketch with the existing clients' sketches and retrieve the nearest neighbor client with minimum distance. Finally, the server retrieves the personalized model on the nearest neighbor client and sends the new client.
We want to highlight that Algorithm| 2 formulates the cold start problem in personalized FL as a nearest neighbor search problem. By measuring data heterogeneity via Euclidean distance between distribution sketches, we provide an intuitive way for selecting a personalized client model.
5 Experiment
In this section, we conduct empirical evaluations to answer three questions: (1) Does the one-pass sketch distance reflect the differences between distributions? (2) Does the sketch distance help convergence in FL. and (3) Does the sketch distance retrieve the best-personalized models? To answer these three questions, we conducted three sets of experiments.
5.1 Dataset, Model and Implementation
Dataset: We evaluate Algorithm 3 and Algorithm@ on both vision and language datasets. For visual classification, federated dataset where each client is generated by random sampling from MNIST under Dirichlet
the MNIST dataset [51] and FEMNIST [52]. For MNIST, introduce distribution with
0.5. We also introduce MNIST Uniform + Dirichlet, dataset with half of the client data generated from random sampling from MNIST under Dirichlet distribution with = 0.5 and the other half of the client generated from uniform sampling from MNIST. We also use the Shakespeare next-character prediction dataset 6 for language-basee FL. The clients in both FEMNIST and Shakespeare datasets are fixed.
For the MNIST dataset, we use fully-connecter network with a single 512-dimension hidden layer. For the FEMNIST dataset,
Model:
LeNet-5 convolutional network (53 51] For the Shakespeare dataset, we build a model with one embedding layer, one long short-term memory (LSTM) (54) block, and output layer. The input and output vocabulary size is 80, and the hidden dimension of the model is 256.
Implementation: Our FL codebase, including FL workflow, LSH functions, and proposed algorithms, implemented on PyTorch [55]. We test Algorithm Algorithm on a server with 8 Nvidia Tesla V100 GPU and 48-core/96-thren processor (Intel(R) Xeon(R) Gold 5220R CPU @ 2.20GHz).
5.2 One-Pass Sketch Distance
In this section, investigate the correlation between one-pass sketch distance and distribution differences.
We generate two groups of subsets from MNIST, Group A and Group B. Datasets from Group are drawn randomly from MNIST under Dirichlet distribution with 0.5. Datasets from Group B are drawn randomly from MNIST under the uniform distribution. For every dataset, we construct one-pass sketch. Then, we calculate the pairwise one- sketch distances between every dataset.
Setting:
Result: We present the result in Figure [ We annotate the one-pass sketch distance in each cell. First, we observe symmetry along the diagonal, which satisfies the basic principle of distance, d(x, y) = d(y,z). Second, we observe group separation. The pairwise distances within group B are smaller, which is expected as they were drawn uniformly. The pairwise distances within Group are larger and more diverse, which expected as each is drawn randomly from the Dirichlet distribution. Lastly, for any dataset from Group A, its distance from any dataset from Group B similar. Our results answer the first question: the distance between the distribution sketches constructed by Algorithm can reflect their divergence in the distribution.
7

---

A
0.12
0.07
A
0.09
0.11
0.11
0.05
0.051
0.051
A,2
0.16
0.11
0.14
0.18
0.1
0.1
0.11
0.1
0.1
A3
0.12
0.11
0.14
0
0.14
0.075
0.076
0.077
0.075
0.075
A,4
0.1
0.0
0.073
0.075
0.018
0.017
0.017
0.016
0,1
0.07
0.05
0.1
0.076
0.1
0.018
0.017
0.016
0.016
0.2
0.07
0.05
0.11
0.077
0.098
0.017
0.017
0
0.017 0.017
0,3
0.071
0.051
0.1
0.075
0.1
0.017
0.016
0.017
0.015
8,4
A,0
A.1
A.2
A.3
A.4
B.O
B.1
8.2
8.3 8.4
Figure 1: Pairwise one-pass sketch distance visualization. Ap. A1. ... An are five datasets drawn randomly under the Dirichlet distribution. Bg. B1, ... Bn are five datasets drawn randomly under the uniform distribution. We mark each cell with the one-pase sketch distance shown in Algorithm [
5.3 Distribution-Aware Client Sampling
In this section, we apply the one-pass distribution sketch (see Algorithm T in the client selection setting (See Section 6) and investigate whether data heterogeneity measurement helps with FL towards faster convergence. Setting: As shown in Section 4.2 in every round of iterative federated
1.30
0.25
320
2015
Todkes Fedlug Cranncs
Toste2
Fasing
Poolig Shouch Obilanc
Tookeg Shich Delanos
Fadlig Shoich Dedance
5.00
Corennacaion Feuse
Communction Mound
Commucctioon..c. Nouve
Coneectte teand
(a) MNIST
(b) MNIST Uniform Dirichlet
(c) FEMNIST
(d) Shakespeare
Figure 2: Result: Convergence plot for MNIST, MNIST Uniform+Dirichlet, FEMNIST, and Shakespeare.
optimization, there would be L activated clients. The standard federated optimization approach would uniformly randomly sample K clients as participants 66. We denote this approach Fedavg. However, due to the data heterogeneity, we often observe model divergence during aggregation. As a result, we propose Algorithm 3 to select the clients with distribution closer to the global distribution. In this experiment, we compare Algorithm witt Fedavg to showcase the importance of data beterogeneity in client selection. We set L 3K. For the MNIST dataset (both MNIST and MNIST Uniform Direchlet), both Algorithm 3 and Fedavg trained by 200 rounds. In each round, K = 3 clients are selected from active clients. Next, each client trained for 20 epochs with batch size 32 and learning rate 1 0.0001. The FEMNIST uses the same training rounds, epochs, batch size, and learning rate as MNIST, except K = 10. For the Shakespeare dataset, both Algorithm and Fedavg active clients. Next, each client trained for 20 epochs with batch size 10 and learning rate r = 0.01. We introduce early stopping prevent overfitting.
trained by 20 rounds. Each round, K 3 clients selected from L
Result: We present the test accuracy plot in Figure As shown in the figure, on four FL datasets/settings, Algorithm converges faster than Fedavg. In almost every round of optimization,

---

Table 1: This table summarizes the test accuracy result for the cold client setting,
Warm Client
Cold Client
Dataset
Global Model
Local Model
Global Model
Neighbor Model
MNIST
0.8764 + 0.0084
0.9784 + 0.0016
0.8791 = 0.0136
0.9702 0.0004
FEMNIST
0.6593 0.0039
0.7336 + 0.0064
0.6667 + 0.0044
0.6937 = 0.0117
Shakespeare
0.3581 0.0250
0.3645 + 0.0249
0.3237 + 0.0186
0.3455 = 0.0082
Algorithm 3) achieves better accuracy than Fedavg. This experiment's results answer the second question. The distances between distribution sketches be used to sample helpful clients towards faster convergence.
5,4 Cold Start in Personalized FL
In this section, would like to evaluate the performance of Algorithm 2 in the cold start problem of personalized FL (see Section 4.2).
For every dataset, we uniformly ran- domly divided them into warm clients and cold clients. The warm clients will participate in global training using Fedavg. Once we have trained the global model. Each client locally fine-tunes the global model. As a result, we have
Setting:
1.0
I
I
Test Accuracy
0.8
0.6
global model the server, and each client has its personalized model. For cold clients, only its label for evaluation. We use MNIST, FEMNIST, and Shakespeare datasets as introduced in Section [.T] We generate MNIST clients following the by-class format shown in [13. There are 100 clients in total. We half of them as cold clients and the rest of them as warm clients. For MNIST, to get the global model, we train 200 rounds, and each round uni- formly selects 10 clients. Each client is trained for 10 epochs on its local data with batch size 32 and learning rate 0.0001. Next, we finetune each client for epochs with the same batch size and learning rate global model training for a local model. For FEMNIST, we set half of the clients as cold clients and the rest as warm clients. We use the same parameters as MNIST, except each round uses 50 clients. For Shakespeare, clients. To get the global model, we train rounds, and each round uniformly selects clients. Each client is trained for 20 epochs on its local data with batch size 10 and learning rate 0.01. Next, finetune each client for 20 epochs for a local model. We set the batch size and learning rate as the same global model training.
0.4
0.2
0.0
Global
Random
Neighbor
Best
Figure 3: Test accuracy cold clients for MNIST. Global: the model trained by Fedavg, Random: uniformly random sample a personalized client model, Neighbor: model selected by Algorithm 2 Best: the upper bound of test accuracy.
use 20 clients cold start clients, and the rest of them are warm
Result: We present the result Table 5.4 For warm clients, we observe that the local models generally have better test accuracy due to local finetuning. For cold clients, using the model from its nearest neighbor always achieves high accuracy compared to the global model. Moreover, we provide an ablation study on MNIST. We obtain the test accuracy each cold client with the global model, the model retrieved by uniformly random sampling, and the model retrieved by Algorithm 2 We also brutal force search the existing personalized models and find the best model with the highest test accuracy. Note that this search impractical since we use client data label information. But is upper bound for Algorithm As shown in Figure we observe that Algorithm Seporithm@Joupperfooo outperforms both the global model and random sampling with higher test accuracy on clients. Moreover, Algorithm approximates the upper bound well. The experiment in this section answers the third question: Algorithm 2 can retrieve suitable personalized models for new clients with unlabeled data.
Related Work
Client Selection in FL. Iterative federated optimization is the de-facto standard for FL.. In this optimization strategy, server aggregates the model updates from each device to create better

---

serves as an essential component of FL in the common situation where many clients are connected the server. each round of iterative federated optimization, the server identifies subset of clients as participants [60]. asks participants to train on client data, and finally aggregates the updates. Client selection necessary for two reasons. First, large-cohort training can cause optimization algorithms to fail. If we use all the clients, local objective misalignment can lead to catastrophic training failures [61]. This observation is backed up by theoretical analysis, which suggests that proper client selection can stabilize convergence [33 62 63 Second, there are systems challenges. The communication required to coordinate a large number of clients imposes heavy burden network bandwidth [64 651, reducing the efficiency of the FL system. Recent work on client selection strategies in FL focuses on the model's s performance on the client data. For instance, [66] and [67 use the training loss of the client model on its data utility score for client selection. [68] and (69 use client model updates and build importance scores to select clients. In this work, we focus on using data heterogeneity for client selection without model information. As a result, we can avoid the computation and communication on unselected clients.
global model. Client selection [56 57 58
59
Personalized FL. Personalization is natural way of improving local predictive performance client 770 71 27 72). Given the global model trained across diverse clients, is a common practice to fine-tune the global model on every client's 's data to obtain better local model [44 73 In this paper, we focus on a new challenging setting from practice: the cold start problem in personalized FL. In this setting, there are plenty of clients with personalized models on their devices. Next, new client comes with unlabeled data. We want to cold start this client with a suitable local model. In this way, we can still provide decent performance on this client without training. Our work focuses retrieving a local model from a client with a similar data distribution to the new client. We aim to achieve this strategy with effective data heterogeneity measurement.
7 Conclusion
We study the problem of data heterogeneity in federated learning (FL), where different clients contain data generated from different distributions. We argue that data heterogeneity should be effectively measured for better global model training and FL personalization. However, there is no explicit and efficient way of measuring data heterogeneity. To address this challenge, we propose one-pass distribution sketch that represents the data distribution while preserving c-differential privacy. We show that the sketch distance can approximate the total variation distance between clients and introduce this technique to the client selection and cold start personalization tasks in FL. Our experiments showcase our sketch's practical impact in effective and efficient FL.
8 Acknowledgement
We would like to thank Tian Li, Jianyu Wang, and the anonymous reviewers for helpful discussions and feedback. Zichang Liu, Benjamin Coleman, and Anshumali Shrivastava supported by NSF- CCS-2211815, ONR-DURIP and NST-BIGDATA-1838111 Zhaozhuo Xu is supported by startup funding from the Stevens Institute of Technology.
References
[1] Andrew Hard, Kanishka Rao, Rajiv Mathews, Swaroop Ramaswamy, Françoise Beaufays, Sean Augenstein, Hubert Eichner, Chloé Kiddon, and Daniel Ramage. Federated learning for mobile keyboard prediction. arXiv preprint arXiv: 1811:030040 2018
[2] Bill Yuchen Lin, Chaoyang He, Zihang Zeng, Hulin Wang, Yufen Huang, Mahdi Soltanolkotabi, Xiang Ren, and Salman Avestimehr. Fednip: A research platform for federated learning in natural language processing. arXiv preprint arXiv:2 104.004445555 2021.
[3] Peyman Passban, Tanya Roosta, Rahul Gupta, Ankit Chadha, and Clement Chung. Training mixed-domain translation models via federated learning. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 2576-2586 2022.
10

---

[4] Wensi Yang, Yuhang Zhang, Kejiang Ye, Li, and Cheng-Zhong Xu, Fíd: A federated learning based method for credit card fraud detection. In International conference on big data, pages 18-32. Springer, 2019.
[5] Guodong Long, Yue Tan, Jing Jiang, and Chengqi Zhang. Federated learning for open banking. In Federated learning, pages 240-254. Springer, 2020.
[6] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efiee
learning of deep networks from decentralized data. In Artificial intelligence and statistics, pages 1273-1282. PMLR, 2017.
[7] Colin Tankard. What the gdpr means for businesses. Network Security, 2016(6):5-8, 2016.
[8] Jacquelyn K O'herrin, Norman Fost, and Kenneth A Kudsk. Health insurance portability accountability act (hipaa) regulations: effect medical record research. Annals of surgery, 233(6):772,, 2004.
[9] Nguyen Truong, Kai Sun, Siyao Wang, Florian Guitton, and YiKe Guo. Privacy preservation federated learning: An insightful survey from the gdpr perspective. Computers & Security, 1 10: 102402, 2021.
[10] Chaoyang He, Songze Li, Jinhyun So, Xiao Zeng, Mi Zhang, Hongyi Wang, Xiaoyang Wang, Praneeth Vepakomma, Abhishek Singh, Hang Qiu, et al. Fedml: A research library and benchmark for federated machine learning. arXiv preprint wrXiv:2007.33518 2020.
[11] Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Ar- jun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, al. Advances and open problems in federated learning. Foundations and Trends@ in Machine Learning, 14(1-2):1 (1112)::12100 1-00 2021.
[12] Jianyu Wang, Zachary Charles, Zheng Xu, Gauri Joshi, H Brendan McMahan, Maruan Al- Shedivat, Galen Andrew, Salman Avestimehr, Katharine Daly, Deepesh Data, al. A field guide to federated optimization. arXiv preprint arXiv:2107.066) 2021.
[13] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Systems, 2:429-450, 2020.
[14] Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data. arXiv preprint arXiv: 00582, 2018.
[15] Shenglai Zeng, Zonghang Li, Hongfang Yu, Yihong He, Zenglin Xu, Dusit Niyato, and Han Yu. Heterogeneous federated learning via grouped sanpential-to-prraaa training. In International Conference on Database Systems Advanced Applicaation,, pages 455-471 Springer, 2022.
[16] Riccardo Zaccone, Andrea Rizzardi, Debora Caldarola, Marco Ciccone, and Barbara Caputo. Speeding up heterogeneous federated learning with sequentially trained superelients. In 2022 26th International Conference on Pattern Recognition (ICPR), pages 3376-3382. IEEE, 2022.
[17] Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International conference on machine learning, pages 5132-5143. PMLR, 2020.
[18] Alp Emre Durmus, Zhao Yue, Matas Ramon, Mattina Matthew, Whatmough Paul, and Saligrama Venkatesh. Federated learning based on dynamic regularization. In International Conference on Learning Representations, 2021.
[19] Sashank J Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Hugh Brendan McMahan. Adaptive federated optimization. International Conference on Learning Representations, 2020.
[20] Zhe Qu, Xingyu Li, Rui Duan, Yao Liu, Bo Tang, and Zhuo Lu. Generalized federated learning via sharpness aware minimization. In International Conference on Machine Learning, pages 19250-18280 PMLR, 2022.

---

[21] Debora Caldarola, Barbara Caputo, and Marco Ciccone. Improving generalization in federated learning by seeking flat minima. In European Conference on Computer Vision, pages 654-672. Springer, 2022.
[22] Jianyu Wang, Vinayak Tantia, Nicolas Ballas, and Michael Rabbat. proving communication arXiv:1910.006433 2019.
Slowmo: Im-
distributed sgd with slow momentum. arXiv preprint
[23] Zheng Chai, Hannan Fayyaz, Zeshan Fayyaz, Ali Anwar, Yi Zhou, Nathalie Baracaldo, Heiko Ludwig, and Yue Cheng. Towards taming the resource and data heterogeneity in federated learning. In OpML, pages 19-21, 2019.
[24] Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Chal- lenges, methods, and future directions. IEEE signal processing magazine, 37(3):50-66) (00000000000000 2020.
[25] Chen Zhang, Yu Xie, Hang Bai, Bin Yu. Weihong Li, and Yuan Gao. A survey federated learning. Knowledge-Baud Systems, 116:106775, 2021.
[26] Canh T Dinh, Nguyen Tran, and Josh Nguyen. Personalized federated learning with moreau envelopes. Advances in Neural Information Processing Systems, 33:21394-21405, 2020.
[27] Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In International Conference on Machine Learning, pages 6357-6368. PMLR, 2021.
[28] Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang. Towards personalized federated learning. IEEE Transactions Neural Networks and Learning Systems, 2022.
[29] Alexandr Andoni, Piotr Indyk, and Robert Krauthgamer. Earth mover distance over high- dimensional spaces. In SODA, volume 8, pages 343-352,20 2008
[30] Wenzhe Yang, Sheng Wang, Yuan Sun, and Zhiyong Peng. Fast dataset search with earth mover's distance. Proceedings of the VLDB Endowment, 15(11):5517-259) 2022.
[31] DC Dowson and BV666017 Landau. The fréchet distance between multivariate normal distribu- tions. Journal of multivariate analysis, 12(3):450-455, 1982.
[32] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
[33] Piotr Indyk and Rajeev Motwani. Approximate nearest neighbors: Towards removing the curse of dimensionality. In Proceedings of the Thirtieth Annual ACM Symposium on the Theory of Computing (STOC), pages 604-613, Dallas, TX, 1998.
[34] Mayur Datar, Nicole Immorlica, Piotr Indyk, and Vahab S. Mirrokni. Locality hashing scheme based on p-stable distributions. In Proceedings of the 20th ACM Symposium on Computational Geometry (SoCG), pages
Brooklyn, NY, 2004.
[35] Alexandr Andoni, Piotr Indyk, Huy L Nguyen, and Ilya Razenshteyn. Beyond locality-sensitive hashing. In Proceedings of the twenty-fifth annual ACM-SIAM symposium Discrete algo- rithms, pages 1018-1028. SIAM, 2014.
[36] Alexandr Andoni and Ilya Razenshteyn. Optimal data-dependent hashing for approximate near neighbors. In Proceedings of the ferty-seventh annual ACM symposium on Theory of computing (STOC), pages 793-801, 2015.
[37] Alexandr Andoni, Thijs Laarhoven, llya Razenshteyn, and Erik Waingarten. Optimal hashing- based time-space trade-offss for approximate near neighbors. In Proceedings of the Twenty-Eight Eightt Annual ACM- SIAM Symposium on Discrete Algorithms (SODA), pages 47-66. SIAM, 2017.
[38] Chen Luo and Anshumali Shrivastava. Arrays of docality-sensitive count estimators (ace) anomaly detection on the edge. Proceedings of the 2018 World Wide Web Conference, pages 1449, 2018
12

---

[39] Benjamin Coleman, Richard Baraniuk, and Anshumali Shrivastava. Sub-linear memory sketches for near neighbor search on streaming data. In International Conference on Machine Learning. pages 2089-2099. PMLR, R, 2020.
[40] Benjamin Coleman, Benito Geordie, Li Chou, RA Leo Elworth, Todd Treangen, and Anshumali Shrivastava. One-pass diversified sampling with application to serabyte-scale genomic sequence streams. In International Conference Machine Learning, pages 4202 112242155 4218. PMLR, 2022.
[41] Zichang Liu, Benjamin Coleman, Tianyi Zhang, and Anshumali Shrivastava. Retaining knowl- edge for learning with dynamic definition. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 14944-14958. Curran Associates, Inc., 2022.
[42] Zellig S Harris. Distributional structure. Word, 10(2-3):146110 162, 1954.
[43] Luc P Devroye and TJ Wagner. The (_1 convergence of kernel density estimates. The Annals of Statistics, 7(5):1 133-1199, 1979.
[44] Dominik Wied and Rafael Weißbach. Consistency of the kernel density estimator: a survey. Statistical Papers, 53(1):1-21, 2012.
[45] Luc Devroye. The equivalence of weak, strong and complete convergence in 11 for kernel density estimates. The Annals of Statistics, pages 896-904, 1983.
[46] Benjamin Coleman and Anshumali Shrivastava. Sub-linear race sketches for approximate kernel density estimation on streaming data. In Proceedings of The Web Conference 2020, pages 1739-1749, 2020.
[47] Indradyumna Roy, Rishi Agarwal, Soumen Chakrabarti, Anirban Dasgupta, and Abir De. Locality sensitive hashing in fourier frequency domain for soft set containment search. Advances in Neural Information Processing Systems, 2023.
[48] Benjamin Coleman and Anshumali Shrivastava. A one-pass distributed and private sketch for kernel sums with applications to machine learning at scale. In Proceedings of the 2021 ACM SIGSAC Conference on Computer and Communications Security, CCS '21, page 3252-3265, New York, NY, USA, 2021. Association for Computing Machinery.
[49] Bharath K Sriperumbudur, Kenji Fukumizu, Arthur Gretton, Bernhard Schölkopf, and Gert RG Lanckriet. On the empirical estimation of integral probability metrics. 2012.
[50] Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep learning. MIT press, 2016.
[51] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(1 2):278--2324 1998.
[52] Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Konečný, H Brendan McMahan, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for federated settings. arXiv preprint arXiv: 22022. 2018
[53] Yann LeCun, Bernhard Boser, John S Denker, Donnie Henderson, Richard E Howard, Wayne Hubbard, and Lawrence D Jackel. Backpropagation applied to handwritten zip code recognition. Neural computation, 1(4):541-551, 1989.
[54] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 98):1133517800 1997.
[55] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[56] Yae Jee Cho, Jianyu Wang, and Gauri Joshi. Towards understanding biased client selection in federated learning. In International Conference on Artificial Intelligence and Statistics, pages 10351-10375. PMLR, 2022.
13

---

[57] Wenkai Yang. Yankai Lin, Guangxiang Zhao, Peng Li, Jie Zhou, and Xu Sun. When to trust aggregated gradients: Addressing negative client sampling federated learning. arXiv preprint arXiv:2301.10400, 2023.
[58] Carl Smestad and Jingyue Li. A systematic literature review on client selection in federated learning. arXiv preprini arXiv:2306.04483 2023.
[59] Yae Jee Cho, Jianyu Wang, and Gauri Joshi. Client selection in federated learning: Convergence analysis and power-of-choice selection strategies. arXiv preprint arXiv:2010.01243, 2020.
[60] Keith Bonawitz, Hubert Eichner, Wolfgang Grieskamp, Dzmitry Huba, Alex Ingerman, Vladimir Ivanov, Chloe Kiddon, Jakub Konečný, Stefano Mazzoechi, Brendan McMahan, al. Towards federated learning at scale: System design. Proceedings of Machine Learning and Systems, 1:374-388, 2019.
[61] Zachary Charles, Zachary Garrett, Zhouyuan Huo, Sergei Shmulyian, and Virginia Smith. On large-cohort training for federated learning. Advances in neural information processing systems, 34:20461-20475, 2021.
[62] Jianyu Wang and Gauri Joshi. Cooperative sgd: A unified framework for the design and analysis of local-update sgd algorithms. Journal of Machine Learning Research, 22, 2021.
[63] Yann Fraboni, Richard Vidal, Laetitia Kameni, and Marco Lorenzi. A general theory for client sampling in federated learning. In IJCAI 2022-31st International joint conférence on artificial intellignce, 2022.
[64] Jie Xu and Heqiang Wang. Client selection and bandwidth allocation in wireless federated learning networks: A long-term perspective. IEEE Transactions on Wireless Communications, 20(2):1 188- 1880-10000 2020.
[65] Osama Shahid, Seyedamin Pouriyeh, Reza M Parizi, Quan Z Sheng, Gautam Srivastava, and Liang Zhao Communication efficiency in federated learning: Achievements and challenges. arXiv preprint arXiv: 100. 10996, 2021.
[66] Fan Lai, Xiangfeng Zhu, Harsha V Madhyastha, and Mosharaf Chowdhury. Oort: Efficient federated learning via guided participant selection. In 15th {USENIX) Symposium on Operating Systems Design and Implementation ({OSDI) 21), pages 19-35, 2021.
[67] Chenning Li, Xiao Zeng, Mi Zhang, and Zhichao Cao. Pyramidfl: A fine-grained client selection framework for efficient federated learning. In Proceedings of the 28th Annual International Conference Mobile Computing And Networking, pages 158-171, 2022.
[68] Georgios Kollias, Theodoros Salonidis, and Shiqiang Wang. Sketch to skip and select: Commu- nication efficient federated learning using locality sensitive hashing. In International Workshop on Trustworthy Federated Learning Conjunction with IJCAI 2022, 2022.
[69] Takayuki Nishio and Ryo Yonetani. Client selection for federated learning with heterogeneous resources in mobile edge. In ICC 2019-2019 IEEE international conference on communications (ICC), pages 1-7. IEEE, 2019.
[70] Virginia Smith, Chao- Kai Chiang, Maziar Sanjabi, and Ameet S Talwalkar, Federated multi-task learning. Advances in neural information processing systems, 30, 2017.
[71] Avishek Ghosh, Jichan Chung, Dong Yin, and Kannan Ramchandran. An efficient framework for clustered federated learning. Advances in Neural Information Processing Systems, 33:1958: 19597, 2020.
[72] Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach. Advances in Neural Information Processing Systems, 33:3557-3568, 2020.
[73] Tao Yu, Fugene Bagdasaryan, and Vitaly Shmatikov. Salvaging federated learning by local adaptation. arXiv preprint arXiv:2002.04758. 2020.
14

---

[74] Moses S. Charikar. Similarity estimation techniques from rounding algorithms. In Proceedings on 34th Annual ACM Symposium on Theory of Computing (STOC), pages 380-388, Montreal, Canada, 2002.
[75] Ping Li. Sign-full random projections. In Proceedings of the Thirty-Third AAAI Conference on Artificial Intelligence (AAAT), pages 4205-4212, Honolulu, HI, 2019.
[76] Xiaoyun Li and Ping Li. Random projections with asymmetric quantization. Advances in Neural Information Processing Systems, 32, 2019.
[77] Xiaoyun Li and Ping Li. Quantization algorithms for random fourier features. In International Conference on Machine Learning, pages 2369-6380 PMLR, 2021.
[78] Allan H Murphy. The finley affair: A signal event in the history of forecast verification. Weather and forecasting, 1 1(())33-0,, 1996.
[79] Joshua Engels, Benjamin Coleman, and Anshumali Shrivastava. Practical near neighbor search via group testing. Advances in Neural Information Processing Systems, 84:9950-9962, 2021.
[80] Xiaoyun Li and Ping Li. Rejection sampling for weighted jaccard similarity revisited. In Proceedings of the AAAI Conference on Artificial Intelligence (AAAI), 2021.
[81] Joshua Engels, Benjamin Coleman, Vihan Lakshman, and Anshumali Shrivastava. Dessert: An efficient algorithm for vector set search with vector set queries. arXiv preprint arXiv:22100 15.44888, 2022.
[82] Ping Li, Anshumali Shrivastava, Joshua Moore, and Arnd König. Hashing algorithms for large-scale learning. Advances in neural information processing systems, 24, 2011.
[83] Weijie Zhao, Deping Xie, Ronglai Jia, Yulci Qian, Ruiquan Ding, Mingming Sun, and Ping
Li. Distributed hierarchical gpu parameter server for massive scale deep learning ads systems. In Proceedings of the 3rd Conference on Third Conference on Machine Learning and Systems (MLSys), Huston, TX, 2020.
[84] Andrei Z. Broder, Moses Charikar, Alan M. Frieze, and Michael Mitzenmacher. Min-wise independent permutations. Proceedings of the Thirtieth Annual ACM Symposium on the Theory of Computing (STOC), pages 327-336, Dallas, TX, 1998.
[85] Ping Li, Art Owen, and Cun-Hui Zhang. One permutation hashing. Advances in Neural Information Processing Systems, 25, 2012.
[86] Ping Li, Xiaoyun Li, and Cun-Hui Zhang, Re-randomized densification for permutation hashing and bin-wise consistent weighted sampling. Advances Neural Information Processing Systems, 32, 2019.
[87] Xiaoyun Li and Ping Li. C-MinHash: Improving minwise hashing with circulant permutation. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 12857-12887. PMLR, 17-23 Jul 2022.
[88] Ping Li, Gennady Samorodnitsky, and John E. Hopcroft. Sign cauchy projections and chi-square kernel. In Advances in Neural Information Processing Systems (NIPS), pages 257 1-2579, Lake Tahoe, NV, 2013.
[89] Ping Li and Cun-Hui Zhang. Theory of the GMM kernel. In Proceedings of the 26th Interna- tional Conference on World Wide Web (WWW), pages 1053-1062, Perth, Australia, 2017.
[90] Ping Li. Several tunable GMM kernels. arXiv preprint arXiv: 2018.
[91] Amir Zandieh, Navid Nouri, Ameya Velingker, Michael Kapralov, and Ilya Razenshteyn. Scaling up kernel ridge regression via locality sensitive hashing. In International Conference on Artificial Intelligence and Statistics, pages 4088-4097. PMLR, 2020.
[92] SS Barsov and Vladimir Ul'yanov. Estimates of the proximity of gaussian measures. Sov. Math., Dokl, volume 34, pages 462-466, 1987.
15

---

Appendix
A Design of LSH Functions:
In practice, we could realize d(r, y) with cosine similarity for dense vectors. In this case, the LSH function family H. should have the following form.
Definition A.1 (Sign Random Projection (SRP) Hash (4)) define a family H as follow:
Given vectors x E R_, any h : E + {0,1} 1 that is in the family H. we have
M(x) Ai))(Ax))
where A E TRdx1 is random matrix that every entry of A is sampled from normal distribution (0,, 1), sign is the sign function that set positive value to and others to O. Moreover, for two vectors z, Id we have
Pr(h(x) h(w)) n." aarcoss
|
As shown in Definition A.I the SRP hash an LSH function built upon random Gaussian projections.
which is a fundamental technique in ML [75 76 (77 If two vectors are close in terms of cosine similarity, their collision probability would also be high.
The SRP hash is usually designed for dense vectors. both x E {0. 1] and E {0, 1)d are high dimensional binary vectors with large d. Their Jaccard similarity (78) is also important measure for search (79 80 81) and learning tasks [2] [83.. There also exists a family of LSH functions for Jacaerd similarity. We define this LSH function as:
Definition A.2 (MinHash 84
A comnctionffmit) H. a MinHash family for any h H. given
vectors E f0, 13d, we have
M(x) argmmn(((()))
where Il is a permutation function on the binary vector x. The arg min operation takes the index of the first non-zero value in 11(x). Moreover, given two binary vectors x, y E {0, 1)d, have
Ea
min( y:) max(
Pr[h(x) h(y))
E
where the right term represents the Jaccard similarity of binary vectors x and y.
Following Definition A.2 MinHash serves as a powerful tool for Jaccard similarity estimation 85 86 87]. We will use both SRP hash and MinHash in the following section to build a sketch for data distribution.
In this paper, we take a kernel view of the collision probability of L.H (see Definition 3.2 This
view aligns with series of research in efficient kernel decomposition [8 89 90 kernel density estimation [66 and kernel learning [9]]
B More Algorithms
In this section, we introduce the client selection algorithm with one-pass sketch.
16

---

Algorithm 3 Client Selection with Distribution Sketch
Clients C = (c1,-.., Co. Number of rounds T' Step size 7), LSH function family H. Hash range B, Rows Rt Number of active clients L, Number of Selected Clients K, Epochs E, Random Seed
Input:
Output: Global model w" Initialize: Global model w', global sketch Sg. random seed s.
for e E C do
Get client data D. from client
Compute sketch S. using Algorithm | with parameter De, H. B, R and 8.
Send Se to server
S + Se
end for
Sg Sg/n
Generate global sketch on server
for E [T] do
L clients (c1....
cL} C C are activated at random.
for j E L do
Pj 1/11S- Syll2
SJ is the sketch for client dj
end for
for j [L. do
exp Ps Σlin exp
end for Pj.
Sample K clients out (c...., without replacement, client ej is selected with probability
cz.
Server send w' to the selected clients. Each selected client updates w' to w by training on its data for E epochs with step size Each selected client sends wf. back to the server.
w -E 1
end for return =
C More Experiments
In this section, we provide two-fold of more experiments. We start with simulation of our sketch (see Algorithm Th for total variation distance (TVD) estimation. Next, perform client selection (see Algorithm CIFAR-10 datasets to observe the convergence improvements of our approach over FedAvg.
C.1 Simulation for Total Variation Distance Estimation
We computed the total variation distance (TVD) for two Gaussians with the same covariance and different means, one the few TVDs between nigh-diimnsional distributions for which we have an exact formula. This formula TVD = (Φ(22) (>) - where 2 = V1n (22) C-1( - U22) (see [4 92 for reference). To get the value of X(A) (see Section 3.2.2).
integrated the high- dimensional collision probability over the d-ball. We drew 1000 samples from two 16-dimensional
Gaussians with random means, computed their sketch and TVD distances, and compared the results.
C.2 CIFAR-10 Experiments
We conducted CIFAR- 10 experiments using 100 clients, where L=3K and K=2. The client dataset follows non-iid distribution, as specified in the split provided at witttt:/// // com/yjlee22/ FedShare We visualized the iteration-wig wise convergence of FedAvg with and without the proposed client selection (see Fi. The results demonstrate that our distribution sketch improves the test accuracy convergence rate, leading faster iterations.
17

---

1.3
Approx TVD
0.5
0.6
0.7
0.8
0.9
1.0
TVD
Figure 4: Visualization of TVD Estimation via Sketch.
44
Test accuracy
43
42
41
40
FedAvg
FedAvg + Sketch Distance
39
2
8
10
12
Communication Round
Figure 5: CIFAR-10 Experiments: We visualized the iteration- convergence of FedAvg with and without the proposed client selection.
D Proofs
To formally prove the Theorems in the paper. We start with introducing a query algorithm to the one-pass distribution sketch.
Algorithm 4 Query to the Distribution Sketch
Input: Query g E Rd LSH functions lu,.... hR. Dataset D, Sketch RR built on with
Algorithm
Output: STQ) E
Initialize: S(g) f- 0
A
+++ 0
for i 1 + do
AU (San.()))
end for S(g) + median(A) return S(q)
Next, we introduce the formal statements in the paper below.
Theorem D.1 (Formal version of Theorem 3.3). Let p(x:) denote a probability density function. Let
D p(x) denote dataset. Let k(x, y) be LSH kernel (see Definition 3.2). Let define the
18

---

function implemented by Algorithm
We show that
S(x) i.p N
Σ k(
with convergence rale OOVlog R/ R))
Proof. Let k(z) Ea.eD 12 x4) be the (non-normalized) kerel density estimate. Theorem 3.4
of 48 provides the following inequality for any ò, where xi(x) = ΣzEED vk(x,
32 K (z)
1/27
Pr [NS(2) x(=)] >
log 1/6)
<
R
This is equivalent to the following inequality, which we can obtain by dividing both sides of the inequality inside the probability by N.
1 )
(32 N2R
32 n (x( log 1/8
1/2
Pr S(x)
<ó
We want to show that the error (1()) proves the main claim of the theorem. To do this, we must show that for any A >
N x<(x)) converges in probability to zero, because this directly
1 k(x)
lim Pr R-000
S(x)
A
N
This can be done by setting 5 R. which yields the following inequality:
(x)
1
Pr S(x)
>
32
log
<
N
Because R < N, the following (simpler, but somewhat looser) inequality also holds:
$
log R
Pr S(x) N
k(0))
32
<
R
R
This implies that S(x)) - i.p. N x(x) by considering R large enough that /log R/ A.
Theorem D.2 (Formal version of Theorem 3.4p. Let S be an e-differenialll private distribution sketch of a dataset D = (II, with D p(z) and let k(æ, y) be an LSH kernel (see Definition 3.2
rid Let define the function implemented by Algorithm Then
S(x)
+
Σ k(r,,)))
Lp.
N
DDED
with convergence rate Oo R/R+ + RRlog (Ne)) when R = w(1) (e.g. R = log.N.
Σ EEED V k(x 2i) be the (non-normalized) kernel density estimate and let ((x)) = Ex.ep K(x, For the e-differentially private version algorithm, Theorem 3.4
Proof. As before. let n(x))= =
of 48 provides the following inequality for any > O.
Pr NS(c) - x(x))<(x>>>>>))))>>))))))))
((
K
(x) +
2 E
32 log 1/8
1/2 <
R
Again, we divide both sides of the inner inequality by N, and we also loosen (and simplify) the inequality with the observation that k(x)/N <.
(()
1/2
Pr S(x)
N k(())
+
32 log 1/6
<
R
2
19

---

We want to show that the error ((()) proves the main claim of the theorem. To do this, we must show that for any A
N x(x) converges in probability to zero, because this directly
lim Pr S(x)
1 | k(z)
>
A =0
20
N
As before, we choose = R. which yields the following inequality:
Pr S(x) _ >
1
((2 +
R
1/2
32 log R
No
R
This is the same as:
1 Pr S(x) N xn(z))
log R
R
1/2
1
32
+
64
log R
R
N2
R
Here, we will loosen the inequality again (also for the sake of presentation). Because Va Va+b S va + Vb, the following inequality is also satisfied:
Pr ((x) Ma(a)) 4
log
Volog R
1
+
<
R
Ne
R
Here, the convergence rate is O log R/R + Vllog R/ (N(N).. For us to have the error converge in probability to zero, we need for
log R R
Volgo R
v2
Ne
as + cc. A simple way to achieve this is for R to be weakly dependent on N (i.e. choose R = w (1)). For example, choosing R = log N satisfies the conditions.
20


Reviews:
[{"review_id": "VAQnrylNhl", "reviewer": "Reviewer_Jfur", "summary": "The authors propose a private strategy for measuring data heterogeneity using a RACE matrix. The RACE matrix essentially measures the normalized frequency of data points being assigned to particular LSH codes. These LSH count matrices are sent up to a server where it is applied into two different regimes:\\\\\n1. For distribution-aware client selection, the server samples clients with probability directly associated to closeness between the local RACE and global RACE matrix. \n2. For cold-start learning, a new client receives a personalized client model whose distribution matrix is closest to the new client's distribution matrix. \n\nThe approach is assessed over federated training on several datasets. Local accuracies are compared by using the global or locally closest model across both warm and cold clients. ", "questions": "1. I can't seem to easily find the settings for the LSH setup. How many LSH functions were used per sketch? Were they refreshed after every averaging round? Does the server distribute the hash functions? \n2. Were there any measurement regarding how the global model accuracy is affected by personalized model selection in the cold client setting? I imagine there should be a classical tradeoff between global and personalized accuracy. \n\n", "limitations": "No limitations discussed. Please consider including drawbacks. ", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 1, "strengths": "-The RACE matrix is theoretically a good estimator of DTV between distributions. \n-Scheme for developing client selection probabilities is simple. \n-There are apparently not many existing algorithms for addressing cold start learning. \n-Tested over multiple datasets with promising local accuracies. ", "weaknesses": "1. The RACE matrix is apparently already well-known as an excellent kernel density estimator. The contribution seems to be application of it to the federated setting, which has limited novelty. \n2. Under the suggested client scheme, the global model is favoring the largest set of homogeneous clients (popularity bias). Clients with atypical data distributions will have a low chance of being selected for global averaging, so there is no benefit for them to participate. If a client with an atypical distribution were to be randomly selected, their model would be washed out with a global model which is, on average, favoring distributions unlike theirs. \n3. The RACE estimation could be heavily affected by quantity skew (which will accompany label-distribution skew in practice). Many entries will be under-filled if a local dataset simply isn't large enough and the output space of the hash function larger than 2. \n4. In the cold start setting, the central server needs to be storing every possible client model. This would not be scalable. Such approach might better fit decentralized learning. \n5. MNIST and FEMNIST are too easy to classify. CIFAR-10 training, at the minimum would have been nice to see, which suffers far more from label-distribution heterogeneity, see: \"Federated Learning on Non-IID Data Silos: An Experimental Study,\" Li et al. (2021). ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are glad that you find the method simple and promising, and we have done our best to address your questions carefully. We hope that you might consider raising the score in light of our response. \n\n**Q1: Our contribution with respect to the RACE matrix**\n\nThe RACE matrix, indeed, proves to be an exceptional KDE technique. However, its application has been restricted to single probability distributions, both in theory and practice. This research aims to extend the capabilities of the RACE matrix by leveraging it as a distribution sketch, enabling the assessment of divergence across multiple distributions. It is important to emphasize that efficient divergence estimation is a nontrivial task. For instance, most sampling-based density estimators would require high-dimensional integrals involving all pairwise interactions between samples from the two clients. As a result, there exists a notable distinction between the conventional use of RACE for kernel density queries, as seen in the existing literature, and its novel application for approximating statistical distances, which is one of our key contributions.\n\n**Q2: Clients with atypical data distributions will have a low chance of being selected for global averaging, so there is no benefit for them to participate. If a client with an atypical distribution were to be randomly selected, their model would be washed out with a global model which is, on average, favoring distributions unlike theirs.**\n\nAtypical distribution clients generally destabilize the training process due to data heterogeneity [1,2]. We agree that sampling clients with distributions close to the goal distribution might enhance the popularity bias and is related to the important topic of fairness in federated learning. However, our distance measure can be used in several different ways, one of which could be to identify/sample atypical clients by looking at the largest distance with a global sketch. We want to emphasize that while we focus on personalized FL, our sketching algorithm provides an efficient and effective way to measure distribution differences in a private way, which is helpful for many different goals: stabilize training, fairness, personalized FL, etc.\n \n\n**Q3: Effect of quantity skew**\n\nThe RACE approximation of the density is an unbiased estimator whose variance (error) decreases with the client data size. RACE can be thought of as a quantization (histogram) of the data distribution, where the quantization’s granularity reduces and accuracy improves as more data arrive. While the accuracy will be better for high-quantity clients than low-quantity ones, we do not expect the under-filling of the matrix to pose serious problems because we are not computing density point estimates. We use RACE to approximate the statistical distance (a robust global quantity that depends on the full support of the distribution) rather than the point density (a more sensitive local quantity that depends on the presence/absence of nearby data). The overall RACE distance will still be meaningful (and unbiased), even if the underlying density estimates are only coarsely quantized.\n\n\n\n**Q4: In the cold start setting, the central server needs to store every possible client model**\n\nIn our approach, the server only stores the distribution sketch of existing clients, which is small in memory. The cold-start client will find its nearest neighbor in terms of distribution sketch in the server. Next, the cold-start client can retrieve the personalized model from its nearest neighbor client. As a result, the server doesn’t have to maintain all the client models.\n\n\n\n**Q5: CIFAR-10 results**\n\nWe conducted CIFAR-10 experiments using 100 clients, where L=3K and K=2. The client dataset follows a non-iid distribution, as specified in the split provided at https://github.com/yjlee22/FedShare. We visualized the iteration-wise convergence of FedAvg with and without the proposed client selection (see the submitted PDF). The results demonstrate that our distribution sketch significantly improves the test accuracy convergence rate, leading to faster iterations.\n\n\n**Q6: LSH setup**\n\nWe suggest using $R=O(\\sqrt{n})$ LSH functions globally for sketching, where $n$ is the maximum possible client data size. We set $B$ to be $O(\\log(n))$ following theoretical guarantees in [3].  We assume the server distributes the same set of hash functions to all clients by sharing the random seed. We don’t need to refresh the hash functions. In our experiment, we use sign random projections (Image) and MinHash (Text) as LSH functions. The server does not need to distribute the hashes, as their parameters can be generated from a single random seed.\n\n**Q7: Global model accuracy in the cold client setting.**\n\nIn this setting, we have trained a global model based on a set of clients. Each client participating in the training has labeled data. Next, for every trained client, we fine-tune the global model on its own data and get a personalized model associated with this client. Given a cold-start client with no label information on its client data, we generate its data distribution sketch and identify the trained client with the closest sketch distance. Finally, the cold-start client retrieves the model from the selected trained client as the personalized cold-start model. We note that the global model on the server is not affected by this procedure.\n\n[1] Li, Tian, et al. \"Federated optimization in heterogeneous networks.\" Proceedings of Machine learning and systems 2 (2020): 429-450.\n\n[2] Gao, D., Yao, X., & Yang, Q. (2022). A survey on heterogeneous federated learning. arXiv preprint arXiv:2210.04505.\n\n[3] Locality-Sensitive Hashing Scheme Based on p-Stable Distributions"}, {"review_id": "ipE0qAkpTo", "reviewer": "Reviewer_QJNz", "summary": "This paper proposed to use sketching for measuring heterogeneity in user data distributions in federated learning settings. The sketching is based on locality sensitive hashing and the authors have shown that in theory the measurement from the count sketch converges to the true kernel density estimation. The authors showed that the heterogeneity measurement and the count sketch can be used to 1) improve convergence by sampling users smartly according to their density in the distribution and 2) improve personalization by warming up the new client with a known nearest client’s model. The effectiveness of the approaches have been empirically tested on simple benchmark FL datasets.\n", "questions": "* What are the values for B and R chosen in this paper?\n* How should one choose these values for different datasets?\n", "limitations": "There is no or minor potential negative societal impact of this paper. The authors have acknowledged limitations in the appendix of the paper.\n", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* Measuring data heterogeneity has been an important problem to solve and the authors proposed a simple, intuitive yet effective idea. The implementation is easy to extend to existing FL frameworks. \n* The experiments demonstrated not only the effectiveness of heterogeneity measurement but also how to apply the measurement for improving federated learning convergence and personalization, showing the general applicability of the proposed methods.\n", "weaknesses": "* For cross-device setting, some algorithms (e.g. with differential privacy) might incur large communication costs depending on the size of the sketch matrix. \n* The datasets and models used in the experiments are very small scale and some do not reflect the true heterogeneity in the wild.\n* Figure 1 can be explained better with the label distribution for users in group A. E.g. Do users with higher dissimilarity have dissimilar digits in MNIST and vice versa?\n* Minor points about writing:\n1) RACE introduced in section 2 has matrix $A \\in {R}^{B\\times R}$ while the sketch matrix in the following paragraph has transposed the axes. There would be less friction of reading if they are consistent.\n2) Line 140, algorithm 0 should be algorithm 4.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments and suggestions to improve our paper! Please see the following clarifications.\n\n**Choices of parameters**\n\nWe suggest choosing $B$ and $R$ as follows: $R$ should be chosen proportional to $O(\\sqrt{n})$, where $n$ where \n is the maximum possible client data size. We set $B$ to be $O(\\log(n))$ following theoretical guarantees in [1].  As a result, our sketch size should be significantly smaller than the dataset size.\n\n\n**Figure 1 can be explained better with the label distribution for users in group A. E.g. Do users with higher dissimilarity have dissimilar digits in MNIST and vice versa?**\n\nThank you for the suggestion. Datasets in group A are drawn randomly under a Dirichlet distribution, and datasets in group B are drawn randomly under a uniform distribution. As a result, the label distribution among the datasets in Group A is more heterogeneous than that of Group B. As shown in the histogram, the pairwise distances among group A are smaller because they include all the digits. On the other hand, the pairwise distances among group B are larger and more inconsistent because they are likely to include different digits.\n\n**CIFAR-10 results**\n\nWe conducted CIFAR-10 experiments using 100 clients, where L=3K and K =2. The client dataset follows a non-iid distribution, as specified in the split provided at https://github.com/yjlee22/FedShare. We visualized the iteration-wise convergence of FedAvg with and without the proposed client selection. The results demonstrate that our distribution sketch significantly improves the test accuracy convergence rate, leading to faster iterations.\n\n\n[1] Locality-Sensitive Hashing Scheme Based on p-Stable Distributions\n\nWe thank the reviewer for the comments on the writing. We have updated the paper accordingly for a better presentation.\n"}, {"review_id": "aZOhuus0DM", "reviewer": "Reviewer_iXQ3", "summary": "The presence of **data heterogeneity** poses a critical challenge in **federated learning** (FL), as it leads to slower and noisier collaborative learning. To address this issue and improve the learning process with faster convergence, it is here proposed a more efficient client selection strategy based on distribution estimation. By utilizing an **efficient distribution sketch**, the server gains access to estimated distributions of the clients' local data, enabling the construction of a sketch representing the underlying global distribution while maintaining privacy constraints. With this information at hand, clients with data distributions closer to the global one can be more actively involved in the training process, thereby mitigating the negative effects of data heterogeneity and facilitating the development of a better global model. In addition, this approach offers the advantage of providing new clients - potentially holding unlabeled data - with initial models that are better aligned with their local distribution, thus enhancing their local training. To be more specific, new clients can receive the model trained by other users with similar local distributions, thus benefiting from their previously acquired knowledge.", "questions": "**Questions**\n1. Regarding the client selection strategy,\n    * How does the model behave on out-of-distribution (OOD) clients, *i.e.*, the clients having maximum distance from the global distribution sketch?\n    * How often are clients having larger distances from the global distribution sketch selected for training with respect to clients having more similar distributions?\n    * Referring to 3) in Weaknesses, what happens when all local distributions are very distant from the global one, i.e. if $p_i - p_j \\approx 0 \\forall i,j\\in[L]$? (referring to the notation in Algorithm 3)\n    * How does this approach compare with other smart client selection strategies? Please see 1) in Weaknesses for examples on related works.\n2. Regarding experiments in Figure 2:\n    * The proposed method seems to be particularly effective on the Shakespeare dataset rather than on MNIST/FEMNIST. Is this behavior more related to the dataset  itself, or to the task? Is the one-pass distribution sketch more effective for NLP? Experiments on an additional vision or NLP dataset could help clarify this behavior. Examples are federated CIFAR10/100, or StackOverflow.\n    * The text mentions the experiments were run for 200 rounds (line 315), but in the plots the X axes show a different number of rounds for each experiment. Which is the correct information?\n3. How are the values of $B$ and $R$ (Alg. 1) selected? Are they hyperparameters? Which values were chosen? How sensitive is the method to such a choice?\n4. Is there a risk of privacy leakage from the distribution sketch if no Laplace noise is added? \n\n**Minor questions** (did not affect my rating)\n1. How many clients were used in the experiments? \n2. What is the average number of samples per client?\n\n**Suggestions**\n1. I would suggest revising the abstract (lines 10-14) to obtain a more effective result. For instance, it is not clear how the distribution sketch helps in the client selection, or how this method helps training on new unlabeled clients. A few additional details might send the message better.\n\n**Minor suggestions** (did not affect my rating)\n* Typos: \n    * line 45: *is* computational ... --> \"are\"\n    * In Sec. 3.2, Algorithm **3** from Sec. D appears referenced as Algorithm **0**.\n    * lines 332-334 probably needs rephrasing (specifically referring to the sentence \"Once we have the global model\")\n* Def 2.1 introduces $p_1$, $p_2$ and $D$, whose set of belonging is not formally defined - *e.g.*, $\\mathbb{R}$.", "limitations": "Main limitations were addressed by the authors. \n\nAdditional limitation: potential lack of fairness towards users belonging to less represented distributions.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The paper addresses the issue of data heterogeneity in FL, relevant to this field of research.\n2. A new method for estimating local and global data distributions is introduced. This method is both efficient and privacy preserving. It can also be adapted to be $\\epsilon$-differentially private. The distribution sketch enables convergence speed up by allowing a smarter client selection, thus reducing the communication rounds between clients and server, significant for the evaluation of FL algorithms.\n3. Both theoretical and complexity analyses are presented for the proposed method. Claims are supported by proofs.\n4. Empirical validation is presented on both vision and NLP datasets.\n5.The paper is well written, easy to follow and understand.\n6. Limitations and potential negative impacts are discussed. Computational requirements are identified.", "weaknesses": "1. Some relevant related works are missing. \n    * Client selection: [1,2]\n    * Personalized FL: [3]\n    * I believe discussing other works proposing different data distribution estimation techniques is relevant to this paper, but is currently missing. Examples are [4,5,6,7].\n    * Algorithms addressing statistical heterogeneity in FL. Examples are unlimited: the most well-known baselines are [10-15]. Recently works such as [16,17] are gaining visibility, as well as [18].\n    * Lastly, since the paper highlights the positive effect of having new clients starting from pretrained models, current literature addressing the impact of pretraining in FL might be of interest [8,9].\n2. The efficacy of the one-pass distribution sketch in efficiently addressing the cold start of new *unlabeled* clients is emphasized. However, it is important to note that these clients still require labeling in order to effectively use the model and perform local training. As no alternative solution is proposed in this context, the necessity for labeled data on the client side remains unchanged. Consequently, the ability for new clients to initially possess unlabeled data does not yield a substantial difference in practical real-world scenarios.\n3. In realistic cross-device [19] federated scenarios, billions of clients with very small datasets are involved in the training. This implies $|\\cup_{k\\in[K]} \\mathcal{D}_k| \\gg |\\mathcal{D}_k|$, with $K$ being the number of clients with local datasets $\\mathcal{D}_k, k\\in[K]$. As a consequence, it is very unlikely that distribution of single local datasets may resemble the global one. Referring to the notation introduced in Algorithm 3, in the real-world scenario, $p_i - p_j \\approx 0 \\forall i,j\\in[L]$. It is not clear how this method would be effective in this case.\n4. To my understanding, in order to perform the warm start on new clients, the server has to keep memory of the last model trained by each client, and store it, together with the distribution sketch. This becomes unfeasible with the increasing number of clients in realistic settings. If the server retrieves the personalized model from the nearest neighbor client instead, this introduces additional communication, and such cost increases with the constant arrival of new clients. This limitation is nowhere addressed.\n5. By favoring clients having local distributions similar to the global one, the resulting learned model might potentially be unfair towards out-of-distribution users. I believe this behavior arises issues in terms of **fairness** towards less represented distributions.\n6. The method is only compared with FedAvg, the first algorithm proposed for FL. Missing comparison with methods i) using different techniques for estimating the local distributions, ii) performing smarter client selection, iii) addressing statistical heterogeneity in FL, iv) using pretrained models in FL.\n\n**Minor weaknesses** (did not affect my rating)\n1. Missing hyperparameters ranges used for tuning the baselines and some details on the federated settings. Please refer to Questions.\n\n**References**\n\n[1] Cho, Yae Jee, Jianyu Wang, and Gauri Joshi. \"Towards understanding biased client selection in federated learning.\" International Conference on Artificial Intelligence and Statistics. PMLR, 2022.\n\n[2] Yang, Wenkai, et al. \"When to Trust Aggregated Gradients: Addressing Negative Client Sampling in Federated Learning.\" TMLR (2023).\n\n[3] Fallah, Alireza, Aryan Mokhtari, and Asuman Ozdaglar. \"Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach.\" Advances in Neural Information Processing Systems 33 (2020): 3557-3568.\n\n[4] Zeng, Shenglai, et al. \"Heterogeneous federated learning via grouped sequential-to-parallel training.\" International Conference on Database Systems for Advanced Applications. Cham: Springer International Publishing, 2022.\n\n[5] Zaccone, Riccardo, et al. \"Speeding up heterogeneous federated learning with sequentially trained superclients.\" 2022 26th International Conference on Pattern Recognition (ICPR). IEEE, 2022.\n\n[6] Lu, You-Ru, Xiaoqian Wang, and Dengfeng Sun. \"Tackling Imbalanced Class in Federated Learning via Class Distribution Estimation.\" (2022).\n\n[7] Chen, Dawei, et al. \"Digital twin for federated analytics using a Bayesian approach.\" IEEE Internet of Things Journal 8.22 (2021): 16301-16312.\n\n[8] Nguyen, John, et al. \"Where to begin? Exploring the impact of pre-training and initialization in federated learning.\" NeurIPSW (2022).\n\n[9] Chen, Hong-You, et al. \"On the importance and applicability of pre-training for federated learning.\" The Eleventh International Conference on Learning Representations. 2022.\n\n[10]  Li, Tian, et al. \"Federated optimization in heterogeneous networks.\" Proceedings of Machine learning and systems 2 (2020): 429-450.\n\n[11] Karimireddy, Sai Praneeth, et al. \"Scaffold: Stochastic controlled averaging for federated learning.\" International conference on machine learning. PMLR, 2020.\n\n[12] Durmus Alp Emre Acar, Yue Zhao, Ramon Matas Navarro, Matthew Mattina, Paul N Whatmough, and Venkatesh Saligrama. Federated learning based on dynamic regularization. International Conference on Learning Representations, 2021. \n\n[13] Reddi, Sashank, et al. \"Adaptive federated optimization.\" ICLR (2021).\n\n[14] Li, Qinbin, Bingsheng He, and Dawn Song. \"Model-contrastive federated learning.\" Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2021.\n\n[15] Karimireddy, Sai Praneeth, et al. \"Mime: Mimicking centralized stochastic algorithms in federated learning.\" Advances in Neural Information Processing Systems, 2021 .\n\n[16] Qu, Zhe, et al. \"Generalized federated learning via sharpness aware minimization.\" International Conference on Machine Learning. PMLR, 2022.\n\n[17] Caldarola, Debora, Barbara Caputo, and Marco Ciccone. \"Improving generalization in federated learning by seeking flat minima.\" European Conference on Computer Vision. Cham: Springer Nature Switzerland, 2022.\n\n[18] Wang, Jianyu, et al. \"Slowmo: Improving communication-efficient distributed sgd with slow momentum.\" ICLR (2020).\n\n[19] Kairouz, Peter, et al. \"Advances and open problems in federated learning.\" Foundations and Trends® in Machine Learning 14.1–2 (2021): 1-210.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments and suggestions to improve our paper! Please see the following clarifications.\n\n**Reference**\n\nWe thank the reviewer for a great summary of the literature. We have updated the paper by including the mentioned reference in the Introduction and Related Work section.\n\n**Client selection strategy**\n\nA major difference between our approach with other smart client selection strategies is that we do not need the loss/gradient of the candidate clients. In other words, we avoid the transmission of the global model to non-selected clients and loss/gradient computation on them.\n\n\n**Probability of selecting distant client**\n\nConsider the following example based on Figure 1. If we set the sketch of A_2 to be the global sketch. Then the distance of A_0, A_1, A_3, and A_4 to the global sketch would be [0.16,0.11,0.14,0.18]. Following Algorithm 3, we mark their selection probability as [0.04745458, 0.8129602 , 0.11588869, 0.02369653]. The distant client, A_3, will be sampled with probability 0.02369653. \n\n\n**All distant clients**\n\nAs shown in Algorithm 3, if all clients are far away from the global sketch, then $p_i$ remains similar across clients. In this case, after the normalization step, the client selection process is equivalent to uniform sampling.\n\n\n**More Experiments**\n\nWe conducted CIFAR-10 experiments using 100 clients, where L=3K and K=2. The client dataset follows a non-iid distribution, as specified in the split provided at https://github.com/yjlee22/FedShare. We visualized the iteration-wise convergence of FedAvg with and without the proposed client selection (please refer to the rebuttal PDF). The results demonstrate that our distribution sketch improves the test accuracy convergence rate, leading to faster iterations.\n\n**Rounds in experiments**\n\nWe introduce early stopping to prevent overfitting. Feel free to refer to line 321 of the paper for the details. \n\n**Choice of Parameters**\n\nWe suggest choosing $B$ and $R$ as follows: $R$ should be chosen proportional to $O(\\sqrt{n})$, where $n$ is the maximum possible client data size. We set $B$ to be $O(\\log(n))$ following theoretical guarantees in [1].  As a result, our sketch size should be significantly smaller than the dataset size.\n\n\n**Privacy leakage from the distribution sketch if no Laplace noise is added**\n\nYes. To understand why, consider a client with an extreme outlier in its dataset. This outlier will likely be the only point in one of the RACE partitions, meaning that we could potentially localize it from the sketch alone, given access to the sketch's randomness. There are several ways to deal with this (e.g. using the distributed private sketch framework of [2]), but the addition of Laplace noise is the simplest method. \n\n\n\n**Experiments settings**\n\nFor FEMNIST and shakespeare, we follow the exact settings in the LEAF https://leaf.cmu.edu/ project.  For MNIST, there are 30 clients in total. For the client selection, we set our parameters following Section 5.3 (line 315 to 320).\n\nWe thank the reviewer for the comments on the writing. We have updated the paper accordingly for a better presentation.\n\n[1] Locality-Sensitive Hashing Scheme Based on p-Stable Distributions\n\n[2] Differentially-Private Multi-Party Sketching for Large-Scale Statistics"}, {"review_id": "BVKuh0KNBH", "reviewer": "Reviewer_c2Su", "summary": "The paper proposes to use a locally sensitive hashing based sketch to measure the similarity between two datasets in a privacy-preserving manner. This sketch can be computed efficiently and differential privacy can be guaranteed via the Laplace mechanism. The paper concludes with presenting two potential applications of their similarity measure: selecting participating clients in federated learning based on the similarity of local datasets, and selecting an initial model for a new client in personalized federated learning.", "questions": "- Would the definition 3.1 of TVD not be more appropriately placed at the beginning of Sec 3.2.2?\n- Would it be possible to quantify the difference of the upper bound (using the Euklidean distance between sketches) to the true TVD, based on properties of the sketches? E.g., assuming sufficiently smooth data distributions.\n- Could you add synthetic experiments where data is drawn from known distributions, so that the accuracy of the TVD estimation can be more clearly assessed?", "limitations": "The paper is clear about the assumptions and limitations.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The idea is clear and sound, a privacy-preserving dataset similartiy measure is useful, in particular in federated and collaborative learning scenarios.\n- Theoretically sound definition of the sketch from which an estimator of the total variation distance can be derived.\n- The paper presents two applications of the similarity measure for federated learning.", "weaknesses": "- The distribution-aware client sampling is interesting, but not well-motivated. Sampling clients with differing local data from the average distribution can be helpful in some settings, while being detrimental to the learning process in others. While the empricial results look promising, the analysis is far too shallow to permit conclusions. \n--- after rebuttal ---\nThe authors have addressed this weakness in the rebuttal.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nWe thank the reviewer for the comments and suggestions to improve our paper! Please see the following clarifications.\n\n**Q: Would the definition 3.1 of TVD not be more appropriately placed at the beginning of Sec 3.2.2?**\n\nWe introduce the TVD early to motivate the theory and provide a high-level sketch of the ideas. However, we agree that equation (2) would be better placed in Section 3.2.2 and plan to clean this up in the revision.\n\n**Q: Would it be possible to quantify the difference of the upper bound (using the Euklidean distance between sketches) to the true TVD, based on properties of the sketches? E.g., assuming sufficiently smooth data distributions.**\n\nThis is difficult to do rigorously, even if we assume a sufficiently smooth data distribution. We can think of two ways to prove rigorous estimation bounds for the TVD, but both have technical difficulties. The first way is to analyze the sketch as a high-dimensional histogram and formalize the arguments of Section 3.2.2. However, we cannot calculate the error of the piecewise integral approximation without upper and lower bounds on $\\lambda(A)$ (which are not generally known for LSH partitions). The second way is to analyze the sketch as a density estimator, and then apply a known technique for high-dimensional integration (such as Monte Carlo or sparse-grids). We briefly describe this option in the first paragraph of Section 3.2.2, but unfortunately we are not aware of a high-dimensional integration method that comes with the computational performance we require in this application.\n\n**Q: Could you add synthetic experiments where data is drawn from known distributions, so that the accuracy of the TVD estimation can be more clearly assessed?**\n\nThis is an excellent suggestion. We computed the TVD for two Gaussians with the same covariance and different means, one of the few TVDs between high-dimensional distributions for which we have an exact formula. This formula is $\\mathrm{TVD} = 2\\Phi(z) - 1$ where $z = \\sqrt{\\mu_1 - \\mu_2)^{\\top}C^{-1}(\\mu_1 - \\mu_2)}$ (see Theorem 1 of [1] and page 5 of [2]). To get the value of $\\lambda(A)$ from line 178, we integrated the high-dimensional collision probability over the d-ball. We drew 1000 samples from two 16-dimensional Gaussians with random means, computed their sketch and TVD distances, and compared the results. We plot the results in the attached PDF. We will add this discussion to the paper.\n\n[1]: \"Estimates of the proximity of Gaussian measures,\" S. S. Barsov and V. V. Ulyanov, Soviet Mathematics, Doklady, 34:462–466, 1987.\n\n[2]: “The total variation distance between high-dimensional Gaussians with the same mean,” Devroye et. al., arXiv 2022.\n\n\t\n"}]

Author rebuttal:
We thank reviewers [R1(Jfur), R2(QJNz), R3(iXQ3), R4(c2Su)] for their thoughtful and highly supportive feedback! We were glad that the reviewers found the problem significant and interesting [R2 R3], the observations and theoretical analysis insightful and valuable [R3, R4], the methods novel and clear [R2, R3, R4], and the experimental results convincing and promising [R2 R3].

We have updated the paper to incorporate constructive suggestions. We summarize the attached additional experiments:

[ R1,  R2, R3] Addition CIFAR-10 experiments on client sampling.

[R4] Simulation on TVD estimation.

Meta-review: This paper an interesting method to measure heterogeneity across client data in federated learning. The reviews are generally positive and agree that this will add value to the FL literature.
Decision: Accept (poster)