Heterogeneity-Guided Client Sampling: Towards Fast and Efficient Non-IID Federated Learning
Authors: Huancheng Chen
Haris Vikalo
Venue: NeurIPS 2024 poster (2024)
Keywords: Federated Learning
Client Selection
Non-IID data
Primary area: other

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.

Helerogeneity-Guidee Client Sampling: Towards Fast and Efficient Non-IID Federated Learning
Huancheng Chen University of Texas at Austin huanchngchhuttexas. edu
Haris Vikalo University of Texas at Austin hvikalodece utexas.ee..... edu
Abstract
Statistical heterogeneity of data present at client devices in federated learning (FL) system renders the training of a global model in such systems difficult. Particu- larly challenging are the settings where due communication resource constraints only small fraction clients can participate in any given round of FL. Recent approaches to training a global model in FL systems with non-IID data have fo- cused 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 heterogene- ity of client's data using the client's update of the network's 's output layer and relies on this information 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 set- tings HiCS-FL achieves faster convergence than state-of-the-ar FL client selection schemes. Notably, HiCS-FL drastically reduces computation cost compared existing selection schemes and adaptable to different heterogeneity scenarios.
1 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 inde- pendent and identically distributed (IID) data. However, compared 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 reduce the variance and improve convergence of FL on non-IID data [6. 9, 14, 17, 19, 30].
On another note, constraints communication resources and therefore the number of clients that may participate in training additionally complicate implementation of FL schemes. It would be particularly unrealistic to require regular contributions training from all the clients in large- scale cross dross-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
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

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 ss 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. 111. 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 learming and propose novel adaptive clustering 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 the clients' data and, based on that, decide which clients to sample. In particular, HiCS-FL enables heterogeneity-guun -giidedd 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, 1O, Mini-ImageNo 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.
sampling method for identifying and selecting
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) novel clustering-based algorithm for neetrooeneityygguidee client selection; (3) extensive simulation results demonstrating HiCS-FL pro- vides significant improvement in terms of convergence speed and variance over competing approaches; and (4) theoretical analysis of the proposed schemes.
2 Background and Related Work
Assume the cross-device federated learning setting with N clients, where client ls owns private local dataset B& with [Bx] samples. The plain vanilla FL considers the objective
N
min F(O) 2
Σ
px (o)).
k=1
where 9 denotes parameters of the global model, FA (0) the loss (empirical risk) of model on
B, and pk denotes the weight assigned to client k. Σx= Pk: = 1. In FedAvg, the weights are set to Pk= = Ba E [B.]- In training round t, the server collects clients' model updates of formed by training on local data and aggregates them to update global model as gt+1 =
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 f by St. In departure from FedAvg, FedProx [19] proposes an alternative strategy for sampling clients based on 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 gr-1 = K Ekest 0g. 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 1..
AFL [12] is the first study to utilize local validation loss as a value function for computing client sampling probabilities; 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-Uhlenbedd process and proposes a strategy, optimal under that assumption, for selecting clients with significant weight updates.
Sewer-of-Choice [8] takes a step further to propose a greedy approach to
2

---

The clustering- basedd 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 the global model (high loss, large update or low class- baassimmallance. Efficient and effective client selection in FL remains an open challenge, motivating the heterogeneity-guin adaptive client selection method presented next.
3 HiCS-FL:: Federated Learning via Hierarchical Clustered Sampling
Existing client sampling methods including Clustered Sam- pling [11] and DivFL [2] aim 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 participa- tion) while minimizing the variance
Modie A.
N
Σ VF)(0') N
1
E (O) RES
(2)
signals
prediction
one hot tabel
K
1=1
Figure 1: The last two network layers.
Clustered Sampling, for instance, groups N clients into K clusters based 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, number of studies in centralized learning has shown that iabalammbledeee 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-FI. is more likely to select diverse clients (i.e., sample uniformly across clusters, thus reducing to the conventional clustered sampling strategy).
3.1
Class-imallance Causes Objective Drift
number of studies explored detrimental effects of non-IID training data on the performance of global model learned via FedAvg. An example is SCAFFOLD [14] which demonstrates objective drift in non-III FL. manifested through large differences between local models 0% 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 not necessarily leading towards the optimal global model 0*. The optimal model 0-, 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(()) denote the gradient of F(O") given the global model 0² at round l; the difference between VF((') and the local gradient V FF 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.
Assumption 3.1 (Bounded Dissimilarity.) Gradient V Fx((0') of the k-th local model at global round is such that
-F S <K- - ( = =
where (D()) is the data label distribution of client k, Do denotes uniform distribution, H(-) is Shannon's 's entropy of stochastic vector, and 3 0, p 0.
3

---

The assumption commonly encountered in literature is recovered by setting the right-hand side of (3) of max; of Intuitively, if the data label distribution of client is highly imbalanced (i.e., (D(() is small), the local gradient F (0*) may significantly differ from the global gradient VF(0') (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.
3.2 Estimating Client's Data Heterogeneity
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]. this section, we demonstrate how to efficiently and accurately estimate data heterogeneity using local updates of the output layer of neural network in a classifica- tion task. Figure illustrates the last two layers in typical neural network. The prediction q E IR is computed by forming weighted average of signals 2 E R. utilizing the weight matrix W E RCXL and bias b E R.C
3.2.1 Local updates of the output layer
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 between local updates of the output layer's '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 widely used multi-class classification framework. The gradient is computed by averaging contributions of the samples in mini-batches, i.e., Vb = E/=1 En=1 VbCli (x((.)), y(d""),, where B denotes the batch size, / is the number of mini- batches, x(i.) is the n-th point in the j-th mini- and y(d.m) E [C] is its label. The contribution of x x(im)) 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)
the local update. To pursue the latter, we focus the correlation
(xllin), y(dm)) = lii = =y(nn) Eaye exp(1ein)) Vo, c(im) (
expla. (imn) Ec exp(gl (im))
+ Lii # y(m))
Ee exp(g! (im)))
where I[-} is an indicator, q(inn) = W - z(i.n + b is the output logit for signals z(in) E R- corresponding to training point (x (xlix), ((Ui)),((nn) y(im)) (see Fig. 1), and where C' denotes the number of classes. We make the following observations: (1) the sign of -(. component of Vb.li is opposite of the sign of other components; and (2) the (( component of ((, is 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 of FedAvg, the selected client k starts from the global model 0 and proceeds to compute local update in R local epochs employing an SGD optimizer with learning rate 17. According to Eq. 4, the i-th component of local update (b()) is computed
B
Ab (A)
b. (j.m.r)
BI
where Vo, ((...r) denotes the gradient of bias at local epoch Note that the local update of client k
Ab(A) (b), is dependent on the label distribution of client k's data, (Ak) = [D(*)] " (TT and the label-specific components of q(í.)) 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 be the subset of local data B that excludes points with label i. Lets " (x) e [0, 1]c be the softmax output of a trained neural network for a training point (x. y) B The i-th component sx)) S; (x), indicates the level of confidence in (erroneously) classifying x as having label i. For convenience, we define E, ((xx))=B-1 [x)] Vi E [C].
In an untraineed/mittaalzzed neural network where classifier makes random predictions, & = 1/C; as training proceeds, & decreases. By taking expectation and simplifying. we obtain (details provided in Appendix A.4)
C -
E
n/t D
Σ c=
E
&
(6)
where (* denotes the true fraction of samples with label i in client k's data, Σ D( 1.
3.2.2 Estimating local data heterogeneity
We quantify the heterogeneity of clients' data by an entropy-like measure defined below. Let ((k) denote the label distribution of client k's data; its entropy is defined as ((((k) a
( In D (k) < In C. Recall that more balanced data results in higher entropy, and that (D) takes the maximal value when ((A) is uniform. The server does not know D(k) and therefore cannot compute H (D(k)) directly. We define
((()())) _ H (softmax (atttmaxx(b())) (Ab()
here T' is scaling syper-parameter (so-called temperature). Note that even though compute H (D(K)) to characterize heterogeneity, (A))) and & 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 classifica- tion task over C classes. Let D(")) and D(k) denote data label distributions of an arbitrary pair of clients and k, respectively. Moreover, let U denote the uniform distribution, and let nj and R be the learning rate and the number of local epochs, respectively. Then
2
(D(W))) [ ((((()) 2
R CT
&g
|DII] - U
2
R T
p
U
C8.
(8)
where C = III and å max; CST2
E & - The proof is provided in Appendix A.6.
E.j.
As illustration, consider the scenario where client has balanced dataset while the dataset of client & is imbalanced; then ([((k) U is relatively large compared to |]D(") Ull The bound in (8) also depends 8, which is reflective of how misleading on average can a class be; small 5 suggests that no class is universally misleading. As shown in Appendix A.4, during training gradually decreases to 0 as E & decreases to 0.
3.2.3 Generalizing beyond FedAvg and SGD
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 appendix A.8 and A.9.
3.3 Neeterogeneity-guide Clustering
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
5

---

Algorithm HiCS-FL
Input:
sample group G'n. according to n';
Datasets distributed across N clients, the number of clients sample K, total global rounds T.
10:
sample client & in G based on pmn S* 8" U k:
11:
1: Initialize updates of bias (b < 0 vk E [M].
12:
end while
global model 8 <- ( So [N]-
13:
end
2: for f 1...., T do
14:
for E do
3: i [N/K] then
15: 0% + LocalUpdate(f)) (b(") E 8¹
4: S <- randomly sample min( K. |5
16:
end for
clients from So. update So <- So S"
17:
0f+1 + π Enes 0i: Ab Ab(R) wk
5: else
18:
6:
estimate (D(()) and cluster N clients into
19:
end for
M groups based on Eq. 9;
Output:
7:
S +- 0:
The global model (T+1
8:
while (S < K do
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 hetero- geneity H (D))). In particular, the proposed measure of distance between clients u and & that we use to form clusters is defined as
Áb (w) (b( Abb()) A ) + - -
Distance(u, k) = arc cos
(9)
where 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-paramen our experiments). For large X, 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.
(set to 10 in all
3.4 Hierarchical Clustered Sampling
To select K out of clients in an FL system, we first organize the clients into M 2 K groups via the proposed Hierarchical Clustered Sampling (HiCS) technique. In particular, during the first [N/K K training rounds the server randomly (without replacement) selects clients and collects from them local updates of (b); the server then estimates H' (D(+)) for each selected client k and clusters the clients using the distance measure defined in Eq. 9. Let Gi..... GM denote the resulting M clusters global round t. and let H TG EkEG.. H' (D(k)) characterize the average heterogeneity of clients in cluster probability vector T", and then from the selected cluster selects a client according to the probability vector Pt.. The two probability vectors T' and pt, defined as
E [M]. Having computed I1 HiCS selects cluster according to the
exp(*+Hi)
exp(>' (H(r)
Pki
G.
Σ Em exp(1 H
Em= exp(1* H
Etec.,. Pk
ΣREG Pk
(10)
where ki...., ki... are the indices of clients in cluster Gm, -0=1000=110000 +) denotes an annealing syper-parameter, and T is the number of global rounds. The annealing parameter is scheduled such that 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 7 as O, when the server samples the clusters uniformly. The described procedure is formalized as Algorithm 1.
3.5 Convergence Analysis
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 Fx(.) is L-smooth for all k E [N]. Let Oe denote parameters of the global model and let F-)) be defined as in Eq. I. Furthermore, assume the stochastic gradient estimator gt. (0') is unbiased and the variance is bounded such that EE (0') F& (0*)) < o². Let 1 and It be the learning rate and the number of local epochs, respectively. If the learning rate is such that
n S SLR R 2 2, then
( -((0)
1 (F(00)-
T-1 N
min
|= &
A1 + A2
Σ
1
+
(11)
t=0
where A1, A2. T are positive constants, and w/ is the probability of sampling client is at round t.
Note that only the second term in the parenthesis on the right-hand side of the bound in Theorem 3.4 related to the sampling method II. Under Assumption 3.1,
N
N
exp (3H(D())) w exp (DH(
woOk SK
HII.
(12)
If the server samples clients with weights proportional to pk. the statistical heterogeneity of the (X
entire FL system may be characterized by Hs = -h Pk exp(,(H((D)) p. If all clients have class- imbalanced data, Hs 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
Ph. exp(y' (tn(("''''(0 (D())) entropy, the proposed HiCS-FL leads to w - exp(y"
When training starts, Hn is large because the server tends to sample clients with higher pk exp(?" (C(/""H((HHH)) (D(())); as 7 decreases, Hn eventually approaches Hs. Further details and the proof of the theorem are in Appendix A.7.
Experiments
Setup. We evaluate the proposed HiCS-FL algorithm on four benchmark datasets (FMNIST, CI- FARIO, 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 a to generate data partitions for a single experiment, leading realistic scenario of varied data heterogeneity across different clients. To quantify the performance of the tested methods, we use two metrics: (I) average training loss, and (2) test accuracy of the learned global model. For better visualization, data points in the results are smoothened by Savitzky-Gole filter with window length 13 and the polynomial order set 3. Further details of the experimental setting and a visualization of data partitions are in Appendix A.1 and A.10.
4.1 Comparison on Test Accuracy and Training Loss
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 a 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 HM monotonically decreases as we go through settings (1) to (3). For 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 shows that HiCS-FL outperforms other methods across different settings, exhibiting the fastest convergence rates and the least amount of variance. Particularly significant the acceleration of convergence in setting (1) where 20% of the participating clients have balanced data. Figure 3 shows that HiCS-FL. helping achieve significant reduction of training variations (as expected, see Section 3.5) as evident by smooth loss trajectory.
7

---

Barving SARPO CA
gola
(a) FMNIST (1)
(b) FMNIST (2)
(c) FMNIST (3)
(d) CIFARIO (1)
(C) CIFARIO (2)
(f) CIFARIO (3)
Figure 2: Test accuracy for the global model groups of data partitions of FMNIST and CIFARIO.
Clabte Sar
Dualtrad Samoodey
(a) FMNIST (1)
(b) CIFARIO (1)
(c) Mini-ImageNeet (1)
Figure 3: Training loss of HiCS-FL compared to four baselines for setting (1) on the three datasets.
CIFARIO. Here we compare the performance of HiCS-FL to FedProx [19] running CNN model with Adam optimizer on the task of training an FL system with 50 clients, where 20% of clients are selected to participate in each training round. Similar to the experiments on FMNIST, sets of the concentration parameter a are considered: (1) {0.001, 0.01, 0.1, 0.5, 1}: (2) {0.001, 0.002, 0.005, 0.01, 0.5): (3) (0.001, 0.002, 0.005, 0.01, 0.1). The interpretation of the scenarios emulated by these setting is as same as the FMNIST experiments. Figure demonstrates improvement of HiCS-FL over all the other methods. HiCS-FL exhibits partieularly significant improvements in settings (2) and (3), where 80% of the clients with extremely imbalanced data benefit from 20% of the clients with either balanced or mildly imbalanced data. The advantage of HiCS-FL in setting (1) where all clients have relatively high data heterogeneity is relatively modest (see Fig.2.(d)) because the system's 's Hs is relatively large (see discussion in Section 3.5).
Mini- Imageem As in the Mini- Mini-ImageNee experiments, we compare HiCS-FL to FedProx running 18 with Adam optimizer but now consider training of an FL system with 100 clients, where 20% of the clients are selected to participate in each round of training. We consider two settings of the concentration parameter a: (1) (0.001, 0.01, 0.1, 0.5. } and (2) {0.001, 0.005, 0.01, 0.1, 1). Setting (1) emulates the scenario where clients have a range of heterogeneity profiles, from extremely imbalanced, through mildly imbalanced, 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 H for setting (1) is larger than H2 for setting (2), which is reflected in 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 different domain, we conduct experiments involv- ing text classification on the THUC news dataset Chinese language (10 labels). Similar to the afore- mentioned experiments, we allocate data to 50 clients by emulating heterogeneous data distributions scenarios with parameter a set to: (1) {0.001, 0.01, 0.1, 0.2,1); (2) (0.001, 0.002, 0.01, 0.1, 0.1.0.):: and (3) (0.001, 0.002, 0.005, 0.01, 0.1). We trained TextRNNs [20] with BiLSTM architecture as
8

---

Table Test accuracy (%) for the global model on 3 groups of data partitions of THUC news dataset.
Schemes
Random
Pow-of-Choice
CS
DivFL
FedCor
HiCS-FL
settin (1)
78.9
80.0
80.6
73.0
81.2
83.2
setting (2)
74.9
75.4
82.8
68.9
81.3
83.9
settin (3)
72.7
66.5
79.4
72.1
76.4
79.7
Table 2:
The number of communication rounds needed to reach a certain test accuracy in the experiments on FMNIST, CIFARIO, Mini-ImageNet and THUC News. All results are for the concentration parameter setting (2).
FMNIST
CIFARTO
Mini- ImagelmageNett
THUC news
Schemes
acc = 0.75
speedup
acc 0.6
speedup
acc = 0.5
speedup
acc = 0.8
speedup
Random
149
1.0x
898
I.0x
191
1.0x
83
1.xx
pow-d
1.8f
1037
0.9
432
0.44
109
0.80
CS
114
1.3f
748
1.2f
186
1.0x
74
1.1f
DivFL
478
0.31
1417
0.6
726
0.3 4
289
0.3
FedCor
88
1.71
711
1.31
229
0.81
100
0.8
HiCS-FL
60
2.51
123
7.31
86
2.21
27
3.11
the classifiers using Adam optimizer. The test accuracy of the global model trained with different schemes for 100 global rounds, reported in Table I, show that our method outperforms baselines in all the settings, demonstrating efficacy of our proposed algorithm in simple NLP task.
4.2 Accelerating the Training Convergence
au tbs mddes)
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 them in Table 2. As can be seen from the table, HiCS-FL significantly reduces the number of communication rounds needed reach target test accuracy. On FMNIST, HiCS-FL needs 60 rounds to reach test accuracy 0. 0.75, achieving 2.5 times faster than the random sampling scheme. On CIFARIO, HiCS-FL requires only 123 rounds to reach 0.6 test accuracy, which 7.3 times faster than random sampling. Significant speedup appears on THUC dataset, in which HiCS-FL only needs 27 rounds to achieve 0.8 test accuracy, 3.1 times faster than the baseline. Acceleration on Mini Mini-Imager ImageNet is relatively modest HiCS-I FL still outperforms other methods, and does so up to times faster than random sampling.
the considered four datasets, and display
pon dn
Der
Lettigg
fouunnno// de pas
Lander Serpr
Lettinng (1)
Figure 4: MinilmageNet acc.
Table also shows that HiCS-FL provides the reported improvements without introducing major computational and communication overhead. The only additional computation is due to estimating data heterogeneity and performing clustering utilizing bias updates, which scales with the total number of classes but does not increase with the size of the neural network model |0*|, Remarkably, HiCS-FL outperforms pow-d, Clustered Sampling, DivFL and FedCor in terms of convergence speed, variance and test accuracy while requiring significantly less computations. More details are provided in Appendix A.11.
4.3 Number of Clustering Groups
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 statistical heterogeneity. despite potential differences in data distribution. Under these cireumstances, 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-FI. does not require M > K but adheres to the CS settings ensure fair comparison. To elucidate the impact of the number of clusters, conducted supplementary experiments with

---

Table 3: Additional experimental results (accuracy in %) on HiCS with the number of clusters
M I K, where K is the number of selected clients each global round.
M
CIFARIO (2 (I)
Mini ImageNet (1)
Mini Mini--maaeeee (2)
M = 0.3K
61.3
57.0
47.5
50,4
50.1
M 0.5K
65.1
61.5
46.2
49.8
50.4
M 0.7K
62.8
59.2
51.2
51.1
49.9
65.5
59.8
50.5
51.2
Table 4: In experiments CIFARIO, only 20 out of 50 clients are available in the beginning; additional 10 clients join each 100 global rounds. The initial 20 clients leave the system after 400 global rounds.
Scheme
Random
pow-d
DivFL.
CS
FedCor
HiCS-FL.
CIFARIO (1)
85.6
86.7
84.0
86.2
80.8
87.4
CIFAR10 (2)
93.7
93.3
91.6
93.7
93.7
94.7
CIFAR1O (3)
94.5
94.7
93.9
94.5
95.0
95.8
Mini- Mini-ImageNee ImageNet
67.3
67.2
67.5
67.8
68.7
69.0
Mini- Emaeeett (2)
71.2
71.8
72.1
72.1
72.7
72.5
HiCS-FL using varying numbers of clusters M and compared these results 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-FI. can perform well with smaller M K as long as M is not too small, such as M = 3.
4.4 Dynamic Availability of Clients
The purpose the warm-up phase (t < [N/K]) shown Alg. 1 to collect updates of output layer from all the available clients in the system in order to facilitate clustering. Although we conduct all the experiments all the clients are available in the warm-up phase and can be adapted to more practical scenarios where clients have dynamic availability.
the setting where clients have fixed availability, our HiCS-FL does not assume
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 f, the server can obtain the information of availability and selects these new clients at round f + to approximate their data heterogeneity. To provide more insights, we conduct additional experiments on CIFAR 10 dataset; the results are reported in Table. can seen there, HiCS-FL outperforms baselines that consider clients' availability.
5 Conclusion
In this paper, we studied federated learning systems where clients that own non-lID data collabora- tively train global model; the system operates under communication constraints and thus only fraction 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 pro- posed heterogeneity estimation method, and the convergence of training 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 minor computational overhead.
Acknowledgement
This work was funded in part by the NSF grant 2148224.
10

---

References
[1] Durmus Alp Emre Acar, Yue Zhao, Ramon Matas Navarro, Matthew Mattina, Paul N Whatmough, and Venkatesh Saligrama. 2021. Federated learning based on dynamic regularization. arXiv preprint arXiv:2111.0426
[2] Ravikumar Balakrishnan, Tian Li. Tianyi Zhou, Nageen Himayat, Virginia Smith, and Jeff Bilmes. 2022.
Diverse client selection for federated learning via submodular maximization. International Conference on Learning Representations.
[3] Mateusz Buda, Atsuto Maki, and Maciej A Mazurowski. 2018. A systematic study of the class imbalance problem convolutional neural networks. Neural networks, 106:249-259.
[4] Nitesh V Chawla, Kevin W Bowyer, Lawrence Hall, and W Philip Kegelmeyer. 2002. Smote: synthetic minority over-sampling technique. Journal of artificial intelligence research, 16:321-355
[5] Huancheng Chen and Haris Vikalo. 2024. Recovering labels from local updates in federated learning. arXiv preprint arXiv:2405.00 00950
[6] Huancheng Chen, Chaining Wang, and Haris Vikalo. 2023. The best of both worlds: Accurate global and personalized models through federated learning with data-free free knoweedeo distillation. In The Eleventh International Conference on Learning Representations.
[7] Wenlin Chen, Samuel Horvath, and Peter Richtarik. 2020. Optimal client sampling for federated learning. arXiv preprint arXiv:2010.1317 137233.
[8] Yae Jee Cho. Jianyu Wang. and Gauri Joshi. 2020. Client selection in federated learning: Convergence analysis and power-of-choice selection strategies. arXiv preprint arXiv:2010.01243.
[9] Liam Collins, Hamed Hassami, Aryan Mokhtari, and Sanjay Shakkottai. 2021. Exploiting shared repre- sentations for personalized federated learning. In International Conference Machine Learning, pages 2089-2099. PML.R.
[10] Sever S Dragomir, Marcel L Scholz, and Jadranka Sunde. 2000. Some upper bounds for relative entropy and applications. Computers & Mathematics with Applications, 1909-10):911111111111111)01 10):911000
[11] Yann Fraboni, Richard Vidal, Lactitia Kameni, and Marco Lorenzi. 2021. Clustered sampling: Low- variance and improved representativity for clients selection in federated Icarning. In International Confer- ence on Machine Learning, pages 3407-3416. PMLR.
[12] Jack Goetz. Kshitiz Malik, Due Bui, Seungwhan Moon, Honglei Liu, and Anuj Kumar. 2019. Active federated learning. arXiv preprint arXiv:1909 1tt.:1999.126. 12641.
[13] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. 2016. Deep residual learning for image recognition. In Proceedings of the IEFF conference computer vision and pattern recognition, pages 770-778.
[14] Sai Pranceth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. 2020. Scaffold: Stochastic controlled averaging for federated learning. In International Conference on Machine Learning, pages 5132 5143. PMLR
[15] Diederik P Kingma and Jimmy Ba. 2014. Adam: A method for stochastic optimization. arXiv preprint arXiv: arXiv:1412.6980 1412,6980.
[16] Qinbin Li, Bingsheng He, and Dawn Song. 2021. Model-contrastive federated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10713-10722.
[17] Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. 2021. Ditto: Fair and robust federated learning through personalization. In International Conference on Machine Learning, pages 6357-6368. PMI.R
[18] Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. 2020. Federated learning: Challenges, methods, and future directions. IEEE signal processing magazine, 37(3):50-4)) 60.
[19] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith, 2020. Federated optimization in heterogeneous networks. Pruceedings of Machine Learning and Systems. 2:429-450.
[20] Pengfei Liu, Xipeng Qiu, and Xuanjing Huang. 2016. Recurrent neural network for text classification with multi-tas task learning, arXiv preprint

---

[21] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. 2017. Communication--ffiiiee learning of deep networks from decentralized data. In Artificial intelligence and statistics. pages 1273-1282. PMLR.
[22] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen. Zeming Lin, Natalia Gimelshein, Luca Antiga, 2019. Pytorch: An imperative style, performere deep learning library. Advances in neural information processing xystems, 32.
[23] Monica Ribero and Haris Vikalo. 2020. Communication--fffcci federated learning optimal client sampling. arXiv preprint 200.
[24] Felix Sattler, Klaus- Robert Müller, and Wojciech Samek. 2020. Clustered federated learning: Model- agnostic distributed multitask optimization under privacy constraints. JEEE transactions on neural networks and learning systems, (28))3710-372
[25] Ronald W Schafer. 2011. What is savitzky-gol golay filter?[lecture notes). IEEE Signal processing magazine, 28(4):111-117.
[26] Li Shen, Zhouchen Lin, and Qingming Huang. 2016. Relay backpropagation for effective learning of deep convolutional neural networks. In Computer Vision ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11 14, 2016, Proceedings, Part VII 14, pages 467 -422. Springer.
[27] Fang Shi, Weiwei Lin, Lisheng Fan, Xiazhi Lai, and Xiumin Wang. 2023. Efficient client selection based contextual combinatorial multi-arn bandits. IEEE Transactions Wireless Communications.
[28] Minxue Tang, Xuefei Ning. Yitu Wang, Jingwei Sun, Yu Wang. Hai Li, and Yiran Chen. 2022. Fedcor: Correlation based active client selection strategy for heterogeneous federated learning. In Proceedings the IEEE/CVF Conference on Computer Vision and Pattenn Recognition, pages 10102-10111.
[29] Aidmar Wainakh, Fabrizio Ventola, Till Müßig, Jens Keim, Carlos Garcia Cordero, Ephraim Zimmer, Tim Grube, Kristian Kersting, and Max Mühlhäuser. 2021. User label leakage from gradients federated learning. arXiv preprint arXiv:2105.09366
[30] Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoolos, and Yasaman Khazaeni. 2020.
Federated leamning with matched averaging. arXiv preprint arXiv:2002.06440
[31] Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, Vincent Poor. 2020. Tackling the objective inconsistency problem in heterogeneous federated optimization. Advances in neural information processing systems, 33:7611-7723. 1-7623.
[32] Joel Wolfrath, Nikhil Sreekumar, Dhruv Kumar, Yuanli Wang, and Abhishek Chandra. 2022. Haces: Hcterogencity-awee aware clustered client selection for accelcrated federated learning. In 2022 IEEE Interna- tional Parallel and Distributed Processing Symposhum (IPDPS), pages 985 995. IEEE.
[33] Haibo Yang, Minghong Fang, and Jia Liu. 2021. Achieving linear specdup with partial worker participation in non-iid federated learning. arXiv preprint arXiv: 201. /203.
[34] Miao Yang, Ximin Wang, Hongbin Zhu, Haifeng Wang, and Hua Qian. 2021. Federated learning with class imbalance reduction. In 2021 29th European Signal Processing Conference (EUSIPCO), pages 2174-21140 2178. TEFE.
[35] Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Nghia Hoang, and Yasaman Khazaeni. 2019. Bayesian nonparametric federated learning of neural networks. International conference on machine learning, pages 7252-7261. PMLR.
[36] Jianyi Zhang, Ang Li, Minxue Tang, Jingwei Sun, Xiang Chen, Fan Zhang, Changyou Chen, Yiran Chen, and Hai Li. 2022. Fed-cbs: A heterogeneity- client sampling mechanism for federated learning via class-imbalance reduction, arXiv preprint arXiv:229.1552:
12

---

A Appendix
A.1
Details of the Experiments
A.1.1 General Settings
The 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 x 2-axxpoliing layers (with a stride of 2). followed by fully-connected layer. In the experiments involving CIFAR 10, each client used CNN- classifier with three 3 -oo-oovolutional layers and two 2x2-mxxpooinng layers (with a stride of 2), followed by two fully-connected layers; dimension of the hidden layer was 64. In the experiments involving Mini-ImaeeNet 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 CIFAR wt/Minii-mageeeee/////e
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, CIFAR 10, Mini-ImageNet and THUC news, respectively. In all the experiments, the number of local epochs R was set to 2 and the size of 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 presented graphs were smoothened by a Savitzky-olla filter [25] with window length 13 and the polynomial order set to 3.
CIFAR MAAARii/iiii- 10/Mini-Immgeeeet For the sake of visualization, data points in the
A.1.2 Hyper-parameters
In all experiments, the hyper-parameter μ of the regularization term in FedProx [19] was set to 0.1. In the Power-of-Chie -hericcc (pow-d) [8] selection strategy, d was set to the total number of clients: 50 in the experiments on FMNIST, CIFAR 10 and THUC news, 100 in the experiments on Mini- Mini--mageNet. When running DivFL [2], used the ideal setting where 1-step gradients were requested from all client in each round (regardless of their participation status), similar to the of-Choice settings. For FedCor [28], we followed all settings in the paper and set the annealing coefficient B controlling the sampling strategy to 0.9 as suggested in the paper. For HiCS-FI. (our method), the scaling parameter T (temperature) used in data heterogeneity estimation was set 0.0025 in the experiments on FMNIST and to 0.0015 in the experiments on CIFAR I//MnniImmageNet. 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 the number of selected clients K. The coefficient 7" was set to 4 in the experiments on FMNIST and CIFAR10 10 while set to 2 in the experiments on Mini- ImageNmt To group clients, both Clustered Sampling [11] and HiCS-FL (our method) utilized an off-the-shte clustering algorithm performing hierarchical clustering with Ward's 's Method.
A.2 Empirical Validation of Assumption 3.1
To illustrate and empirically validate Assumption 3.1, we conducted extensive experiments on FMNIST and CIFARIO with the same model mentioned in Section A. In particular, we varied a over 250 values in the interval (0.01,50 50) generate data partitions allocated to 250 clients; entropy of the generated label distributions ranged from 0 to In 10 (maximum). In these experiments, we allowed all clients 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 x 500 12500 data points that correspond to 250 data partitions. For better visualization, merged adjacent points.
13

---

The results obtained by following these steps in experiments on FMNIST and CIFARIO are shown in Figure For more informative visualization, the horizontal coordinate of point in the scatter plot H (D(()), while the vertical coordinate is || F (0') - 9 V F((*)) )I( The dashed lines correspond to the curves y
cxp(3 {a: (Do)))) p that envelop the majority of the generated points. In the case of FMNIST, the blue dashed lime is parametrized by 3 = 1.0, p = 0.13, and x: 0.14 1.5, p 0.025, and x 0.022; these two lines
while the green dashed line is parametrized by envelop 95% of the generated points. In the case of CIFAR I0, the blue dashed line is parametrized by = 2.0, p 0.30, and 0.36 while the green dashed line is parametrized by B 1.8, p 0.15, and K = 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 (D())) decreases, implying that the local gradient computed by client with more balanced data is closer to the true gradient.
Fitting simdlation points on FMNIST with negative exponential functions
Fitting simulation points
CIFAR10 wit negative exporential functions
4 bcteoon local grad 01 and 1000 9220et.
teett and toa gruduee
0.30
0.22-
0.04
0:35-
0.82
ponns
0.20
6m5
1.5
2.0
23
0.2
D5
HUM
(a) FMNIST
(b) CIFARIO
Figure 5: Visualization of the difference between local gradients and the global gradient (evaluated all the data is centrally collected).
A.3 Gradient of the output (fully connected) layer's bias
Given a batch of samples (x(5.)), yli.n)), the cross-entropy loss is readily computed as
(1m) (gissms
exp
LCE
log
CE
cl..).
g(an) E [C]
(13)
Bi
Ee == exp
f=
gf (im)) Σ
(i.m) Wa + bez
(14)
where B is the batchsize; is the number of mini-batches; is the number of classes; d is the (j.z) denotes the d-th feature in the hidden space given sample ((.))
dimension of the hidden space; 24 the j-th batch; wa,e and be denote the weight of z (j.m) and the bias for the neuron that outputs the
probability of the class c, respectively; and g.2 is the corresponding output logit class c. The
gradient of the bias b, given sample x(j.)) can be computed by the chain rule as
aclim) ;b;
ac(
aQ
0q. (j.m)
(15)
aQ
0q. (..)
bb,
where
(..)
exp
(ysiin)
Q
(16)
Σg exp
(,0)
Then
ac(..n)
(n(
1
1.
(17)
8Q
Q'
ab:
14

---

If = ylim),
exp gyains
2 exp
(g6mm))) - cxp
(dn)
2
Q Q rzt ylsin) exp (gbmm)
aQ
gyltims
0g: (.m)
Σe= exp (gb.mi)) 2
E exp ge (
(18)
If i + ylim),
g (ylsi))
exp (g.)m))
exp (g.5.m)
aQ
exp
(19)
0q, (im)
> exp g6 (((mm)))
C
exp ge ((i)))
By plugging Eq. 18 and 19 in Eq. 15, we obtain
ac(in) bb;
E Eofglsm) exp (gllm) Ee_1 exp de (e.m))
ac(
exp
(
if y(mm)
if * y y
;b;
Ec=1 exp [r ((6mm))
(20)
A.4 Expectation of the local update (b)
By combining Eq. 4 and 5 and taking expectation, we obtain
R
E
7
Σ
E
V
CF
(jn..r)
Bi
j=1
c oxp(ye (hmrrr)
Pii y (1 ((m)))
1
17
BI
3=1
R
expla (hiirr)
Pli y Jilt # y(m)
1)
B1
exp(gl (hirrr))
r=1
R
E (xx,))~B-
c= exp(g (.n..r)
1=1
(21)
R
exply (ini)))
1)
(1 D (k)
I (
0 exp(y (im.))
R
R
D
E E(x.))~B-- (x)] - > (1 D. Exx,))BBB1 [s, (xx)
1=1
1=1
R
D (k)
Σ
E (1 )
C
= =nR
((())
Σ
E, - Ei
15

---

Note that
0
8
(x.9)~B-1
[ (x)
i=1
B
c} exp(o
1x (ym)
BIDe (k)
E exp( (
ofi
==1 m=1
(22)
B
Enei expeeep (im))
P P(ylin)
= Ees exp(e (()m))
C
BID
(k)
j=1==1
(Lin) 2e exp(ge (1n))
1 C - 1 BI
explogisi))
C
+
C
C-1
comparison to LcE Eq. 13 reveals that as LCe decreases during training, so does 5C & Given an untrained/initialized neural network model, E 1/C for Vi E [C], i... Σ E At global round T', if L:v then E E
C + --c-T
A.5 Privacy of (D))
According to Eq. 6, the server is able to obtain C' linear equations from each client,
Ab (k)
nR
D
(8)
Σ E
Eg
for vi E [C],
(23)
D (K)
1,
(24)
where denotes the number of classes. Suppose E[Ab! [ are known by the server. Then D(A)) the variables in the aforementioned equations, cannot be determined uniquely since there are C' variables and C + 1 equations. Therefore, the server unable to infer clients' true data label distribution and the privacy of ((k) is protected.
A.6 Proof of Theorem 3.3
In Section A.3 derived an expression for the gradient of the bias in the output layer given single sample (x (x(mm) y) in the mini-batch. It is worthwhile making the following two observations:
the sign of the ((r*)-th component of V.c(s.") (x(im)),((mm) opposite of the sign of the other components; and
(xlin), ,((in)) is equal in magnitude to all other com- the y(óm)).t thh component of V Vbc(An) (x(ii))))))))) ponents combined.
C denote the local update (made by client k) of the bias in the output layer of the neural network model, and let ((k) [D(*) ..... D(t)) be the (unknown) true data label distribution, D (D(*)
Proof: Let (b = (
Assuming the learning rate n and R local epochs, the
expectation the local update of (b()) ()) is
C n (())
E, - E e.).
Ab
Σ
(25)
Data heterogeneity can be captured via entropy. H(D())) = De x(a² In (1,) where higher H (D(()) indicates that client l has more balanced data. However, since we do not have access to
16

---

the client's 's data distribution, we instead define and use measure of heterogeneity (D(()) _
H (softmax (b(*), where
exp( Ab, /T)
(k)
softmax Ab T)
(26)
=== exp( Abe /T)
and where denotes the temperature of the softmax operator. Suppose there are two clients, u and with class-balanced and class-imbalanced data; let D(") and D(k) denote their data label distributions, respectively, while D(") and (AA) are computed by softmax (oftmax(Ab()) ( ,)) and softmax(A) (Ab(*) T).
Without loss of generality, can re-parameterize Dix))
C Dla) = cU + Σ Z
(27)
==1
where = lt. J denotes uniform distribution; i-th component of Z, while the remaining
components are 0; and all non-negative such that Σ E 1. We can always set ming Ej 0; otherwise, let é c+ miny ey and 6, = c; miny Ejz Vi E [C]: quantifies how close
(a) to U. Due to the concavity of entropy.
H(D())) 2 2 CH(U) Σ (H() clnC.
(28)
We will find the following lemma useful.
Lemma A.1 For probability vectors p and q with dimension C, the Kullback-Leibe Leibler divergence between and satisfies
1 |p p 1
KLD(pllq)
2 2
(29)
where ||p 1---qp 1pp--11- -i E %/l.
For the proof of the lemma, please see [10]. Applying it, we obtain 1
KLD(DO)) H(U) - H(D(()) 2 2 |
- > D - - ul
(30)
Combining Eq. 28 and Eq. 30, we obtain
H(D())) H(D(M)) (6-1)) 1 B(
(31)
By taking expectations of both sides,
EE (((()))
[ H(D()))) > (B[e) 1)
Cn C - ((N)) U 2
(32)
Since |
U 2 is convex (composition of the Fuclidean norm and softmax), according to
Jensen's 's inequality
((D())) H(((M)) 2 (E(c) - 1) C ((k) (E(Ab()))) - Ul
(33)
where
exp nR (D.(0 Ec EEi /T)
Es
(EAA
(34)
Σ exp n De E. - E,) /T)
Selecting 7 such that nR
D (k)
DC & - & / is sufficiently small and applying the first-order
Taylor's expansion of e around 0, we obtain
Σ exp
nR
(k)
Σ E.
Eg
IT
Σ
+ n Σ
( Σ E -E/IT= CT= (35)
1=1
c=c
17

---

where Σj= D/a)) = 1. This leads to simplified B(A))
(E (B|Ab(()))
1/ (h Ec E - E,) /T
D 1 (E(Ab ()))
(36)
C
Substituting Eq. 36 for the second term on the right-hand side of ineq. 33 leads to
C
C
2
(B D - U 2
nR
D
(k)
Σ E
Eg E
(37)
CT
1=1
c=1
Now, consider
-
U = (c 1-) +
(38)
Taking expectations of both sides,
E (c- Σ
aZ =E ( - uj] 2 D(a)) (E[Ab(")) - UU.
(39)
E=1
The above inequality holds component-wise, so for the j-component (ej = 0)
+e, = 1)1² D(n)) (BAb("") -U, 1) 1 (1) De E - -gG
1 lllee- 1)
E
(40)
CT
Therefore,
E/e] 12 R (D.0 Ee- & Ej
1/R (D(a))
Ee Ee - Ei
min
(41)
T
T
Taking absolute value of both sides yields
Eid - -1 = n
T
max
D ((") Σ - &,
= R
T i
max
1 (D(")
C
Σ
C
E
-
+.+
C
1
EE
C
E
(42)
c=1
c=1
By applying the triangle inequality we obtain
Ele] - 1| - nR
max ((u) - - Σ E +
C
RR max
Σ
C
Ee
- E.
(43)
T
T
c=1
Let 0 = max; C EC & - E.}. Since D & < δ
= 1, it holds that
[B(c] - 1| <
R
max D(u)
+
n
8.
(44)
T
C
T
Furthermore, since [[c] - 1 <<<< 0.
E(c]
1
2
R
max
D(u)
1
n
(45)
T
C
T
Note that
C
2
c
2
D
E
E
(D
1
Σ
E.
+
1
E
0
-
1
&
+ 7
Σ
E E
+2
Σ
C
Eg
(D(
d)
Σ
E
E E
c=1
==
(D!!)) - -d)
)022
c
2
(D())
Σ
2 Σ Ee
Σ
E - E
(46)
18

---

Therefore,
C
c
C
D (K) 1 2
Σ
D
88)
E 3
2
Σ
3
C
c=1
c
D (M) - :
1 c
+2
Σ
Σ
(
Σ
& - &
C=1
c=1
C
2
C
Σ
3
3
D(K)
1
C
+2
Σ 3
Σ
D
(()
Σ
c
E
1
Σ
c
&
+
3 C
-
C
1=1
6=1
2
C
Σ
33
Σ
(k)
1
C
c
c
22
Σ
3
&
1
Σ
E
+
1
Σ
&
Σ
D
&
C
C
6=1
c=
1=1
i=1
2
c
C
D
(M)
1
+2
Σ
E
Σ
E
max y
Eg
C
0
c=1
c=1
2
E
(k)
1
26.
(47)
Substituting the above expression in Eq. 33, we obtain
E ((((w)) - -H(D())]
2 nRIn C
((u)
RRInC
max
(48)
T
C
T
RR CT
2
C
(((0) - 1 2
)
-
<n
2
+ 2
E
Σ
-
CT
8, (49)
1=1
and, therefore,
2
E [(((())) - H(D()))] -H(D)))
2
nR CT
Σ E
|ID(])
-
n In C T
D(u) -
Co.
2
(50)
CRInR CTT
where C-
A.7
Convergence Analysis
Here we present the convergence analysis of an FL system deploying FedAvg with SGD wherein only 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.
Assumption A.2 (Smoothness) Each local objective Fx(-) is L-smooth,
-
- | - 8/112
(51)
Assumption A.3 (Gradient oracle) The stochastic gradient estimator gx((",)) -VF) for each global round t and local epoch is such that
+ Sh."
Bick = 0
(52)
19

---

and
I [|stt" 12
[ 0²
(53)
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 1
). We provide generalization that holds for any sampling strategy, resulting in E k EXES: Σm=1 (x=]")
R Eves' Em=1 9x (0%"" - Σ PA (O
Σx Σ F&& where w denotes the probability of sampling client k: in round & under sampling strategy II. Note that Σx= w 1. We assume that all clients deploy the same number of local epochs R and use learning rate 77 at round t.
A.7.1 key lemma
Lemma A.4 (Lemma in [33]) Instate Assumptions 3.1, A.2 and A.3. For any step size n such that 7) S BLR: for any client k it holds that
E
8- 001
S 5Ri + 30 |FF())]
(54)
Proof of Lemma A.4: For any client A E [AV] and r E [1]]
|10 - 091"]] = E -100-- - 0r - ng (0-.--1))) 1)
-B(100%-11 - 0r - no (0 (0%--)) 1 )+ VP((OH
(F' + (F'( - VF(0') + [F(
1
E 186--1 0¹1 +n²E -
11?
1 2R 1R--1
+ 6R
FE
(Ot
+ B 19 - -VF0)))
+ 6R
'E
VF(0")
1 1 2R -
1
E 0hr-11 -
-0010 0²(² + 1 1001²
+ o
6 |F
1
E 1 0r(
1
+ BRIpL
+ +100² + +
2R 1
+ 6R2r2
1
g¹ 1/
1 +
E
+ +1z00² 6R 2 o + 6
R-1
R
(55)
Unrolling the recursion yields
1 r-1
E 100 -091²
(12002 + +6Ry + 6
1 R R-1)
(1+ 11
)1)²²
1
< RR 1)
-1
+ GRI (n'''+
GRI
R R-
s 5Ro² (02 + +6 + 30 |FF(8)) ||2
(56)
20

---

A.7.2 Proof of Theorem 3.4
The model update at global round è is formed as
R
g++1 H¹ -7 "
1
M REST
Σ
94 (0%"))
(57)
where G and 6e denote parameters of the global model at rounds t + 1 and i,, respectively, and 0 denotes parameters of the local model of client k after local training epochs. Let
R
A
K
9k
(0(").
(58)
K
KES
Taking the expectations (conditioned on 06) of both sides, we obtain E [F( = E [F(0' - n]']]
- F(0') 17 (VF(O)) B( + L En'N [|A'
((())+((VV()))) RV (0)) RVF Ar)) + L " || [2]
- - Rm VF(O' B (B)(() -タ E ||
A
A₂
(59)
Inequality (a) in the expression above holds due to the smoothness of F(-) (see Assumption A.2). Note that the term A, can be bounded as
A1 (VF(()) ),E E F(O) - A')>
R Ih
VF(O") RVF(O') K
KES
r=1
R
VF(O') [RVF(O))
Σ
w
(
k=1
R
Σ w/ VRVP(O).
E
(VF.(0t))) -
VR
N
N
R
(
R
|F((")
112
+++
1
Σ E
(F -VF(())
1
w
E
Σ V
2
2R
2R
k=1
k=
(b)
R
|F((') 1|²
+
1
wfE
Σ
R
(VFL(O))) k
VFL(O'))
2
R
km
n
n
1
Σ
(F') VF(O')
1 Σ w
Σ
V
Fu (0
R 1=1
2R
1=1
(e)
R
|F((0)
E
V(A - (1)
2
E
(0') VF(0')12
1
of
Σ
2R
(d)
R
|FF
)|1
2 Σ
0.
0t 12
RRR
w o
1
Σ
Σ
2
2R
(60)
21

---

where equality (a) follows from (x.y) = 2/ (1x= + ||vll - |x - yyll))), inequality (b) is due to
||x + yl2² < 2 ||x]|| + 2 llyl". inequality (c) holds because |Σi=1 zell" S n Σi= 1|2,11 and
incquality (d) follows from Assumptions 3.1 and A.2. By selecting 1 < SLR and applying Lemma A.4 we obtain
R
A
S
+ |F(0)]1
L
R + +6Ro + 30 [V(
2
k=1
N
+
R
2
w E
V (2(
2R
k=
(61)
R
+ 300 1.0L2
|FF(0))|
+ 3L2 1222222222220 30L
Σ
of
2
k=1
N Σ
+
wiok
w
Σ V
2R
k=1
k=1
Furthermore,
A2 E
ge (B
K
KES p=l
11k E st1 N
B
E
Ih (0t2²)
K
T(ks 8t)
R
Σ (x (Fa( + (F
E
K
I[k E Sq
R
Σ 1lke
St)
R
9h (0%" VP((O)))
+
(b)
E
N
I(k E 8t)
Σ
n
(0x)
(8"") |12
+E
NNN
I(k E Sy)
Σ (FL(
n1
K
K
(0) 2
W
Σ V Fu (0
(62)
where equation (a) holds because - Jensen's inequality, and inequality (c) is due to Assumption A.3. Substituting inequalities (61) and (62) into inequality (59) yields
inequality (b) stems from the
[(0++))] & F(0") Rn [FF()) (VF(O') E RVF(8") Ar) + +22
L n² 122 E
||
Ai
A2
< F(80) Rn (1 300 R2n²
|F((0))
6L"R2nn²
+
LR
2 11 0²
N
L
+ (3013 Rp + Rm) Σ w
Σ wfB
Σ V
2
2R
(63)
22

---

If < SLR.
1. it must be that 1/2 -30022R27 > 0 and 11² 2R < 0, leading to
B [(( - F(6") -Rn ( R
|VF(0)||
(64)
+ 5 5 + LR
N
0² + (30L" + Rn)
Σ
2
By rearranging and summing from = 1o T
we obtain
1 2 - 3012R2
E [F((T]]] - F(00) v RRy
V |
1=
T-1
+ 5L 51211112 + 2 11²
LR
To² - -3002 + Ra)
w
1= k=
(65)
s <Pny 1
300.2 1211)
T
min :E(T]
|
T-1
+ (512RR2 2 "1
LR
To² - (30T.2 + Roj)
Σ
Σ
o
Let H" denote the optimal model's parameters, i.e., F(80) S F(B')Vt E [T]. Then
F(00) F()))
min TET
|F
A w + +
(66)
A
(10L" By +
where A = Ry ( 30T, A2 Fo1[1600- Rn
and +
1-GOL" ""
A.8 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 regularization term to combat overfitting, including FedProx [19]. [edDyn[1] I] 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.
A.8.1 FedProx
The objective function used by FedProx [19] is
11 1100"" 001]2
Sprox Loe
(67)
where 0%." is the vector of client k's local model parameters in the r-th local epoch at global round t. Therefore, contribution of sample (xli.)), y(imn) the gradient of Eprox local epoch r
aclion))
ac(j..r) + p (B. - b))
(68)
Ob
Oh.
where b bb. ..... b b' = li..... bl denotes parameters the global model round We assume the model is trained by SGD as the optimizer, and hence
(. denotes parameters of bias in the output layer of the local model, and
B5t bf = b actiar-
2 mjm. -T}
bb
+ (1 mP)) -b6).
(69)
ab;
bb,
23

---

Therefore,
m1pp" 1-s ac(in)
6,"
by
1]t
CE
+ + (1-
6,)
bb;
(70)
7/110) 1-- ac(.mm.)
1/t
0b;
and thus
0L 8c
Oc(mm)))
0C mm.s)
1/111
7114)
CE
(71)
bb.
Taking expectation of both sides yields
B
Oc acloa
E
- ((m))))
Eii
Σ Ogfi
Ee ETi(i + g
Bl
db;
11
R
17:111 1=11)~1-
==1
(1- (1
(((e)
E + (1 D.
17231
oFi
f
C
DD (R)
Σ
E + &
(72)
where c" = (1 - (01--pp)---- > 0 provided πe and μ are sufficiently small. Therefore, the expectation of the local update of bias in the output layer satisfies
C &
| = Onn
D
(k)
Σ
E,
(73)
where C' C- c². Eq. (73) is similar to the expression for the expectation of the local updates of bias when applying FedAvg presented the main paper; clearly, the analysis of HiCS-FL done the context of FedAvg extends to FedProx.
A.8.2 FedDyn
For FedDyn [1], the objective function in local epoch at global round /
Ciyn
pt.r (V.t 1, o%.
t 100 0* 0²1|².
(74)
where R denotes the total number of local epochs. The first order condition for local optima implies
V Ctr
V VC + u(0 0²) 0.
(75)
24

---

and, therefore,
① @. 1
1 (bir b;)
ab;
bb;
i 2t-2,R
dyn
11 (b,-1.11 02-1) 14 (b.- bt)
db.
Σ
1/ -0--bb) b/) 14 (b - -b!)
(76)
Ab! (6,-" - -))
-M
2 dcbr-
AbG
-111-
+ 1, b/
Ohi
Ab +
pilt
bb.
T=1
where b ...",...; ] denotes the bias parameters in the output layer of the local model at local epoch r, and where Åb" = =AAb A is the local update of the bias at round т. Since
11
AbT
(77)
bb.
it holds that
a 2
t-1
(-M
1-1
Abح + pilt
Ab
=-p(1 + pM)E))
Ab
(78)
bb,
and
ac
dyn
1-1 Ab! + pil
(-M Σ
1-1
إoح - (n + 1mm1))
Σ
1-1
= =-M(1 + (1(1))
Σ
L-1
Ab7.
ab;
r=1
r=1
(79)
By induction,
cchr Ub;
-M(1 + pm))
Σ Ab.
(80)
Therefore, the expectation of the local update of bias in the output layer at round / can be computed
R ()) = Σ (1 + 1
E
Ab
E
(81)
1=1
T=1
= ((1 + pmm) R 1)
Σ E Ab! (1)77]
(82)
Since the objective function of E |b [1). coincides with that of FedAvg.
C
Ab [+]
= mi (D:0)
Σ
Es E
E
(83)
o=1
where 7/ is the learning rate at global round = Then,
C
E
A (M),2]
m R ((1 + 1112) -1
(k) Σ E - E
(84)
c=
a102
(A) Σ E, - E;
(85)
cml
25

---

where R and (1 + pmp)) R 1. Furthermore.
C
E
Ab Ab (1.3
a1as(1 + 02)
D.AA)
Σ
E - &
(86)
C
E
Ab (
=aja11+ 02 as - agas)
D
Σ
& E E
(87)
c=1
and
Ab [6),5]
1 Σ E -E:
E
aias 02 + a3 +, + 0203 + 0za1 + a2a3a1)
(88)
By induction,
C
f-31-1 aidh 1+ {1+
T+i
E
(() Σ & &
I+II(T+ +(<+
II
(89)
1=0==2
a (K) & E
(90)
where de (1 + pmme R and aiab
IIT IT 0g 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.
A.8.3 Model-Contrastive Federated Learning (Moon)
Moon [16] relies on the objective function with a contrastive term
Limon 1
(j.)) CE
exp(sim(z (2.m)) (..)) glob
glob
)/T)
1a log
(91)
Bi
exp (eppsim(z((nn
)/T) + exp (xp(sim(z(dtn) Zprcv //T)))/T)
(j.))
n=
where z(im) denotes the output of the feature extractor of the local model ff. the feature extractor of the global model 0², and (j.n) is the output of the feature extractor of the local model in the previous round 0--1_ Since the contrastive term does not depend on the parameters of bias in the output layer, it holds that
(d.n) glob is the output of
achiol Ob.
achi
(92)
ab.
Since 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.
A.9
Optimization Algorithms Beyond SGD
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.
A.9.1 SGD with momentum
In each local epoch r, SGD with momentum updates the model according
um 1,r-1 + (1 (1)
pl. VLr
(93)
9b
m t.r
(94)
Tt. 1.r
(95)
26

---

where mg f.1 V VL
1. denotes the momentum in the r-th local epoch, is the weight for the momentum, and
mi
Then
A0% -7
9k
(96)
where
m
411 = V C
(97)
" t.2 = /V c + (1 - -m
(98)
m = VV (1 p1) V C
(99)
= VVCC + +1(11 μ1) V 222
V + (1- P) V
Therefore,
m m/r = (1 n1) Σ pt V
(100)
and thus we have
r
T
Σ
V
p-
ch +1 P Σ p -
Vi'a
(101)
Similar to the discussion in the previous section,
R
c (k)
E
Ab
(k)
178
Σ
+
(1- ()) Σ nー +1
Σ E - Ei
(102)
r=2
c=1
D (k)
Σ & -
(103)
E> (w- (1 1) Σイ 2) +1)
where a The
Similar result is obtained when SGD applies Nesterov acceleration as long as the optimizers are not using second-orde momentum.
A.9.2 Adam Optimizer
Recall that the two observations regarding the gradient of LcE 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 adaptive learning rate tied its magnitude. With an SGD optimizer, the magnitude of the i-th entry of the local update of bias Ab is approximately proportional to the fraction the samples with label i, (D(
(if is small),
C
Ab!a [k) = n.R
D
E. -
(104)
However, this observation does not hold when using the Adam optimizer for the local update because each entry has different learning rate Th, and thus
E
Ab (K)
D(N) Σ E.
E.
Es
(105)
Although the magnitude o [b is longer approximately proportional to D(*a), we can utilize
i.e. (k)
the sign of TE Ab (k)
if D. > D(K)) then P (B |ba 0) TP ( E | (2) 0).
(106)
Suppose client & has highly imbalanced data, i.e., (D(+)) is small. Then the maximal component max; (*)) is much larger than the other components; in fact, it is likely to have only one positive
27

---

component in the local update of bias (b On the contrary, suppose client has balanced data and thus H (D(w)) is large. The maximal component max; D( is then very close the other components, and it is likely to observe larger number of positive components in the local update of Ab(") While characterizing P(E[Ab; (1)) > 0) appears challenging,
empirically infer that client with more balanced data has a local update of bias Ab (4) with more positive components. With
(D( H(softmaxAAb (Ab (4) T));
(107)
(D( H(softmaxAA T));
(108)
(D(w)) is more likely to be larger than (D())). The examples of estimated entropy when utilizing Adam as the optimizer are provided in Section. A.12.
A.10 Visualization of Data Partitions
To generate non-IID data partitions follow the strategy in [35], utilizing Dirichlet distribution with different concentration parameters a to control the heterogeneity levels. In particular, the number of
x! N. samples with label owned by client & is set to Σ x x where x!".... ""....".""""
(N) x!!') are drawn from
Dir(a) and Nj denotes the total number of samples with label in the overall dataset. For the setting with multiple a, we divide the overall training set into la] equal parts and generate data partitions according the method above. Figures and illustrate the class distribution of local clients by displaying the number of samples with different labels; colors distinguish between magnitudes the darker the color, the more samples are the class.
Heatrrap data parttirn
destrop data partioon
Hastrap of data part tona
(a)
(b)
Figure 6: Results on CIFARIO. Training data is split into 50 partitions according to a Dirichlet distri- bution (50 clients). The concentration parameter is as follows: (1) E (0.001, 0.01, 0.1, 0.5, 1.0}; (2) a f0.001, 0.002, 0.005, 0.01, 0.5): (3) a E (0.001, 0.002, 0.005, 0.01, 0.1). The figures (a), (b) and (c) correspond to settings (1), (2) and (3), respectively.
Heathag of daTa partillrrs
Heatmap data partitons
Gan
(a)
(b)
Figure Results on Mini-ImageNNe.. Training data split into 100 partitions according to Dirichlet distribution (100 clients). The concentration parameter is varied as follows: (1) E (0.001, 0.01, 0.1, 0.5, 1.0); (2) a (0.001, 0.005, 0.01, 0.1, 1.0). The figures (a) and (b) corre- spond to settings (1) and (2), respectively.
28

---

Table 5: The columns "Extra Computation" and "Extra Communication" denote the computation and communication complexity of additional operations each sampling scheme compared to random sampling.
Method
Extra Computation
Extra Communication
Random
pow-d
O [
0(101)
CS
O [
DivFL
O [0)]
0(101)
FedCor
O ((0))
HiCS-FL
(C)
A.11 Computational and Communication Complexity
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 Table. 5. In its ideal setting, pow-d selects K clients with the largest local validation loss among all N clients. To compute the local validation loss the beginning of global training round the server must send the global model to all clients. Compared to the random sampling strategy where the global model is sent to only K clients, pow-d must transmit additional (N K) [0'] model parameters. Moreover, pow-d requires all clients to compute validation loss of the global model on local datasets, which incurs additional O(N
computations. While communication requirements of Clustered Sampling do not exceed those of random sampling, the server must run a clustering algorithm on the local updates of dimension [0') (the same gradients). DivFL relies on maximizing submodular function to select the most diverse clients based on all clients' gradients, leading to a transmission overhead and additional computation involving |0'| -dimensional gradients. 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 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 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 aoneel-agnottcc. Overall, HiCS-FL requires negligible computational overhead to significantly improve the performance of non iid Federated Learning.
A.12 Examples of Estimated Entropy
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 ImageN,,,,,,,, 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 and compare the performance of estimating entropy with SGD and Adam optimizers for the same setting of a. Notably, as shown in the figures, the method capable of distinguishing clients with extremely imbalanced data from those with balanced data.
29

---

EPORUCT Caara sur Siginc 2501, S.D. UA,,
(a)
(b)
Figure 8: The estimated entropy of data label distribution in experiments FMNIST with SGD as the optimizer. The parameter a for the two figures: (a) a {0.01, 0.02, 0.1, (b) {0.001, 0.002, 0.005, 0.01, 0.5}
Mun on wom
mmm
mmm
mny
42-
(a)
(b)
Figure 9: The estimated entropy of data label distribution in experiments Mini-ImageNet with Adam as the optimizer. The parameter a for the two figures: (a) a E {0.001, 0.01, 0.1, 0.5, 1.0): (b) (0.001, 0.005, 0.01, 0.1, 1.0).
CIANTO Data 10000, 8.0. 8.J. 2.5.2.00
CIAATE Data 8.3. 23..0.3.0
sppe
(a)
(b)
Figure 10: The estimated entropy of data label distribution in experiments on CIFAR10 with a E (0.001, 0.01, 0.5, 1.0). (a) The result of the experiments using SGD as the optimizer. (b) The result of the experiments using Adam as the optimizer.
30

---

CPA122 Dutaiet aipias 1000L, 8000, 3.00 raccca
CRANCO DURAIE 1O00), 8.02. 2.005
(a)
(b)
Figure 11: The estimated entropy of data label distribution in experiments on CIFARIO with a E {0.001, 0.002, 0.005, 0.01, 0.5). (a) The result of the experiments using SGD as the optimizer. (b)
The result of the experiments using Adam as the optimizer.
31

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: [NA]
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 perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
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 not discussed in the paper.
The authors are encouraged to create 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 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might 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 address problems of privacy and faimess.
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 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
32

---

Justification: [NA]
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
33

---

Answer: [Yes]
Justification: [NA]
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm cc/
ilyy for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: /CodeSubmic / //iips.co.. //nip..c wttpccccpppccccceeececceccccccct /ttpa////iiieeeeeeeuuubiiiiiii ssissssilllll onPolicy))) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
34

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources Question:
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of exNoccution) needed to repro- duce the experiments?
Answer: [No]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that NeurIPS Code of Ethics https: : /ttp///eerr///ee/ee//ee cc cttpp/bccccccchhccsteeeeeeeeeeee
didn't make into the paper).
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the
Answer: [Yes]
Justification: [NA]
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: [NA]
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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to
35

---

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 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.
there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators 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?
Question:
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 dataset.
The authors should state which version of the asset is used and, if possible, include 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 (c.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, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's is creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
36

---

Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/codeemoodll as part 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 used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere 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.
37


Reviews:
[{"review_id": "GGEOvQgGX3", "reviewer": "Reviewer_uaYT", "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.", "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?\n\n2. Why wasn't an experimental analysis conducted on different values of $ \\lambda $ in Eq. 9?\n\n3. 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": "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.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. This paper proposes HiCS-FL, a new federated learning method designed to adaptively handle clients with varying label distributions. It achieves better convergence and greater efficiency compared to baseline methods in experiments.\n\n2. HiCS-FL utilizes the gradients of the output layer's bias to estimate clients' statistical data heterogeneity. This estimation is then used to cluster and sample clients, enabling efficient training of the global model with relatively data-balanced clients.\n\n3. This paper provides a detailed theoretical analysis of HiCS-FL.", "weaknesses": "See questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the informative comments and valuable questions. Our answers addressing unique questions are below; responses to concerns raised by multiple reviewers are in the Author Rebuttal block for sake of brevity. Note that **Tables 1-3** can be found in the **pdf** file at the **bottom** of the **Author Rebuttal** block.\n\n**Q1.1** \nIs 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 still pose privacy concerns for clients? \n\n**A1.1**\nPlease note that other reviewers raised the same question. For our response, please see **R.1** in the **Author Rebuttal** block.\n\n**Q1.2** \nAdditionally, 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?\n\n**A1.2**\nThanks for the suggestion! To address it, we conducted additional experiments in the oracle settings where the server has access to the true entropy of the clients’ data distribution. As expected, providing the true entropy information further improves the accuracy and speeds up the convergence of the proposed method. The results are shown in **Table 1**; for convenience, the variant of HiCS that relies on the true entropy values (provided by the clients) is labeled as ‘Oracle’.\n\n**Q1.3**\nWhy wasn't an experimental analysis conducted on different values of $\\lambda$ in Eq. 9?\n\n**A1.3**\nFollowing up on the reviewer’s question, in **Table 2** we show the results of experiments for different values of $\\lambda$. As can be seen there, our method achieves the best performance with $\\lambda = 10$.\n\nAs a general comment (and as discussed at the end of Section 3.3), $\\lambda$ has to be set to a value that ensures the second term in equation (9) dominates the first term; if $\\lambda$ is too small, the clustering prioritizes the cosine similarity of gradients over data heterogeneity. In the extreme case of $\\lambda=0$, our algorithm collapses to the prior work Clustered Sampling [3]. \n\n**Q1.4**\nWill 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?\n\n**A1.4**\nSince the original motivation of Federated Learning is to collaboratively train a global model with strong generalization capability, the experimental results shown in the manuscript are performed on the global test dataset with balanced data. Nevertheless, having obtained such a global model, the clients certainly may personalize the global model by fine-tuning it on local data. To investigate this more closely, we conducted supplementary experiments where the global model is personalized by fine-tuning on local data and compared our HiCS with several other client sampling methods. The results of inference tests on the local data from the same distribution as in fine-tuning are shown in **Table 3**. The results demonstrate the value of HiCS training strategy – namely, aiming at achieving strong generalization ability on all classes – as it sets up the fine-tuned model for success. Indeed, the fine-tuned HiCS generally outperforms the competing methods on personalized/local inference tasks.\n\n**Q1.5**\nThe 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.\n\n**A1.5**\nTo partly address the reviewer’s comment, we conducted supplementary experiments where we investigated the impact of clients’ availability on the proposed method. Please note that another reviewer raised the same point -- we provided the same response as **R.2** in the Author Rebuttal block.\n\nAs a general comment, please note that we consider system heterogeneity and data heterogeneity to be complementary challenges in federated learning. System heterogeneity encompasses variations in computational resources, communication bandwidth, and fault tolerance, which are unrelated to the statistical heterogeneity of client data. The former may cause difficulties when clients with valuable balanced data cannot attend a training round due to poor resources; in such scenarios, our HiCS-FL must manage the trade-off between data heterogeneity and training efficiency. We will add this brief discussion to the conclusion section of the manuscript.\n\nIn this paper, the primary focus is on enhancing the generalization capability of a global model trained within the federated learning framework, while the system heterogeneity remains an exciting topic for future extensions of our work.\n\n\nReference:\n\n[1] Li T, Sahu A K, Zaheer M, et al. Federated optimization in heterogeneous networks[J]. Proceedings of Machine learning and systems, 2020\n\n[2] Kingma D P, Ba J. Adam: A method for stochastic optimization[J], ICLR, 2015\n\n[3] Fraboni Y, Vidal R, Kameni L, et al. Clustered sampling: Low-variance and improved representativity for clients selection in federated learning, ICML, 2021", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TpybSTpQCY", "reviewer": "Reviewer_yLvJ", "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.", "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": "See weakness", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The introduction of a hierarchical clustering-based method that considers data heterogeneity is a significant advancement over existing client selection techniques. The solution considers the balanced and imbalanced client data to train the model, which is a very interesting perspective. The extensive experiments on multiple datasets and comparison with several baselines demonstrate the robustness and generalizability of HiCS-FL. This paper is well-written.", "weaknesses": "- 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.\n\n- 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.\n\n- 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.\n\n[1] Oort: Efficient Federated Learning via Guided Participant Selection \n[2] Auxo: Heterogeneity-Mitigating Federated Learning via Scalable Client Clustering\n\n- 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.\n\nThere 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thoughtful comments. Our answers addressing unique questions are below; responses to concerns raised by multiple reviewers are in the Author Rebuttal block for sake of brevity. Note that **Tables 4-5** can be found in the **pdf** file at the bottom of the **Author Rebuttal** block.\n\n**Q2.1**\n 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. \n\n**A2.1**\nWe appreciate the reviewer's point regarding feature skew, wherein clients may possess data with analogous features but divergent labels; this indeed differs from the label skew addressed in our paper. The inconsistency between local models trained on data exhibiting feature skew certainly hinders the convergence of the global model. In scenarios characterized by pronounced feature skew, personalized federated learning (pFL) [1], which aims to train multiple personalized models based on clients' preferences, may help mitigate this challenge. While it appears difficult to simultaneously address both label and feature skew due to their inherently opposing nature, it would certainly be of interest to deploy our method within a pFL framework that allows the server to train multiple global models catering to varied client preferences.\n\n**Q2.2**\nThe paper does not compare HiCS-FL with other relevant algorithms, such as Oort and Auxo, which also focus on client selection and clustering.\n\n**A2.2**\nWe thank the reviewer for pointing out these two relevant papers. Oort relies on the idea of assigning higher importance to the clients with higher training loss and applies the UCB algorithm with the training loss as the reward function to balance exploration-exploitation. Auxo is an interesting paper attempting to address the challenge of data heterogeneity via personalized models aggregated by clustering. While this presents an exciting line of work, it conflicts with the motivation of training a global model with strong generalization ability pursued by our paper, rendering a fair comparison between two methods rather difficult. Therefore, we only implemented Oort and compared it with our HiCS (see **Table 4**). Nevertheless, Auxo introduced an interesting idea of online clustering without assigning a fixed number of clusters, which may potentially be adopted by our method. In any case, we are committed to citing and discussing these two highly relevant works in our revised manuscript.\n\n**Q2.3**\nThe 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. \n\n**A2.3**\nAs 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 can recover the performance of the previously established Clustered Sampling (CS) approach [2]. 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 5**. As shown there, HiCS performs well with smaller $M < K$ as long as M is not too small (e.g., $M > 0.3K$).\n\n**Q2.4**\nThe 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.\n\n**A2.4**\nThe optimal approach would indeed involve collecting updates from all clients within the same round and using those for clustering. However, this may not be feasible due to clients’ varied availability and would undermine the goal of client selection aimed at reducing communication overhead. Consequently, we adopt a compromise by gathering updates from clients during the initial $t \\leq \\lceil N/K \\rceil$ rounds, thereby adhering to communication constraints. As demonstrated in Appendix A.4, update of the output layer $\\Delta \\mathbf{b}^{(k)}$ is influenced by the data distribution of client $k$ and the global model's capability. During the early stages of federated learning (FL) training, the global model's capability is relatively consistent, allowing us to treat updates collected in these rounds as equivalent. This \"warm-up\" strategy is also employed by CS [2] and FedCor [3] to avoid requiring all clients to attend local training in the same round. Even though the gradients computed from different global rounds could be different, the pattern corresponding to the data distribution is contained in the gradients and can be utilized for clustering as Auxo (the paper brought up by the reviewer) does to generate different cohorts in online clustering.\n\nReference:\n\n[1] Tan A Z, Yu H, Cui L, et al. Towards personalized federated learning[J]. IEEE transactions on neural networks and learning systems, 2022\n\n[2] Fraboni Y, Vidal R, Kameni L, et al. Clustered sampling: Low-variance and improved representativity for clients selection in federated learning, ICML2021\n\n[3] Tang M, Ning X, Wang Y, et al. FedCor: Correlation-based active client selection strategy for heterogeneous federated learning, CVPR 2022", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cuXtuXykkM", "reviewer": "Reviewer_7UbW", "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.", "questions": "Data heterogeneity can be mitigated with pre-training and personalization. Would pre-training make clustering unnecessary?", "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.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is very well written. The paper did a good job of summarizing what others have done and how HiCS-FL differs. \n\nFigure 2, and Figure 3 shows considerable gain over the baselines on common FL datasets. \n\nThe method doesn't incur computation or communication overhead and leads to 2x speed up over other methods.", "weaknesses": "The approach of clustering to address heterogeneity might be practical as there are simpler methods that can achieve the same goal. \nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for the valuable comments. Please find our responses below.\n\n**Q3.1**\nThe 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.\n\n**A3.1**\nAs we also stated in our response **A2.2** to reviewer **yLvJ**, clustered federated learning is an alternative approach to tackling the challenges of data heterogeneity. There, the server clusters clients into multiple cohorts and trains personalized models reflecting these cohorts’ preference. However, this approach contradicts our aim at training a single global model with strong generalization capabilities and requires additional resources to store multiple personalized models in the server.\n\nPlease note that our method can readily be applied to the cross-silo federated learning scenarios where a server attempts to train a strong model that can be adapted to downstream tasks efficiently, involving multiple data centers such as hospitals or banks. In the cross-device federated learning settings, HiCS can prove valuable even if thousands of clients are involved as HiCS can recognize the valuable clients with class-balanced data.\n\n\n**Q3.2**\nThe paper doesn't list limitations 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.\n\n**A3.2**\nWe rely on Dirichlet distribution since using different concentration parameters helps model heterogeneity in real-world datasets. The limitation of previous FL papers is in using only one concentration parameter to generate data partitions, which is indeed unrealistic. In this paper, however, we take a step further and use multiple concentration parameters to generate data partitions. Admittedly, the number of concentration parameters alpha is still not quite sufficient to simulate full richness of real-world settings. We appreciate the reviewer pointing out this limitation – we are committed to exploring a wider range of real-world datasets in our future work.\n\n**Q3.3**\nData heterogeneity can be mitigated with pre-training and personalization. Would pre-training make clustering unnecessary?\n\n**A3.3**\nAs 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.\n\nOn another note, personalized federated learning (pFL) has a fundamentally different goal than general federated learning. In pFL, the server does not attempt to train a global model with strong generalization capacity but various personalized models based on clients’ preference, which may perform well on local data distribution. However, these personalized models cannot work when the inference data distribution is unknown or shifts from training, which requires strong generalization ability of the global model.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AOJESjuf4g", "reviewer": "Reviewer_Rb9m", "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.", "questions": "- Why the first training rounds set to be $\\lceil N / K\\rceil$ in line 235?", "limitations": "See Weaknesses.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper found an interesting relationship between the last layer's bias and the label distributions of a given client data.\n- The proposed distance metric can be used to estimate the local data heterogeneity and distinguish which one's data is more balanced.\n- The paper is well-written and easy to follow.", "weaknesses": "- 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]. \n\n\n[1] Evaluating Gradient Inversion Attacks and Defenses in Federated Learning, NeurIPS 2021.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for the valuable comments. Please find our responses below.\n\n**Q4.1**\nA 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.\n\n**A4.1**\nIn the classical federated learning, a trustworthy server is allowed to access clients’ model updates. In the proposed HiCS, no additional information is requested from the clients and thus HiCS does not cause more privacy risk than other client selection methods. Moreover, as we also discussed in response **R.1** in the **Author Rebuttal** block, HiCS can perform well in settings where clients have different objective functions such as in **FedProx** framework [1] and use adaptive optimizers such as **Adam** [2]. In the gradient inversion attacks mentioned by the reviewer, the victim model is typically trained with cross-entropy loss and SGD optimizer. The information contained in the gradients is eliminated by the Adam optimizer because varying step sizes are assigned to different components of the gradients to compute model updates.\n\n**Q4.2**\nWhy the first training rounds set to be $t \\leq \\lceil N/K \\rceil$?\n\n**A4.2**\nGlobal rounds $t = 1$ to $t =\\lceil N/K \\rceil$ are a warm-up phase for the server to collect local updates from clients and perform clustering based on the estimated data heterogeneity.  We do not require all the clients to attend local training and upload model updates because that would violate the assumption of partial participation. After $\\lceil N/K \\rceil$ rounds, all available clients in the system are selected at least once so the server is able to perform clustered sampling in the following global rounds. More details can be found in **R.2** in the **Author Rebuttal** block.\n\n\nReference:\n\n[1] Li T, Sahu A K, Zaheer M, et al. Federated optimization in heterogeneous networks[J]. Proceedings of Machine learning and systems, 2020\n\n[2] Kingma D P, Ba J. Adam: A method for stochastic optimization[J], ICLR2015", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We would like to thank the reviewers for their time and valuable comments. We have attempted to address all the points they raised. Excerpts from their most significant/repeated questions, followed by our responses, are below. Please note that the **tables** reporting new experimental results are in the **pdf** file at the **bottom** of this block.

**W.1** :
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.

**R.1**
While the entropy of a client's label distribution alone does not reveal critical information about the client’s data (i.e., does not provide opportunities for privacy attacks), if combined with other information (e.g., which classes are in the dataset),  it may indeed cause privacy concerns. This, in turn, may compel the client to withhold the true entropy values, which motivated us to develop the method in the manuscript. Please note that, as emphasized in the title of Section 3.2.2, our method characterizes the level of data heterogeneity by relying on a proposed proxy rather than the actual entropy – this proxy, $\hat{H}(\mathcal{D}^{(k)})$, is computed by relying on the label distribution of client $k$ but is not identical to the true entropy. As stated in Theorem 3.3, if client $u$ has a more balanced label distribution than client $k$, $\hat{H}(\mathcal{D}^{(u)})$ is expected to be larger than  $\hat{H}(\mathcal{D}^{(k)})$. Building on this observation, we proposed a method capable of identifying the clients with large $\hat{H}(\cdot)$, implying more balanced data, without exposing the true entropy values. In fact, establishing a precise relationship between $\hat{H}(\mathcal{D}^{(k)})$ and the true entropy of client $k$ appears rather challenging.

Note that the analysis involving expressions (6)-(8) in the paper is predicated upon the assumption that an FL system is running vanilla FedAvg and using SGD optimizer, which is a setting indeed potentially vulnerable to privacy attacks due to the linear relationship between model updates and gradients. However, we also demonstrated that our method, HiCS-FL, can readily be implemented within **FedProx** [1] framework trained by **Adam** [2] optimizer, and outperform competing client selection methods (experiments on CIFAR10 and Mini-ImageNet). Adam is an optimizer which adaptively assigns various stepsizes to different gradient components according to their magnitude so the true information of label distribution contained in the gradients cannot be recovered.

Finally, please note that our method **does not** require any information beyond what is already collected by the servers in federated learning systems, where local model updates (including the output layers) are typically collected by the server for aggregation. We certainly appreciate the reviewer’s comment – further close examination of privacy issues is a major potential extension of the presented work.

**W.2**
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.

**R.2**
The purpose of the “warm-up” phase ($t \leq \lceil N/K \rceil$) is to collect updates of the output layer from all the available clients in the system in order to facilitate clustering. Although we conducted 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, as explained and illustrated by the experiments below.

In such a scenario, the warm-up phase can be implemented by the available clients at the beginning of training. The proposed hierarchical clustered sampling is then implemented only among the available clients; the available clients with more balanced data are preferred. When new clients join the system at global round $t$, the server selects these new clients at round $t+1$ to approximate their data heterogeneity. To provide further insights, we conducted additional experiments on CIFAR10; the results are reported in **Table 6**. As can be seen there, HiCS outperforms baselines that account for variations in clients’ availability.

**W.3**
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.

**A.3**
We discussed the computation and communication aspects of HiCS-FL in Appendix A.11. Compared to the random sampling, HiCS-FL’s only additional computations are in the one-shot clustering of the updates of the clients’ fully-connected layers. Since the dimension of bias in the fully-connected layers is typically $C$, the total number of classes, the computation needed for clustering is negligible. On another note, the competing client sample methods require much more computation and communication resources than HiCS-FL. Specifically, Power-of-Choice and DivFL require all clients in the systems to attend training each global round; FedCor requires all clients to conduct local training in the “warm-up” phase; Clustered Sampling conducts clustering on the entire updates of local models. These additional operations present significant computation overhead as compared to random sampling.

Reference:

[1] Li T, Sahu A K, Zaheer M, et al. Federated optimization in heterogeneous networks[J]. Proceedings of Machine learning and systems, 2020

[2] Kingma D P, Ba J. Adam: A method for stochastic optimization[J], ICLR2015

Meta-review: 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.
Decision: Accept (poster)