Heterogeneity-Guided Client Sampling: Towards Fast and Efficient Non-IID Federated Learning

Authors: huancheng  chen
haris  vikalo
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-04 03:02:42

Abstract: Statistical heterogeneity of data present at client devices in a federated learning (FL) system renders the training of a global model in such systems difficult. Particularly challenging are the settings where due to communication resource constraints only a small fraction of clients can participate in any given round of FL. Recent approaches to training a global model in FL systems with non-IID data have focused on developing client selection methods that aim to sample clients with more informative updates of the model. However, existing client selection techniques either introduce significant computation overhead or perform well only in the scenarios where clients have data with similar heterogeneity profiles. In this paper, we propose HiCS-FL (Federated Learning via Hierarchical Clustered Sampling), a novel client selection method in which the server estimates statistical heterogeneity of a client's data using the client’s update of the network’s output layer and relies on this information to cluster and sample the clients. We analyze the ability of the proposed techniques to compare heterogeneity of different datasets, and characterize convergence of the training process that deploys the introduced client selection method. Extensive experimental results demonstrate that in non-IID settings HiCS-FL achieves faster convergence than state-of-the-art FL client selection schemes. Notably, HiCS-FL drastically reduces computation cost compared to existing selection schemes and is adaptable to different heterogeneity scenarios.

Main paper: 
 introduction: 
 The federated learning (FL) framework enables privacy-preserving collaborative training of machine learning (ML) models across a number of devices (clients) by avoiding the need to collect private data stored at those devices. The participating clients typically experience both the system as well as statistical heterogeneity [18] . The former describes settings where client devices have varying degree of computational resources, communication bandwidth and fault tolerance, while the latter refers to the fact that the data owned by the clients may be drawn from different distributions. In this paper, we focus on FL under statistical heterogeneity and leave studies of system heterogeneity to future work.An early FL method, FedAvg [21] , performs well in the settings where the devices train on independent and identically distributed (IID) data. However, compared to the IID scenario, training on non-IID data is detrimental to the convergence speed, variance and accuracy of the learned model. This has motivated numerous studies aiming to reduce the variance and improve convergence of FL on non-IID data [6, 9, 14, 17, 19, 30] .On another note, constraints on communication resources and therefore on the number of clients that may participate in training additionally complicate implementation of FL schemes. It would be particularly unrealistic to require regular contributions to training from all the clients in a largescale cross-device FL system. Instead, only a fraction of clients participate in any given training round; unfortunately, this further aggravates detrimental effects of statistical heterogeneity. Selecting informative clients in non-IID FL settings is an open problem that has received considerable attention from the research community [8, 11, 12] . Since privacy concerns typically prohibit clients from sharing their local data label distributions, existing studies focus on estimating informativeness of a client's update by analyzing the update itself. This motivated a family of methods that rely on the norms of local updates to assign probabilities of sampling the clients [7, 23] . Aiming to enable efficient use of the available communication and computation resources, another set of methods groups clients with similar data distributions into clusters based on the similarity between clients' model updates [2, 11] . Across the board, the existing methods still struggle to deliver desired performance in an efficient manner and cannot distinguish clients with balanced data from the clients with imbalanced data.In this paper, we consider training a neural network model for classification tasks via federated learning and propose a novel adaptive clustering-based sampling method for identifying and selecting informative clients. The method, referred to as Federated Learning via Hierarchical Clustered Sampling (HiCS-FL), relies on the updates of the (fully connected) output layer in the network to determine how diverse is the clients' data and, based on that, decide which clients to sample. In particular, HiCS-FL enables heterogeneity-guided client selection by utilizing general properties of the gradients of the output layer to distinguish between clients with balanced from those with imbalanced data. Unlike the Clustered Sampling strategies [11] where the clusters of clients are sampled uniformly, HiCS-FL allocates different probabilities (importance) to the clusters according to their average estimated data heterogeneity. Numerous experiments conducted on vision datasets FMNIST, CIFAR10, Mini-ImageNet and a NLP dataset THUC news demonstrate that HiCS-FL achieves significantly faster training convergence and lower variance than the competing methods. Finally, we conduct convergence analysis of HiCS-FL and discuss implications of the results.In summary, the contributions of the paper include: (1) Analytical characterization of the correlation between local updates of the output layer and the FL clients' data label distribution, along with an efficient method for estimating data heterogeneity; (2) a novel clustering-based algorithm for heterogeneity-guided client selection; (3) extensive simulation results demonstrating HiCS-FL provides significant improvement in terms of convergence speed and variance over competing approaches; and (4) theoretical analysis of the proposed schemes. 
background and related work: 
 Assume the cross-device federated learning setting with N clients, where client k owns private local dataset B k with |B k | samples. The plain vanilla FL considers the objectiveEQUATIONwhere θ denotes parameters of the global model, F k (θ) is the loss (empirical risk) of model θ on B k , and p k denotes the weight assigned to client k, N k=1 p k = 1. In FedAvg, the weights are set top k = |B k | / N i=1 |B i |.In training round t, the server collects clients' model updates θ t k formed by training on local data and aggregates them to update global model as θ t+1 = N k=1 p k θ t k . When an FL system operates under resource constraints, typically only K ≪ N clients are selected to participate in any given round of training; denote the set of clients selected in round t by S t . In departure from FedAvg, FedProx [19] proposes an alternative strategy for sampling clients based on a multinomial distribution where the probability of selecting a client is proportional to the size of its local dataset; the global model is then formed as the average of the collected local models θ t+1 = 1 K k∈S t θ t k . This sampling strategy is unbiased since the the updated global model is on expectation equal to the one obtained by the framework with full client participation as Eq.1.AFL [12] is the first study to utilize local validation loss as a value function for computing client sampling probabilities; Power-of-Choice [8] takes a step further to propose a greedy approach to sampling clients with the largest local loss. Both of these methods require all clients to compute the local validation loss, which is often unrealistic. To address this problem, FedCor [28] models the local loss by a Gaussian Process (GP), estimates the GP parameters from experiments, and uses the GP model to predict clients' local losses without requiring them to perform validation. In [7] , Optimal Client Sampling scheme aiming to minimize the variance of local updates by assigning sampling probabilities proportional to the Euclidean norm of the updates is proposed. The study in [23] models the progression of model's weights by an Ornstein-Uhlenbeck process and proposes a strategy, optimal under that assumption, for selecting clients with significant weight updates.The clustering-based sampling method proposed in [11] uses cosine similarity [24] to group together clients with similar local updates, and proceeds to sample one client per cluster in attempt to avoid redundant gradient information. DivFL [2] follows the same principle of identifying representative clients but does so by constructing a submodular set and greedily selecting diverse clients. Both of these techniques are computationally expensive due to the high dimension of the gradients that they need to process.In general, the overviewed methods either: (1) select diverse clients to reduce redundant information; or (2) select clients with a perceived significant contributions to the global model (high loss, large update or low class-imbalance). Efficient and effective client selection in FL remains an open challenge, motivating the heterogeneity-guided adaptive client selection method presented next.3 HiCS-FL: Federated Learning via Hierarchical Clustered Sampling Existing client sampling methods including Clustered Sampling [11] and DivFL [2] aim to select clients such that the resulting model update is an unbiased estimate of the true update (i.e., the update in the case of full client participation) while minimizing the varianceEQUATION)Clustered Sampling, for instance, groups N clients into K clusters based on representative gradients [24] , and randomly selects one client from each cluster to contribute to the global model update. Such an approach unfortunately fails to differentiate between model updates formed on data with balanced and those formed on data with imbalanced label distributions -indeed, in either case the updates are treated as being equally important. However, a number of studies in centralized learning has shown that class-imbalanced datasets have significant detrimental effect on the performance of learning classification tasks [3, 4, 26] . This intuition carries over to the FL settings where one expects the updates from clients training on relatively more balanced local data to have a more beneficial impact on the performance of the system. The Federated Learning via Hierarchical Clustered Sampling (HiCS-FL) framework described in this section adapts to the clients' data heterogeneity in the following way: if the levels of heterogeneity (as quantified by the entropy of data label distribution) vary from one cluster to another, HiCS-FL is more likely to sample clusters containing clients with more balanced data; if the clients grouped in different clusters have similar heterogeneity levels, HiCS-FL is more likely to select diverse clients (i.e., sample uniformly across clusters, thus reducing to the conventional clustered sampling strategy).Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: [NA] Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
class-imbalance causes objective drift: 
 A number of studies explored detrimental effects of non-IID training data on the performance of a global model learned via FedAvg. An example is SCAFFOLD [14] which demonstrates objective drift in non-IID FL manifested through large differences between local models θ * k trained on substantially different data distributions. The drift is due to FedAvg updating the global model in the direction of the weighted average of local optimal models, which is not necessarily leading towards the optimal global model θ * . The optimal model θ * , in principle obtained by solving optimization in Eq. 1, achieves minimal empirical error on the data with uniform label distribution and is intuitively closer to the local optimal models trained on balanced data. Recent work [36] empirically verified this conjecture through extensive experiments. Let ∇F (θ t ) denote the gradient of F (θ t ) given the global model θ t at round t; the difference between ∇F (θ t ) and the local gradient ∇F k (θ t ) computed on client k's data is typically assumed to be bounded [7, 11, 31] . To proceed, we formalize the assumption about the relationship between gradients and data label distributions.EQUATIONwhere D (k) is the data label distribution of client k, D 0 denotes uniform distribution, H(•) is Shannon's entropy of a stochastic vector, and β > 0, κ > ρ > 0.The assumption commonly encountered in literature is recovered by setting the right-hand side of (3) to σ 2 m = max k σ 2 k . Intuitively, if the data label distribution of client k is highly imbalanced (i.e., H(D (k) ) is small), the local gradient ∇F k (θ t ) may significantly differ from the global gradient ∇F (θ t ) (as reflected by the bound above). Analytically, connecting the gradients to the local data label distributions allows one to characterize the effects of client selection on the variance and the rate of convergence. The results of extensive experiments that empirically verify the above assumption are reported in Appendix A.2.If the server were given access to clients' data label distributions, selecting clients would be relatively straightforward [32] . However, privacy concerns typically discourage clients from sharing such information. Previous studies have explored the use of multi-arm bandits for inferring clients' data heterogeneity from local model parameters, or have utilized a validation dataset at the server to accomplish the same [27, 34, 36] . In this section, we demonstrate how to efficiently and accurately estimate data heterogeneity using local updates of the output layer of a neural network in a classification task. Figure 1 illustrates the last two layers in a typical neural network. The prediction q ∈ R C is computed by forming a weighted average of signals z ∈ R L utilizing the weight matrix W ∈ R C×L and bias b ∈ R C .An empirical investigation of the gradients of the output layer's weights while training with FedAvg using mini-batch stochastic gradient descent (SGD) as an optimizer is reported in [5, 29] . There, the focus is on detecting the presence of specific labels in a batch rather than on exploring the effects of class imbalance on the local update. To pursue the latter, we focus on the correlation between local updates of the output layer's bias and the client's data label distribution; we start by analyzing the training via FedAvg that employs SGD and then extend the results to other FL algorithms that utilize optimizers beyond SGD. We assume that the model is trained by minimizing the cross-entropy (CE) loss over one-hot labels -a widely used multi-class classification framework. The gradient is computed by averaging contributions of the samples in mini-batches, i.e., (j,n) , y (j,n) ), where B denotes the batch size, l is the number of minibatches, x (j,n) is the n-th point in the j-th mini-batch and y (j,n) ∈ [C] is its label. The contribution of x (j,n) to the i-th component of the gradient of the output layer's bias b can be found as (details provided in Appendix A.3)∇ b L ce = 1 Bl l j=1 B n=1 ∇ b L (j,n) ce (xEQUATIONwhere I{•} is an indicator, q (j,n) = W • z (j,n) + b is the output logit for signals z (j,n) ∈ R L corresponding to training point (x (j,n) , y (j,n) ) (see Fig. 1 ), and where C denotes the number of classes. We make the following observations: (1) the sign of y(j,n) -th component of ∇ b L (j,n) ce is opposite of the sign of other components; and (2) the y (j,n) -th component of ∇ b L (j,n) ceis equal in magnitude to all the other components combined. Note that the above two observations are standard for neural networks using CE loss for supervised multi-class classification tasks.In each global round t of FedAvg, the selected client k starts from the global model θ t and proceeds to compute local update in R local epochs employing an SGD optimizer with learning rate η. According to Eq. 4, the i-th component of local update ∆b (k) is computed asEQUATIONwhere ∇ bi L (j,n,r) ce denotes the gradient of bias at local epoch r. Note that the local update of client k, ∆b (k) , is dependent on the label distribution of client k's data,EQUATIONC ] T and the label-specific components of q (j,n) which change during training. We proceed by relating expected local updates to the label distributions; for convenience, we first introduce the following definition. Definition 3.2 Let B -i be the subset of local data B that excludes points with label i. Let s -i (x) ∈ [0, 1] C be the softmax output of a trained neural network for a training point (x, y) ∈ B -i . The i-th component of s -i (x), s -i i (x), indicates the level of confidence in (erroneously) classifying x as having label i. For convenience, we defineE i = E (x,y)∼B -i s -i i (x) , ∀i ∈ [C].In an untrained/initialized neural network where classifier makes random predictions, E i = 1/C; as training proceeds, E i decreases. By taking expectation and simplifying, we obtain (details provided in Appendix A.4)EQUATIONwhereD (k) idenotes the true fraction of samples with label i in client k's data,C i=1 D (k) i = 1.We quantify the heterogeneity of clients' data by an entropy-like measure defined below. Let D (k) denote the label distribution of client k's data; its entropy is defined asH(D (k) ) ≜ - C i=1 D (k) i ln D (k) i ≤ ln C.Recall that more balanced data results in higher entropy, and that H(D (k) ) takes the maximal value when D (k) is uniform. The server does not know D (k) and therefore cannot compute H(D (k) ) directly. We defineEQUATIONhere T is a scaling hyper-parameter (so-called temperature). Note that even though we can compute Ĥ(D (k) ) to characterize heterogeneity, Dand E i remain unknown to the server (details in A.5).Theorem 3.3 Consider an FL system in which clients collaboratively train a model for a classification task over C classes. Let D (u) and D (k) denote data label distributions of an arbitrary pair of clients u and k, respectively. Moreover, let U denote the uniform distribution, and let η and R be the learning rate and the number of local epochs, respectively. ThenEQUATIONwhere C = ηR(ηR+C 2 T ln C)C 2 T 2 and δ = max i C c=1 Ec C -E i .The proof is provided in Appendix A.6.As an illustration, consider the scenario where client u has a balanced dataset while the dataset of client k is imbalanced; then ∥D (k) -U∥ 2 2 is relatively large compared to ∥D (u) -U∥ ∞ . The bound in (8) also depends on δ, which is reflective of how misleading on average can a class be; small δ suggests that no class is universally misleading. As shown in Appendix A.4, during training δ gradually decreases to 0 as C i=1 E i decreases to 0.The proposed method for estimating clients' data heterogeneity relies on the properties of the gradient for the cross-entropy loss objective discussed in Section 3.2.1. However, for FL algorithms other than FedAvg, such as FedProx [19] , FedDyn [1] and Moon [16] , which add regularization to combat overfitting, the aforementioned properties may not hold. Moreover, optimization algorithms using second-order momentum such as Adam [15] deploy update rules different from SGD, making the local updates no longer proportional to the gradients. Nevertheless, HiCS-FL remains capable of distinguishing between clients with imbalanced and balanced data, which will be demonstrated in our experiments. Further theoretical discussion of various FL algorithms with optimizers beyond SGD are in appendix A.8 and A.9.Clustered Sampling [11] uses cosine similarity [24] between gradients to quantify proximity between clients' data distributions and subsequently group them into clusters. However, cosine similarity The global model θ T +1 cannot help distinguish between clients with balanced and those with imbalanced datasets. Motivated by this observation, we introduce a new distance measure that incorporates estimates of data heterogeneity Ĥ(D (k) ). In particular, the proposed measure of distance between clients u and k that we use to form clusters is defined asEQUATIONwhere the first term is akin to the cosine similarity used by CS with the major difference that we compute it using only the updates of the bias in the output layer, which is much more efficient than using the weights of the entire network; λ is a pre-defined hyper-parameter (set to 10 in all our experiments). For large λ, the second term dominates when there are clients with different levels of statistical heterogeneity; this allows emergence of clusters that group together clients with balanced datasets. The second term is small when clients have data with similar levels of statistical heterogeneity; in that case, the distance measure reduces to the conventional cosine similarity.To select K out of N clients in an FL system, we first organize the clients into M ≥ K groups via the proposed Hierarchical Clustered Sampling (HiCS) technique. In particular, during the first ⌈N/K⌉ training rounds the server randomly (without replacement) selects clients and collects from them local updates of ∆b (k) ; the server then estimates Ĥt (D (k) ) for each selected client k and clusters the clients using the distance measure defined in Eq. 9. Let G where k 1 , . . . , k |Gm| are the indices of clients in cluster G m , γ t = γ 0 (1 -t T ) denotes an annealing hyper-parameter, and T is the number of global rounds. The annealing parameter is scheduled such that at first it promotes sampling clients with balanced data, thus accelerating and stabilizing the convergence of the global model. To avoid overfitting potentially caused by repeatedly selecting a small subset of clients, the annealing parameter is gradually reduced to γ t ≈ 0, when the server samples the clusters uniformly. The described procedure is formalized as Algorithm 1.Adopting the standard assumptions of smoothness, unbiased gradients and bounded variance [7] , the following theorem holds for FedAvg with SGD optimizer.Theorem 3.4 Assume F k (•) is L-smooth for all k ∈ [N ]. Let θ t denote parameters of the global model and let F (•) be defined as in Eq. 1. Furthermore, assume the stochastic gradient estimator g k (θ t ) is unbiased and the variance is bounded such that E ∥g k (θ t ) -∇F k (θ t )∥ 2 ≤ σ 2 . Let η and R be the learning rate and the number of local epochs, respectively. If the learning rate is such that ηEQUATIONwhere A 1 , A 2 , Φ are positive constants, and ω t k is the probability of sampling client k at round t.Note that only the second term in the parenthesis on the right-hand side of the bound in Theorem 3.4 is related to the sampling method Π. Under Assumption 3.1,EQUATIONIf the server samples clients with weights proportional to p k , the statistical heterogeneity of the entire FL system may be characterized byH S = N k=1 p k exp(β(H(D (k) ))exp(β(H(D0)) ρ. If all clients have classimbalanced data, H S is small and thus random sampling leads to unsatisfactory convergence rate (as indicated by Theorem 3.4). On the other hand, since the clients sharing a cluster have similar data entropy, the proposed HiCS-FL leads toω t k = p k exp(γ t Ĥt (D (k) )) N j=1 pj exp(γ t Ĥt (D (j) )) .When training starts, H Π is large because the server tends to sample clients with higher p k exp(γ t H(D (k) )); as γ t decreases, H Π eventually approaches H S . Further details and the proof of the theorem are in Appendix A.7.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
experiments: 
 Setup. We evaluate the proposed HiCS-FL algorithm on four benchmark datasets (FMNIST, CI-FAR10, Mini-ImageNet and THUC news) using different model architectures. We use four baselines: random sampling, pow-d [8] , clustered sampling (CS) [11] , DivFL [2] and FedCor [28] . To generate non-IID data partitions, we follow the strategy in [35] , utilizing Dirichlet distribution with different concentration parameters α which controls the level of heterogeneity (smaller α leads to generating less balanced data). In a departure from previous works we utilize several different α to generate data partitions for a single experiment, leading to a realistic scenario of varied data heterogeneity across different clients. To quantify the performance of the tested methods, we use two metrics: (1) average training loss, and (2) test accuracy of the learned global model. For better visualization, data points in the results are smoothened by a Savitzky-Golay filter with window length 13 and the polynomial order set to 3. Further details of the experimental setting and a visualization of data partitions are in Appendix A.1 and A.10.FMNIST. We run FedAvg with SGD to train a global model which has CNN architecture in an FL system with 50 clients, where 10% of clients are selected to participate in each round of training. The data partitions are generated using one of 3 sets of the concentration parameter α values: (1) {0.001, 0.002, 0.005, 0.01, 0.5}; (2) {0.001, 0.002, 0.005, 0.01, 0.2}; (3) {0.001}. These are used to generate clients' data so as to emulate the following scenarios: (1) 80% of clients have severely imbalanced data while the remaining 20% have balanced data; (2) 80% clients have severely imbalanced data while the remaining 20% have mildly imbalanced data; (3) all clients have severely imbalanced data. Note that H M monotonically decreases as we go through settings (1) to (3). For a fair comparison, pow-d and DivFL are deployed with their ideal settings where the server requires all clients to precompute in each round a metric that is then used for client selection. Figure 2 shows that HiCS-FL outperforms other methods across different settings, exhibiting the fastest convergence rates and the least amount of variance. Particularly significant is the acceleration of convergence in setting (1) where 20% of the participating clients have balanced data. Figure 3 shows that HiCS-FL is helping achieve significant reduction of training variations (as expected, see Section 3.5) as evident by a smooth loss trajectory. Setting (1) emulates the scenario where clients have a range of heterogeneity profiles, from extremely imbalanced, through mildly imbalanced, to balanced, while setting (2) corresponds to the scenario where 80% of the clients have extremely imbalanced data while the remaining 20% have balanced data. The system's HS for setting (1) is larger than HS for setting (2) , which is reflected in a more significant improvements achieved by HiCS-FL in the latter setting, as shown in Figure 4 .THUC news. To evaluate our method on data from a different domain, we conduct experiments involving text classification on the THUC news dataset in Chinese language (10 labels). Similar to the aforementioned experiments, we allocate data to 50 clients by emulating heterogeneous data distributions scenarios with parameter α set to: (1) {0.001, 0.01, 0.1, 0.2,1}; (2) {0.001, 0.002, 0.01, 0.1, 0.5}; and (3) {0.001, 0.002, 0.005, 0.01, 0.1}. We trained TextRNNs [20] with BiLSTM architecture as 1 , show that our method outperforms baselines in all the settings, demonstrating efficacy of our proposed algorithm in a simple NLP task. In this section we report the communication costs required to achieve convergence when using HiCS-FL, and compare those results with the competing schemes. For brevity, we select one result from each experiment conducted on the considered four datasets, and display them in Table 2 . As can be seen from the As discussed at the end of Section 3.3, the distance function in Equation 9 can be reduced to the conventional cosine similarity when clients exhibit similar levels of statistical heterogeneity, despite potential differences in data distribution. Under these circumstances, our HiCS-FL method can recover the performance of the previously established CS approach [11] . While CS suggests that the number of clusters M should be greater than or equal to the number of selected clients K, our HiCS-FL does not require M > K but adheres to the CS settings to ensure a fair comparison.To elucidate the impact of the number of clusters, we conducted supplementary experiments with HiCS-FL using varying numbers of clusters M and compared these results to those obtained with M = K as presented in the paper. The results of those experiments can be found in Table . 3. As shown there, HiCS-FL can perform well with smaller M < K as long as M is not too small, such as M = 3.The purpose of the warm-up phase (t < ⌈N/K⌉) shown in Alg. 1 is to collect updates of the output layer from all the available clients in the system in order to facilitate clustering. Although we conduct all the experiments in the setting where clients have fixed availability, our HiCS-FL does not assume all the clients are available in the warm-up phase and can be adapted to more practical scenarios where clients have dynamic availability.In such a scenario, the warm-up phase can be implemented by the available clients at the beginning of training. The proposed HiCS-FL is then implemented only among the available clients; the available clients with more balanced data are preferred. When new clients join the system at the global round t, the server can obtain the information of availability and selects these new clients at round t + 1 to approximate their data heterogeneity. To provide more insights, we conduct additional experiments on CIFAR10 dataset; the results are reported in Table . 4. As can be seen there, HiCS-FL outperforms baselines that consider clients' availability. 
conclusion: 
 In this paper, we studied federated learning systems where clients that own non-IID data collaboratively train a global model; the system operates under communication constraints and thus only a fraction of clients participates in any given round of training. We developed HiCS-FL, a hierarchical clustered sampling method which estimates clients' data heterogeneity and uses this information to cluster and select clients to participate in training. We analyzed the performance of the proposed heterogeneity estimation method, and the convergence of training a FL system that deploys HiCS-FL. Extensive benchmarking experiments on four datasets demonstrated significant benefits of the proposed method, including improvement in convergence speed, variance and test accuracy, accomplished with only a minor computational overhead. 
convergence analysis: 
 Here we present the convergence analysis of an FL system deploying FedAvg with SGD wherein only a small fraction of clients participates in any given round of training. Recall that the objective function that comes up when training a neural network model is generally non-convex; we make the standard assumptions of smoothness, unbiased gradient estimate, and bounded variance.EQUATION)Assumption A.3 (Gradient oracle) The stochastic gradient estimator g k (θ t,r k ) = ∇F k (θ t,r k ) + ζ t,r kfor each global round t and local epoch r is such thatE[ζ t,r k ] = 0 (52)andEQUATION)With these three assumptions in place, we provide the proof of Theorem 3.4 stated in the main paper. The proof relies on the technique previously used in [7, 33] , where the sampling method is unbiased and thus E 1K k∈S t R r=1 g k (θ t,r k ) = N k=1 R r=1 p k ∇F k (θ t,r k ).We provide a generalization that holds for any sampling strategy, resulting in E 1K k∈S t R r=1 g k (θ t,r k ) = N k=1 R r=1 ω t k ∇F k (θ t,r k ), where ω t k denotes the probability of sampling client k in round t under sampling strategy Π. Note that N k=1 ω t k = 1. We assume that all clients deploy the same number of local epochs R and use learning rate η at round t. 
regularization terms in the objective function: 
 The proposed method for estimating clients' data heterogeneity relies on the properties of gradient computed for the cross-entropy loss objective. However, the method also applies to the FL algorithms other than FedAvg, in particular those that add a regularization term to combat overfitting, including FedProx [19] , FedDyn [1] and Moon [16] . In the following discussion, we demonstrate that HiCS-FL remains capable of distinguishing between clients with imbalanced and balanced data when using these other FL algorithms. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: [NA] Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: [NA] Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: [NA] Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: A.1 Details of the ExperimentsThe experimental results were obtained using Pytorch [22] . In the experiments involving FMNIST, each client used a CNN-based classifier with two 5×5-convolutional layers and two 2×2-maxpooling layers (with a stride of 2), followed by a fully-connected layer. In the experiments involving CIFAR10, each client used a CNN-based classifier with three 3×3-convolutional layers and two 2×2-maxpooling layers (with a stride of 2), followed by two fully-connected layers; dimension of the hidden layer was 64. In the experiments involving Mini-ImageNet and THUC news, each client fine-tuned a pretrained ResNet18 [13] and learned a TextRNNs [20] , respectively. The optimizers used for model training in the experiments on FMNIST and CIFAR10/Mini-ImageNet/THUC news were the mini-batch stochastic gradient descent (SGD) and Adam [15] , respectively. The learning rate was initially set to 0.001 and then decreased every 10 iterations, with a decay factor 0.5. The number of global communication rounds was set to 200, 500, 100 and 100 for the experiments on FMNIST, CIFAR10, Mini-ImageNet and THUC news, respectively. In all the experiments, the number of local epochs R was set to 2 and the size of a mini-batch was set to 64. The sampling rate (fraction of the clients participating in a training round) was set to 0.1 for the experiments on FMNIST/THUC news, and to 0.2 for the experiments on CIFAR10/Mini-ImageNet. For the sake of visualization, data points in the presented graphs were smoothened by a Savitzky-Golay filter [25] with window length 13 and the polynomial order set to 3.In all experiments, the hyper-parameter µ of the regularization term in FedProx [19] was set to 0.1.In the Power-of-Choice (pow-d) [8] selection strategy, d was set to the total number of clients: 50 in the experiments on FMNIST, CIFAR10 and THUC news, 100 in the experiments on Mini-ImageNet. When running DivFL [2] , we used the ideal setting where 1-step gradients were requested from all client in each round (regardless of their participation status), similar to the Power-of-Choice settings.For FedCor [28] , we followed all settings in the paper and set the annealing coefficient β controlling the sampling strategy to 0.9 as suggested in the paper. For HiCS-FL (our method), the scaling parameter T (temperature) used in data heterogeneity estimation was set to 0.0025 in the experiments on FMNIST and to 0.0015 in the experiments on CIFAR10/Mini-ImageNet. In all experiments, parameter λ which multiplies the difference between clients' estimated data heterogeneity (used in clustering) was set to 10. In all experiments, the number of clusters m was for convenience set to be equal to the number of selected clients K. The coefficient γ 0 was set to 4 in the experiments on FMNIST and CIFAR10 while set to 2 in the experiments on Mini-ImageNet. To group clients, both Clustered Sampling [11] and HiCS-FL (our method) utilized an off-the-shelf clustering algorithm performing hierarchical clustering with Ward's Method.To illustrate and empirically validate Assumption 3.1, we conducted extensive experiments on FMNIST and CIFAR10 with the same model mentioned in Section A.1. In particular, we varied α over 250 values in the interval [0.01, 50] to generate data partitions allocated to 250 clients; entropy of the generated label distributions ranged from 0 to ln 10 (maximum). In these experiments, we allowed all clients to participate in each of 500 training rounds. To facilitate the desired study, in addition to these 250 clients we also simulated a super-client which owns a data set aggregating the data from all the clients (the set of labels in the aggregated dataset is uniformly distributed).In each round, clients start from the initialized global model and compute local gradients on their datasets; the super-client does the same on the aggregated dataset. The server computes and records squared Euclidean norm of the difference between the local gradients and the "true" gradient (i.e., the super-client's gradient). In each round, the difference between the local gradient and the true gradient changes in a pattern similar to what is stated in Assumption 3.1. As an illustration, we plot all such gradient differences computed during the entire training process of a client. Specifically, the server computes the difference between local gradient and the true gradient in each round of training, obtaining 250 × 500 = 12500 data points that correspond to 250 data partitions. For better visualization, we merged adjacent points.The results obtained by following these steps in experiments on FMNIST and CIFAR10 are shown in Figure 5 . For a more informative visualization, the horizontal coordinate of a point in the scatter plot is H(D (k) ), while the vertical coordinate is ∥η t ∇F k (θ t ) -η t ∇F (θ t )∥ 2 . The dashed lines correspond to the curves y = -exp(β [x -H(D 0 )])ρ + κ that envelop the majority of the generated points. In the case of FMNIST, the blue dashed line is parametrized by β = 1.0, ρ = 0.13, and κ = 0.14 while the green dashed line is parametrized by β = 1.5, ρ = 0.025, and κ = 0.022; these two lines envelop 95% of the generated points. In the case of CIFAR10, the blue dashed line is parametrized by β = 2.0, ρ = 0.30, and κ = 0.36 while the green dashed line is parametrized by β = 1.8, ρ = 0.15, and κ = 0.20; as in the other plot, these two lines envelop 95% of the generated points. As the plots indicate, the difference between the local gradient and the true gradient increases as H(D (k) ) decreases, implying that the local gradient computed by a client with more balanced data is closer to the true gradient. Given a batch of samples (x (j,n) , y (j,n) ), the cross-entropy loss is readily computed asEQUATIONEQUATION)where B is the batchsize; l is the number of mini-batches; C is the number of classes; d is the dimension of the hidden space; z(j,n) ddenotes the d-th feature in the hidden space given sample x (j,n) in the j-th batch; w d,c and b c denote the weight of z (j,n) d and the bias for the neuron that outputs the probability of the class c, respectively; and q(j,n) cis the corresponding output logit on class c. The gradient of the bias b i given sample x (j,n) can be computed by the chain rule asEQUATIONwhereEQUATION)Then ∂LEQUATIONEQUATIONEQUATION)By plugging Eq. 18 and 19 in Eq. 15, we obtainEQUATIONA.4 Expectation of the local update ∆b (k) By combining Eq. 4 and 5 and taking expectation, we obtainE ∆b (k) i = - η Bl l j=1 B n=1 R r=1 E ∇ bi L (j,n,r) CE = η R r=1 P{i = y (j,n) }I{i = y (j,n) } 1 Bl l j=1 B n=1 c̸ =i exp(q (j,n,r) c ) C c=1 exp(q (j,n,r) c ) -η R r=1 P{i ̸ = y (j,n) }I{i ̸ = y (j,n) } 1 Bl l j=1 B n=1 exp(q (j,n,r) i ) C c=1 exp(q (j,n,r) c ) = η R r=1 D (k) i c̸ =i E (x,y)∼B -cexp(qEQUATIONNote thatEQUATIONA comparison to L CE in Eq. 13 reveals that as L CE decreases during training, so doesC i=1 E i . Given an untrained/initialized neural network model, E 0 i = 1/C for ∀i ∈ [C], i.e., C i=1 E 0 i = -1 C-1 + C C-1 = 1. At global round T , if L * CE = 0, then C i=1 E T i = -C C-1 + C C-1 = 0. A.5 Privacy of D (k)According to Eq. 6, the server is able to obtain C linear equations from each client,EQUATIONEQUATIONwhere C denotes the number of classes. Suppose E[∆b i , the variables in the aforementioned equations, cannot be determined uniquely since there are C variables and C + 1 equations. Therefore, the server is unable to infer clients' true data label distribution and the privacy of D (k) is protected.In Section A.3 we derived an expression for the gradient of the bias in the output layer given a single sample (x (j,n) , y) in the mini-batch. It is worthwhile making the following two observations:• the sign of the y (j,n) -th component of ∇ b L (j,n) CE (x (j,n) , y (j,n) ) is opposite of the sign of the other components; and• the y (j,n) -th component of ∇ b L (j,n) CE (x (j,n) , y (j,n)) is equal in magnitude to all other components combined. EQUATIONData heterogeneity can be captured via entropy,H(D (k) ) = - C c=1 D (k) i ln D (k)i , where higher H(D (k) ) indicates that client k has more balanced data. However, since we do not have access to the client's data distribution, we instead define and use as a measure of heterogeneity Ĥ(D (k) ) ≜ H(softmax(∆b (k) , T )), where softmax(∆b (k) , T ) i = exp(∆bEQUATIONand where T denotes the temperature of the softmax operator. Suppose there are two clients, u and k, with class-balanced and class-imbalanced data; let D (u) and D (k) denote their data label distributions, respectively, while D(u) and D(k) are computed by softmax(∆b (u) , T ) and softmax(∆b (k) , T ). Without a loss of generality, we can re-parameterize D(u) asEQUATIONwhere U = [ 1 C , . . . , 1 C ] denotes uniform distribution; i-th component of Z i is 1 while the remaining components are 0; ϵ and ϵ i are all non-negative such that ϵ + C i=1 ϵ i = 1. We can always set min j ϵ j = 0; otherwise, let ϵ ′ = ϵ + min j ϵ j and ϵEQUATION)We will find the following lemma useful.Lemma A.1 For two probability vectors p and q with dimension C, the Kullback-Leibler divergence between p and q satisfies KLD(p||qEQUATIONwhere ∥p -q∥ 1 = C i=1 |p i -q i |.For the proof of the lemma, please see [10] . Applying it, we obtainEQUATION)Combining Eq. 28 and Eq. 30, we obtainEQUATIONBy taking expectations of both sides,EQUATION) Since D(k) -U 2 2is convex (composition of the Euclidean norm and softmax), according to Jensen's inequalityEQUATIONwhereEQUATION)Selecting T such that ηR D(k) i C c=1 E c -E i /T is sufficiently small and applying the first-order Taylor's expansion of e x around 0, we obtain (35) whereC j exp ηR D (k) j C c=1 E c -E j /T = C j 1 + ηR C j D (k) j C c=1 E c -E j /T = C,EQUATIONSubstituting Eq. 36 for the second term on the right-hand side of ineq. 33 leads toEQUATION)Now, considerEQUATIONTaking expectations of both sides,EQUATIONThe above inequality holds component-wise, so for the j-component (ϵ j = 0)EQUATIONTherefore,EQUATIONTaking absolute value of both sides yieldsEQUATIONBy applying the triangle inequality we obtainEQUATIONEQUATIONNote thatEQUATIONTherefore,C i=1 D (k) i C c=1 E c -E i 2 ≥ C c=1 E c 2 C i=1 D (k) i - 1 C 2 + 2 C c=1 E c C i=1 D (k) i - 1 C 1 C C c=1 E c -E i = C c=1 E c 2 C i=1 D (k) i - 1 C 2 + 2 C c=1 E c C i=1 D (k) i C C c=1 E c - 1 C 2 C c=1 E c + E i C -D (k) i E i = C c=1 E c 2 C i=1 D (k) i - 1 C 2 + 2 C c=1 E c 1 C C c=1 E c - 1 C C c=1 E c + 1 C C i=1 E i - C i=1 D (k) i E i ≥ C c=1 E c 2 C i=1 D (k) i - 1 C 2 + 2 C c=1 E c 1 C C c=1 E c -max j E j ≥ C c=1 E c 2 C i=1 D (k) i - 1 C 2 -2δ.(47) Substituting the above expression in Eq. 33, we obtainEQUATIONand, therefore,EQUATIONwhereC = ηR(ηR+C 2 T ln C) C 2 T 2 . ■ A.Lemma A.4 (Lemma 2 in [33] ) Instate Assumptions 3.1, A.2 and A.3. For any step size η such that η ≤ 1 8LR , for any client k it holds thatEQUATION)Proof of Lemma A.4:EQUATION)Unrolling the recursion yieldsE θ t,r k -θ t 2 ≤ R r=1 1 + 1 R -1 r-1 η 2 σ 2 + 6Rη 2 σ 2 k + 6Rη 2 E ∇F (θ t ) 2 ≤ (R -1) 1 + 1 R -1 R -1 η 2 σ 2 + 6Rη 2 σ 2 k + 6Rη 2 E ∇F (θ t ) 2 ≤ 5Rη 2 σ 2 + 6Rσ 2 k + 30R 2 η 2 ∇F (θ t ) 2 .(56) ■ A.7.2 Proof of Theorem 3.4The model update at global round t is formed asEQUATIONwhere θ t+1 and θ t denote parameters of the global model at rounds t + 1 and t, respectively, and θ t,r k denotes parameters of the local model of client k after r local training epochs. LetEQUATION)Taking the expectations (conditioned on θ t ) of both sides, we obtainE F (θ t+1 ) = E F (θ t -η∆ t ) (a) ≤ F (θ t ) -η ∇F (θ t ), E ∆ t + L 2 η 2 E ∆ t 2 = F (θ t ) + η ∇F (θ t ), E R∇F (θ t ) -R∇F (θ t ) -∆ t + L 2 η 2 E ∆ t 2 = F (θ t ) -Rη ∇F (θ t ) 2 + η ∇F (θ t ), E R∇F (θ t ) -∆ t A1 + L 2 η 2 E ∆ t 2 A2 .(59) Inequality (a) in the expression above holds due to the smoothness of F (•) (see Assumption A.2). Note that the term A 1 can be bounded asA 1 = ∇F (θ t ), E R∇F (θ t ) -∆ t = ∇F (θ t ), E R∇F (θ t ) - 1 K k∈S t R r=1 g k (θ t,r k ) = ∇F (θ t ), E R∇F (θ t ) - N k=1 R r=1 ω t k ∇F k (θ t,r k ) = N k=1 ω t k √ R∇F (θ t ), - 1 √ R E R r=1 ∇F k (θ t,r k ) -∇F (θ t ) (a) = R 2 ∇F (θ t ) 2 + 1 2R N k=1 ω t k E R r=1 ∇F k (θ t,r k ) -∇F (θ t ) 2 - 1 2R N k=1 ω t k E R r=1 ∇F k (θ t,r k ) 2 (b) ≤ R 2 ∇F (θ t ) 2 + 1 R N k=1 ω t k E R r=1 ∇F k (θ t,r k ) -∇F k (θ t ) 2 + 1 R N k=1 ω t k E R r=1 ∇F k (θ t ) -∇F (θ t ) 2 - 1 2R N k=1 ω t k E R r=1 ∇F k (θ t,r k ) 2 (c) ≤ R 2 ∇F (θ t ) 2 + N k=1 ω t k R r=1 E ∇F k (θ t,r k ) -∇F k (θ t ) 2 + N k=1 ω t k R r=1 E ∇F k (θ t ) -∇F (θ t ) 2 - 1 2R N k=1 ω t k E R r=1 ∇F k (θ t,r k ) 2 (d) ≤ R 2 ∇F (θ t ) 2 + L 2 N k=1 ω t k R r=1 E θ t,r k -θ t 2 + R N k=1 ω t k σ 2 k - 1 2R N k=1 ω t k E R r=1 ∇F k (θ t,r k )where equality (a) follows from ⟨x, y⟩ = EQUATION)Furthermore,A 2 = E   1 K k∈S t R r=1 g k (θ t,r k ) 2   = E   N k=1 I{k ∈ S t } K R r=1 g k (θ t,r k ) 2   = E   N k=1 I{k ∈ S t } K R r=1 g k (θ t,r k ) -∇F k (θ t,r k ) + ∇F k (θ t,r k ) 2   (a) = E   N k=1 I{k ∈ S t } K R r=1 g k (θ t,r k ) -∇F k (θ t,r k ) 2   + E   N k=1 I{k ∈ S t } K R r=1 ∇F k (θ t,r k ) 2   (b) ≤ E N k=1 I{k ∈ S t } K R r=1 g k (θ t,r k ) -∇F k (θ t,r k ) 2 + E   N k=1 I{k ∈ S t } K R r=1 ∇F k (θ t,r k ) 2   (c) ≤ Rσ 2 + N k=1 ω t k E R r=1 ∇F k (θ t,r k ) 2 ,(62) where equation (a) holds because E g k (θ t,r k ) -∇F k (θ t,r k ) = 0, inequality (b) stems from the Jensen's inequality, and inequality (c) is due to Assumption A.3.EQUATIONEQUATIONBy rearranging and summing from t = 0 to t = T -1 we obtainEQUATION)Let θ * denote the optimal model's parameters, i.e., F (θEQUATIONwhereA 1 = Rη 1 2 -30L 2 R 2 η 2 , A 2 = 60L 2 R 3 η 3 +2Rη Rη(1-60L 2 R 2 η 2 ) and Φ = (10L 2 Rη 2 +Lη)σ 2 1-60L 2 R 2 η 2 . ■ A.The objective function used by FedProx [19] isEQUATIONwhere θ t,r k is the vector of client k's local model parameters in the r-th local epoch at global round t. Therefore, contribution of sample (x (j,n) , y (j,n) ) to the gradient of L prox in local epoch r is EQUATIONEQUATIONTherefore,EQUATIONand thusEQUATIONTaking expectation of both sides yieldsEQUATIONwhere c r = (1 -η t µ) r-1 > 0 provided η t and µ are sufficiently small. Therefore, the expectation of the local update of bias in the output layer satisfiesEQUATIONwhere C = R r=1 c r . Eq. ( 73) is similar to the expression for the expectation of the local updates of bias when applying FedAvg presented in the main paper; clearly, the analysis of HiCS-FL done in the context of FedAvg extends to FedProx.For FedDyn [1] , the objective function in local epoch r at global round t isEQUATIONwhere R denotes the total number of local epochs. The first order condition for local optima impliesEQUATIONand, therefore, EQUATIONEQUATIONit holds that ∂L t,2EQUATIONand ∂L t,3dyn ∂b i = -µ t-1 τ =1 ∆b τ i + µη t -µ t-1 τ =1 ∆b τ i -(µ + µ 2 η t ) t-1 τ =1 ∆b τ i = -µ(1 + µη t ) 2 t-1 τ =1 ∆b τ i .(79) By induction, ∂L t,rEQUATIONTherefore, the expectation of the local update of bias in the output layer at round t can be computed asEQUATIONEQUATION)Since the objective function of E ∆bEQUATIONwhere η 1 is the learning rate at global round t = 1. Then,EQUATIONwhereEQUATIONEQUATIONandEQUATIONBy induction,EQUATIONwherea t = (1 + µη t ) R -1 and a = a 1 a t 1 + t-3 i=0 t-1 τ =2 I(τ + i < t)τ +i i=τ a s > 0. After comparing Eq. ( 89) with its counterpart in the case of FedAvg, we conclude that the previously presented analysis of HiCS-FL extends to FedDyn.Moon [16] relies on the objective function with a contrastive termEQUATIONwhere z (j,n) denotes the output of the feature extractor of the local model θ t k , z(j,n)glob is the output of the feature extractor of the global model θ t , and z (j,n) prev is the output of the feature extractor of the local model in the previous round θ t-1 k . Since the contrastive term does not depend on the parameters of bias in the output layer, it holds thatEQUATIONSince the expectation of the local updates of bias in the output layer coincides with the one in case of FedAvg, previously presented analysis of HiCS-FL extends to Moon.Optimizers beyond SGD utilize different model update rules which in principle may lead to different properties of the local update of the bias in the output layer. However, for several variants of SGD, the properties of the local update of the bias remain such that our presented analysis still applies.In each local epoch r, SGD with momentum updates the model according toEQUATIONEQUATIONEQUATIONwhere m t,r k denotes the momentum in the r-th local epoch, µ is the weight for the momentum, and m t,1 k = ∇L t,1 CE . ThenEQUATIONwhereEQUATIONEQUATIONEQUATIONTherefore,EQUATIONand thus we haveEQUATIONSimilar to the discussion in the previous section,EQUATIONwherea = η t R r=2 µ r-1 + (1 -µ) r τ =2 µ r-τ + 1 > 0.Similar result is obtained when SGD applies Nesterov acceleration as long as the optimizers are not using second-order momentum.Recall that the two observations regarding the gradient of L CE still hold when training the model with an adaptive optimizer such as Adam [15] . However, Adam updates the model differently from SGD. In particular, each entry of the gradient has an adaptive learning rate tied to its magnitude. With an SGD optimizer, the magnitude of the i-th entry of the local update of bias ∆b (k) is approximately proportional to the fraction of the samples with label i, DEQUATIONHowever, this observation does not hold when using the Adam optimizer for the local update because each entry has a different learning rate η t,i and thusEQUATIONAlthough the magnitude of E ∆b EQUATION)Suppose client k has highly imbalanced data, i.e., H(D (k) ) is small. Then the maximal componentmax i D (k) iis much larger than the other components; in fact, it is likely to have only one positive component in the local update of bias ∆b (k) . On the contrary, suppose client u has balanced data and thus H(D (u) ) is large. The maximal component max i D(u) i is then very close to the other components, and it is likely to observe larger number of positive components in the local update of ∆b (u) . While characterizing P(E[∆b (k) i ] > 0) appears challenging, we can empirically infer that client u with more balanced data has a local update of bias ∆b (u) with more positive components. With Ĥ(D (u) ) ≜ H(softmax(∆b (u) , T )),EQUATIONĤ(D (u) ) is more likely to be larger than Ĥ(D (k) ). The examples of estimated entropy when utilizing Adam as the optimizer are provided in Section. A.12.To generate non-IID data partitions we follow the strategy in [35] , utilizing Dirichlet distribution with different concentration parameters α to control the heterogeneity levels. In particular, the number of samples with label i owned by client k is set toX (k) i Ni N j=1 X (j) i,where X We compare the communication and computational costs of HiCS-FL with those of the competing methods, including Power of Choice (pow-d) [8] , Clustered Sampling [11] and DivFL [2] , and map them against random sampling, as shown in In our experiments, DivFL has consistently required the longest training time and memory usage due to its dependence on the submodularity maximizer. FedCor [28] cliams that only partial clients participating in the global update after warm-up stage but still needs all clients to perform inference for computing validation loss in the warm-up stage. Our proposed method, HiCS-FL, does not require any additional transmission of model parameters; furthermore, in HiCS-FL the server clusters clients based on their local updates of the bias in the output layer, which is low-dimensional and model-agnostic. Overall, HiCS-FL requires negligible computational overhead to significantly improve the performance of non-iid Federated Learning.To further illustrate the proposed framework, here we show a comparison between the estimated entropy of data label distribution and the true entropy. Specifically, Figures 8 and 9 show that the entropy estimated by the proposed method is close to the true entropy; the experiments were conducted on FMNIST and Mini-ImageNet, using SGD and Adam as optimizers, respectively. As stated in Theorem 3.3, the clients with larger true entropy are likely to have lager estimated entropy.In case where the model is trained with Adam, estimated entropy of data label distribution is not as accurate as in the case of using SGD. Figures 10 and 11 compare the performance of estimating entropy with SGD and Adam optimizers for the same setting of α. Notably, as shown in the figures, the method is capable of distinguishing clients with extremely imbalanced data from those with balanced data. Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.38th Conference on Neural Information Processing Systems (NeurIPS 2024).,(60) 


Decision: The paper introduces a client selection method that operates over  non-i.i.d. datasets in FL systems under communication constraints. The proposed algorithm estimates the heterogeneity of a client using updates from the network’s output layer.  Clients are clustered based on this information, and those with more balanced data are preferentially sampled. Experimental results demonstrate the scheme's superior performance compared to state-of-the-art methods.

Reviewers were generally satisfied with responses provided by the authors, even if novelty is somewhat limited. The authors should surface some of their responses in their submission (e.g., additional experiments reported, the relationship to Auxo, discussion on overhead, and the several clarification questions asked by the reviewers). Limitations should be explicitly discussed.

Reviews: review: paper_summary: paper_summary: The authors propose a novel client selection method to address federated learning scenarios where clients exhibit varying degrees of data imbalance. The authors estimate the label distribution entropy of clients on the server side using the gradient of the output layer's bias. Based on this estimation, they cluster and sample clients, utilizing those with more balanced data to train the global model. Experiments in the paper demonstrate that the proposed method, HiCS-FL, achieves better and faster convergence compared to baselines. The authors provide detailed theoretical proofs to support their approach.
strength_weakness: strength_weakness: See questions.
questions: questions: 1. Is it due to privacy concerns that clients cannot directly compute and upload label distribution entropy to the server? If so, does estimating label distribution entropy through $ \Delta \mathbf{b}^{(k)} $ still pose privacy concerns for clients? Additionally, could you consider conducting an experiment comparing the federated learning performance (accuracy, convergence speed, etc.) using the true label distribution entropy versus the gradient-based label distribution entropy estimation proposed in this paper?

2. Why wasn't an experimental analysis conducted on different values of $ \lambda $ in Eq. 9?

3. Will the final trained global model be used for inference on local clients? If so, wouldn't the client selection method that prefers more data-balanced clients be suboptimal for local clients with severe data imbalance?
limitations: limitations: The authors mention leaving studies of system heterogeneity to future work in the first paragraph of the Introduction, but there is no further discussion on the limitations of the proposed approach. There are no concerns raised about the societal impacts of the research.
paper_summary: paper_summary: The paper addresses the challenges posed by non-IID data in FL systems, particularly under communication constraints where only a small fraction of clients can participate in each training round. It introduces HiCS-FL, a novel client selection method. HiCS-FL estimates the statistical heterogeneity of a client’s data using updates from the network’s output layer to cluster and sample clients more effectively. Based on the heterogeneity information, clients are clustered, and those with more balanced data are preferentially sampled. Experimental results demonstrate HiCS-FL's superior performance compared to state-of-the-art methods.
strength_weakness: strength_weakness: - The paper primarily focuses on label imbalance when discussing data heterogeneity. However, other types of imbalances, such as feature imbalances within the data, are not considered. Addressing these other forms of data imbalance could provide a more comprehensive solution.

- The method involves collecting updates from the fully connected layer, which may inadvertently reveal the label distribution of a client’s data. This could raise privacy concerns, as it potentially compromises the confidentiality of the client's data distribution.

- The paper does not compare HiCS-FL with other relevant algorithms, such as Oort and Auxo, which also focus on client selection and clustering. Including these comparisons or discussions would provide a clearer picture of HiCS-FL's relative performance and effectiveness.

[1] Oort: Efficient Federated Learning via Guided Participant Selection 
[2] Auxo: Heterogeneity-Mitigating Federated Learning via Scalable Client Clustering

- The assumption that clients can be clustered into M groups, where M is larger than the number of selected clients (K) per round, may not be practical in large-scale systems. Specifically, the paper assumes that all clients are available during the initial clustering phase and that no new clients join the system later, which is unrealistic in dynamic FL environments.

There might be additional overhead associated with the proposed clustering method. The paper should ensure that the experimental comparisons account for this overhead to maintain fairness in the evaluation of HiCS-FL’s performance.
questions: questions: - The rationale behind using gradients from different rounds for clustering clients is not clearly explained. Given that gradients can vary significantly across rounds, further clarification is needed to justify this approach and its impact on the effectiveness of the clustering process.
limitations: limitations: See weakness
paper_summary: paper_summary: The paper address data heterogeneity by clustering clients via the gradients of the output layer to distinguish between clients with balanced from those with imbalanced data. HiCS-FL assigns different importance to the clusters according to their average estimated data heterogeneity. The paper found that there is a correlation between the gradient of the output layer and the label distribution.
strength_weakness: strength_weakness: The approach of clustering to address heterogeneity might be practical as there are simpler methods that can achieve the same goal. 
The experiments were done on a small set of clients, 50. As the number of clients increases, the more pronounced the effect of heterogeneity can have.
questions: questions: Data heterogeneity can be mitigated with pre-training and personalization. Would pre-training make clustering unnecessary?
limitations: limitations: The paper doesn't list limitation but I think a limitation there is the experimental setting. Using  Dirichlet to generate clients is unrealistic (I know this is common in FL papers), and the paper focuses only on classification tasks.
paper_summary: paper_summary: This paper provides an interesting approach HiCS-FL to investigate the client sampling problem in federated learning, especially for the non-iid setting. This paper estimates the clients' data statistical heterogeneity (label distributions) via the client-updated gradients of the output layer's weights. By using this distance information, the server can distinguish which client owns more balanced data.
strength_weakness: strength_weakness: - A potential privacy issue may be raised since it needs to access the individual client update (i.e., gradient) information.  More client's data information may be leaked from the gradient inversion attacks [1]. 


[1] Evaluating Gradient Inversion Attacks and Defenses in Federated Learning, NeurIPS 2021.
questions: questions: - Why the first training rounds set to be $\lceil N / K\rceil$ in line 235?
limitations: limitations: See Weaknesses.

Comments: title: Thank you for your contribution
Experiments with pretrained models
comment: We’re happy to clear up any confusion the reviewer may have had. We’d like to express our gratitude once again for the reviewer’s invaluable assistance in improving our work.
Thanks  for the clarification. I've updated my score.
We’re happy to clear up any confusion the reviewer may have had. We’d like to express our gratitude once again for the reviewer’s invaluable assistance in improving our work.
I appreciate the authors' responses to my questions. I would like to keep my current score.
We’re happy to clear up any confusion the reviewer may have had. We’d like to express our gratitude once again for the reviewer’s invaluable assistance in improving our work, and we are open to incorporating additional tasks, such as regression, in future work.
We appreciate the reviewer for pointing out this potential confusion, which may also affect other reviewers, and for helping us improve our paper. As stated in the General Settings in Appendix A1.1 (line 456), all experiments were conducted on Mini-ImageNet using a pretrained ResNet18. To clarify, we kept the feature extractor of the pretrained model (on ImageNet1K) frozen and only fine-tuned the fully connected layers on the target dataset. As a result, the test accuracy on the target dataset improves from the initial, similar to training from scratch.

The experimental results on Mini-ImageNet indicate performance degradation under severe data heterogeneity. We apologize for the unclear description of our experimental setup and will clarify the pretrained settings in the main paper in the revised version.
> A3.3 As shown by the experimental results on Mini-ImageNet dataset in the paper and the supplementary experiments in this rebuttal, the pretrained models on ImageNet still encounter performance degradation when faced with severe data heterogeneity.

I don't understand the setup for " experimental results on Mini-ImageNet dataset in the paper and the supplementary experiments in this rebuttal", which table should I be looking at? I checked the paper and don't see which dataset did you pretrain with.
Thank you for the authors efforts in addressing the review comments in the rebuttal.

Overall, the authors' responses and new experiments have resolved my concerns.  I am willing to increase the rating.

It is interesting that the authors found relevance between the output layer bias gradients and class distribution entropy, and applied it to federated learning. I hope that in the future, they might consider how to use this in other tasks like regression, not only classification.