SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zichang Liu*
zichangliu@rice.edu
Zhaozhuo Xu*
zxu79@stevens.edu
Rice University†
Rice University & ThirdAI Corp.
anshumali@rice.edu
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.
Federated learning (FL) has recently become the go-to solution for challenging on-device applications such as typing prediction [1], machine translation [2, 3] and fraud detection [4, 5]. Initially proposed by McMahan et. al. [6], the FL setting defines a learning problem where multiple clients collaborate to construct a global model from local client data. Unlike in 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 [7, 8, 9], and it has correspondingly seen a 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 distributions [6, 13, 14, 12, 15, 16, 17, 18, 19, 20, 21, 22]. For instance, in the typing prediction task, each client is a cellphone user. Differences in user interest, language, and geography will produce extremely diverse word distributions. Data heterogeneity is a severe problem for FL because the models trained on each client may have many different behaviors. It is challenging to aggregate such different client models into a coherent and performant global model.
*Equal contribution. The order of authors is determined by flipping a coin.
†Now with Google DeepMind.
Code is available at https://github.com/lzepsilon/RACE_Distance
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 quality of FL [23, 24, 25]. 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 a 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, 28]. 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 a new client arrives with unlabeled data. We can search for a 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.
Challenges in Measuring Data Heterogeneity in FL. 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 in high-dimensional space. Existing approaches, including earth mover distance [29, 30] or Fréchet distance [31, 32] are computationally 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 a significant efficiency issue in practical deployment. Further, to keep with the privacy advantage of FL, data heterogeneity measurements must be 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 sketch-based sampling strategy for federated optimization. Given sketches for each FL client, we find a global distribution sketch by merging all client sketches. We propose a form of client importance sampling where, in each training round, we preferentially select clients that have a 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 an efficient cold start approach for personalized FL where a new client (with unlabeled data) selects the most appropriate model from the personalized finetuned models of other clients. To do this, we compare the distribution sketch of the new client with each existing client. We perform a nearest neighbor search over clients under the sketch distance to identify the client that has similar data distribution. Our cold-start approach provides a simple yet effective way to find a suitable model for new clients.
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 a hash value, usually binary code or integer. If two vectors are close to each other, with a high probability they will share the same hash value. Formally, we define the LSH function as:
Definition 2.1 (( $D, cD, p_1, p_2$ -sensitive hash family). We define a function family  $\mathcal{H}$  to be  $(D, cD, p_1, p_2)$ -sensitive with respect to distance function  $d: \mathbb{R}^d \times \mathbb{R}^d \to \mathbb{R}$  if for any two vector  $x, y \in \mathbb{R}^d$ , any  $h \in \mathcal{H}$  chosen uniformly at random satisfies:
- If  $d(x, y) \leq D$  then  $Pr[h(x) = h(y)] \geq p_1$ .
- If  $d(x, y) \geq cD$  then  $Pr[h(x) = h(y)] \leq p_2$ .
Here  $D \in \mathbb{R}$  and  $c > 1$ . We denote  $Pr[h(x) = h(y)]$  as the collision probability of vector  $x$  and  $y$ . Here the term "collision" means  $x$  and  $y$  have the same hash value. In practice, we will use LSH functions where the collision probability between  $x$  and  $y$  is a monotonically increasing function with respect to a distance measure  $d: \mathbb{R}^d \times \mathbb{R}^d \to \mathbb{R}$ . Formally, we denote this relationship as:
$$
P r [ h (x) = h (y) ] \propto f (d (x, y)). \tag {1}
$$
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 [41]. RACE builds a two-dimensional matrix  $A \in \mathbb{R}^{R \times B}$  using  $R$  LSH functions (see Definition 2.1), each LSH function  $h_i$  maps a vector in  $\mathbb{R}^d$  to an integer in  $[B]$ . Given a dataset  $\mathcal{D} \subset \mathbb{R}^d$ , for every  $x \in \mathcal{D}$ , RACE maps it to  $R$  different hash values, denoted as  $\{h_i(x) | i \in [R]\}$ . Then we increment entry  $A_{i,h_i(x)}$  with 1 for every  $i \in [R]$ . Following this manner, RACE maintains a matrix where each column can be regarded as a histogram of the dataset  $\mathcal{D}$ . As a 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 a data distribution. Given a dataset  $\mathcal{D} \subset \mathbb{R}^d$ , if we build a RACE matrix  $S$  by inserting every element in  $\mathcal{D}$ , we build a histogram of data samples with hashing-based projection. Next, if we normalize the counts in the RACE matrix  $S$  with the dataset size  $|\mathcal{D}|$ , we will get a representation of  $\mathcal{D}$ .
Following this intuition, we introduce our distribution sketch algorithm in Algorithm 1. Our distribution sketch  $S$  is a two-dimensional matrix with  $R$  rows and  $B$  columns. Given a dataset  $\mathcal{D}$ , we first initialize a zero matrix  $S\in \mathbb{R}^{R\times B}$  and  $R$  LSH functions  $\{h_1,\dots ,h_R\}$ . Here, each hash function  $h_i$  generates an integer hash value within  $[B]$ . We suggest  $R$  to be  $O(\sqrt{|\mathcal{D}|})$ . The choice of  $B$  depends on the LSH functions we use. Next, for every  $x\in \mathcal{D}$ , we compute  $R$  hash values  $\{h_1(x),\dots ,h_R(x)\}$ . Finally, we increment  $S_{i,h_i(x)}$  for every  $i\in [R]$ . In practice, for vision tasks with dense feature vectors, we use SRP hash (See Definition A.1) as our LSH function. For language modeling tasks, we tokenize every sentence in the dataset into a bag of
Algorithm 1 One-Pass Distribution Sketch
Input: Dataset  $\mathcal{D} \subset \mathbb{R}^d$ , LSH function family
$\mathcal{H}$  , Hash range  $B$  , Rows  $R$  , Random Seed  $s$
Output: Sketch  $S\in \mathbb{R}^{R\times B}$
Initialize:  $S\gets 0^{R\times B}$
Generated  $R$  independent LSH functions  $h_1, \ldots, h_R$  from  $\mathcal{H}$  with range  $B$  and random seed  $s$ .
for  $x\in \mathcal{D}$  do
$$
\begin{array}{l} \mathbf {f o r} i = 1 \rightarrow R \mathbf {d o} \\ S _ {i, h _ {i} (x)} + = 1 \\ \end{array}
$$
$$
S = S / | \mathcal {D} |
$$
return  $S$
words or characters [42]. Next, a sentence can be formed as a binary vector with a nonzero index representing a token. In this case, we can use MinHash (see Definition A.2) as an LSH function.
Complexity Analysis: Given a dataset with  $n$  vectors, the running time of Algorithm 1 is  $O(nRd + RB)$ . Since the number of LSH functions  $R$  and hash range  $B$  are usually smaller than the dataset size  $n$ , Algorithm 1 is a linear passing over the dataset. Moreover, the space complexity of Algorithm 1 is  $O(RB)$ , which is a much smaller than the memory cost of storing the dataset.
In this section, we describe how to measure the data heterogeneity using the one-pass distribution sketch in Algorithm 1. 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  $P$  and  $Q$  on a space  $(\Omega, \mathcal{F})$ ,
$$
\operatorname {T V D} (P, Q) = \sup  _ {A \in \mathcal {F}} | P (A) - Q (A) | \tag {2}
$$
If  $P$  and  $Q$  have density functions  $p$  and  $q$ , this expression is equivalent to the following  $\ell_1$  norm.
$$
\operatorname {T V D} (p, q) = \frac {1}{2} \int_ {\Omega} | p - q | d v \tag {3}
$$
Our goal is to describe a method for consistent estimation of the TVD that is efficient to implement in the distributed and private FL setting. At a high level, our strategy is as follows.
1. Given a dataset  $D$  drawn from a distribution with probability density function  $p(x)$ , construct a sketch-based consistent estimator  $S_{D}(x)$  for the kernel density estimate (KDE) over  $D$ .
2. Argue that  $S_{D}$  is a consistent estimator of  $p$ , because the KDE is a consistent estimator of  $p$  and  $S_{D}$  is a consistent estimator of the KDE.
3. Use  $S_{D}$  as a plug-in estimator for  $p$  in the TVD integral (Equation 3).
Assumptions: We consider probability measures on  $\mathbb{R}^d$  that have a distribution function. We assume that the probability densities  $p$  and  $q$  are continuous almost everywhere, so that the  $\mathrm{KDE} \to p$  in step 2 [43]. We also assume that the support of  $P$  and  $Q$  is a bounded subset of  $\mathbb{R}^d$ , 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 sure uniform convergence.
Notation: In a slight abuse of notation, we will use  $O(\cdot)$  to denote both stochastic bounds and the more typical (deterministic) asymptotic bounds. More concretely,  $X_{N} = O(f(N))$  means that for any  $\epsilon >0$ , we can find constants  $c$  and  $N_0$  such that for all  $N > N_0$ ,  $\operatorname *{Pr}[|X|\geq cf(N)]\leq \epsilon$ . In the deterministic case, this still holds with  $\epsilon = 0$ . We will assume two probability measures  $P,Q$  having densities  $p,q$  and two datasets  $\mathcal{D}_p\sim p$  and  $\mathcal{D}_q\sim q$ . We denote the sketch of  $\mathcal{D}_p$  with  $S_{p}$  and  $\mathcal{D}_q$  with  $S_{q}$ . We write the distribution sketch  $S$  generated by Algorithm 1 as a functional  $S(x)$ . We explicitly define this functional in Algorithm 4, Section D.
3.2.1 Consistent Estimation with Sketches
This section shows that our sketches provide a consistent estimator of the data-generating distribution. We introduce the distribution sketch's query algorithm in Algorithm 4 (See Section D in the appendix). For a sketch constructed via Algorithm 1, given a query  $q \in \mathbb{R}^d$ , we use the same LSH functions in construction sketch and get  $R$  hash values. The  $i$ th hash value corresponds to a location in the  $i$ th row. Next, we take the median over the average counts in the hashed locations. We show that the median-of-means count  $S(q)$  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 a kernel  $k(x,y)$  to be an LSH kernel if for every  $x,y$ ,  $k(x,y)$  represents the LSH collision probability (See Eq. (1)) of  $x$  and  $y$ , which is monotonically decreasing with distance  $d(x,y) \geq 0$ . An LSH kernel  $k(x,y)$  is a positive semi-definite kernel.
Next, we have the following statement.
Theorem 3.3 (Informal version of Theorem D.1). Let  $P(x)$  denote a probability density function. Let  $\mathcal{D} \sim P(x)$  denote a dataset. Let  $k(x, y)$  be an LSH kernel (see Definition 3.2). Let  $S$  define the function implemented by Algorithm 4. We show that
$$
S (x) \underset {\text {i . p .}} {\rightarrow} \frac {1}{N} \sum_ {x _ {i} \in \mathcal {D}} k \left(x _ {i}, q\right)
$$
with convergence rate  $O(\sqrt{\log R / R})$
Theorem 3.3 shows the consistent estimation of the KDE as  $R \to \infty$ , which is itself a consistent estimator of the distribution under some assumptions about  $k(x,y)$  and  $P$ . The convergence of the KDE to the distribution as  $N \to \infty$  is a 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 to 1, for the kernel bandwidth to decay with  $N$ , and for the bandwidth to be bounded above by  $N^{1 / d}$  [45].
We argue that these conditions hold for the LSH kernels described by [46]. The bandwidth requirements are easy to satisfy because we can 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 kernels, such as the  $p$ -stable LSH probability, which is not absolutely integrable over  $\mathbb{R}^d$ . For this reason, some LSH kernels cannot be normalized to have a unit integral. However, this is not a problem because we assume that  $P$  and  $Q$  have bounded support in  $\mathbb{R}^d$ . Note that this type of input domain truncation is a 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.e.  $R = \omega(1)$ ).
Private Sketches: When the sketch is required to preserve  $\epsilon$ -differential privacy, we must add independent Laplace noise to the sketch so that each counter is private [48]. This constrains our choice of  $R$  because the noise is distributed  $Z \stackrel{\mathrm{iid}}{\sim} \mathrm{Laplace}(R\epsilon^{-1})$ .
Theorem 3.4 (Informal version of Theorem D.2). Let  $S$  be an  $\epsilon$ -differentially private distribution sketch of a dataset  $\mathcal{D} = \{x_1, \ldots, x_N\}$  with  $\mathcal{D} \sim P(x)$  and let  $k(x, y)$  be an LSH kernel (see Definition 3.2). Let  $S$  define the function implemented by Algorithm 4. Then
$$
S (x) \underset {\text {i . p .}} {\rightarrow} \frac {1}{N} \sum_ {x _ {i} \in \mathcal {D}} k (x _ {i}, x)
$$
with convergence rate  $O\left(\sqrt{\log R / R} + \sqrt{R \log R} / (N \epsilon)\right)$  when  $R = \omega(1)$ , e.g.  $R = \log N$ .
In exchange for privacy, we pay a price in the convergence rate, which now has an additional  $\epsilon^{-1}$  term. While this dependence can be improved to  $\epsilon^{-1/2}$  (by optimizing  $R$  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 on  $\mathbb{R}^d$ ), the TVD between KDEs is 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 3. The integral itself can be calculated using standard numerical integration techniques (e.g., Monte Carlo estimation or sparse grids), giving us a consistent sketch-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 justify a simpler method of computing the distance, albeit without formal guarantees.
As a collection of histograms, our sketch is a piecewise-constant approximation to the KDE over a partition  $\pi$  of events. This permits a simpler expression for Equation 3, where we compute the value in each partition, weighed by the partition size. Here, we suppose that the partition  $\pi$  consists of cells  $\{A_1,\ldots \}$ . We let  $x\in A$  be any point inside  $A$  and use  $\lambda (A)$  to denote the Lebesgue measure of  $A$ .
$$
\operatorname {T V D} (p, q) = \frac {1}{2} \lim _ {\max \lambda (A) \to 0} \sum_ {x \in A \in \pi} | S _ {p} (x) - S _ {q} (x) | \lambda (A)
$$
To further simplify this expression, we make the following observations. Each cell of the partition, which we denote as  $A \in \pi$ , 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_{\pi} \lambda(A)$  is not generally known for most LSH functions, it is clear that  $\lambda(A)$  is non-increasing with  $R$  because each cell of the partition  $\pi$  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 LSH functions is a zero-probability event. Also,  $\mathbb{E}[\lambda(A)]$  is the same for all the cells in the partition by symmetry arguments.
Therefore, it is reasonable to consider an approximation where we compute entry-wise differences between sketches. Let  $\Delta_{r,b}$  be the difference between row  $r$  and column (or bucket)  $b$  of the matrices  $S_{p}$  and  $S_{q}$ . We can construct a TVD estimator as a sum over all possible intersections of hash buckets.
$$
\operatorname {T V D} (p, q) \approx B ^ {- R} \sum_ {b _ {1} = 0} ^ {B} \dots \sum_ {b _ {R} = 0} ^ {B} | \Delta_ {1, b _ {1}} + \dots + \Delta_ {R, b _ {R}} |
$$
There are  $B$  buckets in each repetition and, therefore,  $B^{R}$  total "intersection paths" through the sketch. To circumvent the  $O(B^{R})$  complexity of computing each path, we consider the sum of absolute deviations  $|\Delta_{1,b_1}| + \ldots + |\Delta_{R,b_R}|$  instead of  $|\Delta_{1,b_1} + \ldots + \Delta_{R,b_R}|$ . This gives an upper bound on the TVD that can be computed in  $O(BR)$  time because we can now distribute the sums. In particular, we can compute the upper bound by taking the Frobenius norm (Euclidean distance) of  $S_p - S_q$ . We denote  $\| S_p - S_q\|_2$  as sketch distance. Our empirical results show that this heuristic is an excellent indicator of distribution similarity.
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-Aware Client Selection
In our work, we study the setting where there are  $n$  clients checking into the server. In every round of iterative federated optimization,  $L$  clients are activated and ready for training. We assume that  $L$  clients are activated uniformly random out of  $n$  clients. Our client selection algorithm is presented in Algorithm 3 in the supplementary material. For each one of the  $n$  clients, we compute its sketch with Algorithm 1 and send it to the server. As introduced in Theorem 3.4, our sketch can be  $\epsilon$ -differentially private with respect to the client data. As a result, we are able to communicate with the sketch between the server and clients. Next, the server takes the average of all the client sketches towards a global sketch  $S_{g}$ . In each round of iterative federated optimization, the server looks up the sketch for every activated client. Next, we compute the Euclidean distance between  $S_{g}$  and every activated client sketch  $S_{j}$ . Next, we use the reciprocal of the distance as sample probability  $p_{j}$ . Next, we normalize the  $p_{j}$  across the  $L$  clients with Softmax function [50]. Next, we sample  $K$  clients out of the  $L$  activated clients. Each client is sampled with normalized probability  $p_{j}$ . Finally, we perform local training on the  $K$  selected clients and aggregate their model updates.
Algorithm 2 Cold Start with Distribution Sketch
Input: Clients  $\mathcal{C} = \{c_1,\dots ,c_n\}$  , New Client  $C_q$  , LSH function family  $\mathcal{H}$  , Hash range  $B$  Rows  $R$  Random seed  $s$
Output: Personalized model  $w$ . for  $i \in [n]$  do
Get client data  $\mathcal{D}_i$  from client  $c_{i}$
Get client model  $w_{i}$  from client  $c_{i}$ .
Compute sketch  $S_{i}$  using Algorithm 1 with parameter  $\mathcal{D}_i$ ,  $\mathcal{H}$ ,  $B$ ,  $R$  and  $s$ .
Send  $S_{i}$  to server
Get client data  $\mathcal{D}_q$  from client  $c_{q}$
Compute sketch  $S_{q}$  using Algorithm 1 with parameter  $\mathcal{D}_q$ ,  $\mathcal{H}$ ,  $B$ ,  $R$  and  $s$ .
Send  $S_{q}$  to the server.
$$
d _ {\min } = - \infty
$$
$$
i d = n u l l
$$
$$
\mathbf {f o r} i \in [ n ] \mathbf {d o}
$$
$$
d _ {i} = \left\| S _ {q} - S _ {i} \right\| _ {2}
$$
$$
\text {i f} d _ {i} <   d _ {\min } \text {t h e n}
$$
$$
d _ {\mathrm {m i n}} \gets d _ {i}
$$
$$
\operatorname {i d} \leftarrow i
$$
return  $w_{\mathrm{id}}$