FedSampling: A Better Sampling Strategy for Federated Learning

1 Introduction

Federated learning aims to utilize decentralized data to train machine learning models [McMahan et al., 2017], and has become a popular privacy-preserving machine learning paradigm [Dayan et al., 2021; Bai et al., 2021; Jiang et al., 2021; Rothchild et al., 2020]. The mainstream federated learning framework [Reddi et al., 2020; Yang et al., 2018; Sun et al., 2021] usually uniformly samples some clients for local model training in each round. However, sizes of samples on different clients may largely differ [Fraboni et al., 2021b]. The uniform client sampling may prevent the effective exploitation of the clients with more data for model training and lead to a suboptimal performance [Fraboni et al., 2021b]. For instance, consider a scenario with 3 clients (Fig. 1), in which the client c0, c1 and c2 keep 2, 1 and 1 samples respectively, and we need to select a client for training in each round. For uniform client sampling methods, although the sampling probability of each client is uniform and identical to 13 , the weights of samples in different clients for model updating are biased. For the client c0, two samples in it will be used for training at the same time and thus their weights are identical to 12 , while the weights of samples in the client c1 and c2 are identical to 1. Thus, the model training may pay more attention to samples in the client c1 and c2, which leads to a biased data exploitation.

∗The corresponding author.

To address this issue, some recent methods have enabled the server to track the local sample size in each client to perform weighted client sampling [Li et al., 2019; Wang et al., 2020; Fraboni et al., 2021a]. In these methods, the sampling probability of a client is the ratio of its local sample size to the sample size of all clients. However, these methods usually over-emphasize clients with more data, thus fail to satisfy an unbiased exploitation of samples, leading to sub-optimal performance. In the example of Fig. 1, for weighted client sampling methods, the sampling probabilities of data in the client c0 are identical to 12 while the sampling probabilities of data in the client c1 and c2 are identical to 1 4 . Thus, half of the total samples in the client c0 will be sampled for training more frequently than the remaining half of the total samples, which is still biased in data exploitation. Besides, in some certain scenarios (e.g., financial transactions), the local sample size represents the frequency of some privacy-sensitive behaviors. Thus, allowing the server to track local sample sizes may also arouse privacy concerns [Gerber et al., 2018].

In short, existing client-level sampling methods are usually biased in the data exploitation due to the dependent or nonidentical data sampling. In this paper, we propose a uniform data sampling framework (named FedSampling) to improve the effectiveness of federated learning. At each training round of FedSampling, each sample in each client is independently sampled with an identical probability, where the probability is the ratio of the server desired sample size to the total sample size in all available clients. Considering the privacy of local sample size in each client, we propose a privacy-preserving method to estimate the total sample size based on the local differential privacy (LDP) technique. Each client randomly chooses to send the server a true or a randomly-generated local sample size to protect personal privacy, and the server can estimate an unbiased total sample size from the randomized responses. Besides, the analysis on the utility and privacy of FedSampling is also provided. We conduct extensive experiments on four benchmark datasets across different domains. Results show that FedSampling can outperform many recent FL methods, especially under imbalanced data size distribution and non-IID data distribution, and meanwhile achieve an effective trade-off between utility and privacy.

2 RelatedWork

Due to the importance of user privacy, privacy-persevering machine learning techniques have attracted increasing attention [Avdiukhin and Kasiviswanathan, 2021; Achituve et al., 2021; Jumper et al., 2021]. Federated learning can utilize massive private data distributed on local clients to benefit the training of a shared ML model [Hamer et al., 2020; Wang et al., 2021; Li et al., 2020], and thereby has become a popular privacy-aware machine learning framework [Dayan et al., 2021; Yoon et al., 2021; Huang et al., 2021; Bai et al., 2021]. Existing methods usually follow a similar paradigm, where a server uniformly samples some clients and collects their local model parameters to update a global model, iteratively [Liang et al., 2021; Yuan et al., 2021; Fraboni et al., 2021a]. For example, McMahan et al. [2017] proposed to locally train models on clients for multiple steps, and further average local model parameters from sampled clients to update the global model. Reddi et al. [2020] proposed a federated version of the Adam optimization algorithm to smooth the model training, where the current local parameters are combined with historical local parameters to learn the global model. However, the sample size of different clients may be highly imbalanced, and clients with more data cannot be sampled more frequently to contribute more to the model training in these methods, which usually results in a sub-optimal performance.

Some methods explored weighted client sampling to handle this problem [Li et al., 2019; Wang et al., 2020; Fraboni et al., 2021a]. For example, Li et al. [2019] proposed to first sample clients for local model training based on the ratios between their local samples and the total samples in all clients, then average local model updates in the server to update the global model. However, client-level weighting usually overemphasizes the clients with more data and also cannot achieve a uniform data exploitation. Besides, these methods need the server to track the local sample size of each client, which may be impracticable in some scenarios due to privacy concerns. Different from these methods, we propose a data-level uniform sampling strategy, which can achieve a similar data exploitation like centralized learning to improve the performance of federated model training.

3 FedSampling

3.1 ProblemFormulation

In our work, we assume that there are H available clients that can participate in the federated learning, and the set of available clients is denoted as C. Each client c locally keeps a private dataset Dc and never shares it with the outside, where Dc = {si|i = 1, 2, ..., |Dc|} and si denotes the i-th local sample. Moreover, the local sample size |Dc| of each client c is also assumed to be privacy-sensitive and cannot be disclosed. Besides, there is a server that takes charge of maintaining an ML model and scheduling the workflow of the federated model training. The core problem of this work is how to effectively sample decentralized data for model training to improve the effectiveness of federated learning.

3.2 UniformData Sampling

Uniform data sampling has the potential to improve the effectiveness of federated learning since it has a similar data exploitation pattern like centralized learning. Intuitively, by allowing the server to track the local sample size of each client, we can allocate quotas for each client according to the size of its local samples to achieve a uniform data sampling. However, in many scenarios (e.g., medicine and financial transactions) the size of local data is privacy-sensitive [Gerber et al., 2018], and thereby cannot be tracked by the server. For instance, the size of medical records in a client represents the frequency of medical activities, which can be highly privacysensitive for many users. In our work, we propose a unified method that can uniformly and independently sample decentralized data for federated training without the collection of local sample sizes.

Without loss of generalization, we assume the server needs K samples to collaboratively update the model in a training round. Besides, we assume the server can obtain an estimated number ˜ N of total samples in all available clients 1 . In FedSampling, the server first broadcasts K and ˜ N to each client and then each client can locally sample local data for training according to its local sample size and the server desire. Take a client c as example, we first independently draw a random variable zc i from a Bernoulli distribution B( K˜ N ) for each local sample in Dc, where the ratio K˜ N is the probability of assigning 1 to zc i , which means the i-th local sample sci will be sampled for local training. In this way, each sample in each client can be independently sampled for training with identical probabilities. In addition, the expected number of samples for participating in this training round is NE[ K˜ N ], which can asymptotically converge to K and meet the server demand (the convergence is discussed in the following section). After the client c obtains a set Sc of locally sampled data, the client c will employ the current model Θ to calculate model gradients gs on each sample s ∈ Sc, and build the local model update based on normalized local gradients: Gc = 1K ∑s∈Sc gs . Then the local update is uploaded to the server for the global model updating.

1 We will introduce how we obtain ˜ N in the next section.

3.3 Privacy-Preserving Ratio Estimation

Next, we will introduce how to estimate the ratio of the server desired sample number K to the total sample number N in a privacy-preserving way. Intuitively, we may bypass this problem if we sample data based on a fixed ratio r (e.g., 1%) instead of a total sample number-aware ratio KN . However, this will arise new challenges that need to be carefully addressed. Specifically, in this naive method, the local update Gc in the client c cannot be locally normalized by the number of sampled data in this round due to the lack of knowledge of it (i.e., r × N). The client c can only obtain a local model update like mainstream FL frameworks by averaging local gradients: Gc = 1 |Sc| ∑s∈Sc gs . On one hand, to obtain an unbiased global model updating, the server in this naive method needs to employ the number of locally sampled data in each client to weight the aggregation of the corresponding local updates. However, this will cause privacy leakage due to the disclosure of the number of locally sampled data in each client. On the other hand, the naive method can uniformly average the local model updates to avoid the potential privacy leakage. However, this will lead to a biased model update and result in a sub-optimal performance. Thus, the naive method mentioned above is not optimal for uniform data sampling.

Next, we will introduce a carefully designed method to estimate this ratio in a privacy-preserving way. Specifically, the server first queries clients for their local sample sizes. When a client c receives the query, it will first generate a true response nc by clipping the local sample size |Dc| via a size threshold M: nc = min(|Dc|,M - 1). Then the client c will generate a fake response by drawing a random variable ˆnc from a uniform multinomial distribution: ˆnc ∼ P(M), where we randomly select an integer from 1 to M - 1 with identical probabilities and assign it to ˆnc. Furthermore, the client c will draw a random variable xc from a Bernoulli distribution for the response selection:

where α and 1 - α denotes the probability of assigning 1 and 0 to xc respectively, and rc is the selected response. Under the protection of this method, it is difficult for the server to obtain the accurate knowledge on the local sample size of any client. The privacy protection ability of our privacy-preserving ratio estimation method can satisfy the ϵ-LDP, which will be discussed in the following section. Furthermore, after receiving responses from clients, the server further aggregates them to estimate the total sample size N:

where ˜ N is the estimated total sample size and is proved to be unbiased in the following section. Then the server distributes the ratio of K to ˜ N to clients for uniform data sampling.

3.4 Federated Training Framework

Next, we will introduce the workflow of federated training in FedSampling, which is summarized in Fig. 2 and Algorithm 1. The workflow of FedSampling is similar to the mainstream FL methods in local model training and global model updating, while different in data sampling. The t-th training round in FedSampling includes the following steps.

First, the server broadcasts the current model Θt-1 for model aligning and queries each client for the local sample size, where Θt-1 is the model in the t - 1-th round. Then clients respond to the server based on Eq. 1 and the server

Algorithm 1 Workflow of FedSampling

1: for all t ← 1 to T do 2: Broadcast model Θt-1 and query sample sizes 3: for c in C do 4: Draw xc fromB(α), ˆnc fromP(M) 5: Obtain rc via Eq. 1 6: Send rc to the server 7: end for 8: Calculate ˜ N via Eq. 2, and broadcast ˜ N, K 9: for c in C do 10: Initialize an empty data queue Sc 11: for i ← 1 to |Dc| do 12: Draw zc i ∼ B( K˜ N ) 13: Push sci into Sc if zc i > 0 14: end for 15: Build local model update based on samples in Sc 16: Send local model update to the server 17: end for 18: Θt = Θt-1 + η ∑G∈G G 19: end for

estimates the total sample size ˜ N via Eq. 2. Second, the server broadcasts K and ˜ N to clients. Then clients can further uniformly sample data to obtain the local updates and upload them to the server. Third, until receiving all local updates, the server aggregates them to update the global model: Θt = Θt-1 +η ∑G∈G G, where G is the set of received local updates, and η is the learning rate.

3.5 Discussions on Utility and Privacy

Next, we will analyze the utility and privacy of FedSampling in theory and compare it with mainstream FL methods.

Utility Analysis on FedSampling

Definition 3.1. An FL algorithm is data-level unbiased, iff any data can be independently sampled for a training step with identical probabilities to centralized training.

Based on this definition, data-level unbiased FL methods have a similar data exploitation pattern with centralized learning. Thus, data-level unbiased FL methods have the potential to more effectively exploit decentralized data under imbalanced data size distribution and even the non-IID data distribution. We remark that FL methods based on the uniform or weighted client sampling are not data-level unbiased, due to the violation of independent or identical sampling conditions. Furthermore, based on Lemma 3.1, in FedSampling data can be sampled with probabilities identical to centralized learning in a training round. Besides, since FedSampling controls the sampling of data via independent random variables, and thereby it is data-level unbiased.

Lemma 3.1. Let p(x) and ˆ p(x) denote the probability of a sample x that can participate in a training step in the centralized learning and FedSampling. The mean square error between p(·) and ˆ p(·) asymptotically converges to 0.

Proof. Denote the data size distribution as PD, then we can view the selected response rc as a combination of three random variables (Eq. 4), and further obtain its expectation:

where n is the averaged sample size of different clients. Furthermore, we prove ˜ N is an unbiased estimation of N:

Furthermore, without the loss of generalization, we assume that there are K samples selected for model training in a single step. Based on the uniform data sampling strategy in the centralized learning, and the sampling strategy in FedSampling, we can obtain the following equations:

Let f(x) = ( 1x - 1N )2 , based on the Taylor series, we have:

Let x = ˜ N, and note that f(N) = 0 and f l (N) = (-1)l (l-1)l! N l+2 , then we have the following equation:

Furthermore, the following inequation holds:

Further, based on tight version of the Cr inequation [Bahr and Esseen, 1965], we have:

Note that E[|rc -E[rc]|l ] is a constant determined by the data size distribution and the setting of our size anonymization method (i.e., the value of α and M), and we denote it as δl. Note that |C| < N, then we obtain this inequation:

Note that δl, and α are constants and irrelevant with |C|, then we can obtain our expected conclusion:

Corollary 3.1. The mean square error between NK ˜ N and K can asymptotically converge to 0:

Privacy Analysis on FedSampling

For privacy protection, most FL methods disclose sample sizes and may violate privacy constrains. We are the first to propose a privacy-preserving ratio estimation method that can protect privacy in the local sample sizes. Our method can achieve ϵ-LDP in a training round via Lemma 3.2.

Definition 3.2. A randomized mechanism M(·) satisfies ϵLDP, iff for two arbitrary input x and x′ , and any output y ∈ range(M), the following inequation holds:

The privacy budget ϵ quantifies the privacy protection ability, where a smaller ϵ means better privacy protection.

Lemma 3.2. Given an arbitrary size threshold M, FedSampling can meet ϵ-LDP in protecting the privacy of local sample size in each client, when α = eϵ-1 eϵ+M-2 .

Proof. Let M(·) denote the privacy preserving ratio estimation method. Consider an arbitrary client c with a true response nc, and a protected response rc, then we obtain:

where we can find the Eq. 17 always larger or equal than Eq. 18. Furthermore, for two arbitrary c and c′ in C and any output y ∈ {y|y = 1, 2, ..., M - 1}, we have:

Furthermore, according to the definition of ϵ-LDP, we have:

Thus, when M and ϵ is fixed, we can obtain:

We can further de-link responses from the same client at different iterations via client anonymity [Sun et al., 2021] to avoid the accumulation of privacy budget. Thus, FedSampling can meet ϵ-LDP in the whole training process. Besides, the disclose of local updates may also leak user privacy. We remark that this concern is out of the scope of this paper, as there are already many works [Lyu et al., 2020] that protect local updates via LDP. It is straightforward to apply many of them to FedSampling for further privacy protection.

4 Experiment

4.1 Datasets and Experimental Settings

Experiments are conducted on four benchmark datasets: a text dataset MIND [Wu et al., 2020], two Amazon review datasets (i.e., Toys and Beauty) [McAuley et al., 2015], and an image dataset EMNIST [Cohen et al., 2017]. First, MIND, Toys and Beauty are used to evaluate model effectiveness under imbalanced data size distribution. These three datasets are used for the text classification task, where Text-CNN [Kim, 2014] and the mainstream NLP model Transformer [Vaswani et al., 2017] are used as the basic model. Data in Toys and Beauty can be naturally partitioned into clients based on user IDs, which follows a long-tail distribution (shown in Appendix). Besides, since texts in MIND dot not include user information, we partition MIND based on a typical long-tail distribution, i.e., log-normal distribution. The average size of samples in each client is set to 2 and we adjust the variance σ of the distribution to control the data size imbalance. σ is set to 4 in experiments. Second, we use MIND and EMNIST to verify how non-IID data distribution affects different methods. Motivated by McMahan et al. [2017], we sort training data in each dataset by their labels, and partition them into different clients of size 300 and 6000 respectively. Thus, distributions of local data in different clients are varied. We use the ResNet network [He et al., 2016] for the image classification on EMNIST. Marco-F1 and Accuracy are used for classification task evaluation. In our FedSampling method, the size threshold M and privacy budget ϵ are set to 300 and 3 respectively. The learning rate η is set to 0.05, and the number K of samples for participating in a training round is set to 2048. All hyper-parameters are selected on the validation set. More model details are in the Appendix and Codes (https://github.com/taoqi98/FedSampling).

4.2 Performance Evaluation

We compare FedSampling with several recent federated learning methods. We first compare the standard FedAvg algorithm and three adaptive federated optimization methods proposed to improve model effectiveness under heterogeneous data distribution, including: (1) FedAvg [McMahan et al., 2017]: uniformly selecting clients for federated model training, and averaging local model updates to train the global model. (2) FedYogi [Reddi et al., 2020]: a federated version of the Yogi optimization algorithm that smooths global model updates based on momentum strategy [Reddi et al., 2018]. (3) FedAdagrad [Reddi et al., 2020], a federated version of the Adagrad optimization algorithm [Lydia and Francis, 2019]. (4) FedAdam [Reddi et al., 2020], a federated version of the Adam optimization algorithm [Kingma and Ba, 2015]. We remark all that these baseline FL methods are based on uniform client sampling according to their original settings. We also apply the centralized model training (Centralization) as a baseline method to quantify the performance decline of different federated learning methods.

We repeat each experiment five times and show average performance and standard deviations in Table 1. First, compared with Centralization, baseline FL methods have serious performance declines on all datasets. This is because the size distributions of many real-world user data are usually imbalanced. However, baseline FL methods usually uniformly select some clients to participate in a training round, and thereby fail to effectively exploit massive training data in some long-tail clients. Second, our FedSampling method can consistently achieve comparable performance with Centralization. This is because our FedSampling method has a sample-level sampling strategy, which can uniformly and independently sample data to achieve a uniform and unbiased data exploitation like centralized learning.

To more effectively exploit decentralized data for model training, some existing federated learning methods [Li et al., 2019; Wang et al., 2020] sample clients based on local sample sizes. Thus, we apply the weighted client sampling strategy to baseline FL methods for further comparisons. Due to space limitations, we only show results on MIND with the Transformer model in the following sections. Results are presented in Fig. 3, and we have several findings. First, we can find that weighted client sampling significantly improves the performance of baseline FL methods. This is because based on weighted client sampling long-tail clients have more opportunities to participate in a training round, and thereby massive data in them can be more effectively exploited. Second, compared with Centralization and FedSampling, performance of baseline FL methods with weighted sampling are still inferior. This is because weighted client sampling usually emphasizes Distribution Variance:

Figure 4: Influence of data size imbalance degree.

long-tail clients during model training while head clients can also accumulate non-neglect samples. These methods still cannot achieve uniform data exploitation, which leads to suboptimal performance. Different from these methods, we propose a sample-level sampling strategy which can independently and uniformly sample decentralized data for training. In addition, the weighted client sampling strategy requires the server to track local samples sizes of each client, which may arouse privacy concerns due to the privacy sensitivity of local sample sizes. Different from them, FedSampling protects the local sample size of each client via a privacy-preserving ratio estimation method with a theoretical privacy guarantee. Besides, we also compare the naive method introduced in approach section with FedSampling. Results show that performance of the naive method is inferior to FedSampling. This is because to avoid the privacy leakage on sample sizes, the naive method needs to uniformly average local updates to learn the global model, which is biased in data exploitation.

4.3 Influence of Data Size Distribution Imbalance

Next, we will analyze how the imbalance of data size distribution affects different methods. According to our experimental settings, the data partition of MIND is based on a log-normal distribution. Thus, we can control the size imbalance by adjusting the variance σ of the log-normal distribution, where larger σ leads to more imbalanced sample sizes. Results in Fig. 4 show the influence of σ on the effectiveness of different methods, from which can we summarize two major phenomenons. First, with the increase of σ, the performance of baseline methods consistently declines. This is because the increase of σ makes sample sizes more imbalanced, where long-tail clients can keep more data. Thus, there will be more data that fails to be effectively exploited by baseline FL methods. Second, FedSampling has comparable performance with Centralization under different σ. This is because FedSampling achieves uniform sampling in data level and can effectively learn models under imbalanced data size distribution.

4.4 Comparisons under Non-IIDDistribution

Since FedSampling can uniformly sample decentralized data, it has the potential to improve the effectiveness of FL under non-IID data distribution. Next, we compare different methods under the non-IID data distribution. We set clients to have equal sample size, i.e., 300, which makes no differences for baseline FL methods to perform uniform or weighted client sampling. Thus, in this section, we only show results of baseline FL methods with uniform client sampling. Results are shown in Fig. 5, from which we can find the convergence of FedSampling is faster and smoother. This is because distributions of data on different clients are highly heterogeneous, where most clients only keep data with the same classification category. However, baseline FL methods sample data for training at the client level, which results in heterogeneous local model updates and hurts the model convergence. Different from these methods, FedSampling can independently select samples in different clients for model training with identical opportunities. Thus, FedSampling can achieve a similar data sampling pattern with the canonical uniform data sampling in centralized learning, and improve the effectiveness of FL models under the non-IID data distribution.

4.5 Utility and Privacy Analysis

Next, we will analyze the trade-off between utility and privacy of our FedSampling method. As shown in Fig. 6, given a size threshold M, we evaluate the performance of FedSampling under different privacy budget ϵ. First, with a constant size threshold M, a smaller privacy budget ϵ usually leads to a worse model performance. This is because a smaller privacy budget ϵ increases the probability of clients for sending a fake response to the server, which hurts the accuracy of the total sample number estimation and the quality of data sampling. Second, when the privacy budget ϵ is fixed, FedSampling can achieve the best performance under a moderate size threshold M, e.g., 102 . This is because when M is too small, many true responses may be truncated by a small value and cause bias on the estimation of total sample size. Besides, according to Lemma 3.2, when the privacy budget ϵ is fixed, a larger size threshold M requires larger probabilities of sending fake responses, which hurts the accuracy of data size distribution estimation. Thus, we select a moderate value of size threshold M (i.e., 100) and privacy budget ϵ (i.e., 3) in experiments.

5 Conclusion

In this paper, we propose a uniform data sampling strategy for federated learning (named FedSampling), which can achieve an uniform data exploitation. In each round of FedSampling, each sample in each client is independently sampled for model training according to an identically probability, which is based on the ratio of the server desired sample size to the total sample size in all available clients. We also propose a privacy-preserving method to estimate an unbiased total sample size with an LDP privacy guarantee on the local sample size of each client. We conduct extensive experiments on four benchmark datasets. Experimental results show that FedSampling can outperform many FL methods especially under imbalanced data size distribution and non-IID data distribution. Although effective, FedSampling may increase the communication cost of federated learning. In our future work we will explore to handle this problem.

Acknowledgments

This work was supported by the National Key Research and Development Program of China 2022YFC3302104, Tsinghua University Initiative Scientific Research Program of Precision Medicine under Grant number 2022ZLA0073, Tsinghua-Toyota Joint Research Fund 20213930033, and China Postdoctoral Science Fundation (No.2022M721892).