SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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 is efficient. We only need one-shot, one-pass computation to generate a one-pass distribution sketch of client data. Moreover, the sketch is mergeable so a linear combination will give us a sketch for global data distribution. Furthermore, we can efficiently estimate the divergence of the client data distribution to the global distribution by taking the Euclidean 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 a global model through iterative federated optimization [12]. Next, each client builds a personalized model based on the global model that achieves decent predictive performance on its local data (See Section 6). Next, a new client connects to the server. The new client has local but unlabeled data. In this setting, it
remains unknown which model, either global or local, should be applied to the new client as a cold starter.
We propose Algorithm 2 for this cold start problem. For the clients that already connect to the server, we compute its sketch with Algorithm 1 and send it to the server. Next, given a new client, we perform the sketching via Algorithm 1 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 it to 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.
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 2 on both vision and language datasets. For visual classification, we use the MNIST dataset [51] and FEMNIST [52]. For MNIST, we introduce a federated dataset where each client is generated by random sampling from MNIST under Dirichlet distribution with  $\alpha = 0.5$ . We also introduce MNIST Uniform + Dirichlet, a dataset with half of the client data generated from random sampling from MNIST under Dirichlet distribution with  $\alpha = 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-based FL. The clients in both FEMNIST and Shakespeare datasets are fixed.
Model: For the MNIST dataset, we use a fully-connected network with a single 512-dimension hidden layer. For the FEMNIST dataset, we use 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 an 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, is implemented on PyTorch [55]. We test Algorithm 3 and Algorithm 2 on a server with 8 Nvidia Tesla V100 GPU and a 48-core/96-thread processor (Intel(R) Xeon(R) Gold 5220R CPU @ 2.20GHz).
5.2 One-Pass Sketch Distance
In this section, we investigate the correlation between one-pass sketch distance and distribution differences.
Setting: We generate two groups of subsets from MNIST, Group A and Group B. Datasets from Group A are drawn randomly from MNIST under Dirichlet distribution with  $\alpha = 0.5$ . Datasets from Group B are drawn randomly from MNIST under the uniform distribution. For every dataset, we construct a one-pass sketch. Then, we calculate the pairwise one-pass sketch distances between every dataset.
Result: We present the result in Figure 1. 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,x)$ . 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 A are larger and more diverse, which is 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 is similar. Our results answer the first question: the distance between the distribution sketches constructed by Algorithm 1 can reflect their divergence in the distribution.
Figure 1: Pairwise one-pass sketch distance visualization.  $A_0, A_1, \ldots, A_n$  are five datasets drawn randomly under the Dirichlet distribution.  $B_0, B_1, \ldots, B_n$  are five datasets drawn randomly under the uniform distribution. We mark each cell with the one-pass sketch distance shown in Algorithm 1.
5.3 Distribution-Aware Client Sampling
In this section, we apply the one-pass distribution sketch (see Algorithm 1) 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
(b) MNIST Uniform + Dirichlet
(c) FEMNIST
Figure 2: Result: Convergence plot for MNIST, MNIST Uniform+Dirichlet, FEMNIST, and Shakespeare.
(d) Shakespeare
optimization, there would be  $L$  activated clients. The standard federated optimization approach would uniformly randomly sample  $K$  clients as participants [6]. We denote this approach as 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 3 with Fedavg to showcase the importance of data heterogeneity in client selection. We set  $L = 3K$ . For the MNIST dataset (both MNIST and MNIST Uniform + Direchlet), both Algorithm 3 and Fedavg are trained by 200 rounds. In each round,  $K = 3$  clients are selected from  $L$  active clients. Next, each client is trained for 20 epochs with batch size 32 and learning rate  $\eta = 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 3 and Fedavg are trained by 20 rounds. Each round,  $K = 3$  clients are selected from  $L$  active clients. Next, each client is trained for 20 epochs with batch size 10 and learning rate  $\eta = 0.01$ . We introduce early stopping to prevent overfitting.
Result: We present the test accuracy plot in Figure 2. As shown in the figure, on four FL datasets/settings, Algorithm 3 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 can be used to sample helpful clients towards faster convergence.
5.4 Cold Start in Personalized FL
In this section, we would like to evaluate the performance of Algorithm 2 in the cold start problem of personalized FL (see Section 4.2).
Setting: For every dataset, we uniformly randomly 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 a global model on the server, and each client has its personalized model. For cold clients, we only use its label for evaluation. We use MNIST, FEMNIST, and Shakespeare datasets as introduced in Section 5.1. We generate MNIST clients following the by-class format shown in [13]. There are 100 clients in total. We set 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 uniformly selects 10 clients. Each client is trained for 10 epochs on its local data with batch size
Figure 3: Test accuracy of 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.
32 and learning rate 0.0001. Next, we finetune each client for 10 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, we use 20 clients as cold start clients, and the rest of them are warm clients. To get the global model, we train 5 rounds, and each round uniformly selects 3 clients. Each client is trained for 20 epochs on its local data with batch size 10 and learning rate 0.01. Next, we finetune each client for 20 epochs for a local model. We set the batch size and learning rate as the same as global model training.
Result: We present the result in 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 on each cold client with the global model, the model retrieved by uniformly random sampling, and the model retrieved by Algorithm 2. We also run a brutal force search on the existing personalized models and find the best model with the highest test accuracy. Note that this search is impractical since we use client data label information. But it is an upper bound for Algorithm 2. As shown in Figure 3, we observe that Algorithm 2 outperforms both the global model and random sampling with higher test accuracy on clients. Moreover, Algorithm 2 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.
Client Selection in FL. Iterative federated optimization is the de-facto standard for FL. In this optimization strategy, a server aggregates the model updates from each device to create a better
global model. Client selection [56, 57, 58, 59] serves as an essential component of FL in the common situation where many clients are connected to the server. In each round of iterative federated optimization, the server identifies a subset of clients as participants [60], asks participants to train on client data, and finally aggregates the updates. Client selection is 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 [13, 62, 63]. Second, there are systems challenges. The communication required to coordinate a large number of clients imposes a heavy burden on network bandwidth [64, 65], reducing the efficiency of the FL system. Recent work on client selection strategies in FL focuses on the model's performance on the client data. For instance, [66] and [67] use the training loss of the client model on its data as a 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.
Personalized FL. Personalization is a natural way of improving local predictive performance on a client [70, 71, 27, 72]. Given the global model trained across diverse clients, it is a common practice to fine-tune the global model on every client's data to obtain a better local model [14, 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, a 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 on retrieving a local model from a client with a similar data distribution to the new client. We aim to achieve this strategy with an effective data heterogeneity measurement.
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 a one-pass distribution sketch that represents the data distribution while preserving  $\epsilon$ -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.
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 are supported by NSF-CCS-2211815, ONR-DURIP and NSF-BIGDATA-1838177. Zhaozhuo Xu is supported by startup funding from the Stevens Institute of Technology.
[1] Andrew Hard, Kanishka Rao, Rajiv Mathews, Swaroop Ramaswamy, Françoise Beaufays, Sean Augenstein, Hubert Eichner, Chloe Kiddon, and Daniel Ramage. Federated learning for mobile keyboard prediction. arXiv preprint arXiv:1811.03604, 2018.
[2] Bill Yuchen Lin, Chaoyang He, Zihang Zeng, Hulin Wang, Yufen Huang, Mahdi Soltanolkotabi, Xiang Ren, and Salman Avestimehr. Fednlp: A research platform for federated learning in natural language processing. arXiv preprint arXiv:2104.08815, 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.
A Design of LSH Functions:
In practice, we could realize  $d(x,y)$  with cosine similarity for dense vectors. In this case, the LSH function family  $\mathcal{H}$  should have the following form.
Definition A.1 (Sign Random Projection (SRP) Hash [74]). We define a function family  $\mathcal{H}$  as follows: Given a vectors  $x\in \mathbb{R}^d$  any  $h:\mathbb{R}^d\to \{0,1\}$  that is in the family  $\mathcal{H}$ , we have
$$
h (x) = \operatorname {s i g n} (A x),
$$
where  $A \in \mathbb{R}^{d \times 1}$  is a random matrix that every entry of  $A$  is sampled from normal distribution  $\mathcal{N}(0,1)$ ,  $\text{sign}$  is the sign function that set positive value to 1 and others to 0. Moreover, for two vectors  $x, y \in \mathbb{R}^d$  we have
$$
\Pr [ h (x) = h (y) ] = 1 - \pi^ {- 1} \operatorname {a r c c o s} \frac {x ^ {\top} y}{\| x \| _ {2} \| y \| _ {2}}.
$$
As shown in Definition A.1, the SRP hash is 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. If both  $x \in \{0,1\}^d$  and  $y \in \{0,1\}^d$  are high dimensional binary vectors with large  $d$ . Their Jaccard similarity [78] is also an important measure for search [79, 80, 81] and learning tasks [82, 83]. There also exists a family of LSH functions for Jacacd similarity. We define this LSH function as:
Definition A.2 (MinHash [84]). A function family  $\mathcal{H}$  is a MinHash family if for any  $h\in \mathcal{H}$ , given a vectors  $x\in \{0,1\}^d$ , we have
$$
h (x) = \arg \min  (\Pi (x))
$$
where  $\Pi$  is a permutation function on the binary vector  $x$ . The arg min operation takes the index of the first non-zero value in  $\Pi(x)$ . Moreover, given two binary vectors  $x, y \in \{0,1\}^d$ , we have
$$
\operatorname * {P r} [ h (x) = h (y) ] = \frac {\sum_ {i = 1} ^ {d} \min (x _ {i} , y _ {i})}{\sum_ {i = 1} ^ {d} \max (x _ {i} , y _ {i})},
$$
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 LSH (see Definition 3.2). This view aligns with a series of research in efficient kernel decomposition [88, 89, 90], kernel density estimation [46] and kernel learning [91].
B More Algorithms
In this section, we introduce the client selection algorithm with our one-pass sketch.
Algorithm 3 Client Selection with Distribution Sketch
Input: Clients  $\mathcal{C} = \{c_1,\dots ,c_n\}$  , Number of rounds  $T$  , Step size  $\eta$  , LSH function family  $\mathcal{H}$  , Hash range  $B$  , Rows  $R$  , Number of active clients  $L$  , Number of Selected Clients  $K$  , Epochs  $E$  , Random Seed s   
Output: Global model  $w^{T}$    
Initialize: Global model  $w^{1}$  , global sketch  $S_{g}$  , random seed s.   
for  $c\in \mathcal{C}$  do Get client data  $\mathcal{D}_c$  from client c. Compute sketch  $S_{c}$  using Algorithm 1 with parameter  $\mathcal{D}_c$  ,  $\mathcal{H}$ $B,R$  and s. Send  $S_{c}$  to server  $S_{g}\leftarrow S_{g} + S_{c}$    
end for   
 $S_{g}\gets S_{g} / n$  Generate global sketch on server   
for  $i\in [T]$  do L clients  $\{c_1,\dots ,c_L\} \subset \mathcal{C}$  are activated at random. for  $j\in [L]$  do  $p_j = 1 / \| S_g - S_j\| _2$ $S_{j}$  is the sketch for client  $c_{j}$  end for   
for  $j\in [L]$  do  $p_j = \frac{\exp p_j}{\sum_{l = 1}^L\exp p_l}$    
end for Sample  $K$  clients out of  $\{c_1,\dots ,c_L\}$  without replacement, client  $c_{j}$  is selected with probability   
 $p_j$  Server send  $w^{i}$  to the selected clients. Each selected client  $c_k$  updates  $w^{i}$  to  $w_{k}^{i}$  by training on its data for  $E$  epochs with step size  $\eta$  Each selected client sends  $w_{k}^{i}$  back to the server.  $w^{i + 1} = \sum_{k = 1}^{K}w_{k}^{i}$    
end for   
return  $w^{T}$
In this section, we provide two-fold of more experiments. We start with a simulation of our sketch (see Algorithm 1) for total variation distance (TVD) estimation. Next, we perform client selection (see Algorithm 3) on 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 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 [34, 92] for reference). To get the value of  $\lambda(A)$  (see Section 3.2.2), 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.
C.2 CIFAR-10 Experiments
We conducted CIFAR-10 experiments using 100 clients, where  $\mathrm{L} = 3\mathrm{K}$  and  $\mathrm{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 Figure 5). The results demonstrate that our distribution sketch improves the test accuracy convergence rate, leading to faster iterations.
Figure 4: Visualization of TVD Estimation via Sketch.
Figure 5: CIFAR-10 Experiments: We visualized the iteration-wise convergence of FedAvg with and without the proposed client selection.
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  $q\in \mathbb{R}^d$  , LSH functions  $h_1,\dots,h_R$  , Dataset  $\mathcal{D}$  , Sketch  $S\in \mathbb{R}^{R\times B}$  built on  $\mathcal{D}$  with Algorithm 1   
Output:  $S(q)\in R$    
Initialize:  $S(q)\gets 0$ $A\gets \emptyset$    
for  $i = 1\rightarrow R$  do   
 $A\gets A\cup \{S_{i,h_i(q)}\}$    
end for   
 $S(q)\leftarrow$  median(A)   
return  $S(q)$
Next, we introduce the formal statements in the paper as below.
Theorem D.1 (Formal version of Theorem 3.3). 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 (x _ {i}, q)
$$
with convergence rate  $O(\sqrt{\log R} / \sqrt{R})$ .
Proof. Let  $\kappa(x) = \sum_{x_i \in \mathcal{D}} \sqrt{k(x, x_i)}$  be the (non-normalized) kernel density estimate. Theorem 3.4 of [48] provides the following inequality for any  $\delta$ , where  $\widetilde{\kappa}(x) = \sum_{x_i \in \mathcal{D}} \sqrt{k(x, x_i)}$ :
$$
\Pr \left[ | N S (x) - \kappa (x) | > \left(3 2 \frac {\widetilde {\kappa} ^ {2} (x)}{R} \log 1 / \delta\right) ^ {1 / 2} \right] <   \delta
$$
This is equivalent to the following inequality, which we can obtain by dividing both sides of the inequality inside the probability by  $N$ .
$$
\Pr \left[ \left| S (x) - \frac {1}{N} \kappa (x) \right| > \left(3 2 \frac {\widetilde {\kappa} ^ {2} (x)}{N ^ {2} R} \log 1 / \delta\right) ^ {1 / 2} \right] <   \delta
$$
We want to show that the error  $\left|S(x) - \frac{1}{N}\kappa (x)\right|$  converges in probability to zero, because this directly proves the main claim of the theorem. To do this, we must show that for any  $\Delta >0$
$$
\lim  _ {R \rightarrow \infty} \Pr \left[\left| S (x) - \frac {1}{N} \kappa (x) \right| > \Delta \right] = 0
$$
This can be done by setting  $\delta = \frac{1}{R}$ , which yields the following inequality:
$$
\Pr \left[ \left| S (x) - \frac {1}{N} \kappa (x) \right| > \left(3 2 \frac {\widetilde {\kappa} ^ {2} (x)}{N ^ {2} R} \log R\right) ^ {\frac {1}{2}} \right] <   \frac {1}{R}
$$
Because  $\widetilde{\kappa} < N$ , the following (simplex, but somewhat looser) inequality also holds:
$$
\Pr \left[ \left| S (x) - \frac {1}{N} \kappa (x) \right| > \left(3 2 \frac {\log R}{R}\right) ^ {\frac {1}{2}} \right] <   \frac {1}{R}
$$
This implies that  $S(x) \underset{\mathrm{i.p.}}{\rightarrow} \frac{1}{N} \kappa(x)$  by considering  $R$  large enough that  $\sqrt{\log R / R} < \Delta$ .